#artificial intelligence #machine learning Preprint

ARM: Attention with Routed-Memory for Learnable Sparse Control

Qiu-Hao Zeng Jerry M. Huang Peng Lu Rui-Yi Fang Ge-Zheng Xu Zi-Hao Jing Yu-Fei Cui Charles X. Ling Gang Niu Bo-Yu Wang

Sep 2026 · 0 citations · 109 references

Computer Science

TL;DR

Attention with Routed Memory is proposed a novel KV caching structure that introduces a fully differentiable, fixed-size memory system organized as a hierarchical router that achieves superior performance and efficiency compared to fixed KV-caching approaches, while remaining efficient and scalable in terms of both memory and generation latency.

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.

DOI arXiv Semantic Scholar

View source

Similar papers

#artificial intelligence Preprint Sep 2026

MemAgent: Learning to Manage Heterogeneous Memory Providers for LLM Agents

MemAgent, featuring a content-aware routing architecture and a training-data synthesis pipeline, is proposed, featuring a content-aware routing architecture and a training-data synthesis pipeline that improves average accuracy and outperforms every individual memory method across all three benchmarks.

Yong-Xian Wei, Yi-Lin Zhao, Run-Xi Cheng et al. · 0 citations

#artificial intelligence Preprint Oct 2026

Memory Type Varies: Empowering LLM Agents for Long-Term Memory with Diverse Strategies

This work proposes a novel memory framework, named MemoType, which can adaptively recognize each memory and query type with the learned router model, and proposes a novel retrieval framework, named MemoType, which can adaptively recognize each memory and query type with the learned router model.

Yi Wen, De-Rong Xu, Peng-Yue Jia et al. · 0 citations

#natural language process... Preprint Sep 2026

On-Demand Attention: Language Models Know When to Recall

These findings show that pretrained decoding states can support both token prediction and decisions about accessing distant information, allowing models to allocate global computation as needed during decoding, allowing models to allocate global computation as needed during decoding.

Hai-Bo Feng, Ruiqi Liang, Dong-Yang Jin et al. · 0 citations

Conference Aug 2026

Learning what to Remember: Reinforcement Learning for Cost-Aware Memory in LLM Agents

Long-term memory is a critical capability for LLM-based agents, yet most current systems rely on RetrievalAugmented Generation (RAG), which stores information indiscriminately and retrieves by similarity alone. This passive approach lacks any notion of utility: it does not distinguish valuable knowledge from noise and...

Cristian-Marius Cringasu, Adriana Tapus · IEEE International Symposium... · 0 citations

#natural language process... Preprint Aug 2026

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

This work provides a new method for fine-tuning models with sparse attention that works for any KV cache policy, runs on a moderate hardware budget, and allows the model to co-adapt with the policy, often outperforming models trained with exact attention (sequence parallelism).

Matthias W. Seeger, Zeyu Zhang, Vihang Patil et al. · 0 citations

#natural language process... Preprint Oct 2026

LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches

Reinforcement learning has greatly advanced the capabilities of large language models, but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations su...

Shao-Kun Zhang, Yi-Fan Zhang, Jian Hu et al. · 0 citations