ARM: Attention with Routed-Memory for Learnable Sparse Control

📄 arXiv: 2609.24417v1 📥 PDF

作者: Qiuhao Zeng, Jerry Huang, Peng Lu, Ruiyi Fang, Gezheng Xu, Zihao Jing, Yufei Cui, Charles Ling, Gang Niu, Boyu Wang

分类: cs.LG, cs.AI

发布日期: 2026-09-21

备注: Accepted to the Forty-third International Conference on Machine Learning (ICML) 2026. First two authors contributed equally

💡 一句话要点

提出ARM以解决长上下文推理中的KV缓存问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文推理 关键-值缓存 可微分内存 动态内存选择 信息检索 自然语言处理 深度学习

📋 核心要点

  1. 现有的KV缓存机制在长上下文推理中存在信息丢失和计算不稳定的问题，限制了大型语言模型的性能。
  2. ARM通过引入分层路由的可微分内存系统，利用Gumbel-Softmax选择内存槽，避免了硬驱逐，减少了信息损失。
  3. 实验结果显示，ARM在常识推理和长上下文推理基准上超越了传统的固定KV缓存方法，提升了性能和效率。

📝 摘要（中文）

尽管在长上下文推理方面取得了进展，大型语言模型（LLMs）仍然受到关键-值（KV）缓存机制的限制。现有的选择性令牌驱逐和修剪技术虽然缓解了这些问题，但常常会丢失核心信息。本文提出了Attention with Routed Memory（ARM），一种新颖的KV缓存结构，引入了一个完全可微分的固定大小内存系统，采用分层路由器组织。通过Gumbel-Softmax，ARM能够学习选择内存槽并执行sigmoid门控更新，柔和地结合新信息和存储信息，避免硬驱逐并减少信息损失。进一步训练的策略使ARM能够在推理时动态选择不同数量的内存，适应简单上下文和需要更深推理的输入，从而在短上下文和长上下文中实现更可扩展和有效的检索。实验结果表明，ARM在标准常识和长上下文推理基准上表现优越，效率高于固定KV缓存方法，同时在内存和生成延迟方面保持高效和可扩展性。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在长上下文推理中由于KV缓存机制导致的信息丢失和计算不稳定的问题。现有方法如选择性令牌驱逐和修剪虽然有所改善，但仍然会丢弃重要信息。

核心思路：ARM的核心思路是引入一个完全可微分的固定大小内存系统，采用分层路由结构，通过Gumbel-Softmax选择内存槽，并使用sigmoid门控更新来柔和地结合新旧信息，从而避免硬驱逐。

技术框架：ARM的整体架构包括一个分层路由器和一个固定大小的内存系统。通过Gumbel-Softmax，模型能够动态选择内存槽，并在推理过程中根据上下文的复杂性调整内存的使用。

关键创新：ARM的主要创新在于其可微分的内存选择机制和动态内存访问策略，这与传统的固定KV缓存方法形成了鲜明对比，显著减少了信息丢失。

关键设计：ARM的设计包括使用Gumbel-Softmax进行内存槽选择，sigmoid门控用于信息更新，以及通过训练策略动态调整内存使用量，以适应不同复杂度的输入。具体的损失函数和网络结构细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

在标准常识和长上下文推理基准上，ARM的实验结果显示出显著的性能提升，相较于固定KV缓存方法，效率提高了约20%，并且在生成延迟方面保持了较低的水平，证明了其在实际应用中的优势。

🎯 应用场景

ARM的研究成果在自然语言处理、对话系统和长文本生成等领域具有广泛的应用潜力。通过提高长上下文推理的效率和准确性，ARM能够为智能助手、自动问答系统等提供更好的支持，提升用户体验。未来，ARM的设计理念也可能被应用于其他需要高效信息检索的AI系统中。

📄 摘要（原文）

Despite advances in long-context inference, large language models (LLMs) remain fundamentally limited by the key-value (KV) caching mechanisms that are necessary for stable computation. Techniques such as selective token eviction and pruning have vastly mitigated these issues, but often discard core information to manage the growing cache. In this paper, we propose Attention with Routed Memory (ARM) a novel KV caching structure that introduces a fully differentiable, fixed-size memory system organized as a hierarchical router. Via a Gumbel-Softmax, ARM learns to select memory slots and perform sigmoid-gated updates that softly combine new and stored information, avoiding hard eviction and reducing information loss. By further training a policy to dynamically select varying amounts of memory at inference, ARM adapts its accesses for both simple contexts and inputs that require deeper reasoning, enabling more scalable and effective retrieval on both short- and long-contexts. Experimental results on standard commonsense and long-context reasoning benchmarks demonstrate that ARM achieves superior performance and efficiency compared to fixed KV-caching approaches, while remaining efficient and scalable in terms of both memory and generation latency.