Scaling Sparse and Dense Retrieval in Decoder-Only LLMs
Hansi Zeng, Julian Killingback, Hamed Zamani
cs.IR
2025-02-21
UMass compares Llama-3 retrievers from 1B to 8B: sparse beats dense everywhere; distillation does not scale, and 8B sparse Lion-SP is SOTA on MS MARCO and BEIR.
Most recent LLM retrievers are dense dual encoders trained with contrastive loss. Llama2-7B bi-encoders already beat BERT retrievers, and Fang et al. documented scaling for dense models. Two pieces were missing. Sparse retrieval, which writes a high-dimensional vocabulary-space vector and can use an inverted index, had no comparable scaling study. Knowledge distillation, which is strong on small retrievers, had not been checked once the student grows past the teacher.
Sparse retrieval also hits a structural snag. Causal attention means the hidden state at position i never sees token i itself, yet sparse models project every position back into the vocabulary for term expansion. Leave the mask causal, and that projection is half-blind.
UMass Amherst ran a 2x3 grid on Llama-3 at 1B, 3B, and 8B under a fixed compute budget: sparse versus dense, and contrastive loss (CL) versus knowledge distillation (KD) versus both. Training data is MS MARCO passages (8.8M docs, 532K queries). In-domain tests are MS MARCO Dev (MRR@10) and TREC DL 19/20 (nDCG@10). Out-of-domain is nDCG@10 on 13 BEIR datasets. Fine-tuning is LoRA (r=16, α=32). Each run takes about 40-44 hours on 4xA100.
They first swap the causal mask for a bidirectional one, then run LLM2Vec-style masked next-token prediction: mask 20% of tokens, but predict the masked token from the previous position so the loss still matches causal pretraining. At most 10,000 steps on MS MARCO; Llama-3-8B finishes in about 17 hours on 2xA100.
Dense retrieval mean-pools token states into a D-dimensional vector. Sparse retrieval multiplies hidden states by the embedding table, then ReLU, max-pool, and log(1+·), with FLOP regularization to keep the vector sparse (0.05 on queries, 0.04 on documents). Relevance is a dot product in both cases.
CL is InfoNCE with hard negatives. Solo KD uses the cross-encoder ms-marco-MiniLM-L-6-v2 and MarginMSE. When CL and KD are mixed, KD switches to listwise KL and the two terms are weighted 0.5 each.
Scaling shows up clearly only under CL. KD is strong in-domain: a 1B KD model already beats an 8B CL model on TREC 19+20. On BEIR, KD overfits. At 8B both sparse-KD and dense-KD fall behind their CL counterparts. Dense+KD even drops on BEIR when going from 1B to 3B/8B.
Sparse beats dense at every objective and every size, with a larger gap out of domain. At 8B, sparse-KD is 10.5% above dense-KD on BEIR, and sparse-CL is 4.3% above dense-CL.
CL+KD is the best trade-off. At 1B and 3B it improves every set over CL-only and KD-only. Sparse CL+KD at 8B:
| Set | 1B | 3B | 8B |
| MS MARCO Dev | 0.410 | 0.417 | 0.417 |
| TREC 19+20 | 0.749 | 0.759 | 0.762 |
| BEIR | 0.535 | 0.544 | 0.552 |
Relative to CL-only, 8B sparse CL+KD is +0.5% on Dev, +3.2% on TREC, and -0.9% on BEIR. Dense CL+KD at 8B scores only 0.501 on BEIR, 6.2% below dense CL.
They also rerank each retriever's BEIR output with the teacher. Sparse-CL already beats the teacher at 3B (+1.3%) and reaches +7.4% at 8B (0.557). Dense-CL needs 8B to pass the teacher (+3.1%). Once KD is in the mix, sparse still beats the teacher; dense CL+KD lags it by 1.8%, 2.4%, and 1.4% at 1B, 3B, and 8B.
The public name is Lion. Against published SOTA:
| Model | MARCO Dev | TREC-19 | TREC-20 | BEIR-13 |
| ColBERTv2 | 0.397 | 0.750 | 0.746 | 0.499 |
| Lion-SP-1B | 0.410 | 0.747 | 0.751 | 0.535 |
| RepLLaMA | 0.412 | 0.743 | 0.721 | 0.551 |
| Lion-SP-8B | 0.417 | 0.758 | 0.766 | 0.552 |
| Lion-DS-8B | 0.417 | 0.755 | 0.759 | 0.501 |
Lion-SP-1B is +3.2% on Dev and +7.2% on BEIR versus ColBERTv2. Lion-SP-8B is +10.6% on BEIR versus ColBERTv2 and +4.1% on TREC 19+20 versus RepLLaMA. Dense 8B keeps up in-domain and then falls to 0.501 on BEIR.
Two takeaways for anyone training retrievers. Sparse retrieval on a decoder-only backbone is not a nostalgia move: it generalizes better, tolerates a noisy teacher, and at 8B can itself outrun that teacher. Distillation is not free soft labels. With a MiniLM-scale cross-encoder, once the student outgrows the teacher, KD pins the model to the in-domain distribution and hurts BEIR. Small models drink KD. Large models need CL. Mixing both covers the two regimes.
In production, sparse vectors still hit an inverted index, so there is no ANN tax. The cost is the bidirectional+MNTP adaptation. You cannot contrastively fine-tune a chat checkpoint and call it sparse retrieval.
This is a controlled comparison, not a new loss. The claims sit on Llama-3 and MS MARCO. The direction is still clear: do not assume that a bigger dense bi-encoder plus distillation is the scaling path for retrieval.
The authors used a single teacher. A larger cross-encoder could rewrite the KD scaling curve. That experiment is missing.
A fixed compute budget means the 8B model sees fewer steps or a smaller batch than the 1B model, so the curves are not "same data, bigger net." Documents are truncated at 128 tokens in training and 512 on BEIR; long-document collections may favor lexical sparse matching. Lion-SP-8B's BEIR 0.552 versus RepLLaMA's 0.551 is a 0.001 gap. The SOTA claim rests more on TREC and on not collapsing anywhere than on a wide out-of-domain lead. Dense 8B at 0.501 BEIR is essentially the 1B dense mixed-loss number (0.500); dense+mix did not scale here. Hard-negative mining is thinly specified and will be a hidden variable in reproductions.