Compression Beyond the Uncompressed: A Two-Stage Training Recipe for Soft Context Compression in RAG
Shuyu Guo, Shuo Zhang, Zhaochun Ren
cs.CL
2026-09-04
DEX-Comp先只蒸馏未压缩RAG答对的题，再在答错的题上做GRPO。16倍压缩、首token最多快24倍，五个开放域QA上平均CEM反超原文，top-30显存降4倍。
RAG 把检索到的文档整段塞进解码器，输入变长，延迟跟着涨，还可能顶破上下文窗口。软压缩把每篇文档编成更短的连续向量序列，信息密度比抽句子、做摘要这类硬压缩更高。现有做法几乎都在蒸馏未压缩 RAG 的输出，学生被钉在老师读离散 token 时养成的计算习惯上。压缩槽位理论上能装下更多信息，实测却长期低于未压缩基线。
先前评测还多停在 top-5。真实系统经常要 top-15、top-30，浅检索上的结论推不过去。
DEX-Comp 是两阶段训练配方，压缩器和解码器共用 Mistral-7B 骨干、各挂一套 LoRA。压缩在离线做：文档后面拼 mi=⌊Li/τ⌋ 个压缩 token，取这些位置的末层隐状态当软向量，整库预计算缓存。推理时解码器只读缓存向量加查询。
第一阶段叫 Pure Distillation。按老师答案是否命中金标，把训练集切成 S⁺ 和 S⁻。只在 S⁺ 上最小化老师与学生的响应 KL，梯度同时流过解码器和压缩器。错题留给下一阶段，避免把老师的失败模式写进初始化。
第二阶段叫 Hard Exploration。只在 S⁻ 上跑 GRPO：每题采一组 rollout，对金标做 0/1 对错奖励，组内标准化 advantage，并用 KL 把策略拴在 PD 检查点上。老师已经答错的题上，模仿没有正确答案可抄，学生只能在压缩表示里另找计算方式。
评测覆盖 NQ、TriviaQA、HotpotQA、ASQA、PopQA，检索深度 top-5 到 top-30，语料是切成 128 token 的 Wikipedia-KILT。检索器 Splade-v3，重排 DeBERTa-v3。指标是 Containment Exact Match 和 Gemini 3 Flash 做裁判。对照包括未压缩 Mistral-7B、硬压缩 LLMLingua-2，以及 xRAG、ICAE、COCOM、PISCO。PISCO 额外按 top-30 重训，方便深检索对比。
训练推理深度对齐、压缩 16 倍时，DEX-Comp 在每个数据集、每个深度上都超过未压缩 RAG，配对置换检验 p<0.05。平均 CEM/LLM：top-5 为 68.38/72.54 对 64.56/70.53，top-15 为 67.83/72.88 对 63.94/70.54，top-30 为 66.03/72.02 对 62.91/69.49。PISCO 在同样 16 倍、top-30 训练下平均 CEM 只有 57.90。单模型按 top-30 训练、跨深度推理，平均分仍高于未压缩基线，PopQA 在 top-5/15 是例外。
效率按 batch=8 测在一块 RTX PRO 6000 上。top-30 的 TTFT 从 4782ms 降到 201ms（23.73 倍），GFLOPs 降 13.61 倍，峰值显存从 59328MB 降到 14766MB（4.02 倍）。top-5 已经有 4.42 倍 TTFT。压缩率 16 最好，32 也能打平未压缩基线。
消融说明两阶段都要。只用全部老师输出做 SFT，平均 LLM 掉到 64.79；PD 单独是 67.75；PD 后再在对的题或混合集上做 RL，分别是 70.60 和 71.47，仍低于完整配方的 72.02。Resilience Rate 的增益更大，压缩顺手把噪声文档滤了一层。跨骨干和 BioASQ、CovidQA、FEVER 也能打平或超过对应未压缩 RAG。
这是第一篇在多个真实检索深度上、查询无关软压缩反超未压缩 RAG 的报告。对要缓存整库、查询时不想再压一遍的系统，16 倍压缩加最高约 24 倍首 token 加速，换来的是平均两到四个点的 CEM，不是精度换延迟。配方本身不绑特定压缩结构，换自动编码压缩器也能套。
未压缩 RAG 没有做任务相关 RL，作者写明是因为长上下文算不起。学生在压缩后的短输入上做了 GRPO，老师没有对等的强化学习，反超里有多少来自「短输入更好训」、多少来自压缩表示本身，分不开。作者把更干净的对照留给后续。
压缩必须离线预计算并存储整库向量，语料频繁更新时增量重压的成本会露出来。评测停在开放域 QA 的答案正确性，忠实度、引用、对抗和长文本生成都没测。软向量对下游不可读，出了错不好追到哪一段原文。