CompRank:用 token 压缩 + decoding-free 打分把 LLM listwise 重排做轻
1. 问题与动机
LLM listwise 重排序在候选列表变长时成本急剧上升,两大瓶颈:一是逐 token 自回归解码生成文档 ID,列表越长越慢;二是把全部文档 token 都做稠密注意力,存在大量冗余。作者的核心观察是"相关性判断往往只依赖一小部分信息量大的文档 token",其余 token 对排序贡献有限——排序信号本质稀疏,却用稠密计算去算,二者不匹配。在 top-100 乃至 top-500 重排场景下冗余尤为致命。更糟的是生成式方法在训练列表长度(30 篇)外推到 500 篇时大面积出错:100% 的 query 出现"漏 ID"(平均每条漏 78.2 个),66% 出现重复 ID,NDCG@10 从 75.8 崩到 59.7。CompRank 要在保住排序质量的同时,把重排成本压到可上线水平。
2. 方法详解
CompRank 由四部分组成。(1) 文档表示解耦:候选排成独立 block [D1;D2;…;Dn;Bq],每个文档 token 用文档内局部位置 pos=t(与文档序号无关),query block 加常数偏移 C+t;注意力上做文档内因果自注意力、禁止文档间互相注意、query 注意所有文档。如此每个文档的 KV 状态 Ki,Vi=fθ(Di) 与其他候选无关、可复用,复杂度从 O((NL)²d) 降到 O(NL²d)。(2) token 级压缩,仅作用于 query→document 注意力:对比了 Last-k(保留末 k 个 token)、Segment-wise(每隔 k 个位置取 1 个,默认 Step-10 保留约 10.2% token)、以及基于文档内注意力 αi,t 的 Intrinsic Importance(top-k + anchor)。文档表示本身不变,只压缩 query 能看到的 token。(3) decoding-free 打分:在打分层 ℓ*(Mistral-7B 取第 20 层)、决策 token u(首个输出 token '[')上,对保留 token 算注意力 logit aᵤ,ₜ=QᵤKₜᵀ/√d+mᵤ,ₜ,softmax 后把文档 i 上的注意力质量聚合成分数 sᵢ,按 sᵢ 降序排序,全程不生成 token。(4) CopyNet 式训练目标:不做生成文档 ID 的 SFT,而优化指针分布 P_copy(i)=exp(sᵢ/τ)/Σexp(sⱼ/τ)(τ=0.05),损失为 -log P_copy(正样本)。训练用 Mistral-7B 全参微调、MS MARCO 5 万样本(每 query 30 篇候选)、8 张 H20、lr 3e-6、bf16、1 epoch、文档截断 256 token。
3. 核心技术贡献
四点贡献各有机制支撑。(1) Block 解耦去掉文档间注意力,使候选数 N 上的复杂度从二次降到线性,并让文档 KV 可缓存复用——这是 RAG/agentic 反复重排同一批文档时省算力的基础。(2) Segment-wise 压缩理论上带来约 1.43×(3/(2+ρ),ρ≈0.1)的注意力提速;实验证明均匀采样(Step-10)优于只取末段(Last-k)或纯按注意力取(Top-k),因为文档内显著性≠对 query 的相关性,纯 Top-4~16 只有 35.8–36.8 NDCG。(3) decoding-free 打分用注意力质量直接出分,避开自回归生成,带来对生成式基线 4.9–9.5× 的提速,且输出空间不随列表长度漂移。(4) CopyNet 目标在"推理时真正用到的注意力分数"上做监督,预测空间是候选级而非词表级,因此 30 篇训练能稳定外推到 500 篇、零解码错误(生成式在 500 篇时 100% query 漏 ID)。
4. 实验结果分析
主结果在 7 个 BEIR 数据集(FiQA、SciFact、TREC-COVID、DBPedia、Climate-FEVER、NFCorpus、SciDocs)上:CompRank 全 token 上界 39.7,Segment-10 用 10.2% token 拿到 39.2,仅掉 0.5,已超过 ICR(38.4)、Corehead(38.1)、RankGPT-Mistral(35.1)与 BM25(35.4)。压缩策略消融显示 Step-10 > Step-5(20% token,39.1,提升微弱)> Step-20 > Last-k > Top-k;Last-32 用 14.1% token 只有 38.0,说明末段偏置不如均匀采样。外推实验最有说服力:30 篇训练、评测到 500 篇,CompRank(Segment-10) 在 500 篇仍有 79.5 NDCG@10,而生成式 Direct-List 从 75.8 跌到 59.7。延迟上,500 篇时 Direct-List 要 114.5s,CompRank 全 token 16.22s,Segment-10 仅 12.0s(4.9–9.5× vs 生成式、1.3× vs 全 token)。延迟拆解更关键:500 篇端到端 12.0s 中文档重算(full-forward)占 11.72s,而走缓存 KV 的 query-forward 只要 0.05s——若上线做 KV 缓存理论上可再降约 234×。层与决策 token 消融显示第 20 层(68.5% 注意力准确率)和首 token '['(68.1%)远好于早层(≤20%)和 'Final'/'Answer' 模板(29–31%)。
5. 局限与未来方向
作者明确列出:(1) 模型多样性有限,主要在 Mistral-7B 上验证,跨架构泛化待测;(2) 当前流水线仍重算文档、未实现生产级 KV 缓存,缓存的存储/加载/显存搬运开销未建模——234× 是模拟值而非端到端实测;(3) Segment-wise 是静态压缩,当证据稀疏或局部时可能次优,未来可做 query 条件的自适应 token 选择;(4) 算力受限下多为单次实验,缺乏多次统计;(5) 仅在 BEIR 上评测,其他领域泛化未知。隐含局限是打分层/决策 token 的选择依赖具体 backbone,需逐模型标定。
6. 工业实践意义
对搜索/推荐精排有几条可落地启示。其一,"排序信号稀疏"这一判断值得借鉴:上线 LLM 重排不必把全文喂进去,按固定步长均匀采样保留约 10% token 往往就够,比只取开头/结尾或按显著性取更稳。其二,decoding-free 打分(用注意力质量或 logits 出分)应作为 LLM 重排上线的默认形态,能省掉自回归生成这一最大延迟源,对社区搜索这种延迟敏感场景尤为关键。其三,文档表示解耦 + KV 缓存是 RAG/多轮搜索里反复重排同一批 UGC 文档时的省算力关键——把文档侧 KV 预计算缓存,在线只算 query 侧(0.05s 级),但要预留缓存存储与加载成本。其四,训练-推理列表长度对齐:生成式重排在训练长度外大面积出错(30→500 篇 100% 漏 ID),用候选级指针目标(CopyNet)训练能稳定外推,对"训练只能用短列表、线上要排长列表"的工业约束很对症。需注意论文未给线上 A/B,最佳层/决策 token 需按自家 backbone 重新标定。