DataScience-UIBK/SmallReason-ColBERT-32M

SmallReason-ColBERT (32M)

An ultra-small late-interaction retriever for reasoning-intensive retrieval. 32M parameters, plus a 129-parameter query-side importance head.

21.41 mean nDCG@10 on BRIGHT — above every ≤33M ColBERT we evaluated, and within 1.21 of the 4.7× larger 150M Reason-ModernColBERT.

---

⚠ Read this before loading

This model is a ColBERT base plus a small importance head stored in importance_head/. The head is not part of modules.json, so a standard PyLate / sentence-transformers load silently ignores it and gives you the un-headed base:

How you load it	What you get	BRIGHT mean
pylate.models.ColBERT(...) — plain load	base only, head ignored, no error	19.61
WeightedColBERT.from_base(...) — see below	full model	21.41

There is no warning when the head is skipped, so if you are reproducing the paper number, use the second path. WeightedColBERT.from_base resolves the head from this repo automatically and raises if it cannot find one, so that path cannot fail silently. Pass require_head=False if you deliberately want the base.

---

Usage

The loader is a single file, weighted_colbert.py, from the companion repository.

from weighted_colbert import WeightedColBERT

model = WeightedColBERT.from_base(
    "DataScience-UIBK/SmallReason-ColBERT-32M",   # auto-detects importance_head/
    query_length=256,
    document_length=2048,
    device="cuda:0",
)

queries = ["What factors affect the number of Hadley cells a planet has, and how?"]
docs = [
    "Hadley cells are driven by differential solar heating; their number scales with "
    "planetary rotation rate and atmospheric depth.",
    "The best pasta recipe uses semolina flour and plenty of salted boiling water.",
]

q_embs, q_weights = model.encode(queries, is_query=True, return_weights=True)
d_embs = model.encode(docs, is_query=False)

for i, d in enumerate(d_embs):
    score = WeightedColBERT.weighted_maxsim(q_embs[0], q_weights[0], d)
    print(i, float(score))

weighted_maxsim implements the evaluation-time score

$$s(q,d)=\frac{\sum _t{w}_t\cdot \underset{j}{\max }{\mathbf{Q}}_t\cdot {\mathbf{D}}_j}{\sum _t{w}_t}s(q,d)\; =\; \backslash frac\{\backslash sum\_t\; w\_t\; \backslash cdot\; \backslash max\_j\; \backslash mathbf\{Q\}\_t\; \backslash cdot\; \backslash mathbf\{D\}\_j\}\{\backslash sum\_t\; w\_t\}$$

where $w_t = \sigma(\mathbf{W}\mathbf{Q}_t + b)$ is the learned per-query-token gate. The 1/\sum_t w_t factor is constant across documents for a fixed query, so it does not change ranking — it only keeps scores comparable across queries of different length.

Base only (no head)

If you want the reasoning-tuned base without the gate (19.61 on BRIGHT), load it as an ordinary PyLate ColBERT — the head files are simply unused:

from pylate import models
base = models.ColBERT("DataScience-UIBK/SmallReason-ColBERT-32M",
                      query_length=256, document_length=2048)

---

Results

BRIGHT (nDCG@10 ×100)

Evaluated with brute-force MaxSim, query_length=256 (Pony: 32), document_length=2048.

Split	upstream 32M	base (no head)	SmallReason-ColBERT
biology	28.70	33.16	34.17
earth_science	42.29	44.28	45.03
economics	17.65	20.25	19.99
psychology	21.93	24.91	24.94
robotics	18.09	18.65	18.14
stackoverflow	16.49	16.66	17.21
sustainable_living	18.64	20.11	21.07
pony	12.90	22.77	19.33
leetcode	16.15	17.40	29.98
aops	9.80	4.89	10.29
theoremqa_questions	12.51	9.04	13.00
theoremqa_theorems	2.76	3.19	3.74
Mean	18.16	19.61	21.41

The head is worth +1.80 mean nDCG@10 over the same base, concentrated in the long, symbol-dense splits: LeetCode +12.58, AoPS +5.40, TheoremQA-questions +3.96.

Reference points

Model	Params	BRIGHT mean
SmallReason-ColBERT	32M	21.41
answerai-colbert-small-v1	33M	18.49
mxbai-edge-colbert-v0-17m	17M	18.60
GTE-ModernColBERT-v1	150M	21.72
Reason-ModernColBERT	150M	21.97 (our protocol) / 22.62 (published)

NanoBEIR sanity (classical IR)

The gate is trained on long reasoning queries, so it is expected to give a little back on short keyword queries. It does, but not much:

