MoSAR: Mixture of Semantic Attention Regimes for Learning Adaptive and Approximable Attention Geometries

Michele Paolicelli; Alessandro Petruzzelli; Alessandro Franceso Maria Martina; Cataldo Musto; Giovanni Semeraro

cs.CLSep 25, 2026

MoSAR: Mixture of Semantic Attention Regimes for Learning Adaptive and Approximable Attention Geometries

Authors: Michele Paolicelli, Alessandro Petruzzelli, Alessandro Franceso Maria Martina, Cataldo Musto, Giovanni Semeraro

Organizations: Department of Computer Science Università degli Studi di Bari Aldo Moro Bari, 70121, Italy

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.

Figures & tables

Figure 1: Illustrative MoSAR decay profiles for representative query–key regime pairs. Each pair induces a distance-dependent regime $g_{mn}(d)$ with an initial unit plateau, a smooth decay region, and a finite floor; pairs involving broader regimes retain interactions over longer causal distances.

Method	Data- learned	Input adaptive	Token- wise	Attention approx.	Cost loss	Learned routing	Distance decay	Learned decay
Efficient and approximate attention
Reformer ( Kitaev et al., 2020 )	$\triangle$	✓	✓	✓	✗	✗	✗	✗
Linformer ( Wang et al., 2020 )	✓	✗	✗	✓	✗	✗	✗	✗
Longformer ( Beltagy et al., 2020 )	✗	✗	✗	✓	✗	✗	✗	✗
Positional and geometric modifications
ALiBi ( Press et al., 2021 )	✗	✗	✗	✗	✗	✗	✓	✗

Table 1: Comparison of methods that modify attention geometry or reduce attention computation. Columns refer to the mechanism introduced beyond standard content-based self-attention. ✓ denotes an explicit property, ✗ its absence, and

\triangle

a partial, indirect, or inference-only property.

Method	LM loss	PPL@2k	PPL@4k	PPL@8k	$\Delta\%$ @2k	$\Delta\%$ @8k	Reach
MoSAR (ours)	2.5071	12.269	11.753	12.208	-1.99	-19.6	0.2210
ALiBi	2.5186	12.411	11.873	12.232	-0.85	-19.4	1.0000
Fixed-S (ours)	2.5107	12.313	11.864	12.453	-1.63	-17.9	0.0625
Fixed-M (ours)	2.5069	12.267	11.827	12.520	-2.00	-17.5	0.2500
RoPE-M-mask	2.5419	12.703	12.212	12.584	+1.49	-17.1	0.2500
MoSAR-Cost (ours)	2.5299	12.553	12.053	13.276	+0.28	-12.5	0.0897

Table 2: Performance and length extrapolation in soft/native inference mode after 610k optimization steps. Expected cost is normalized with respect to full attention. All models are evaluated on the same validation corpus, rebuilt at each context length. Lower LM loss and perplexity are better.

Panel of Figure 2 — Figure 2: Validation perplexity over pre-training checkpoints at the training context length. MoSAR remains competitive throughout training and reaches the final checkpoint within the strongest group of variants, while learning an adaptive controlled-decay geometry.

Model	Length	Soft PPL	Top-1 PPL	Gap	Soft reach	Top-1 reach	$q_{S}$	$q_{G}$	$k_{S}$	$k_{G}$
MoSAR	2k	12.269	12.703	+3.5%	0.2193	0.1544	0.741	0.108	0.681	0.166
MoSAR	4k	11.753	12.216	+3.9%	0.1878	0.1276	0.734	0.114	0.666	0.180
MoSAR	8k	12.208	12.872	+5.4%	0.1690	0.1038	0.734	0.115	0.640	0.181
MoSAR-Cost	2k	12.553	13.123	+4.5%	0.0887	0.0724	0.955	0.010	0.890	0.021
MoSAR-Cost	4k	12.053	12.737	+5.7%	0.0562	0.0424	0.934	0.017	0.878	0.020
MoSAR-Cost	8k	13.276	13.795	+3.9%	0.0858	0.0481	0.812	0.079	0.794	0.051

