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.

Large language model inference is fundamentally limited by a cache of previous computations that needs to be maintained throughout generation. This becomes an issue during long-context inference, thus techniques have been introduced in the past to directly remove elements as the process continues to reduce this burden. In this work, we introduce ARM, a new caching structure shaped as a hierarchy with a fixed number of memory slots in which information can be routed and stored. ARM gradually mixes in information in a soft manner, avoiding issues that come with directly pruning or evicting individual pieces of information. We further augment ARM to dynamically select differing amounts of information from the memories depending on the input context, enabling faster, more effective retrieval on varying types of problems, which is further confirmed by empirical validations.