Adaptive attention geometry improves long context language models
MoSAR: Mixture of Semantic Attention Regimes for Learning Adaptive and Approximable Attention Geometries
Computation and LanguageArtificial Intelligence
Summary
Modeling language with attention is often slow because it looks at everything at once, which takes a lot of computation. The authors suggest learning which words should pay attention to others based on the input itself, rather than fixing a pattern ahead of time. Their approach, MoSAR, learns to focus on different ranges of nearby and distant words in a flexible way, adapting to the context. This lets the model work faster without losing quality, even when dealing with longer texts than it has seen before.
What this means in practice
- •For language model engineers: Build language models that adaptively focus on relevant words to reduce computation while keeping quality for long documents.
- •For natural language processing teams: Improve inference efficiency of transformer models by deploying learned attention geometries that are stable under discrete approximation.
Authors
Michele Paolicelli, Alessandro Petruzzelli, Alessandro Franceso Maria Martina, Cataldo Musto, Giovanni Semeraro
Abstract
The quadratic complexity of dense self-attention remains a central bottleneck for long-context language modeling. Many efficient alternatives address this cost by deciding in advance where attention should be sparse or local. We argue that attention approximation should instead be approached as a geometric problem, with the relevant interaction geometry learned from data: natural-language dependencies are input-dependent and difficult to prescribe in advance, so the model should learn where positional relevance can decay and where broader interactions must be preserved. We introduce Mixture of Semantic Attention Regimes (MoSAR), which learns such an adaptive, controlled-decay geometry over query--key interactions. Input-conditioned query and key routers, applied after positional encoding, select mixtures over short, medium, and global regimes, inducing a continuous distance-dependent attention field rather than a fixed sparsity pattern. This geometry is learned during training and can subsequently be discretized through top-1 routing. In controlled pre-training experiments with matched 500M-parameter models, MoSAR learns a substantially lower-reach attention geometry without degrading language-modeling quality, improving perplexity over dense RoPE at the training context length. Under length extrapolation, MoSAR achieves the best perplexity among all evaluated variants, including strong baselines such as ALiBi. Moreover, the learned geometry remains stable under deterministic top-1 discretization, suggesting that it is not only adaptive, but also amenable to low-cost approximation at inference time.