← 全部深度解读
📖 深度解读 · AI SEARCH

Layer-wise Token Compression for Efficient Document Reranking

Shengyao Zhuang, zhichao Xu, Ivano Lauriola · Amazon · arXiv 2605.20683 · 发表 2026-05-21
8 分钟阅读时长 3096字 4.70综合分 第 1 名05-22 AI Search 日报
目录 · 6 节1问题与动机2方法详解3核心技术贡献4实验结果分析5局限与未来方向6工业实践意义
相关 5新颖 4实用 5证据 4时效 5
双塔检索常用的输入层 token 压缩对 cross-encoder 无效,把压缩挪到中间层是关键发现;更意外的是用短 passage 训练、压缩后的模型在长文档排序上反超未压缩版本,说明压缩起到了鼓励长度无关表征的正则作用。对线上重排降本与跨长度泛化都有借鉴价值。

1. 问题与动机

Transformer cross-encoder 重排序器是现代 IR 系统的核心组件,但要把 query 与 document 拼成长序列一起过 Transformer,推理成本高、QPS 受限。一个已知的提速思路是 token compression——在 embedding 层把相邻 token 聚合,减少有效 token 数。这一招在双塔(bi-encoder)检索器上被证明有效,但作者通过实验发现:直接搬到 cross-encoder 上会失效。证据很硬:沿用 Jasper 的 embedding 层压缩虽然带来 25% 的 QPS 提升,却让 nDCG@10 从基线 0.727 暴跌到 DL19 的 0.5832、DL20 从 0.665 跌到 0.4975,约 20% 的效果损失。

作者给出的根因假设是:cross-encoder 的相关性信号来自 query/document token 之间细粒度的交互,而这些交互从最早的 Transformer 层就开始发生;在 embedding 层就把 token 压掉,等于在交互发生之前抹掉了信号。双塔则是独立编码、只在向量层比相似度,所以对 token 压缩不敏感。这个观察直接引出本文的设计动机:压缩不能在交互之前做,而要等交互充分进行之后再做。

2. 方法详解

LTC(Layer-wise Token Compression)的核心是把 1D 自适应平均池化从 embedding 层挪到中间 Transformer 层。形式化为 H_compressed^(l) = C(H^(l), r),输出 token 数 n' = ⌊n·r⌋,压缩比 r ∈ (0,1]。具体实现:在目标层 l∈{1,…,L} 之前的 1 到 l-1 层处理完整序列;从 l 层起,沿 token 维做自适应平均池化,把相邻 token 表征合并,同时更新 attention mask 与位置索引;l 到 L 层在压缩后的短序列上计算。池化用「相邻 token 平均」而非丢弃,尽量保留语义结构。

对 listwise reranker,作者额外引入 document mask M∈{0,1}^(n×k),只允许文档内池化、禁止跨文档压缩,并保护 instruction prompt 与 query 不被压。

训练设置很关键:pointwise 用 Qwen3-0.6B-Base(28 层 decoder + 分类头),在 MS MARCO passage(约 880 万 passage、50 万训练 query,用 Tevatron 取 hard negative)上训练,损失为 1 正 + 5 hard negative 的交叉熵,3 epoch、batch 32、lr 1e-5、AdamW,且压缩在训练阶段就施加以让模型适应。listwise 用 Mistral-7B-Instruct-v0.3(取首生成 token 的 logits),RankZephyr/FIRST 数据(每 query 20 篇 GPT-4 排序文档),语言建模 + ListMLE 联合损失、LoRA rank 16。QPS 在单张 NVIDIA L40S 上测。

一个非常重要的非显然结论:压缩必须「compression-aware training」。把压缩只在推理时加到未经压缩训练的模型上会灾难性掉点——r=0.2、layer 2 时 nDCG@10 从 0.728 跌到 0.287(相对下降 60%),说明未经训练的中间层表征不具备对 token 缩减的鲁棒性。

3. 核心技术贡献

