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

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 $\sim$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 $\sim$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.

Modern AI language models can struggle when asked to process very long inputs, such as entire books or long documents, because standard methods require the model to compare every part of the input with every other part, which becomes very slow and expensive. In this work, we introduce a more efficient way for models to handle long inputs by allowing them to selectively decide when they actually need to pay attention to distant information. Instead of always looking at the entire context, the model learns when local information is sufficient and when it is necessary to access information from far away in the input. This selective mechanism significantly reduces unnecessary computation while still preserving the model’s ability to understand and reason over long contexts. We show that this approach can extend the effective range of existing models to much longer inputs without requiring expensive retraining from scratch. It also makes training and inference more efficient by avoiding a large fraction of costly long-range computations. Overall, this work moves toward making long-context language models more practical by teaching them to focus their attention only when it is truly needed.