PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.31261v1cs.CLcs.AI

TL;DR

Dense self-attention’s quadratic cost and fixed attention patterns motivate learning which interactions should remain relevant at different distances. MoSAR routes queries and keys over controlled-decay regimes, and achieves the best evaluated perplexity at longer contexts while retaining competitive perplexity under top-1 routing.

  • Problem

    Dense self-attention has quadratic sequence-length complexity, while many efficient alternatives prescribe fixed sparse patterns despite input-dependent language dependencies and locality.

    Source: §1 · p. 1, §2 · p. 4

  • Method

    MoSAR uses input-conditioned query and key routers to select mixtures of controlled distance-decay regimes, trained with soft routing and discretized through deterministic top-1 routing.

    Source: §1 · p. 2, §3 · p. 6

  • Results

    MoSAR achieves the best perplexity among evaluated variants at 4096 and 8192 tokens, and deterministic top-1 routing reduces expected normalized reach while preserving competitive perplexity under extrapolation.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The experiments support controlled decay over abrupt local masking and show that adaptive regime selection can outperform fixed locality in long-context evaluation.

    Source: §7 · p. 10

    The expected attention-cost metric is a proxy rather than measured wall-clock speedup, which would require optimized sparse kernels and hardware-aware implementations.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Dense self-attention’s quadratic sequence-length cost motivates alternatives, but fixed sparsity patterns assume one locality scale fits every token and context. MoSAR instead learns adaptive, controlled-decay attention geometries, and controlled pre-training finds strong language-modeling results alongside lower expected reach.

1 INTRODUCTION

  • Dense self-attention has quadratic complexity in sequence length, while fixed sparsity patterns assume the same locality is appropriate across tokens, layers, and contexts.The introduction frames the alternative as learning from data where local attention suffices and where broader interactions should be retained.Source: §1 · pp. 1–2
  • MoSAR uses post-RoPE query and key routers to select mixtures of smoothly decaying attention regimes, replacing a fixed sparse pattern with geometries selected dynamically by the model.The regimes preserve local interactions, smoothly attenuate them through a transition region, and expose low-relevance regions for approximation.Source: §1 · p. 2
  • In controlled 500M-parameter pre-training, MoSAR achieves the best perplexity among evaluated variants at 4096 and 8192 tokens.Deterministic top-1 routing reduces expected normalized reach while preserving competitive perplexity under length extrapolation.Source: §1 · p. 3

2 BACKGROUND AND RELATED WORK

Prior positional methods prescribe how attention varies with distance, while dynamic selection methods adapt which interactions are computed. MoSAR is motivated by the gap between choosing interactions and learning how their strength decays with distance.

2 BACKGROUND AND RELATED WORK

  • RoPE supplies relative positional geometry but produces oscillatory rather than monotone decay; ALiBi, XPos, YaRN, and p-RoPE prescribe alternative positional laws before training.CoPE and HoPE learn parts of positional geometry, but neither lets each token choose how quickly its own relevance decays.Source: §2 · pp. 3–4
  • Efficient-attention methods reduce the O(n^2d) similarity cost through fixed patterns, hashing, low-rank projection, or input-adaptive interaction selection, but generally do not learn continuous distance-dependent strength within selected interactions.The cited dynamic-selection methods include Selective Attention, FlexPrefill, MoBA, and NSA.Source: §2 · p. 4
  • MoSAR addresses the catch-22 that exact query–key relevance is costly to evaluate by routing tokens over controlled distance regimes before computing pairwise similarities.Unlike Adaptive Attention Span’s per-head scalar span applied to every token, MoSAR routes queries and keys separately in context and combines their regime preferences.Source: §2 · p. 4

3 MOSAR: MIXTURE OF SEMANTIC ATTENTION REGIMES

MoSAR learns token-conditioned attention geometry by routing post-positional queries and keys over controlled distance-decay regimes. Soft routing supports training, while deterministic top-1 routing defines a potential sparse support before query–key similarities are evaluated.

