Tevatron-Elastic: A Unified Abstraction for Training Elastic Retrievers and Rerankers
Yu Wang, Shengyao Zhuang, Xueguang Ma, Zongyu Wu, Jimmy Lin, Vivek Srikumar, Zhichao Xu
cs.CL
2026-08-10
Three ways to shrink a retriever (fewer layers, fewer tokens, shorter embeddings) fold into one abstraction; one checkpoint serves any size at a near-zero tax.
Production retrieval systems face a real constraint: the same model needs different trade-offs in different deployments. Some settings want faster, some want a smaller index, and the right balance shifts with workload. A transformer retriever can be shrunk along three axes: fewer layers (depth), fewer tokens through the upper layers (token), and shorter embeddings (width). Each axis saves a different resource: depth and token save compute, width saves storage and search. The problem is that these three lines have each been studied alone, with their own code and training setup, so combining them or swapping backbones means rewriting code. Tevatron-Elastic pulls them under one abstraction.
The core is a frozen dataclass called Granularity whose four fields pin down any size the model can run at: exit layer (depth), dim (width, how many leading dimensions to keep), and keepratio plus poollayer (token, where to pool and what fraction of positions to keep). A set of sizes is written as a declarative schedule, for example "2:32,6:128,12:768" trains several operating points from layer 2 dim 32 up to layer 12 dim 768 at once.
During training each batch runs the backbone forward only once, then does a cheap readout for every size in the schedule: take the needed layer from the residual-stream hidden states, apply token pooling, truncate the embedding dimension, and compute the contrastive loss at that size. Because the abstraction routes through interfaces Hugging Face transformers already expose, a new backbone is a configuration change rather than new code, so the same framework covers BERT, ModernBERT, Qwen3, Llama3, and Mistral with no modeling changes.
Several prior methods become special cases of this grammar: pure early exit, Matryoshka (MRL), Starbucks-style depth-plus-width 2D Matryoshka, and LTC token compression. The framework also enables a new variant, MLTC (Matryoshka LTC), which jointly trains several token-compression ratios in one retriever checkpoint, something that previously required modifying the HF model class.
The authors trained 20 checkpoints across three backbones (BERT-base, ModernBERT-base, Qwen3-0.6B) and two tasks (retrieval and reranking), evaluated on BEIR-15 average nDCG@10 and MS MARCO MRR@10. Three core findings:
The elasticity tax, the quality a single checkpoint pays to train many sizes at once compared with a model trained for one size, is close to zero on rerankers and small on retrievers:
| Backbone | Single size (plain) | Elastic (2D) | Tax |
| BERT @12:768 | 0.461 | 0.448 | -0.013 |
| ModernBERT @22:768 | 0.476 | 0.450 | -0.026 |
| Qwen3 @28:1024 | 0.513 | 0.521 | +0.008 |
Measured speedups track the analytic cost model. On depth, a Qwen3 reranker exiting at layer 4 measured 6.86x speedup (7.0x predicted), and at layer 16 a 1.75x speedup with no quality loss. On width, cutting Qwen3 embeddings from 1024 to 64 dimensions shrinks the MS MARCO index from 36 GB to 2.3 GB while MRR@10 drops from 0.329 to 0.261. The token axis barely helps short queries (an 8-token query has little to pool) but does help long-document encoding.
For retrieval-system engineers this converges a fragmented set of elastic-compression methods into one maintainable framework. Train one checkpoint and pick the size online per workload, swap in a new backbone without rewriting modeling code, and treat several prior papers as configuration options. The near-zero quality tax means elasticity is no longer a luxury in real deployment. Code and checkpoints are released and usable directly for building production retrieval systems.
The authors are candid: the framework does not beat prior single-axis methods on absolute quality, and they do not claim that turning on all three axes is necessarily better. The training recipe is fixed (SFT from base checkpoints only, without the pre-training some prior systems use), to study elastic checkpoints in isolation. Rerankers admit no width axis (they have no embedding to truncate), the one structural exclusion in the task-by-axis matrix. The token axis currently targets global-attention backbones; extending it to sliding-window models like ModernBERT needs compatible pooling rules. All three axes at once (e.g. "24:[email protected]/20") is expressible in the grammar but the authors report no operating point for it, leaving the configuration to future work.