Table 3: Top-1 approximation of learned MoSAR geometries and Soft MoSAR routing statistics across context lengths (medium is omitted, since

q_{M}=1-q_{S}-q_{G}

, likewise for keys). The gap is the relative perplexity increase of discrete routing with respect to the corresponding soft geometry.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Illustration of the pairwise MoSAR decay kernel $g_{mn}(d)$ as a function of causal distance $d=i-j$ . Interactions remain unattenuated up to the pair-dependent threshold $T_{mn}$ , decay smoothly over a transition region of length $L_{mn}$ , and reach the finite floor $e^{-B}$ thereafter. The global–global regime is exempt from decay, with $g_{GG}(d)=1$ for all $d\geq 0$ .

Panel of Figure 4 — Figure 4: Schematic attention regimes at sequence length $2048$ , rendered as explicit cell grids so that each square corresponds to a single $q_{i}k_{j}$ interaction. Fixed $S$ , $M$ , and $G$ produce contiguous causal bands. In contrast, MoSAR activates only router-selected $(i,j)$ regions: the resulting pattern is sparse and disconnected, with isolated attention islands rather than continuous bridges across the entire past context.

Figure 5: MoSAR-augmented decoder-only Transformer architecture. The model follows a pre-norm decoder-only backbone, while RoPE is applied internally at each layer to queries and keys rather than added once at the input level. The MoSAR router operates only on the query–key pathway and selects the attention regime used by causal self-attention, while values $V$ bypass the router. The overall layout is intentionally designed as a visual adaptation of the iconic Transformer schematic introduced by Vaswani et al. (2017) , so as to make the architectural differences immediately comparable. The present figure is nevertheless a new, task-specific reformulation for Gemma-style decoder-only models with post-RoPE routing.

Mathematical object	Implementation	Compatibility condition
Configuration $(N,\delta,\alpha,d_{r},\tau,B,p,\varepsilon)$	mosar/config.py : MoSARConfig	Values and validation must match the experimental configuration.
$x_{i,Q},x_{j,K}$	mosar/routers.py : PostPositionRouter.forward	Inputs must be post positional; key heads must be unexpanded GQA heads.
$u_{Q},z_{Q},\pi_{Q}$ and $u_{K},z_{K},\pi_{K}$	mosar/routers.py	Separate Q and K parameters; GELU MLP; final temperature scaled softmax.
Router initialization	mosar/routers.py : PostPositionRouter.reset_parameters	$W_{1}$ std $0.02$ , $W_{2}$ std $10^{-3}$ , all biases zero.
$\delta_{mn},T_{mn},L_{mn}$	mosar/regimes.py : pair_geometry	Arithmetic means and transition definitions must match the equations.
$\mathbf{G}(d)$	mosar/regimes.py : RegimeKernel.forward	Non increasing non GG gates in $[e^{-B},1]$ ; exact $GG=1$ .

Appendix

Table 9

	Requirement
$\square$	Query routing is computed from post-positional query representations.
$\square$	Key routing is computed from post-positional, unexpanded key representations.
$\square$	Query and key routers are separately parameterized.
$\square$	Each router ends in a temperature-scaled softmax over regimes.
$\square$	Router MLP biases do not bypass the final softmax.
$\square$	The factorized pair weights are $\pi_{Q}(m)\pi_{K}(n)$ .

Appendix

Table 10

Attention Mechanisms

Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.

See all 513 papers

Sep 19, 2026

Block-Sparse Attention with Semantic-Geometric Decoupled Routing

Long-context inference has become a defining capability of large language models, but exact dense attention remains costly due to its quadratic scaling with sequence length. Block-sparse attention offers a hardware-friendly alternative by routing each query block to a small set of relevant key blocks, yet accurate training-free block routing remains difficult. Existing routers often pool post-RoPE token representations, which entangles semantic aggregation with RoPE-induced geometry and attenuates local positional cues through high-frequency phase cancellation. To resolve this mismatch, we propose Semantic-Geometric Decoupled Routing, a training-free block routing framework that shifts semantic aggregation to the pre-RoPE space and reconstructs geometric bias with an offline structural prior and relative block distances. This decomposition yields an explicit closed-form block routing score without token-level search or post-hoc calibration. Experiments on long-context text and video tasks show that our method approaches full-attention accuracy across 4K--128K contexts, keeps routing overhead below 3.4 ms, and achieves a 5.03

