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
原文摘要
A single model scale challenges the flexibility of a production retrieval system: some settings need it faster, others need a smaller index, and the right trade-off changes with the workload. In the context of information retrieval (IR), a transformer-based model can be made smaller in three ways --- using fewer layers, passing fewer tokens through the upper layers, or producing a shorter embedding --- and each way saves a different resource. These options have been studied in isolation, each as its own method with its own code and training setup, which makes them hard to combine or adapt to a new model. In this work, we present Tevatron-Elastic to bring all three under one unified abstraction: a single object names any size the model can run at, and a short schedule lists the sizes to train. Training then produces one checkpoint that serves all of those sizes, and at deployment the user has the liberty to pick any of them. The same abstraction covers both retrievers and rerankers and both encoder and decoder models, as it works through interfaces that Hugging Face transformers already expose; adding a new backbone can be a configuration change rather than a model rewrite. Prior methods --- Matryoshka embeddings, early exit, 2D Matryoshka (Starbucks), and layerwise token compression (LTC) --- become special cases of our unified abstraction. With Tevatron-Elastic, we also propose a new elastic algorithm Matryoshka~LTC (MLTC) to jointly train several token-compression ratios in one retriever checkpoint. To validate our framework, we train 20 checkpoints across three backbones and two tasks: the quality curves are smooth, full-point quality remains close to that of models trained for a single size, and a controlled study confirms the wall-clock speedups. We release the framework and all checkpoints to support elastic retrieval research.
cs.CL