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

Block-Sparse Attention with Semantic-Geometric Decoupled Routing
Sep 19, 2026Xinwei Long, Weigao Sun, Weibo Gao +6Attention Routing Sparse Attention
Department of Electronic Engineering, Tsinghua University, Beijing, China · Intern at Tongyi Lab, Alibaba Group · Tongyi Lab, Alibaba Group, Shanghai, China +2
Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values
Oct 7, 2026Fang Wan, Xufeng Liu, Fan Li +1Sparse Attention Efficient Attention
Stony Brook University · Duke University
SMat-Attention: Structured Long-Context Sequence Modeling
Sep 28, 2026Emile Anand, Abdullah Ateyeh, Archer Wang +1Long-Context Modeling Linear Attention
Massachusetts Institute of Technology

Sep 19, 2026J/K move · Enter open · S save

Block-Sparse Attention with Semantic-Geometric Decoupled Routing

Xinwei Long, Weigao Sun, Weibo Gao, Pengkun Jiao+5

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

Open paper AskPDF arXiv

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.

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