\times

speedup over FlashAttn at a 128K context length.

Xinwei Long, Weigao Sun, Weibo Gao +6

Department of Electronic Engineering, Tsinghua University, Beijing, China · Intern at Tongyi Lab, Alibaba Group · Tongyi Lab, Alibaba Group, Shanghai, China +2

Attention Routing Sparse Attention Block-Sparse Attention LLM Inference Acceleration Long-Context Language Modeling

Ask PDF

Oct 7, 2026

Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values

Large Language Models (LLMs) achieve strong performance across many domains, but their efficiency is limited by the quadratic cost of attention with respect to prompt length. Sparse attention reduces this cost by retaining only a small fraction of query-key interactions to approximate the full attention matrix. However, existing methods are trapped in a mathematically wrong view: they simply keep large scalar entries or high-mass regions of the attention matrix. This treats the attention matrix as a bag of values, ignoring that it is used as a structured matrix whose entries jointly determine the attention output through multiplication with value vectors. We argue that this is the core conceptual issue: sparse attention should be formulated as matrix approximation, not as blindly choosing the largest values from a bag of entries. Based on this view, we propose Matrix Approximation Sparse Attention (MASA). MASA replaces raw attention-mass ranking with a closed-form score that measures how much each sparse unit reduces matrix-product approximation error. As a theory-grounded plug-in correction, MASA can be added to existing sparse attention frameworks without changing their sparse kernels or budgets. Extensive experiments across multiple sparse attention methods, benchmarks, and LLM backbones show consistent accuracy gains, supporting both MASA and the matrix-approximation view of sparse attention.

Fang Wan, Xufeng Liu, Fan Li +1

Stony Brook University · Duke University

Sparse Attention Efficient Attention Attention Mechanisms Attention Approximation LLM Inference Acceleration

Ask PDF

Sep 28, 2026

SMat-Attention: Structured Long-Context Sequence Modeling

Long-context sequence models face a fundamental tradeoff: softmax attention uses flexible token-level interactions at quadratic cost, whereas linear attention obtains linear-time training and constant-time decoding by compressing history into a fixed-size state. In this work, we ask whether we can connect these regimes through a tunable notion of structure. To this end, we introduce Structured Matrix Attention (SMat-Attention) via a family of causal masks with structured long-range routing whose row supports have VC-dimension

d

. In our construction,

d=1

recovers the standard causal mask, and increasing

d

permits richer subset-routing patterns. We give chunkwise forward and backward algorithms to enable hardware-efficiency. For sequences of length

T

, the hard-routing construction takes

O(T^{2-3/d}+T)

work, despite the mask being dense, for our prescribed family. In fixed-horizon streaming, decoding after the distant prefix takes constant time per token using

O(T^{1-1/d})

cached states. SMat-Attention therefore makes VC-dimension an explicit knob governing access-pattern complexity, prefill cost, and decoding memory. Empirically, subset-routing and rule-assisted multi-key retrieval experiments illustrate the masks' routing expressiveness. Extensions to Mamba-2 and Gated DeltaNet using learned routing with top-

k

query reads retain subquadratic prefill, improve recall accuracy over the backbones in several settings, and achieve comparable small-scale language-modeling performance.

Emile Anand, Abdullah Ateyeh, Archer Wang +1

Massachusetts Institute of Technology

Long-Context Modeling Linear Attention Attention Mechanisms Attention Routing

Ask PDF

Abstract

Attention Mechanisms

Explore similar work

Block-Sparse Attention with Semantic-Geometric Decoupled Routing

Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values

SMat-Attention: Structured Long-Context Sequence Modeling