# ARM: Attention with Routed-Memory for Learnable Sparse Control

arXiv: https://arxiv.org/abs/2609.24417 · HTML: https://mlsanity.com/paper/2609.24417/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 84 of its 109 references that are in the map.
- First posted on arXiv 21 Sep 2026.

## Builds on (most cited first, 30 of 84)

- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs](https://mlsanity.com/paper/1903.00161/) (2019)

## Abstract (arXiv)

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.