第一,定位并解释了「embedding 层压缩对 cross-encoder 失效」这一现象,给出 token 交互时机的机理解释,纠正了从双塔直接迁移的错误直觉。第二,提出中间层自适应池化 LTC,把压缩点作为可调超参 l* 与压缩比 r 的二维设计空间,实验找到 layer 8–14 的「甜区」:既保排序质量又显著提速。第三,把方法无缝扩展到 listwise LLM reranker,用 document mask 实现文档内压缩、保护 prompt/query,长上下文 listwise 收益更大(27–73% QPS)。第四,意外发现压缩在长文档泛化上起到正则作用——用短 passage 训练、压缩后的模型迁移到长文档排序时反超未压缩版本。

4. 实验结果分析

Pointwise passage(Qwen3-0.6B,TREC DL19/DL20):基线 DL19 nDCG@10=0.727、QPS=5.35。最优 LTC(r=0.4、layer 8)DL19 nDCG@10=0.7274 与基线持平,QPS 5.35→6.69,提速 25%。消融揭示规律:早层压缩伤害大(layer 2、r=0.2 跌到 0.638,降 12.2%,p<0.001);中层 8–14 中等压缩保质量且显著提速;后层 14–20 轻压缩(r=0.8)甚至略超基线(0.7314 vs 0.727,不显著),印证正则效应。

Pointwise document(512 token 截断):r=0.4、layer 20 在 DL19 上 nDCG@10=0.663,反超基线 0.651——压缩鼓励长度无关表征。文档排序因序列更长,QPS 提升最高可达 116%。

BEIR 六数据集(NFCorpus、FiQA、ScIDocs、SciFact、TREC-COVID、ArguAna,BM25 top-100 重排)零样本:无压缩 avg nDCG@10=0.422;r=0.8/layer20 为 0.424(+9% QPS);r=0.6/layer14 为 0.420(+18%);r=0.4/layer8 为 0.416(+25%),最激进压缩也仅相对降 1.4%。Listwise(Mistral-7B)passage 在无显著掉点下获 27–73% QPS 提升,document 提速空间更大。

最值得警惕的反面结果是 zero-shot 压缩的崩盘(0.728→0.287),它把「必须压缩感知训练」这条工程红线划得很清楚。

5. 局限与未来方向

作者明确列出的局限:(1) 目标层 l* 与压缩比 r 目前靠人工 sweep,未来可用 NAS 自动选择以降低部署成本;(2) 基线对比偏窄,主要对 Jasper,缺与 MRL-Reranker、prompt compression、KV-cache 压缩、early-exit 等的系统比较;(3) 激进压缩依赖压缩感知微调,完全 training-free 的变体仍开放;(4) 正则化假设还需表征级分析(隐状态范数、attention 稀疏度随长度变化)与 runtime 拆解来确证。除作者所述,评测主要在 MS MARCO/BEIR 学术集,缺真实线上 A/B 与多语种、短文本 UGC 场景验证。

6. 工业实践意义

对搜索重排的直接启示:cross-encoder 想降本,不要在输入层砍 token,而应在中间层(经验上 8–14 层)做自适应池化,并务必在训练阶段就加入压缩。落地收益清晰——passage 重排 +25% QPS、长文档重排最高 +116% QPS,且质量基本无损;这对 query-document 序列偏长、重排算力吃紧的线上系统很有吸引力。压缩比 r 与压缩层 l* 构成「质量-时延」可调旋钮:要更激进提速就早压、要更稳就晚压且轻压。

值得借鉴的还有两点:其一,「压缩即正则」的现象提示,用短文本训练、压缩部署,可能换来对长文档/长 query 更好的泛化,适合 UGC 长短文混合的检索场景;其二,listwise LLM reranker 用 document mask 做文档内压缩、保护 instruction 与 query,是把 LLM 重排搬上线时控成本的实用模式。部署注意:必须重训(compression-aware),不能对存量模型推理期直接加压;层/比例需离线 sweep 选定,建议结合自身延迟预算与 nDCG 容忍度确定甜区。

← 全部深度解读 出现于报告:2026-05-22 · AI Search #1 (4.70) 原文 ↗