Learning When to Attend: Rethinking Memory Access in Long-Context Language Models

Long-context modeling remains one of the persistent bottlenecks in large language model deployment. While frontier models support hundreds of thousands of tokens, extending base models beyond their pretraining context length requires continued pretraining on long sequences, a process whose cost scales quadratically with sequence length due to the attention mechanism. The paper Learning When to Attend: Conditional Memory Access for Long-Context LLMs by Choudhary et al. introduces a fundamental architectural shift that challenges the assumption that every token position requires identical computational resources.

The core observation driving this work is both intuitive and analytically profound: within a 100,000 token document, a token predicting the next word in a concluding sentence typically relies on local context, while a token referencing "Section 2" may need to retrieve information from thousands of tokens prior. Standard attention mechanisms treat both cases identically, computing full quadratic attention for every position. L2A, or Learning To Attend, proposes conditional memory access where a learned router determines, for each individual token, whether expensive global attention is necessary at all.

The Architecture of Conditional Attention

L2A operates through a two-stage attention mechanism that decouples local context processing from long-range memory retrieval. First, all tokens process through Local Attention, attending only to a fixed recent window. Subsequently, a lightweight routing function processes these local representations to produce a binary decision per token: invoke Global Attention over the full context, or bypass it entirely. This creates a dynamic compute graph where only approximately 20% of tokens trigger the full quadratic operation, while the remainder operate strictly within local windows.

This approach fundamentally differs from prior sparse attention methodologies such as Native Sparse Attention or Expansion Span. Those methods address the question of what information to retrieve by sparsifying the key-value cache that each query attends to. However, they still execute attention computations for every token position, merely reducing the number of keys and values considered. L2A instead addresses when to retrieve, applying exact attention over the full context but only for a sparse subset of query tokens. This token-wise conditional routing allows the model to adapt its compute budget to the specific linguistic demands of each position, rather than applying uniform sparsity patterns regardless of local complexity.

The training objective combines standard next-token prediction with a sparsity regularizer that encourages the router to minimize global attention invocations. This formulation treats global attention as a specialized long-term memory retrieval mechanism, invoked only when local context proves insufficient for prediction. Notably, the router learns this behavior through task performance alone, without explicit supervision indicating which tokens require long-range dependencies.

Empirical Validation and System Optimization

The authors implement L2A as a drop-in replacement for standard attention layers, evaluating on Qwen 2.5 and Qwen 3 models at scales of 1.5B and 7B parameters. Through continued pretraining, they extend effective context length from 32K to 128K tokens. On long-context benchmarks including HELMET, BabiLong, and MRCR, L2A matches the performance of standard Continued Long-Context Pre-training within 1.5 to 3% accuracy, despite skipping global attention for 75 to 80% of tokens.

Critically, the paper demonstrates that algorithmic sparsity translates to measurable wall-clock speedups. Custom Triton kernels achieve approximately 2× improvements in training throughput and time-to-first-token compared to FlashAttention-2. This efficiency gain stems from bypassing memory-bound global attention operations for the majority of positions, rather than merely approx