Model	All 13	Excl. Touche-2020
upstream 32M	60.47	65.51
base (no head)	60.93	65.35
SmallReason-ColBERT	60.00	65.00

---

How it works

Three stages, on top of mixedbread-ai/mxbai-edge-colbert-v0-32m:

1. Widen the projection 64 → 128 dims. The first 64 rows are inherited; the new 64 are initialised from N(0, σ²) with σ at 10% of the original weight-matrix std — small enough to leave MaxSim ≈ unchanged at step 0, non-zero so the new channels actually receive gradient.
2. Two-stage base training — a varied-length warmup on ReasonIR-VL, then a hard-negative polish on merged ReasonIR-HQ + BGE-Reasoner. Both stages use PyLate's CachedContrastive loss over in-batch negatives.
3. Importance head — freeze the base, train a single Linear(128, 1) + sigmoid (129 parameters) to weight each query token.

The one non-obvious trick

The head is trained against the un-normalised weighted score Σ w_t · max_j(Q_t · D_j) but evaluated against the length-normalised one.

This asymmetry is the single most consequential choice in the recipe. Train against the normalised score instead and the per-pair score difference is bounded by one token's cosine range, the cross-entropy gradient collapses, the loss stalls near ln 2, the gates never leave their initialisation — and BRIGHT drops by 3.59 nDCG@10.

The head is initialised W = 0, b = 5, so every gate starts at σ(5) ≈ 0.993 and the head is a no-op against the frozen base at step zero.

What the head actually learns

Not soft-IDF. Across ~199K BRIGHT query tokens the gate–IDF Spearman correlation is ρ = −0.02 — statistically detectable, practically zero. Per-split mean gate sits in 0.43–0.47 with std ≈ 0.10: the head is a soft re-weighting, not a selector. A fixed IDF gate on the same base reaches only 20.06, against 21.41 for the learned head.

---

Training

	Warmup	Polish	Head
Data	ReasonIR-VL (~245K)	merged ReasonIR-HQ + BGE-Reasoner (~2.7M)	same merged set
Loss	CachedContrastive	CachedContrastive	CE over [s_pos, s_neg]
LR	1e-5	5e-6	5e-4 (AdamW, wd=0)
Batch	32/GPU × accum 4 × 8 GPU	32/GPU × accum 2 × 8 GPU	16 triples/step, 1 GPU
Steps	1 epoch (~8 h)	1 epoch (~16 h)	3,000 steps (~12 min)
Lengths	q 256 / doc 2048	q 256 / doc 2048	q 256 / doc 2048
Precision	bf16 + FA2	bf16 + FA2	fp32 head, frozen bf16 base

Base training: 8× H100 across two nodes, ~24 h total. Head training: one H100, ~12 min.

---

Limitations

• Scale. The recipe was developed and validated at 32M. It does not transfer for free — the same head at 17M gives no gain.
• Frozen base. The head is trained on a frozen base; joint fine-tuning is unexplored.
• Late-interaction cost. The head is nearly free, but the model still carries multi-vector storage and scoring costs. The efficiency claim is about parameter count, not about matching single-vector retrieval.
• Short queries. Pony (32-token queries) regresses relative to the un-headed base — a per-token gate needs tokens to discriminate between.
• Oblique queries. On OBLIQ-Bench (stance / intent / tip-of-the-tongue) the model is near zero (mean 3.66) and is beaten by every baseline there. Reported as a deliberate negative result; embedding similarity is the wrong tool for that class of query.
• Synthetic teacher data. Training data is synthetic with cross-encoder-mined hard negatives; biases in that mining can propagate.

License

CC-BY-NC-4.0, inherited from the ReasonIR and BGE-Reasoner training data. The upstream base model (mixedbread-ai/mxbai-edge-colbert-v0-32m) is Apache-2.0, and the companion training/inference code is released under Apache-2.0 — but these weights are non-commercial.

Citation

@inproceedings{smallreason-colbert,
  title     = {SmallReason-ColBERT: An Ultra-Small Late-Interaction Retriever
               for Reasoning Intensive Retrieval},
  author    = {Abdallah, Abdelrahman and Ali, Mohammed and Jatowt, Adam},
  booktitle = {Proceedings of the 2026 Conference on Empirical Methods in
               Natural Language Processing (EMNLP)},
  year      = {2026}
}

Acknowledgements

Thanks to Antoine Chaffin (LightOn, Reason-ModernColBERT) for flagging the upstream 2_Dense/use_residual config bug in mxbai-edge-colbert-v0-32m — the base weights were trained with a residual on that layer while the shipped config said otherwise. This model uses the patched config (use_residual: true).