3 MOSAR: MIXTURE OF SEMANTIC ATTENTION REGIMES

  • Post-RoPE query and key routers assign locality preferences separately, whose interaction selects a regime for each pair and can localize attention-matrix regions before pairwise similarities are computed.The router is a lightweight module on the post-RoPE Q–K pathway; its role is to induce a regime-dependent attention geometry.Source: §3 · pp. 4–5
  • Each semantic regime preserves RoPE interactions on a locality plateau, decays smoothly through a transition zone, and holds distant interactions at a finite floor rather than deleting them.The available regimes are fixed before training; MoSAR learns their distributions, and pairwise query–key regimes combine their reaches by averaging.Source: §3 · p. 5
  • Training uses soft routing and dense evaluation; MoSAR-Cost alone adds an auxiliary expected-reach term, which is a cost proxy rather than measured FLOPs or wall-clock acceleration.The nominal normalized reaches for short, medium, and global regimes are 0.0625, 0.25, and 1 at the training context length.Source: §3 · p. 6
  • Deterministic top-1 routing can define computational support before pairwise similarities, but this work’s top-1 evaluation still computes all causal similarities densely.A sparse implementation could omit pairs outside the support; hardware-friendly sparse kernels and their benchmarking remain future work.Source: §3 · p. 7

4 EXPERIMENTAL SETUP

The experiments compare eight from-scratch 500M-parameter decoder-only models with different positional, locality, decay, and routing choices. Shared architecture, data, and training conditions are used to isolate attention geometry, followed by post-training evaluations at longer contexts and with top-1 routing.

4 EXPERIMENTAL SETUP

  • Eight 500M-parameter models compare dense RoPE, ALiBi, 0.75-RoPE, fixed short and medium regimes, hard medium masking, MoSAR, and MoSAR-Cost.These variants distinguish positional encoding, locality, smooth decay, hard truncation, and adaptive routing.Source: §4 · p. 7
  • All variants share an 18-layer architecture, initialization and data seeds, tokenizer, packed-token stream, learning-rate schedule, and training-context evaluation protocol.Models train at sequence length 2048; MoSAR routers are two-layer MLPs with hidden dimension 64.Source: §4 · p. 7
  • Models train for a 610k-step target, and post-hoc evaluations test checkpoints trained at 2048 tokens at 4096 and 8192 tokens, with additional deterministic top-1 evaluations for MoSAR variants.Table 2 reports soft/native performance and length extrapolation after 610k optimization steps; expected cost is normalized to full attention.Source: §4 · pp. 7–8, Table 2 · p. 8

5 RESULTS

MoSAR achieves the best perplexity among evaluated variants at 4k and 8k tokens, while deterministic top-1 routing retains strong performance at lower cost. Its learned routing reach decreases as context grows, and its training-context quality remains close to the best fixed-locality variant.

5 RESULTS

  • At 4k and 8k tokens, MoSAR achieves the best perplexity among evaluated variants, slightly improving over ALiBi at 8k and substantially outperforming dense RoPE.Fixed-S and Fixed-M remain strong, but neither matches MoSAR at longer contexts.Source: §5 · p. 8
  • 0.002 PPL separates MoSAR from Fixed-M at the training context length, while MoSAR uses lower expected cost.MoSAR-Cost reaches the lowest learned expected cost but pays a small training-context perplexity penalty.Source: §5 · p. 8
  • Hard top-1 routing raises perplexity relative to soft MoSAR, but the degradation remains bounded and MoSAR stays substantially better than dense RoPE at 4k and 8k.The top-1 approximation reduces cost increasingly with context length.Source: §5 · p. 8, Table 3 · p. 9
  • Across pre-training checkpoints, MoSAR stays competitive and finishes within the strongest group, remaining below dense RoPE and separated from hard RoPE-M-mask.This behavior is reported throughout training, not only at the final checkpoint.Source: §5 · p. 8, Figure 2 · p. 9
  • Under top-1 routing, MoSAR’s support density ρΩ falls from 0.154 at n=2048 to 0.104 at n=8192, a 33% reduction over a 4× longer context.Query-side routing remains stable while key-side mass shifts toward medium and global regimes as context grows.Source: §5 · p. 9, Table 3 · p. 9

6 LIMITATIONS AND FUTURE WORK

The evaluation is a controlled first study, limited in scale, training budget, seeds, and downstream breadth; the reported cost is an expected-cost proxy rather than measured wall-clock speedup.

6 LIMITATIONS AND FUTURE WORK

  • Experiments are limited in model scale, training budget, number of seeds, and breadth of downstream evaluation, so the study does not establish performance on complex NLP tasks.The authors consider perplexity and training dynamics appropriate first signals, and regard broader conclusions as premature.Source: §6 · p. 9
  • The expected attention-cost metric is a proxy, not a measured wall-clock speedup; real acceleration requires optimized sparse kernels, memory layouts, and hardware-aware implementations.The authors identify scaling MoSAR, evaluating long-context behavior, and developing sparse kernels as future work.Source: §6 · p. 10

