长上下文的真正瓶颈,从来不是模型“看不看得完”,而是 KV cache 撑不撑得住。过去的主流路径是:训练用精确注意力,推理时再套一层稀疏化的 KV 淘汰策略把缓存压小——训练与部署的内存机制并不一致。8 月 20 日提交到 arXiv 的论文《Learning how to Forget》(2608.19920)把这个顺序倒了过来:与其事后打补丁,不如让模型在微调阶段就“带着遗忘”学习。
核心思路:让权重与淘汰策略共同适应
论文由 Matthias Seeger 等五位作者完成,配套代码开源在 AWS 的 awslabs 组织下。方法的关键动作是:微调时让 KV cache 淘汰策略全程生效,权重主动适应“信息会被逐出”这一现实,与策略共同适应(co-adapt)。摘要报告的结果是,这种做法常常反超用精确注意力(序列并行)训练出来的模型。
复现门槛也压得低:README 给出的示例基于 Qwen3-4B-Instruct-2507,单张 A100 40GB 显存即可在 Helmet 基准的 128k token 序列上做 LoRA 微调,缓存策略 h2o-torch-quantized8、槽位 16384;换到 8 张 A100 的 AWS p4d.24xlarge 实例上,全局 batch 可以到 32。
把 H2O 从论文做进内核
论文实验中表现最好的淘汰策略是 H2O(arXiv:2306.14048):按注意力权重的累计和对 KV 条目打分,分数最低者被逐出——README 的说法是,这在强意义上就是缓存世界里的 LRU。
团队还点出一个被内核生态忽视的缺口:主流快速 SDPA 内核都不返回“沿 query 轴求和的注意力权重”,而这恰是 H2O 类策略打分的必需品。库里的解法是给 FlashInfer CUDA 内核补上该返回值,并用 Triton score-sum 内核补齐计算;配合 4/8-bit 缓存量化与 CPU offload,README 明确说 4-bit 量化可把 KV 缓存的 GPU 显存需求降为原来的四分之一。
限制也写得很清楚
README 坦承:纯推理性能不敌 vLLM、SGLang,缺少多设备策略,定位是研究与评估;但对 H2O 这类高级缓存策略的支持比 vLLM 更好。另一个技术亮点在梯度计算:前向阶段把淘汰决策写进“重放日志”,反向阶段用重放缓存逐格回放、配合激活检查点,才让稀疏注意力下的长序列微调在显存上变得可行。
所以呢
这篇论文值得记住的不是某个跑分,而是一次方向修正:长上下文的下一程优化,可能不在“更大的窗口”,而在训练与部署的内存机制对齐——当淘汰策略成为训练的一部分,“遗忘”就从缺陷变成了可学习的能力。对做长文本落地的团队来说,单卡可复现的实验配置加全开源代码,意味着这条路线周末就能上手验证(论文:https://arxiv.org/abs/2608.19920 ,代码:https://github.com/awslabs/keys_values )。