# Learning When to Attend: Conditional Memory Access for Long-Context LLMs

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 68 references that are in the map.
- First posted on arXiv 18 Mar 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [SWE-bench: Can Language Models Resolve Real-World GitHub Issues?](https://mlsanity.com/paper/2310.06770/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Evaluating Very Long-Term Conversational Memory of LLM Agents](https://mlsanity.com/paper/2402.17753/) (2024)
- [Mem0: Building Production-Ready AI Agents with Scalable Long-Term Memory](https://mlsanity.com/paper/2504.19413/) (2025)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://mlsanity.com/paper/2402.02750/) (2024)
- [YaRN: Efficient Context Window Extension of Large Language Models](https://mlsanity.com/paper/2309.00071/) (2023)
- [Enabling Large Language Models to Generate Text with Citations](https://mlsanity.com/paper/2305.14627/) (2023)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity](https://mlsanity.com/paper/2403.14403/) (2024)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens](https://mlsanity.com/paper/2402.13753/) (2024)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [TokenSkip: Controllable Chain-of-Thought Compression in LLMs](https://mlsanity.com/paper/2502.12067/) (2025)
- [ASQA: Factoid Questions Meet Long-Form Answers](https://mlsanity.com/paper/2204.06092/) (2022)
- [BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack](https://mlsanity.com/paper/2406.10149/) (2024)

## Built on by (most cited first, 3 of 3)

- [The Workload-Router-Pool Architecture for LLM Inference Optimization: A Vision Paper from the vLLM Semantic Router Project](https://mlsanity.com/paper/2603.21354/) (2026)
- [Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling](https://mlsanity.com/paper/2604.24715/) (2026)
- [On-Demand Attention: Language Models Know When to Recall](https://mlsanity.com/paper/2609.20734/) (2026)

## Abstract (arXiv)

Language models struggle to generalize beyond pretraining context lengths, limiting long-horizon reasoning and retrieval. Continued pretraining on long-context data can help but is expensive due to the quadratic scaling of Attention. We observe that most tokens do not require (Global) Attention over the entire sequence and can rely on local context. Based on this, we propose L2A (Learning To Attend), a layer that enables conditional (token-wise) long-range memory access by deciding when to invoke global attention. We evaluate L2A on Qwen 2.5 and Qwen 3 models, extending their effective context length from 32K to 128K tokens. L2A matches the performance of standard long-context training to within 3% while skipping Global Attention for ~80% of tokens, outperforming prior baselines. We also design custom Triton kernels to efficiently implement this token-wise conditional Attention on GPUs, achieving up to ~2× improvements in training throughput and time-to-first-token over FlashAttention. Moreover, L2A enables post-training pruning of highly sparse Global Attention layers, reducing KV cache memory by up to 50% with negligible performance loss. Our code is released under Apache 2.0 at https://github.com/awslabs/hybrid-model-factory/tree/main/examples/research/L2A.