7 CONCLUSIONS

MoSAR replaces preset attention sparsity with learned, content-dependent interaction geometry. The conclusion highlights controlled decay, adaptive reach, and top-1 approximation, while the accompanying appendices and statements clarify the rationale and scope of the work.

7 CONCLUSIONS

  • The conclusion reports that controlled decay is more stable than abrupt local masking, adaptive regimes improve over fixed locality at long context, and top-1 routing preserves competitive perplexity while reducing the attention domain.The reported advantage is especially strong under length extrapolation.Source: §7 · p. 10

AI USE STATEMENT

  • The authors report using generative AI for research and writing assistance, while stating that it did not generate training data or results, run experiments, or make autonomous scientific decisions.They say the authors inspected and verified implementations, runs, evaluations, numerical results, tables, and figures.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The reproducibility statement describes shared experimental controls and identifies scripts, configurations, data utilities, evaluation outputs, and checkpoints as available materials.Variants shared seeds, tokenizer, packed-token stream, schedule, backend stack, and validation protocol.Source: REPRODUCIBILITY STATEMENT · p. 10

A FROM THE CATCH-22 OF ATTENTION APPROXIMATION TO MOSAR

  • Appendix A is presented as conceptual and mathematical motivation for MoSAR, not as a substitute for empirical evidence or a claim of theoretical optimality.It is framed as a possible starting point for further work on adaptive and sparse attention.Source: A FROM THE CATCH-22 OF ATTENTION… · p. 13

A.1 THE CIRCULARITY OF SPARSE ATTENTION

  • MoSAR addresses the sparse-attention catch-22 by learning content-dependent interaction geometry instead of predicting exact relevance in advance or imposing a rigid local mask.Appropriate sparse patterns can preserve essential properties of full attention, but discovering dispensable interactions itself requires relevance information.Source: §A.1 · pp. 13–14

A.2 THE INFINITE-DIMENSIONAL NATURE OF THE ATTENTION APPROXIMATION PROBLEM

  • Because language-modeling contexts are open-ended, the appendix argues there is no universal limiting attention matrix that can be sparsified a priori; MoSAR instead learns geometry along the sequence.At each context length, routing selects a computational regime for the current interaction.Source: §A.2 · p. 14

A.3 SEMANTIC SPARSITY

  • MoSAR treats semantic sparsity as distinct from fixed positional sparsity: useful interaction support is unknown beforehand, and some queries or keys may need greater computational reach.The abstract definition does not require locality, block structure, randomness, or low rank.Source: §A.3 · pp. 14–15

A.4 FROM SEMANTIC SPARSITY TO SPARSE COMPUTATION

The appendix frames MoSAR as a trade-off between the resolution of learned support and the cost of discovering and executing it. It distinguishes learned support selection from sparse computation and notes that arithmetic savings do not guarantee faster hardware execution.

A.5 ROUTER-INDUCED SUPPORT AND SPARSE QUERY–KEY COMPUTATION

  • After deterministic routing, the support can serve as a computational stencil: construct admissible regions from router assignments and distance, then evaluate query–key similarities only within them.The authors note that irregular supports may map poorly to GPUs, motivating block, tiled, or other hardware-aligned approximations.Source: §A.4 · p. 15, §A.5 · p. 18

A.4 FROM SEMANTIC SPARSITY TO SPARSE COMPUTATION

  • More regimes can tighten support around useful interactions, but increase router capacity and routing and execution costs; sufficiently fine routing approaches the pairwise selection problem that dense attention already solves.Coarse routing instead accepts over-coverage to avoid evaluating every pair independently.Source: §A.4 · pp. 15–16
  • The useful number and spacing of regimes remain open design choices that may depend on model scale, context length, hardware, and training-distribution structure.The appendix raises, rather than establishes, whether regime count should grow slowly or saturate as context grows.Source: §A.4 · p. 16
  • The proposed long-context scaling picture is a hypothesis, although the appendix says it is consistent with MoSAR's strong language-modeling performance and decreasing normalized routing reach under extrapolation.The suggested decomposition is local semantic construction plus sparse long-range links.Source: §A.4 · p. 16
  • MoSAR takes an intermediate position between hard-local and dense-global attention, but this work focuses on learning support; efficiently constructing and evaluating that support is left as a next step.The appendix separates deciding where pairwise similarity should be computed from computing selected similarities efficiently.Source: §A.4 · p. 17

