论文

ARM 提出可微路由记忆结构,改进 LLM 长上下文 KV 缓存

ARM: Attention with Routed-Memory for Learnable Sparse Control

精选理由

做长上下文推理的可以看看这篇:ARM 用可微路由记忆替代硬淘汰 token,缓存不再丢关键信息,延迟还可控。

arXiv 论文 2609.24417 提出 ARM(Attention with Routed Memory),把传统 KV 缓存改造成固定大小、按层级路由组织的可微记忆系统。ARM 通过 Gumbel-Softmax 学习挑选记忆槽,再用 sigmoid 门控把新信息与已存信息软性融合,避免硬性淘汰 token 带来的信息丢失。它还训练了一个策略,在推理时按输入难度动态决定访问多少记忆,短上下文和需要深度推理的长上下文可以区别对待。在常识推理和长上下文推理基准上,ARM 的表现超过固定 KV 缓存方法,在内存占用和生成延迟上也保持可扩展。

原文 · arXiv cs.LG

ARM: Attention with Routed-Memory for Learnable Sparse Control

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.