AdaMem: Adaptive Memory Token Allocation for Soft Compression in Retrieval-Augmented Generation
Artem Sakhno, Grigorii Davydenko, Omar Zoloev, Julia Belikova, Andrey Savchenko, Maksim Makarenko
cs.CL, cs.IR, cs.LG
2026-08-12
AdaMem 按段落相关性动态分配记忆 token,64 倍压缩下比统一分配基线 OSCAR 高 14.6%,延迟降至四分之一。
RAG 系统把检索到的段落压缩成连续记忆向量(soft compression),用来降低推理成本。已有方法比如 OSCAR 会用相关性分数排序或筛掉低分段落,但每个留下来的段落分到的记忆 token 数量完全一样,不管这段内容对回答问题有多重要。检索池变大、内存预算又紧张时,这种一刀切分配会浪费本该给高价值段落的容量,论文把这个问题定义为「不仅要选哪些文档,还要在异质价值的文档间分配一个固定预算」。
AdaMem 用同一个压缩器在一次前向里同时产出两样东西:每个段落的连续记忆向量,和一个相关性分数,分数取自压缩器给每个段落附加的专用 <RERANK> 位置的最后一层隐状态。分数在检索池内标准化后,经过带温度参数 τ 的 softmax 转成每个段落应得的记忆预算份额,再用最大余数法取整,分数低到一定程度的段落直接分不到 token,相当于被丢弃。
温度 τ 控制分配的锐利程度:τ 趋于无穷退化成 OSCAR 式的均匀分配,τ 趋于 0 则接近排序后硬截断。论文证明这套 softmax 分配规则是一个对数效用最大化问题的精确解(基于 Kelly 提出的比例公平框架),也是经典比特分配理论里逆水填充最优解的一阶近似,理论上预测 τ 应该随压缩率提高而变小,实验里最优 τ 确实从 16 倍压缩下的 1.0 降到 64 倍压缩下的 0.5。训练分三阶段:先用自编码目标预训练压缩器,再在中等规模检索语料上做 query 相关的联合训练,用交叉编码器蒸馏初始化打分头,最后在作者构建的新数据集 KILT-SCR 上联合优化生成损失与排序损失。
在 25 篇检索文档、16 倍压缩的标准设置下,AdaMem 在六个 QA 基准上全面超过同架构的均匀分配基线 OSCAR:
| 压缩率 | 指标 | AdaMem | OSCAR | Full Context |
| 16× | TriviaQA Match | 84.0 | 82.4 | 84.4 |
| 16× | PopQA Match | 61.6 | 58.4 | 60.2 |
| 64× | TriviaQA Match | 83.3 | 77.1 | 84.4 |
| 64× | PopQA Match | 59.6 | 49.8 | 60.2 |
16 倍压缩下平均相对提升 3.4%,64 倍压缩下拉大到 14.6%,PopQA 单项最高涨 9.8 个点(19.7%)。更关键的是稳定性:压缩率从 16 倍加到 128 倍,AdaMem 的效果只掉 2% 左右,OSCAR 同区间掉 7% 以上。32 倍压缩下的 AdaMem 已经能超过 16 倍压缩下的 OSCAR,相当于用一半记忆预算达到更好效果。对比最强的硬压缩基线 LLMLingua-2,AdaMem 在上下文短 8 倍的情况下平均领先 9.7%。速度上,32 倍压缩的 AdaMem 能以 4 倍更低延迟、3.6 倍更少算力匹配未压缩 Full Context 的效果。
对检索池较大、内存预算紧张的 RAG 场景(比如一次要塞进 25 篇以上文档),AdaMem 提供了一个几乎不增加额外训练成本的改法:把原本只用来排序筛选的相关性分数,顺手用来决定每个段落该分到多少记忆容量。作者用同一套数据、同一套训练流程分别训练 PISCO、OSCAR、AdaMem 做对照,差异只在分配策略上,这让 3.4%14.6% 的提升可以直接归因于分配方式而非其他因素。对已经在用 query-conditioned soft compressor 的团队,这个改动可以直接叠加在现有架构上,不需要重新设计压缩流水线。
论文的最优性证明建立在段落价值随 token 数对数增长这个效用模型假设上,现实中这个假设是否成立没有独立验证。实验只覆盖英文开放域 QA,压缩器用 Llama-3.2-1B、解码器用 Mistral-7B-Instruct 这一组固定搭配,换成更大或更小的模型是否保持同样的提升幅度未知。六个基准里 BioASQ 是唯一的例外,硬压缩方法 LLMLingua-2 在这个基准上反而更强,论文没有深入解释为什么生物医学问答场景下软压缩的优势会消失。AdaMem 的峰值显存也比 OSCAR 略高,因为要先编码完整候选池才能做分配,这个开销在低压缩率下更明显。