A.5 ROUTER-INDUCED SUPPORT AND SPARSE QUERY–KEY COMPUTATION

  • When support density is sufficiently small—on the order of 1/d_h—sparse query–key evaluation can approach the arithmetic scale of an O(n^2) matrix operation, but this does not imply comparable wall-clock time.Indexing, support construction, and hardware-alignment overheads remain; the straightforward approach also retains quadratic support construction.Source: §A.5 · pp. 17–18

A.6 A CONDITIONAL SPARSE-APPROXIMATION LEMMA

The lemma gives a conditional basis for sparse attention: small discarded attention mass limits the change caused by masking and renormalizing. Decay assumptions can make a sparse support constructible, but the result does not claim that natural language always has uniform decay.

A.6 A CONDITIONAL SPARSE-APPROXIMATION LEMMA

  • If each row’s discarded attention mass is at most ε, masking and renormalizing changes its attention distribution by at most 2ε in L1; bounded values also bound the output change.The argument is deliberately conditional and uses a row-wise L1 bound.Source: §A.6 · p. 19
  • A monotone decay with a summable tail and suitable bounds on distant logits can make the mass outside a sufficiently large window arbitrarily small, yielding a constructive sparse support.Without an a priori decay estimate, the existence of a low-error sparse approximation does not identify its support.Source: §A.6 · pp. 19–20
  • Norm equivalence at each fixed context length does not ensure uniform approximation bounds as attention matrices grow; generic-norm guarantees require control of constants across lengths.The argument is stated in row-wise L1, which directly tracks removed probability mass.Source: §A.6 · p. 20

A.7 FROM CONTROLLED DECAY TO ADAPTIVE REGIMES

MoSAR combines smooth distance decay with input-dependent routing over short, medium, and global regimes. Its routing and decay parameters balance local computation, broader interactions, and the possibility of sparse execution.

A.7 FROM CONTROLLED DECAY TO ADAPTIVE REGIMES

  • MoSAR’s smooth decay preserves a unit plateau through threshold T_mn, attenuates interactions over transition length L_mn, and reaches a finite floor; global–global interactions do not decay.The stated implementation uses p = 2 for a smooth onset and B = 6 for the finite attenuation floor.Source: Figure 3 · p. 21, §A.7 · p. 21
  • At 8192 tokens, MoSAR reaches perplexity 12.208 versus 12.453 for Fixed-S and 12.520 for Fixed-M; at 2048 tokens, MoSAR reaches 12.269.At 2048 tokens, Fixed-S and Fixed-M reach perplexities of 12.313 and 12.267; the authors caution these comparisons do not directly measure discarded attention mass.Source: §A.7 · p. 22
  • The decay strength B trades smoother geometric regularization against sparse approximability; the experiments use B = 6 and do not establish an optimal value.At long contexts, many individually small floor-weight interactions may still accumulate non-negligible attention probability.Source: §A.7 · p. 22
  • MoSAR uses short, medium, and global regimes, with query- and key-side routers selecting marginal regimes whose factorized combination assigns geometry to each pair.The routers select a computational regime rather than predicting the exact relevant interaction; subsequent attention similarity determines whether the pair is useful.Source: §A.7 · pp. 22–24
  • Pairwise reach is composed symmetrically from the two marginal reaches using their arithmetic mean, placing mixed query–key regimes between their independently selected scales.The authors choose this interpolation over min or max composition, leaving routing behavior to adapt through the language-modeling objective.Source: §A.7 · p. 23

A.8 FULLY EXPANDED FORMULA

  • The expanded formula places router biases inside the MLP logits before regime softmaxes, while the worth gate enters attention through an outer logarithm.Source: §A.8 · p. 24

A.9 MINIMAL PSEUDOCODE

  • The pseudocode applies positional transforms to Q and K, routes their resulting states, adds the log-worth bias to attention scores, then softmaxes and combines attention with V.Source: §A.9 · pp. 24–25

A.10 MOSAR ARCHITECTURE DIAGRAM

The architecture retains a pre-norm decoder-only backbone and modifies attention geometry through MoSAR routing on post-RoPE queries and keys. The values bypass the router, and the rest of the decoder block retains residual and MLP stages.

