Tevatron-Elastic：训练弹性检索器与重排序器的统一抽象

# Tevatron-Elastic：训练弹性检索器和重排序器的统一抽象

Yu Wang[1], Shengyao Zhuang[2], Xueguang Ma[3], Zongyu Wu[4], Jimmy Lin[3], Vivek Srikumar[1], Zhichao Xu[1]  
[1] University of Utah, [2] The University of Queensland, [3] University of Waterloo, [4] Pennsylvania State University

###### 摘要

单一模型规模对生产级检索系统的灵活性构成挑战：有些场景需要更快，另一些需要更小的索引，而恰当的权衡随工作负载变化。在信息检索（IR）背景下，基于 transformer 的模型可以通过三种方式变小——使用更少的层、让更少的 token 通过上层、或生成更短的嵌入——每种方式节省不同的计算资源。这些选项此前曾被逐一研究，各自拥有独立的方法、代码和训练设置，因此难以组合或适配新模型。我们提出 **Tevatron-Elastic**，将三者统一到一个简单抽象之下：一个对象即可命名模型可运行的任意规模，一个简短的 schedule 列出要训练的规模。训练产生的单一 checkpoint 能服务所有这些规模，部署时用户可任选其一。同一抽象同时覆盖检索器与重排序器，以及 encoder 与 decoder 模型，因为它通过 Hugging Face transformers 已经暴露的接口工作；更换骨干网络只是配置更改，而非新的建模代码。先前的方法——[Matryoshka embeddings](https://arxiv.org/html/2608.08809#bib.bib13)（Kusupati et al., 2022）、[early exit](https://arxiv.org/html/2608.08809#bib.bib17)（Liu et al., 2020；[Xin et al., 2020](https://arxiv.org/html/2608.08809#bib.bib25)）、[2D Matryoshka](https://arxiv.org/html/2608.08809#bib.bib38)（如 Starbucks，Zhuang et al., 2026a）和 [layerwise token compression](https://arxiv.org/html/2608.08809#bib.bib39)（LTC，Zhuang et al., 2026b）——成为我们统一抽象的特例。同一接口还支持 Matryoshka LTC（MLTC），它能在单个检索器 checkpoint 中联合训练多个 token 压缩比。为验证我们的框架，我们在三个骨干网络和两个任务上训练了 20 个 checkpoint：质量曲线平滑，一个 checkpoint 相对单规模专用模型只增加少量开销，受控研究证实了墙钟加速。我们发布框架及所有 checkpoint，作为构建弹性检索系统的资源。

## 1 引言

基于 transformer 的信息检索（IR）模型（[Lin et al., 2022](https://arxiv.org/html/2608.08809#bib.bib16)；[Xu et al., 2026b](https://arxiv.org/html/2608.08809#bib.bib31)）有三个在很大程度上可分离的轴，控制推理时的成本：

- **depth**——在读出之前执行多少层 transformer。层数越少，encoder/scorer 计算越少，延迟越低（[Devlin et al., 2019](https://arxiv.org/html/2608.08809#bib.bib2)；[Liu et al., 2020](https://arxiv.org/html/2608.08809#bib.bib17)；[Xin et al., 2020](https://arxiv.org/html/2608.08809#bib.bib25)；[Warner et al., 2024](https://arxiv.org/html/2608.08809#bib.bib23)）。
- **token**——有多少序列位置流经上层。在堆叠中间某个位置对序列做池化，可减少上层浮点运算（FLOPs）（[Goyal et al., 2020](https://arxiv.org/html/2608.08809#bib.bib8)；[Zhang et al., 2025](https://arxiv.org/html/2608.08809#bib.bib37)）。
- **width**——输出嵌入的维度。更小的嵌入会减小存储索引的大小以及最近邻搜索的成本（[Johnson et al., 2017](https://arxiv.org/html/2608.08809#bib.bib12)；[Kusupati et al., 2022](https://arxiv.org/html/2608.08809#bib.bib13)）。

给这些轴分组的一个有用心智模型是看每个轴降低的部署成本：depth 和 token 降低*计算*成本，而 width 降低*存储和搜索*成本。depth 和 token 在计算时起作用，索引保持不变；width 在存储时起作用，而逐输入的计算量不变。

先前关于弹性建模的工作往往只专注于单一轴。Matryoshka 表示学习（MRL）改变 width（[Kusupati et al., 2022](https://arxiv.org/html/2608.08809#bib.bib13)）；跳层和 early-exit 方法改变 depth（[Liu et al., 2020](https://arxiv.org/html/2608.08809#bib.bib17)；[Xin et al., 2020](https://arxiv.org/html/2608.08809#bib.bib25)；[Elhoushi et al., 2024](https://arxiv.org/html/2608.08809#bib.bib4)）；[Zhang et al. (2025)](https://arxiv.org/html/2608.08809#bib.bib37) 和 [Zhuang et al. (2026b)](https://arxiv.org/html/2608.08809#bib.bib39) 在 token 轴上压缩检索器和重排序器；2D Matryoshka 检索器如 2DMSE 和 Starbucks 则同时改变 depth 和 width（[Li et al., 2024](https://arxiv.org/html/2608.08809#bib.bib15)；[Zhuang et al., 2026a](https://arxiv.org/html/2608.08809#bib.bib38)）。上述每种方法在其自身设定下都有效，但每种都是独立实现，带有各自的模型类定义和训练目标，因此组合它们或将其适配到新任务/骨干网络都需要重新实现。

我们认为，这些压缩轴是互补的而非竞争的：由于它们降低不同的成本，因此可以独立应用并组合。然而，现有工作没有提供统一抽象，来统一不同压缩模型类的定义以及相应的训练和推理。

看到这一空白，我们提出 **Tevatron-Elastic**，基于 [Tevatron](https://arxiv.org/html/2608.08809#bib.bib18) 为这些轴提供统一抽象（[Ma et al., 2025](https://arxiv.org/html/2608.08809#bib.bib18)；[Xu et al., 2026a](https://arxiv.org/html/2608.08809#bib.bib30)）。我们首先将压缩模型的动作抽象为一个**工作点**（operating point），例如，在某个层提前退出，就是为该层定义 early exit 工作点。基于这一定义，Tevatron-Elastic 引入了两个核心组件：一个冻结的 dataclass `Granularity`，用每个轴一个可选字段命名任意工作点；以及 `GranularitySchedule`，列出要训练底层弹性模型的工作点。具体来说，训练循环在单次骨干网络前向传播中，对 schedule 中的每个粒度求和损失，从而产生一个能服务所有列出工作点的 checkpoint。该抽象不修改模型内部，而是通过每个 Hugging Face transformer 都暴露的 hidden-states 元组、层模块列表和池化步骤进行路由，因此更换骨干网络只是一个启动参数，而非新代码。

我们的核心贡献是容纳不同类型压缩策略的统一抽象：

1. 一个统一抽象横跨两个任务（检索、重排序）、三个轴和五个骨干家族（BERT、ModernBERT、Qwen3、Llama3、Mistral），每个骨干无需新模型代码（[第 3 节](https://arxiv.org/html/2608.08809#S3)）；
2. 它复现了先前的弹性方法：MRL、early exit、2D Matryoshka、Starbucks 和 LTC 均作为配置出现（[第 4 节](https://arxiv.org/html/2608.08809#S4)）；
3. 联合训练的工作点表现符合预期——我们实证研究中的质量曲线平滑，联合训练相对专用单点模型的开销很小（[第 4 节](https://arxiv.org/html/2608.08809#S4)）。此外，受控研究显示实测推理加速符合解析成本模型（[第 4.3 节](https://arxiv.org/html/2608.08809#S4.SS3)）。
4. 它引入了 **Matryoshka Layerwise Token Compression**（MLTC），作为 [LTC](https://arxiv.org/html/2608.08809#bib.bib39)（[Zhuang et al., 2026b](https://arxiv.org/html/2608.08809#bib.bib39)）和 [Zhang et al. (2025)](https://arxiv.org/html/2608.08809#bib.bib37) 的扩展，在单个 checkpoint 中联合训练多个 token 压缩比。MLTC 使用框架现有的 schedule 和 token 读出路径；若不支持中间堆叠的多处池化点，则需要修改模型类。

我们指出，Tevatron-Elastic 并未在绝对质量上超越先前方法，也不意味着组合全部三个轴一定优于只使用一个轴。在我们的实证实验中，我们固定训练方案（从基础 checkpoint 进行监督微调，而不使用某些先前系统使用的预训练，如 [Gao and Callan, 2022](https://arxiv.org/html/2608.08809#bib.bib7)；[Zhuang et al., 2026a](https://arxiv.org/html/2608.08809#bib.bib38)；[Xu et al., 2026f](https://arxiv.org/html/2608.08809#bib.bib35)），以便隔离地研究训练得到的弹性 checkpoint。我们计划发布框架和所有 checkpoint，帮助研究者复现现有弹性模型方案并开发新方案。

## 2 相关工作与背景

我们先沿三个压缩轴回顾 IR 模型压缩工作（[第 2.1 节](https://arxiv.org/html/2608.08809#S2.SS1)），然后定义本文其余部分使用的记号（[第 2.2 节](https://arxiv.org/html/2608.08809#S2.SS2)）。

### 2.1 既有方法

我们关注使检索或排序模型具有*弹性*——训练后成本可调——的工作，并按其所压缩的轴分组。

#### 宽度。

Matryoshka 表示学习训练一种嵌入，其前导坐标可以单独使用，因此一个模型可服务多种向量维度，在检索中即可服务多种索引大小（[Kusupati et al., 2022](https://arxiv.org/html/2608.08809#bib.bib13)）。截断发生在输出端，因此 width 改变存储和最近邻搜索成本（[Johnson et al., 2017](https://arxiv.org/html/2608.08809#bib.bib12)），但不改变逐输入的计算量。

#### 深度。

Early-exit 和跳层（layer-dropping）方法在中间层读出表示，以质量换取更少的层计算，既可以是固定退出，也可以是输入自适应的退出（[Liu et al., 2020](https://arxiv.org/html/2608.08809#bib.bib17)；[Xin et al., 2020](https://arxiv.org/html/2608.08809#bib.bib25)；[Elhoushi et al., 2024](https://arxiv.org/html/2608.08809#bib.bib4)）。在 IR 中，这会得到更浅的 encoder（检索器）或 scorer（重排序器）。

#### Token。

序列压缩方法在堆叠内部缩短 token 序列，使上层处理更少的位置（[Goyal et al., 2020](https://arxiv.org/html/2608.08809#bib.bib8)）。先前工作 Jasper-Token-Compression-600M（[Zhang et al., 2025](https://arxiv.org/html/2608.08809#bib.bib37)）和 Layer-wise token compression（LTC，[Zhuang et al., 2026b](https://arxiv.org/html/2608.08809#bib.bib39)）分别在所选层对 token 序列做池化，并在缩短后的序列上运行其余层，分别用于检索和重排序。

#### 组合。

少数方法组合了多个轴。2D Matryoshka 嵌入同时改变 depth 和 width（[Li et al., 2024](https://arxiv.org/html/2608.08809#bib.bib15)），Starbucks 训练一个对两者均弹性的检索器，因此一个 checkpoint 可服务 \((layer, dim)\) 的网格（[Zhuang et al., 2026a](https://arxiv.org/html/2608.08809#bib.bib38)）。这些是与我们最接近的系统；但它们各自固定了一对轴和一个任务。

#### 我们的定位。

在现有工作中，每个轴（或固定的一对轴）都作为独立方法实现，带有自己的模型类定义、实现和训练策略，而且几乎总在单一任务上。我们不引入新的池化操作；相反，我们提供一个统一抽象，在其中 depth、token 和 width 被统一命名，要训练的工作点集合被写成数据，而上述方法——MRL、early exit、LTC 和 2D Matryoshka——在检索和重排序中都以配置形式出现。每个工作点上的目标本身是可插拔的，可以是标准的对比损失、排序学习损失（[Cao et al., 2007](https://arxiv.org/html/2608.08809#bib.bib1)；[Xia et al., 2008](https://arxiv.org/html/2608.08809#bib.bib24)；[Oord et al., 2018](https://arxiv.org/html/2608.08809#bib.bib19)），或蒸馏（[Hinton et al., 2015](https://arxiv.org/html/2608.08809#bib.bib10)；[Xu et al., 2025b](https://arxiv.org/html/2608.08809#bib.bib28)）（[第 3 节](https://arxiv.org/html/