arXiv:2605.23200cs.LGcs.AI2026-05被引 1

解决长文本推理中键值缓存膨胀问题,避免逻辑断裂。

Adaptive Mass-Segmented KV Compression for Long-Context Reasoning

论文配图:Adaptive Mass-Segmented KV Compression for Long-Context Reasoning
图 1 · 摘自论文原文
  • 按注意力分布分段分配缓存,避免连续推理块被误删。
  • 在数学推理等任务上提升模型表现,最高提升12.3%准确率。
  • 可无缝接入现有压缩方法,适配主流推理系统。

长序列大模型推理中,键值(KV)缓存的线性增长是关键瓶颈。现有压缩方法依赖全局Top-k重要性评分剔除标记,但导致区域抹除:连续推理块被严重删除,破坏逻辑连贯性。为此,我们提出自适应质量分段(AMS)KV压缩框架,将范式从标记级竞争转向区域感知的配额分配。AMS根据注意力质量的空间分布动态划分KV缓存,确保结构关键的推理片段获得保障性内存。为保证迭代解码稳定性,引入基于指数移动平均(EMA)的平滑机制,防止分段边界抖动。关键的是,AMS作为通用即插即用层,与现有评分器正交,可无缝集成至TOVA、Expected Attention、KeyDiff、R-KV和TriAttention等方法。同时兼容vLLM等现代分页KV服务框架,支持高效聚合并紧凑执行,无额外稳态注意力开销。在数学推理(MATH500、AIME、GSM8K)、代码补全、开放域问答及稀疏检索等多样化任务上的大量实验表明,AMS持续缓解结构碎片化,显著提升模型性能。

原文摘要 · Abstract (English)

The linear growth of the Key-Value (KV) cache is a critical bottleneck in long-form LLM inference. Existing KV compression methods mitigate this by evicting tokens based on importance scores. However, we show that their reliance on global Top-k selection triggers Region Wipe-out: the severe eviction of contiguous reasoning blocks that derails logical coherence. To address this, we propose Adaptive Mass-Segmented (AMS) KV Compression, a framework that shifts the paradigm from token-level competition to region-aware quota allocation. AMS adaptively partitions the KV cache based on the spatial distribution of attention mass, ensuring structurally vital reasoning segments receive guaranteed memory quotas. To ensure stability during iterative decoding, an EMA-based smoothing mechanism is incorporated to prevent jitter in segment boundaries. Crucially, AMS is a universal plug-and-play layer that is orthogonal to existing scorers. It can be seamlessly integrated into representative methods such as TOVA, Expected Attention, KeyDiff, R-KV and TriAttention. AMS is also system-compatible with modern paged-KV serving frameworks such as vLLM, supporting efficient gather-and-compact KV execution without introducing additional steady-state attention overhead. Extensive experiments across a diverse suite of tasks, including mathematical reasoning (MATH500, AIME, GSM8K), code completion, open-domain QA, and sparse retrieval, demonstrate that AMS consistently mitigates structural fragmentation and boosts model performance.

长上下文KV压缩推理优化大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。