Attention with routed memory improves efficiency in large language models
ARM: Attention with Routed-Memory for Learnable Sparse Control
Machine LearningArtificial Intelligence
Summary
Large language models have trouble remembering everything they need during long conversations because they store information in a limited cache that can fill up. The authors introduce ARM, a new way to organize this memory that learns how to choose and update stored information without losing important details. ARM adapts how much memory it uses depending on how complicated the input is, which helps it understand both simple and complex texts better. Tests show that ARM performs better and faster than older methods of handling memory in language models.
What this means in practice
- •For natural language processing engineers: Integrate ARM into language model systems to improve accuracy and efficiency when handling long and complex text contexts.
- •For conversational ai developers: Develop chatbots that maintain better context over extended conversations by using ARM for dynamic memory management.
Authors
Qiuhao Zeng, Jerry Huang, Peng Lu, Ruiyi Fang, Gezheng Xu, Zihao Jing, Yufei Cui, Charles Ling, Gang Niu, Boyu Wang
Abstract
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.