KV-streams for Efficient Compaction in Agentic Reinforcement Learning
作者: Emiliano Penaloza, Dane Malenfant, Dheeraj Vattikonda, Roger Creus Castanyer, Siddarth Venkatraman, Abhay Puri, Jonathan Light, Matthew James Sargent, Augustine N. Mavor-Parker, Massimo Caccia, Lucas Caccia, Glen Berseth, Esmeralda S. Whitammer, Alessandro Sordoni, Minseon Kim, Marc-Alexandre Côté, Laurent Charlin, Guillaume Lajoie
分类: cs.LG, cs.AI
发布日期: 2026-09-28
💡 一句话要点
提出KV-streams以解决长上下文训练效率问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 上下文压缩 流式处理 代理型大语言模型 训练效率 强化学习
📋 核心要点
- 现有的上下文压缩策略在训练过程中需要多次预填充LLM上下文，导致训练效率低下。
- 本文提出KV-streams，通过流式传输KV缓存来实现高效的上下文压缩，避免了频繁的缓存清空。
- 实验结果显示，KV-streams使训练速度提升了2.6到5倍，并且能够有效传递长期信息。
📝 摘要（中文）
随着代理型大语言模型（LLMs）上下文长度的增加，GPU内存的限制成为了一个瓶颈。上下文压缩是缓解这一问题的常用机制，但大多数压缩策略需要多次预填充LLM上下文，影响训练效率。为了解决这一瓶颈并实现高效的可训练压缩，本文提出了KV-streams，这是一种与任何压缩策略兼容的即插即用策略，显著提高了训练的吞吐量，同时未见性能下降。KV-streams通过向前流动KV缓存而非在每次压缩后清空它，实现了可扩展的压缩。实验表明，KV-streams使三种不同的压缩策略实现了2.6到5倍的训练速度提升。此外，流式KV缓存可以作为递归状态，传递早已消失于上下文的信息。总体而言，KV-streams是任何后训练管道的高效且轻量的补充。
🔬 方法详解
问题定义：本文旨在解决代理型大语言模型在训练过程中因上下文长度增加而导致的GPU内存瓶颈问题。现有的压缩策略往往需要多次预填充上下文，影响训练效率。
核心思路：论文提出KV-streams作为一种即插即用的策略，通过流式传输KV缓存而非在每次压缩后清空，从而提高训练的吞吐量和效率。这样的设计使得模型能够在不牺牲性能的情况下，处理更长的上下文。
技术框架：整体架构包括三个主要模块：流式KV缓存管理、上下文压缩策略的集成和训练过程的优化。KV缓存管理模块负责动态流式传输，而压缩策略模块则可以与现有的压缩方法结合使用。
关键创新：KV-streams的最大创新在于其流式处理机制，这与传统的压缩方法形成了鲜明对比，后者通常需要在每次压缩后清空缓存。此创新显著提高了训练效率，并保持了模型性能。
关键设计：在设计中，KV-streams的参数设置经过精细调整，以确保流式缓存的有效性。同时，损失函数和网络结构也经过优化，以支持流式信息传递的需求。
🖼️ 关键图片
📊 实验亮点
实验结果表明，KV-streams使得三种不同的上下文压缩策略实现了2.6到5倍的训练速度提升，显著提高了训练效率。此外，流式KV缓存能够有效传递长期信息，展现出与传统方法的显著差异。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、强化学习和大规模机器学习系统。通过提高训练效率，KV-streams可以加速模型的开发和部署，降低计算资源的需求，推动智能系统的广泛应用与发展。
📄 摘要（原文）
Scaling the horizon of agentic LLMs is bottlenecked by the need to fit ever longer context traces in GPU memory. Context compaction has been the most popular mechanism to alleviate this issue, keeping GPU memory constant for a given trace. Unfortunately, most compaction strategies rely on prefilling the LLM context many times over, hindering training throughput. To alleviate this bottleneck and enable efficient trainable compaction, we propose KV-streams, a plug-and-play strategy compatible with any compaction strategy that substantially increases throughput while showing no evidence of hindering performance. KV-streams enable scalable compaction by streaming the KV cache forward rather than flushing it after each compaction. We show that KV-streams enable three different compaction strategies, achieving a 2.6 to 5x wall-clock speedup in training. Beyond efficiency, we find that the streamed KV cache can act as a recurrent state, carrying forward information that has long since disappeared from the context. Specifically, in a controlled setting we show that, contrary to prior work, RL alone is all that is needed for this behavior to emerge. Overall, we show KV-streams to be an efficient and lightweight plug-and-play addition to any post-training pipeline.