Cache Me If You Can：陈丹琦团队如何「抓住」关键缓存，解放LLM内存？
文章摘要
【关 键 词】 机器学习、长上下文、KV缓存、注意力机制、内存优化
普林斯顿大学陈丹琦团队提出统一框架解决长上下文语言模型中的KV缓存内存瓶颈问题。随着”长思维链”等技术的兴起，模型需要处理数万token的输入，导致基于Transformer架构的语言模型面临KV缓存内存线性增长的挑战。例如，Llama-3-70B模型处理128K token输入时需分配高达42GB内存用于KV缓存存储。
研究团队创新性地提出”KV足迹”作为统一度量标准，该指标聚合所有时间步中未被驱逐键值缓存条目的比例，同时涵盖预填充和解码两阶段的开销。为确保证据比较的实用性，团队进一步定义”关键KV足迹”概念——在保持模型性能不低于全注意力机制90%的前提下，方法能达到的最小KV足迹。这一硬性标准有效区分了真正实用的优化方法与牺牲模型能力的方案。
研究发现，现有KV驱逐方法存在显著缺陷。后填充驱逐方法因与预填充阶段不兼容，导致KV足迹居高不下；而新近度驱逐方法虽能减少内存占用，却可能丢失重要历史信息。基于这些发现，团队提出PruLong（长程精简注意力机制），通过端到端优化实现注意力头专业化分工：部分头保留完整KV缓存处理全局上下文，其余头采用驱逐策略处理局部上下文。
PruLong在三个方面实现创新突破：首先采用下一个token预测损失替代传统重建误差，使训练目标与文本生成任务直接对齐；其次引入离散掩码优化技术，通过伯努利分布参数化消除训练-测试差距；最后利用自然长上下文数据进行训练，相比合成数据更能捕捉复杂的长程依赖关系。实验表明，PruLong的KV足迹比现有方法降低12%，在召回任务中保持原始性能，同时显著减少内存消耗。
研究还建立了KV缓存驱逐的统一分析框架，将推理过程分解为预填充和解码两个阶段。预填充阶段通过单次前向传播处理整个输入序列并存储键值状态，而解码阶段逐个生成token时持续更新KV缓存。团队发现，分块预填充等新兴技术虽然能降低峰值内存，但缺乏统一指标来衡量跨阶段的内存使用效率。为此提出的KV足迹指标能精确跟踪KV条目生命周期，为不同方法提供公平比较基础。
在技术实现层面，PruLong采用混合注意力机制：通过正则化损失约束掩码稀疏度，使用拉格朗日乘子进行min-max优化，实现端到端的注意力头分类训练。相比DuoAttention的合成数据训练，PruLong采用代码仓库和书籍等自然长上下文数据，使模型能学习更丰富的长程依赖模式。该方法不仅适用于传统解码场景，还能无缝适配多轮对话、交错工具调用等复杂推理流程。
这项研究为长上下文语言模型的内存优化提供了系统化解决方案，通过统一度量标准、理论框架和实用算法三者的结合，显著推进了高效推理技术的发展。团队开源的PruLong实现为后续研究奠定基础，其方法论也可拓展至量化、架构设计等正交技术领域，具有广泛的应用前景。
原文和模型
【原文链接】 阅读原文 [ 3628字 | 15分钟 ]
【原文作者】 机器之心
【摘要模型】 deepseek/deepseek-v3-0324
【摘要评分】 ★★★★★
相关文章