A.10 MOSAR ARCHITECTURE DIAGRAM

  • RoPE is applied to Q and K within each layer, after which the MoSAR router selects the regime used by causal self-attention; V bypasses the router.Source: §A.10 · p. 25
  • The decoder block uses pre-normalization, then attention and a residual connection, followed by a second pre-norm, an MLP block, and another residual connection.Source: §A.10 · p. 25
  • MoSAR changes the geometry of attention without altering the overall decoder-only backbone.Source: §A.10 · p. 25

A.11 CHOICE OF BACKBONE AND EXPERIMENTAL SCALE

The study uses matched Gemma2-style 500M-parameter models to compare attention geometries under controlled training conditions. Its scope is a mechanistic study, with dedicated retrieval and long-context reasoning tests left for future work.

A.11 CHOICE OF BACKBONE AND EXPERIMENTAL SCALE

  • All variants share a Gemma2-style backbone and differ only in the positional or attention mechanism under study.The study follows the backbone-specific query pre-attention scaling and tanh soft-capping before applying the MoSAR bias.Source: §A.11 · p. 25
  • The authors reduce the scale and training setting from the 2B-model comparison they had planned, prioritizing matched initialization, data order, optimization, and training budget.The resulting experiments are framed as a controlled mechanistic study, not a scale-matched reproduction of the Gemma-2B experiments.Source: §A.11 · p. 25
  • The controlled study trains 500M-parameter models for approximately 10B tokens and evaluates language modeling, routing, expected normalized reach, and length extrapolation.Source: §A.11 · p. 25
  • Dedicated retrieval and long-context reasoning benchmarks remain future work to test whether learned global regimes preserve task-relevant distant information.Source: §A.11 · pp. 25–26

A.12 SUMMARY

MoSAR frames attention approximation as learning a content-dependent geometry, rather than imposing a fixed sparse mask. Its controlled study finds that regime geometry matters beyond locality, while the reported cost measures expected reach rather than wall-clock speedup.

A.12 SUMMARY

  • Hard local masks are substantially worse than fixed MoSAR regimes at matched nominal cost, supporting a role for regime geometry beyond locality alone.The comparison is reported as an ablation in Table 2.Source: §A.12 · p. 26
  • Explicit compute pressure in MoSAR-Cost induces a quality–cost trade-off, suggesting compute pressure can act as a geometric regularizer.MoSAR-Cost adds an explicit cost term; standard MoSAR is trained with language-modeling loss alone.Source: §A.12 · p. 27
  • The study is a controlled mechanistic test of lower expected attention cost and preserved language-modeling quality, not a broad benchmark or claim of optimized sparse-kernel speedup.Soft training uses strictly positive gates and evaluates similarities densely; exact computational sparsity belongs to a hard policy with a corresponding sparse kernel.Source: §A.12 · p. 27, §F.3 · p. 30
  • MoSAR replaces a universal decay scale with content-dependent regimes that can keep attention local, expand to medium range, or preserve sparse long-range links.The central framing is to learn an approximable geometry rather than impose a sparse mask in advance.Source: §A.12 · p. 27

C.2 FUTURE POSITIONS

  • The causal mask remains authoritative, pure global-global routing adds exactly zero bias, and production attention may use fused or grouped operations if numerically compatible.The dense reference path materializes O(N^2) tensors for clarity.Source: §C.2 · p. 28, §C.3 · p. 28, §F.4 · p. 30

F.1 MODEL AGNOSTICISM

  • For architectures with pairwise positional biases, an adapter must explicitly choose which representation the routers observe; the current post-positional contract does not specify that choice.The operator also requires access to the query and key representations after token-wise positional transformation and a shared bias added before softmax.Source: §F.1 · p. 30

F.5 HARDWARE AND EXPERIMENTAL CONSTRAINTS

Compute constraints limited the breadth and cadence of the experiments, but the main variants shared the same controlled training and evaluation setup. The authors frame the work as a controlled first study rather than a complete efficient-attention evaluation.

F.5 HARDWARE AND EXPERIMENTAL CONSTRAINTS

  • Shared-cluster queueing, reservations, time limits, and segmented training constrained experiment cadence and the number of variants, seeds, and ablations.Additional compute would allow more seeds, cost coefficients, longer extrapolation and hard-inference ablations, and denser milestone curves.Source: §F.5 · p. 31
  • Despite asynchronous training and limited experiment breadth, main variants shared initialization, data, tokenizer, packed-token stream, model scale, sequence length, schedule, and evaluation protocol.These comparisons are intended to isolate attention geometry rather than differences in training setup.Source: §F.5 · p. 31

Related papers

Loading 2609.31261v1…