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
检索模型三种变小法（少层／少 token 过上层／短向量）统一进一个 Granularity 抽象，一个检查点部署成任意大小，弹性训练代价接近零。
生产级检索系统面对一个现实：同一个模型在不同的部署里要的权衡不一样。有的场景要更快，有的要更小的索引，正确的取舍还随负载变化。一个 transformer 检索模型能从三个轴变小：少跑几层（深度）、让上层少处理几个 token（词）、输出更短的向量（宽度），每个轴省的是不同的资源（深度和词省算力，宽度省存储和搜索）。问题是这三条线此前各做各的，各有各的代码和训练流程，想组合或换主干就得重写。Tevatron-Elastic 把它们拢到一套抽象下。
核心是一个 frozen dataclass 叫 Granularity，四个字段就标定了模型能跑的任意一个大小：exit layer（深度）、dim（宽度，保留前多少维）、keepratio 加 poollayer（词，在多少层池化、保留多少比例位置）。一组大小写成一个声明式 schedule，比如 "2:32,6:128,12:768" 表示训练时同时覆盖从第 2 层 32 维到第 12 层 768 维的多个工作点。
训练时每个 batch 只跑一次主干前向，然后对 schedule 里每个大小做廉价 readout：从 residual stream 的 hidden states 里按需截取层、做 token 池化压缩、截取向量维度，得到该大小下的表示再算对比损失。这条设计的好处是抽象全走 Hugging Face transformers 已有的接口，换主干只是配置改动不是新代码，所以同一套框架无改代码地覆盖了 BERT、ModernBERT、Qwen3、Llama3、Mistral。
以往的几种方法都成了这套语法的特例：纯多层退出、Matryoshka（MRL）、Starbucks 那种深度加宽度 2D Matryoshka、LTC 词压缩。框架还顺手做出一个新东西 MLTC（Matryoshka LTC），在一个检索检查点里联合训练多个 token 压缩比，这在没有统一接口前要改 HF 模型类才能实现。
作者训了 20 个检查点，横跨三个主干（BERT-base、ModernBERT-base、Qwen3-0.6B）和检索／重排两类任务，在 BEIR-15 平均 nDCG@10 和 MS MARCO MRR@10 上评。三条核心结论：
弹性税，也就是一个检查点同时训多个大小、跟专门为单一大小训的模型相比要付多少质量代价，在重排器上几乎为零，在检索器上也很小：
| 主干 | 单一大小（plain） | 弹性（2D） | 质量代价 |
| 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 |
实测加速跟解析成本模型对得上。深度方向，Qwen3 重排器在第 4 层退出实测 6.86 倍加速（预测 7.0 倍），第 16 层 1.75 倍且质量不掉。宽度方向，把 Qwen3 向量从 1024 维砍到 64 维，MS MARCO 索引从 36 GB 缩到 2.3 GB，MRR@10 从 0.329 降到 0.261。词方向对短查询基本帮不上忙（8 个 token 的查询没多少可池化），对长文档编码有帮助。
做检索系统工程的人会发现这把碎片化的弹性压缩方法收敛成一套可维护的框架：训一个检查点就能在线上按负载现挑大小，换新主干不用重写模型代码，以往几篇独立工作成了配置项。质量代价小到可以忽略意味着「弹性」在实际部署里不再是奢侈品。代码和检查点都放了，可直接拿来搭生产检索系统。
作者很坦白：这套框架在绝对质量上不会超过以往单独的方法，也不主张三个轴一起开就一定更好。训练 recipe 是固定的（只从 base 检查点做 SFT，不带某些以往系统用的预训练），目的是单独研究弹性检查点。重排器没有宽度轴这个选项（它没有能截断的 embedding），这是任务乘轴矩阵里唯一的结构性排除。词压缩目前只针对全局注意力的主干，要扩到 sliding-window 的 ModernBERT 还得配兼容的池化规则。三个轴同时开（比如 "24:[email protected]/20"）语法上能表达，但作者没给出实测工作点，把配置留给后续。