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:

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{\sumt wt \cdot \maxj \mathbf{Q}t \cdot \mathbf{D}j}{\sumt w_t}$$

where $wt = \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.

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

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:

How it works

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

1. 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. 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. 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

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).