Reason-mxbai-colbert-v0.1-32m

v0.1 of the Reason-mxbai-colbert series — same edge-scale late-interaction retriever as v0, retrained with the correct projection-head architecture (use_residual: true in the 2_Dense layer). Mean BRIGHT nDCG@10 improves from 19.00 → 19.61 (+0.61), with the largest gains on the natural-language and formal-math splits that v0 struggled with.

What changed vs v0

The v0 release inherited use_residual: false in the 2_Dense (768→768) layer from the upstream mixedbread-ai/mxbai-edge-colbert-v0-32m config. That turned out to be a shipped-config bug — the base weights were trained with a residual connection on 2_Dense, but the config flag said otherwise. PyLate respects the flag, so every downstream fine-tune was working with a silently-broken architecture.

v0.1 fixes this by:

1. Copying the base model and flipping 2_Dense/config.json: "use_residual": true.
2. Re-widening the 3_Dense projection from 64 → 128 dims on top of the corrected base (small-random init for the new channels).
3. Retraining the full curriculum (Stage 1 VL warmup → Stage 2 BGE-reasoner + HQ-hn polish) on the corrected architecture.

Under the corrected architecture, stage-2 training loss drops from 38.73 → 31.71 (−7.02) — real learning. The v0 equivalent run under the broken architecture was stuck at loss ~45-51 and never moved the new-dim weights.

BRIGHT results

Evaluated with MTEB BrightRetrieval (brute-force MaxSim, query_length=256, document_length=2048).

8 wins / 4 losses. Highlights:

• Natural-language splits (psychology +2.29, economics +1.54, pony +2.03) — the correct architecture lets the model learn meaningfully richer representations.
• Formal-math splits (theoremqa_theorems +0.94, theoremqa_questions +0.66) — previously stuck near the floor, now moving.
• Small code-split regressions (leetcode −0.28, stackoverflow −0.12) are within noise.

Model Details

• Model Type: PyLate ColBERT (late-interaction, multi-vector)
• Base model: mixedbread-ai/mxbai-edge-colbert-v0-32m (with 2_Dense/use_residual: true patched to match trained weights)
• Parameters: ~32M backbone + widened 128-dim projection
• Document Length (training): 2048 tokens
• Query Length (training): 256 tokens
• Output Dimensionality: 128 per token (widened from base's 64)
• Similarity Function: MaxSim
• Training Data: hanhainebula/bge-reasoner-data + reasonir/reasonir-data (VL split for warmup, HQ with hard negatives for polish)
• Language: en

Model Architecture

ColBERT(
  (0): Transformer({'max_seq_length': 127, 'do_lower_case': True}) with ModernBertModel
      hidden_size=384, num_hidden_layers=10, num_attention_heads=6,
      position_embedding_type='sans_pos', max_position_embeddings=7999
  (1): Dense(384 → 768, bias=False, use_residual=False)
  (2): Dense(768 → 768, bias=False, use_residual=True)     # ← THE FIX (was False in v0)
  (3): Dense(768 → 128, bias=False, use_residual=False)    # widened from 64 → 128 on correct base
)

Usage

from pylate import indexes, models, retrieve

model = models.ColBERT(model_name_or_path="DataScience-UIBK/Reason-mxbai-colbert-v0.1-32m")

# Retrieval (see https://lightonai.github.io/pylate/ for full API)
index = indexes.Voyager(index_folder="pylate-index", index_name="index", override=True)
docs = ["document 1 text", "document 2 text"]
doc_embs = model.encode(docs, is_query=False, batch_size=32, show_progress_bar=True)
index.add_documents(documents_ids=["1","2"], documents_embeddings=doc_embs)

retriever = retrieve.ColBERT(index=index)
q_embs = model.encode(
    ["Given a Psychology post, retrieve relevant passages that help answer the post.\nQuery: why do I procrastinate?"],
    is_query=True,
)
scores = retriever.retrieve(queries_embeddings=q_embs, k=10)

Sentence Transformers

As of Sentence Transformers v6.0.0, this model also loads directly as a multi-vector (ColBERT-style late interaction) retriever via the MultiVectorEncoder:

pip install "sentence-transformers>=6.0.0"

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("DataScience-UIBK/Reason-mxbai-colbert-v0.1-32m")

query = "Given a Psychology post, retrieve relevant passages that help answer the post.\nQuery: why do I procrastinate?"
documents = [
    "Procrastination is often driven by difficulty regulating negative emotions around a task, not laziness, since delaying provides short term relief.",
    "The hippocampus plays a central role in consolidating short term memories into long term storage.",
]

query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings[0].shape)
# torch.Size([28, 128]) torch.Size([31, 128])

# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[25.2429, 24.4635]], device='cuda:0')

Acknowledgements

Thanks to Antoine Chaffin (LightOn / Reason-ModernColBERT author) for flagging the upstream 2_Dense config bug in mxbai-edge-colbert-v0-32m — without that heads-up, v0.1 wouldn't exist.