Gaussian Mixture Attention: Linear-Time Sequence Mixing via Probabilistic Latent Routing

# Gaussian Mixture Attention: Linear-Time Sequence Mixing via Probabilistic Latent Routing
Source: [https://arxiv.org/html/2606.18283](https://arxiv.org/html/2606.18283)
Yongchao Huang111yongchao\.huang@abdn\.ac\.ukHassan Raza222h\.raza\.24@abdn\.ac\.uk

\(16/05/2026\)

###### Abstract

The dense token\-to\-token interaction pattern of standard dot\-product attention remains a central bottleneck in scaling Transformer architectures to long contexts\. We introduceGaussian Mixture Attention \(GMA\), a probabilistic attention\-style sequence mixer that replaces explicit pairwise query–key comparison with routing throughKKlearned Gaussian mixture components\. Queries and keys are mapped to posteriorresponsibilityvectors over a shared latent routing space; their overlap defines an implicit responsibility\-space affinity, while values are written into and read from aKK\-slot latent memory\. By exploiting the associativity of matrix multiplication, GMA avoids materializing the inducedN×NN\\times Naffinity matrix and instead uses two responsibility matrices whose dominant activation storage scales as𝒪​\(N​K\)\\mathcal\{O\}\(NK\)rather than𝒪​\(N2\)\\mathcal\{O\}\(N^\{2\}\)for fixedKK\. We formulate bidirectional and causal variants of GMA, provide an end\-to\-end differentiable parameterization of the Gaussian mixture components, and analyze its responsibility\-modulated gradient structure, constrained non\-negative low\-rank affinity interpretation, and local routing stability\. Empirically, GMA exhibits the intended fixed\-KKlinear memory scaling and is competitive with attention\-style baselines on long\-context classification, while causal GMA improves over tested linear/random\-feature attention variants on WikiText\-103 but remains behind optimized causal SDPA and Mamba in the current implementation\. Analysis of learned responsibilities further shows broad component usage and moderate alignment with surface\-form token categories, supporting GMA as a probabilistic, interpretable, fixed\-KKlinear\-time attention\-style alternative rather than a universal replacement for optimized softmax attention or state\-space models\.

## 1Introduction

TheTransformerarchitecture has achieved strong performance across language, vision, and multimodal learning, driven largely by the representational power and parallelism of dot\-product Multi\-Head Attention \(MHA\)\(Vaswaniet al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib11); Dosovitskiyet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib12); Radfordet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib13)\)\. However, standard self\-attention computes pairwise interactions between allNNtokens in a sequence\. For one attention head, letQ∈ℝN×dkQ\\in\\mathbb\{R\}^\{N\\times d\_\{k\}\}denote the query matrix,Katt∈ℝN×dkK\_\{\\mathrm\{att\}\}\\in\\mathbb\{R\}^\{N\\times d\_\{k\}\}denote the key matrix, andV∈ℝN×dvV\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\}denote the value matrix\. Heredkd\_\{k\}is the query/key channel dimension used to compute dot\-product scores, whiledvd\_\{v\}is the value dimension of the vectors being aggregated\. Scaled dot\-product attention computes

O=softmax⁡\(Q​Katt⊤dk\)​V,O∈ℝN×dv\.O=\\operatorname\{softmax\}\\left\(\\frac\{QK\_\{\\mathrm\{att\}\}^\{\\top\}\}\{\\sqrt\{d\_\{k\}\}\}\\right\)V,\\qquad O\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\}\.\(1\)The intermediate score matrixQ​Katt⊤∈ℝN×NQK\_\{\\mathrm\{att\}\}^\{\\top\}\\in\\mathbb\{R\}^\{N\\times N\}contains token\-to\-token scores between all query and key positions\. Although the final output has dimensionO∈ℝN×dvO\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\}, computing dense attention requires forming or implicitly representing interactions between all pairs of positions\. In the standard explicit formulation, this leads to𝒪​\(N2\)\\mathcal\{O\}\(N^\{2\}\)attention\-score storage and𝒪​\(N2​dk\+N2​dv\)\\mathcal\{O\}\(N^\{2\}d\_\{k\}\+N^\{2\}d\_\{v\}\)arithmetic for the score and value multiplications\. Whendkd\_\{k\}anddvd\_\{v\}are treated as fixed, this gives the familiarquadratic scaling in sequence length\. This quadratic dependence makes long\-context modelling expensive and motivates efficient alternatives for settings such as long\-document processing, byte\-level classification, high\-resolution vision, genomics, and autoregressive language modelling\(Beltagyet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib14); Zaheeret al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib15); Tayet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib6); Wanget al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib5); Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4); Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7); Avsecet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib16)\)\.

A large body of work has attempted to reduce the quadratic attention bottleneck through sparse attention patterns, low\-rank projections, kernel approximations, optimized exact\-attention kernels, and recurrent or state\-space sequence models\(Tayet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib17); Beltagyet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib14); Zaheeret al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib15); Wanget al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib5); Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7); Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4); Daoet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib47); Guet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib18); Gu and Dao,[2024](https://arxiv.org/html/2606.18283#bib.bib8)\)\. Sparse attention methods, such asLongformerandBigBird, reduce computation by restricting each token to attend to a subset of local, global, or structured positions\(Beltagyet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib14); Zaheeret al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib15)\)\. Low\-rank methods such asLinformercompress the attention matrix through learned projections along the sequence dimension\(Wanget al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib5)\)\. Kernel\-based methods replace the softmax attention kernel with feature\-map constructions, including deterministic positive feature maps inLinear Transformerand random\-feature approximations inPerformer\(Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7); Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4)\)\. IO\-aware implementations such asFlashAttentionretain exact softmax attention while reducing memory traffic through hardware\-aware tiling\(Daoet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib47)\)\. More recently, structured state\-space and selective state\-space models, includingS4andMamba, have pursued linear or near\-linear sequence modelling by replacing explicit attention with recurrent or state\-space dynamics\(Guet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib18); Gu and Dao,[2024](https://arxiv.org/html/2606.18283#bib.bib8)\)\. These approaches expose different trade\-offs: they can be highly efficient, but their internal routing structure is often less directly probabilistic or less immediately interpretable as a token\-to\-token distribution than the original softmax attention weights\.

In this work, we proposeGaussian Mixture Attention \(GMA\), a probabilistic alternative to dot\-product attention that reconceptualizes sequence mixing aslatent responsibility\-based routing\. Instead of explicitly computing all pairwise token\-to\-token similarities, GMA introducesKKlearned Gaussian mixture components in a projectedrouting representation space\. Query and key representations are mapped to posterior responsibility vectors over these components\. Key responsibilities write values into a latent memoryV~∈ℝK×dv\\tilde\{V\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\}together with a component\-wise normalizerZ∈ℝKZ\\in\\mathbb\{R\}^\{K\}, while query responsibilities read from this normalized latent memory to produce token\-level outputs\. Thus, GMA replaces the explicitN×NN\\times Ntoken\-to\-token attention matrix with two non\- negativeresponsibility matricesΓQ,ΓK∈ℝN×K\\Gamma^\{Q\},\\Gamma^\{K\}\\in\\mathbb\{R\}^\{N\\times K\}\. Although the induced affinityΓQ​\(ΓK\)⊤\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}is still anN×NN\\times Nmatrix algebraically, GMA does not materialize it\. Instead, it exploits the associativity of matrix multiplication by computing the key–value latent memory\(ΓK\)⊤​VX\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}first and then multiplying byΓQ\\Gamma^\{Q\}\. For fixedKK, this gives linear\-in\-NNdominant activation scaling while retaining a normalized attention\-style routing interpretation\.

A key motivation for GMA is that the responsibility matrices are not merely computational intermediates: they are analyzable probabilistic objects\. The marginal usage of mixture components can diagnose whether the latent routing space is broadly used or collapsed to a small subset of components, while hard assignments derived fromzi=arg⁡maxk⁡γi,kz\_\{i\}=\\arg\\max\_\{k\}\\gamma\_\{i,k\}can be compared with token categories or other annotations\. This provides an interpretability handle that is less direct in random\-feature attention approximations or implicit recurrent/state\-space hidden dynamics\. Our later analysis shows that learned GMA responsibilities use most available components and exhibit moderate alignment with surface\-form token categories, although the components should not be interpreted as clean semantic classes \(Section[5\.4](https://arxiv.org/html/2606.18283#S5.SS4)\)\.

Our contributions and findings are as follows:

1. 1\.We introduceGaussian Mixture Attention\(GMA\), a normalized responsibility\-based sequence mixer that replaces explicit token\-to\-token attention with routing through learned Gaussian mixture components\.
2. 2\.We derive both bidirectional and causal GMA\. The causal variant uses prefix latent memories and prefix normalizers so that autoregressive predictions at positioniidepend only on positionsj≤ij\\leq i, while preserving fixed\-KKlinear\-in\-sequence\-length scaling\.
3. 3\.We analyze GMA’s optimization and representational structure, including responsibility\-modulated gradients, a constrained non\-negative low\-rank affinity interpretation, and local Lipschitz continuity under bounded inputs and variance lower bounds\.
4. 4\.We evaluate GMA in 4 empirical settings: controlled systems profiling \(Table[1](https://arxiv.org/html/2606.18283#S5.T1)\), Long Range Arena \(LRA\) long\-context classification\(Tayet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib6)\)\(Table[2](https://arxiv.org/html/2606.18283#S5.T2)\), WikiText\-103 autoregressive language modelling\(Merityet al\.,[2016](https://arxiv.org/html/2606.18283#bib.bib10)\)\(Table[3](https://arxiv.org/html/2606.18283#S5.T3)\), and latent\-responsibility interpretability analysis \(Table[4](https://arxiv.org/html/2606.18283#S5.T4); Figures[2](https://arxiv.org/html/2606.18283#S5.F2)–[3](https://arxiv.org/html/2606.18283#S5.F3)\)\. The results show that GMA exhibits the intended linear memory scaling and achieves competitive performance among attention\-style baselines evaluated in our pipeline\. On LRA, it gives the strongest average performance among those attention\-style baselines \(Table[2](https://arxiv.org/html/2606.18283#S5.T2)\)\. On WikiText\-103, causal GMA improves over Linear Transformer\(Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7)\)and Performer\(Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4)\), although optimized causal SDPA\(Vaswaniet al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib11)\)and Mamba\(Gu and Dao,[2024](https://arxiv.org/html/2606.18283#bib.bib8)\)remain stronger \(Table[3](https://arxiv.org/html/2606.18283#S5.T3)\)\. Finally, learned GMA responsibilities use most available components and show moderate alignment with surface\-form token categories \(Table[4](https://arxiv.org/html/2606.18283#S5.T4); Figures[2](https://arxiv.org/html/2606.18283#S5.F2)–[3](https://arxiv.org/html/2606.18283#S5.F3)\)\.
5. 5\.We discuss future extensions of the GMA framework, including optimized and hybrid GMA implementations, cross\-attention and multimodal routing, Bayesian and Dirichlet\-process variants for adaptive component weighting, and probabilistic Mixture\-of\-Experts routing\.

## 2Related Work

#### Attention as learned compatibility\.

Attention mechanisms can be viewed broadly as methods for computingcompatibilitybetweenquery representationsandcontext representations, and then using the resulting weights to aggregate values\. Early neural encoder–decoder models used learned alignment mechanisms to focus decoding on relevant source positions\(Bahdanauet al\.,[2016](https://arxiv.org/html/2606.18283#bib.bib35); Luonget al\.,[2015](https://arxiv.org/html/2606.18283#bib.bib49)\)\. TheTransformermadescaled dot\-product attentionthe dominant formulation by computing token\-to\-token scores throughQ​Katt⊤QK\_\{\\mathrm\{att\}\}^\{\\top\}and applying row\-wise softmax normalization\(Vaswaniet al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib11)\), as in Eq\. \([1](https://arxiv.org/html/2606.18283#S1.E1)\)\. This design is highly expressive and parallelizable: every token can form query–key scores against all other tokens using batched matrix multiplication, and the value aggregation can likewise be computed by dense linear algebra on modern accelerators\. This parallel structure is a major reason for the success ofTransformerarchitectures, but it also requires forming or implicitly representing anN×NN\\times Ninteraction pattern, giving the familiar quadratic scaling in sequence length\.

From a broader design perspective, the dot product is only one possible compatibility function\. Attention scores may also be based on learned additive scores, kernels, distances, divergences, sparsity patterns, or latent routing structures\.GMAfollows this broader view by replacing direct pairwise query–key comparison with compatibility in a probabilistic responsibility space\. A more explicit taxonomy of similarity, distance, divergence, and latent\-space attention designs is given in Appendix[G](https://arxiv.org/html/2606.18283#A7)\.

#### Sparse and structured attention\.

One major line of efficient\-Transformer research reduces the number of query–key pairs that are compared\.Sparse Transformeruses structured sparse patterns to generate long sequences more efficiently than dense attention\(Childet al\.,[2019](https://arxiv.org/html/2606.18283#bib.bib36)\)\.Reformerreplaces dense dot\-product attention with locality\-sensitive hashing attention and also uses reversible residual layers to reduce activation storage\(Kitaevet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib50)\)\.Longformercombines a local sliding\-window pattern with task\-motivated global attention, giving linear scaling in sequence length and strong performance on long\-document tasks\(Beltagyet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib14)\)\.BigBirdcombines local, random, and global attention patterns, obtaining linear sparse attention while preserving important theoretical properties of full attention, including universal approximation and Turing\-completeness results under its sparse pattern\(Zaheeret al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib15)\)\. These methods retain token\-to\-token attention but restrict the attention graph\. Their efficiency therefore depends on a designed or sampled sparsity pattern\.GMAdiffers by keeping dense global information available through latent mixture components rather than choosing a sparse set of token pairs\.

#### Low\-rank and landmark approximations\.

A second line approximates the attention matrix by a lower\-dimensional representation\.Linformerargues that self\-attention can be approximated by a low\-rank matrix and projects the sequence dimension to reduce the complexity of attention to linear time and space\(Wanget al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib5)\)\.Nyströmformeruses the Nyström method to approximate self\-attention through a set of landmark points, also targeting linear complexity with favorable long\-sequence performance\(Xionget al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib51)\)\. These approaches are related toGMAin that they replace the fullN×NN\\times Ninteraction structure by a smaller intermediate representation\. However, their intermediate factors are projection\- or landmark\-based approximations of attention\.GMAinstead uses posterior responsibilities under learned Gaussian mixture components, producing a constrained non\-negative low\-rank affinityΓQ​\(ΓK\)⊤\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}with a probabilistic routing interpretation\.

#### Kernelized and associative linear attention\.

Kernelized attentionmethods exploit the associativity of matrix multiplication\.Linear Transformerrewrites attention using feature maps so that the key–value summary can be computed before multiplying by queries, reducing the cost from quadratic to linear in sequence length\(Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7)\)\.Performerapproximates softmax attention using positive orthogonal random features in the FAVOR\+ mechanism, giving linear space and time with theoretical guarantees for the random\-feature approximation\(Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4)\)\. These methods are especially close to GMA at the computational level: all use an ordering of the form “summarize keys and values first, then read with queries”\. The difference lies in the feature map\. Kernelized attention uses deterministic or random feature maps designed to approximate a kernel such as softmax\.GMAuses learned Gaussian\-mixture posterior responsibilities as the feature map, so the intermediate representation is a probability vector over latent components\. This givesGMAan additional diagnostic object, i\.e\. the responsibility matrix, but also introduces extra constant factors from Gaussian log\-density and responsibility computations\.

#### IO\-aware exact attention\.

Another important direction does not change the mathematical attention definition, but optimizes its implementation\.FlashAttentioncomputes exact softmax attention using IO\-aware tiling, reducing reads and writes between GPU high\-bandwidth memory and on\-chip SRAM\(Daoet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib47)\)\. LaterFlashAttentionvariants further improve hardware utilization through more aggressive scheduling and low\-precision support\(Shahet al\.,[2024](https://arxiv.org/html/2606.18283#bib.bib52)\)\. These methods are essential baselines because they show that quadratic attention can be much faster and more memory efficient in practice without changing the underlying attention weights\.GMAaddresses a different question: it changes the attention\-style operator itself to obtain fixed\-KKlinear activation scaling and interpretable latent routing\. Consequently, the relevant comparison is not only asymptotic complexity but also wall\-clock efficiency, kernel maturity, memory traffic, and model quality\.

#### Latent bottleneck attention\.

Latent\-array architectures reduce the cost of processing large inputs by routing information through a smaller set of latent variables\. ThePerceiveruses asymmetric cross\-attention to distill high\-dimensional inputs into a latent bottleneck and then processes the latent array with Transformer blocks\(Jaegleet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib53)\)\.PerceiverIO extends this idea to flexible structured outputs by querying the latent representation with output\-specific queries\(Jaegleet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib54)\)\.GMAshares the high\-level idea that a smaller latent substrate can mediate interactions among many input positions\. However, the latent objects are different\.Perceiveruses learned latent vectors that attend to inputs through cross\-attention\. GMA uses learned Gaussian mixture components in routing space and maps each query/key token to a posterior responsibility vector over these components\. The resulting latent memory is not a separate learned sequence, but a key\-weighted aggregation of values through mixture responsibilities\.

#### State\-space and recurrent alternatives to attention\.

Structured state\-space models provide a separate route to long\-sequence modelling by replacing explicit attention with recurrent or convolutional sequence dynamics\.S4introduced structured state\-space layers for efficient long\-range sequence modelling\(Guet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib18)\)\.Mambabuilds on this line by making state\-space parameters input\-dependent, allowing selective propagation or forgetting of information, and by using hardware\-aware parallel algorithms for efficient training and inference\(Gu and Dao,[2024](https://arxiv.org/html/2606.18283#bib.bib8)\)\. These models often provide strong throughput and strong performance, especially in long\-sequence or discrete\-token settings\. They are not attention mechanisms in the usual token\-to\-token sense, and their internal routing structure is less directly exposed as an attention matrix\.GMAoccupies a different position: it remains an attention\-style sequence mixer with an induced token\-to\-token affinity, while introducing a probabilistic latent routing space that can be inspected through responsibility diagnostics\.

#### Mixture models and probabilistic attention\.

Mixture models have a long history as probabilistic tools for soft assignment, density modelling, and responsibility\-based learning\(Dempsteret al\.,[1977](https://arxiv.org/html/2606.18283#bib.bib27); Bishop,[1994](https://arxiv.org/html/2606.18283#bib.bib30); McLachlan and Peel,[2000](https://arxiv.org/html/2606.18283#bib.bib28); Bishop,[2006](https://arxiv.org/html/2606.18283#bib.bib29)\)\. Several recent works have explored connections between mixture models and attention\.Transformer\-MGK\(Transformer with a Mixture of Gaussian Keys\) replaces redundant attention heads with a mixture of Gaussian keys at each head, allowing each head to focus on different parts of the input while reducing parameters and computation\(yenet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib55)\)\. However, the mixture is local to the key representation inside an attention head, and the model still relies on query–key attention computations unless combined with separate linear\-attention approximations\. In machine translation,GMM\-based cross\-attentionhas been used to model concentrated attention around central source positions, improving alignment quality and long\-sentence translation behaviour\(Zhang and Feng,[2021](https://arxiv.org/html/2606.18283#bib.bib56)\)\. These works show that Gaussian mixtures can be useful inside attention, but they do not use a shared global mixture responsibility space to replace the explicitN×NN\\times Ntoken\-to\-token attention matrix\.

GMAdiffers from these mixture\-based approaches in three ways\. First, it uses a shared global Gaussian mixture over routing vectors, rather than a local mixture attached to each position or attention head\. Second, both queries and keys are mapped into responsibility vectors over the same mixture components, inducing the non\-negative affinityA~GMA=ΓQ​\(ΓK\)⊤\\widetilde\{A\}^\{\\mathrm\{GMA\}\}=\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\. Third, the efficient implementation does not materialize this affinity; it uses the associative ordering\(ΓK\)⊤​VX\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}followed by multiplication withΓQ\\Gamma^\{Q\}\. Thus, GMA is not merely a Gaussian modification of attention scores\. It combines probabilistic responsibility\-space similarity with latent\-memory routing to obtain fixed\-KKlinear activation scaling\.

#### Mixture\-of\-Experts and probabilistic routing\.

GMAis also related to the broader literature on conditional computation and Mixture\-of\-Experts \(MoE\)\. Classical MoE models combine multiple expert predictors through input\-dependent gates\(Jacobset al\.,[1991](https://arxiv.org/html/2606.18283#bib.bib57)\)\. Modern sparsely gated MoE models scale Transformer capacity by routing each token to a small number of feed\-forward experts\(Shazeeret al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib43); Lepikhinet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib45); Feduset al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib44)\)\. Recent work also explores MoE\-style routing inside the attention mechanism itself, for example by switching among attention heads or computing fewer attention matrices\(Csordáset al\.,[2024](https://arxiv.org/html/2606.18283#bib.bib58)\)\. These methods typically use learned softmax or noisy top\-kkgates and are motivated by conditional computation\.GMA’s routing is different: the gating weights are posterior responsibilities under a learned Gaussian mixture model\. This makes the routing probabilities density\-based and interpretable as soft assignments to latent components\. Although our present work uses GMA for sequence mixing rather than sparse expert selection, the same responsibility view suggests future probabilistic MoE extensions, as discussed in Section[6](https://arxiv.org/html/2606.18283#S6.SSx1)\.

#### Positioning of GMA\.

Overall, prior work has reduced the quadratic attention bottleneck by sparsifying the attention graph, approximating the attention matrix, kernelizing the softmax, optimizing exact attention kernels, introducing latent bottlenecks, or replacing attention with state\-space dynamics\.GMAcontributes a complementary design point\. It treats attention as routing through a learned probabilistic latent space: queries and keys are first converted into Gaussian\-mixture responsibilities, values are written into a latent memory by key responsibilities, and queries read from that memory through normalized responsibility overlap\. This givesGMAa fixed\-KKlinear\-time attention\-style operator with an analyzable latent responsibility structure\. Its purpose is not to replace all optimized attention or state\-space models, but to expose a probabilistic and interpretable route to efficient sequence mixing\.

## 3Methodology: Gaussian Mixture Attention

We present GMA primarily in theself\-attentionsetting, where queries, keys, and values are all projected from the same sequenceXX\. For clarity, the main equations are written for a single attention head; a multi\-head GMA layer applies the same computation independently across heads and then combines the resulting head outputs in the usual way\. This formulation covers the bidirectional sequence\-mixing experiments and the causal language\-modelling experiments considered later in this paper\.Cross\-attentionfollows analogously by forming queries from a query\-side sequenceXXand keys/values from a separate context sequenceYY; the same responsibility\-space routing mechanism can then be applied withΓXQ\\Gamma\_\{X\}^\{Q\}andΓYK\\Gamma\_\{Y\}^\{K\}\. A detailed distinction between standard self\-attention, standard cross\-attention, and their GMA counterparts is given in Appendix[E](https://arxiv.org/html/2606.18283#A5)\. We leave a systematic empirical study of GMA cross\-attention to future work\.

### 3\.1Method Overview

Standard self\-attention relies on direct pairwise token\-to\-token comparisons, which gives the familiar quadratic interaction pattern in sequence length\. GMA changes this computation by introducing a latent memory architecture based onKKlearned Gaussian routing components\. Instead of comparing each query with each key directly, GMA maps projected queries and keys to posterior responsibility vectors over a shared Gaussian mixture\.

Intuitively, the projected value matrixVXV\_\{X\}stores token\-level information\. The key responsibility matrixΓK\\Gamma^\{K\}acts as a soft writing mechanism, routing value vectors intoKKlatent memory slots\. This produces a latent memoryV~\\tilde\{V\}together with a component\-wise normalizerZZ\. Query responsibilitiesΓQ\\Gamma^\{Q\}then act as a soft reading mechanism, determining how each output token reads from the normalized latent memory\. Thus, queries do not explicitly compare against all keys in the efficient implementation; instead, query–key compatibility is mediated through shared responsibility vectors over the latent Gaussian components\. A consolidated notation table for the methodology and analysis is provided in Appendix[A](https://arxiv.org/html/2606.18283#A1)\.

### 3\.2Probabilistic Latent Routing

Instead of computing direct token\-to\-token similarities, we introduceKKlearned latent centers, representing Gaussian routing components\. We model these centers as the components of a Gaussian Mixture Model \(GMM\)Huang \([2025](https://arxiv.org/html/2606.18283#bib.bib23),[2026a](https://arxiv.org/html/2606.18283#bib.bib1)\)\. Thekk\-th component is parameterized by a meanμk∈ℝdr\\mu\_\{k\}\\in\\mathbb\{R\}^\{d\_\{r\}\}and a covariance matrixΣk∈𝕊\+\+dr\\Sigma\_\{k\}\\in\\mathbb\{S\}\_\{\+\+\}^\{d\_\{r\}\}, wheredrd\_\{r\}is the routing dimension\. For computational tractability, we use diagonal covariances,Σk=diag⁡\(σk,12,…,σk,dr2\)\\Sigma\_\{k\}=\\operatorname\{diag\}\(\\sigma\_\{k,1\}^\{2\},\\ldots,\\sigma\_\{k,d\_\{r\}\}^\{2\}\)\. The mixture prior of componentkkisπk\\pi\_\{k\}, with∑k=1Kπk=1\\sum\_\{k=1\}^\{K\}\\pi\_\{k\}=1\.

For a routing vectorxi∈ℝdrx\_\{i\}\\in\\mathbb\{R\}^\{d\_\{r\}\}, which may be either a query projection or a key projection, the density under componentkkis

𝒩​\(xi∣μk,Σk\)=1\(2​π\)dr/2​\|Σk\|1/2​exp⁡\(−12​\(xi−μk\)⊤​Σk−1​\(xi−μk\)\)\.\\mathcal\{N\}\(x\_\{i\}\\mid\\mu\_\{k\},\\Sigma\_\{k\}\)=\\frac\{1\}\{\(2\\pi\)^\{d\_\{r\}/2\}\|\\Sigma\_\{k\}\|^\{1/2\}\}\\exp\\left\(\-\\frac\{1\}\{2\}\(x\_\{i\}\-\\mu\_\{k\}\)^\{\\top\}\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\)\\right\)\.\(2\)Letzi∈\{1,…,K\}z\_\{i\}\\in\\\{1,\\ldots,K\\\}denote the latent component assignment forxix\_\{i\}\. The responsibility of componentkkforxix\_\{i\}is the posterior probabilityp​\(zi=k∣xi\)p\(z\_\{i\}=k\\mid x\_\{i\}\):

γi,k≡p​\(zi=k∣xi\)=p​\(zi=k\)​p​\(xi∣zi=k\)p​\(xi\)=πk​𝒩​\(xi∣μk,Σk\)∑ℓ=1Kπℓ​𝒩​\(xi∣μℓ,Σℓ\)\.\\gamma\_\{i,k\}\\equiv p\(z\_\{i\}=k\\mid x\_\{i\}\)=\\frac\{p\(z\_\{i\}=k\)p\(x\_\{i\}\\mid z\_\{i\}=k\)\}\{p\(x\_\{i\}\)\}=\\frac\{\\pi\_\{k\}\\mathcal\{N\}\(x\_\{i\}\\mid\\mu\_\{k\},\\Sigma\_\{k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\pi\_\{\\ell\}\\mathcal\{N\}\(x\_\{i\}\\mid\\mu\_\{\\ell\},\\Sigma\_\{\\ell\}\)\}\.\(3\)This operation maps a sequence of routing vectors into a responsibility matrixΓ∈ℝN×K\\Gamma\\in\\mathbb\{R\}^\{N\\times K\}\. Since∑k=1Kγi,k=1\\sum\_\{k=1\}^\{K\}\\gamma\_\{i,k\}=1for every tokenii, each row ofΓ\\Gammais a probability distribution over the latent Gaussian components\.

### 3\.3Linear\-Time Normalized Attention via GMA

LetX∈ℝN×dmodelX\\in\\mathbb\{R\}^\{N\\times d\_\{\\mathrm\{model\}\}\}denote the input sequence representations to a GMA layer, whereNNis the sequence length anddmodeld\_\{\\mathrm\{model\}\}is the model hidden dimension, i\.e\., the dimensionality of each token representation before it is projected into queries, keys, and values\. As in standard attention, we first form query, key, and value projections:

QX=X​WQ∈ℝN×dr,KX=X​WK∈ℝN×dr,VX=X​WV∈ℝN×dv,Q\_\{X\}=XW\_\{Q\}\\in\\mathbb\{R\}^\{N\\times d\_\{r\}\},\\qquad K\_\{X\}=XW\_\{K\}\\in\\mathbb\{R\}^\{N\\times d\_\{r\}\},\\qquad V\_\{X\}=XW\_\{V\}\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\},\(4\)whereWQ,WK∈ℝdmodel×drW\_\{Q\},W\_\{K\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{model\}\}\\times d\_\{r\}\}andWV∈ℝdmodel×dvW\_\{V\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{model\}\}\\times d\_\{v\}\}\. Heredrd\_\{r\}is the routing dimension used to compute Gaussian responsibilities, anddvd\_\{v\}is the value dimension\. We reserve scalarKKfor the number of Gaussian mixture components\.

GivenQXQ\_\{X\}andKXK\_\{X\}, GMA computes two responsibility matrices:

ΓQ∈ℝN×K,ΓK∈ℝN×K\.\\Gamma^\{Q\}\\in\\mathbb\{R\}^\{N\\times K\},\\qquad\\Gamma^\{K\}\\in\\mathbb\{R\}^\{N\\times K\}\.The entryγi,kQ\\gamma^\{Q\}\_\{i,k\}is the posterior responsibility of componentkkfor the query vectorqi=\(QX\)i,:∈ℝdrq\_\{i\}=\(Q\_\{X\}\)\_\{i,:\}\\in\\mathbb\{R\}^\{d\_\{r\}\}, whileγj,kK\\gamma^\{K\}\_\{j,k\}is the posterior responsibility of componentkkfor the key vectorkj=\(KX\)j,:∈ℝdrk\_\{j\}=\(K\_\{X\}\)\_\{j,:\}\\in\\mathbb\{R\}^\{d\_\{r\}\}\. These responsibilities are calculated using the Gaussian\-mixture posterior in Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\), applied separately to the query and key projections\. The superscriptKKinΓK\\Gamma^\{K\}denotes key responsibilities, not the number of mixture components\.

To construct an end\-to\-end𝒪​\(N​K\)\\mathcal\{O\}\(NK\)attention mechanism, GMA avoids forming the fullN×NN\\times Ntoken\-to\-token attention matrix\. Instead, it uses the key responsibilities to write the projected valuesVXV\_\{X\}intoKKlatent memory slots, and then uses the query responsibilities to read from those slots\. To make the induced attention weights stochastic over the sequence dimension, we normalize the latent routing operation\. The computation proceeds in two associative steps\.

1. 1\.Latent aggregation\.The key responsibilities aggregate the value vectors intoKKlatent memory slots and compute the total key mass assigned to each component: V~=\(ΓK\)⊤​VX∈ℝK×dv,Z=\(ΓK\)⊤​𝟏N∈ℝK,\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\},\\qquad Z=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\}\\in\\mathbb\{R\}^\{K\},\(5\)where𝟏N\\mathbf\{1\}\_\{N\}denotes the all\-ones vector\. Component\-wise333From a linear algebra perspective, thekk\-th row ofV~\\tilde\{V\}, i\.e\.V~k\\tilde\{V\}\_\{k\}, is a linear combination of all rows ofVXV\_\{X\}, with weights being thekk\-th column ofΓK\\Gamma^\{K\}\., V~k=∑j=1Nγj,kK​VX,j∈ℝdv,Zk=∑j=1Nγj,kK\.\\tilde\{V\}\_\{k\}=\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\\in\\mathbb\{R\}^\{d\_\{v\}\},\\qquad Z\_\{k\}=\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\.\(6\)Thus,V~k\\tilde\{V\}\_\{k\}is the responsibility\-weighted sum of all value vectors assigned to componentkk, andZkZ\_\{k\}is the total key\-responsibility mass assigned to that component\.
2. 2\.Normalized broadcasting\.The query responsibilities read from the latent memory and normalize by the corresponding routed key mass: O=ΓQ​V~ΓQ​Z\+ϵ=𝚪𝑸​\(𝚪𝑲\)⊤​𝑽𝑿𝚪𝑸​𝒁\+ϵ=AGMA​VX∈ℝN×dv,O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}=\\boldsymbol\{\\frac\{\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}\}=A^\{\\mathrm\{GMA\}\}V\_\{X\}\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\},\(7\)where we have defined the implicit normalized affinityAGMA=ΓQ​\(ΓK\)⊤ΓQ​Z\+ϵA^\{\\mathrm\{GMA\}\}=\\frac\{\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}, which is useful for interpretation and analysis, but not materialized in the efficient implementation\. The denominator is broadcast across the value dimension, the division is appliedrow\-wise, andϵ\>0\\epsilon\>0is used for numerical stability444Equivalently, one may write the row\-wise normalization in matrix form\. Defined=ΓQ​Z\+ϵ​𝟏N∈ℝN\.d=\\Gamma^\{Q\}Z\+\\epsilon\\mathbf\{1\}\_\{N\}\\in\\mathbb\{R\}^\{N\}\.ThenO=diag\(d\)−1ΓQV~=diag\(d\)−1ΓQ\(ΓK\)⊤VX\.O=\\operatorname\{diag\}\(d\)^\{\-1\}\\Gamma^\{Q\}\\tilde\{V\}=\\operatorname\{diag\}\(d\)^\{\-1\}\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\.The corresponding implicit normalized affinity matrix isAGMA=diag\(d\)−1ΓQ\(ΓK\)⊤\.A^\{\\mathrm\{GMA\}\}=\\operatorname\{diag\}\(d\)^\{\-1\}\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\.This matrix is useful for interpretation and analysis, but it is not materialized in the efficient implementation\.\.

Specifically, for tokenii, the output can be written as

Oi=∑k=1Kγi,kQ​V~k∑k=1Kγi,kQ​Zk\+ϵ=∑k=1Kγi,kQ​∑j=1Nγj,kK​VX,j∑k=1Kγi,kQ​∑j=1Nγj,kK\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z\_\{k\}\+\\epsilon\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\+\\epsilon\}\.\(8\)Thus, GMA induces the normalized attention weights

Ai​jGMA=∑k=1Kγi,kQ​γj,kK∑ℓ=1N∑k=1Kγi,kQ​γℓ,kK\+ϵ,A^\{\\mathrm\{GMA\}\}\_\{ij\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\}\{\\sum\_\{\\ell=1\}^\{N\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{\\ell,k\}\+\\epsilon\},\(9\)so that

Oi=∑j=1NAi​jGMA​VX,jO\_\{i\}=\\sum\_\{j=1\}^\{N\}A^\{\\mathrm\{GMA\}\}\_\{ij\}V\_\{X,j\}\(10\)up to the numerical stabilizer\. Thus, GMA replaces the explicitN×NN\\times Nattention matrix by two non\-negative responsibility matricesΓQ,ΓK∈ℝN×K\\Gamma^\{Q\},\\Gamma^\{K\}\\in\\mathbb\{R\}^\{N\\times K\}, while still admitting the implicit normalized affinity in Eq\. \([9](https://arxiv.org/html/2606.18283#S3.E9)\)\.

The whole GMA workflow is visually represented in Fig\.[1](https://arxiv.org/html/2606.18283#S3.F1)\. This ordering reflects a different routing philosophy from standard dot\-product attention in Eq\. \([1](https://arxiv.org/html/2606.18283#S1.E1)\)\. Standard attention first evaluates query–key similarity by forming the token\-to\-token interaction matrixQ​K⊤QK^\{\\top\}, normalizes this matrix into attention weights, and then uses those weights to collect information from the value matrixVV\. Informally, each query first “knocks on” all keys to decide where relevant content is located, and then retrieves the corresponding values\. GMA reverses this order\. It first uses the key projectionKXK\_\{X\}to compute key responsibilitiesΓK\\Gamma^\{K\}, which write the value projectionVXV\_\{X\}into a compact latent memoryV~=\(ΓK\)⊤​VX\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\. The query responsibilitiesΓQ\\Gamma^\{Q\}then read directly from thisKK\-slot memory through normalized broadcasting\. Thus, standard attention follows a “query–key matching then value collection” principle, whereas GMA follows a “key\-induced memory construction then query\-based memory reading” principle\.

Input Sequence\(X\)\(X\)ℝN×dmodel\\mathbb\{R\}^\{N\\times d\_\{\\mathrm\{model\}\}\}Query Projection\(QX\)\(Q\_\{X\}\)QX=X​WQQ\_\{X\}=XW\_\{Q\}ℝN×dr\\mathbb\{R\}^\{N\\times d\_\{r\}\}Read Responsibilities\(ΓQ\)\(\\Gamma^\{Q\}\)via Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\)ℝN×K\\mathbb\{R\}^\{N\\times K\}Output\(O\)\(O\)O=ΓQ​V~ΓQ​Z\+ϵ\\displaystyle O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}ℝN×dv\\mathbb\{R\}^\{N\\times d\_\{v\}\}Key Projection\(KX\)\(K\_\{X\}\)KX=X​WKK\_\{X\}=XW\_\{K\}ℝN×dr\\mathbb\{R\}^\{N\\times d\_\{r\}\}Shared GMM Parameters\{πk,μk,Σk\}k=1K\\\{\\pi\_\{k\},\\mu\_\{k\},\\Sigma\_\{k\}\\\}\_\{k=1\}^\{K\}Value Projection\(VX\)\(V\_\{X\}\)VX=X​WVV\_\{X\}=XW\_\{V\}ℝN×dv\\mathbb\{R\}^\{N\\times d\_\{v\}\}Write Responsibilities\(ΓK\)\(\\Gamma^\{K\}\)via Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\)ℝN×K\\mathbb\{R\}^\{N\\times K\}Latent Memory\(V~\)\(\\tilde\{V\}\)V~=\(ΓK\)⊤​VX\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}ℝK×dv\\mathbb\{R\}^\{K\\times d\_\{v\}\}Normalizer\(Z\)\(Z\)Z=\(ΓK\)⊤​𝟏NZ=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\}ℝK\\mathbb\{R\}^\{K\}writereadFigure 1:End\-to\-end workflow of Gaussian Mixture Attention \(GMA\)\. The input sequenceXXis first projected into query, key, and value representationsQXQ\_\{X\},KXK\_\{X\}, andVXV\_\{X\}\. A shared Gaussian Mixture Model with parameters\{πk,μk,Σk\}k=1K\\\{\\pi\_\{k\},\\mu\_\{k\},\\Sigma\_\{k\}\\\}\_\{k=1\}^\{K\}mapsQXQ\_\{X\}andKXK\_\{X\}to query and key responsibility matrices,ΓQ\\Gamma^\{Q\}andΓK\\Gamma^\{K\}, using the posterior responsibility formula in Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\)\. The key responsibilities write the projected valuesVXV\_\{X\}into a latent memoryV~\\tilde\{V\}and a component\-wise normalizerZZ\. The query responsibilities then read from this normalized latent memory to produce the output sequenceOO\. This replaces explicitN×NN\\times Ntoken\-to\-token attention with routing throughKKlatent Gaussian components\.
### Understanding the GMA Workflow

The GMA computation can be summarized as*responsibility\-space affinity plus latent\-memory routing*\. Unlike standard dot\-product attention, GMA does not directly score a query and key by the dot productqi⊤​kjq\_\{i\}^\{\\top\}k\_\{j\}\. Instead, it first maps both query and key representations into posterior responsibility vectors over the sameKKlearned Gaussian mixture components\.

For a routing vectorxx, the responsibility of componentkkis \(Eq\.[3](https://arxiv.org/html/2606.18283#S3.E3)\)

γk​\(x\)=p​\(z=k∣x\)=πk​𝒩​\(x∣μk,Σk\)∑ℓ=1Kπℓ​𝒩​\(x∣μℓ,Σℓ\)\.\\gamma\_\{k\}\(x\)=p\(z=k\\mid x\)=\\frac\{\\pi\_\{k\}\\mathcal\{N\}\(x\\mid\\mu\_\{k\},\\Sigma\_\{k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\pi\_\{\\ell\}\\mathcal\{N\}\(x\\mid\\mu\_\{\\ell\},\\Sigma\_\{\\ell\}\)\}\.For queries and keys, this gives

γi,kQ=p​\(z=k∣qi\),γj,kK=p​\(z=k∣kj\)\.\\gamma^\{Q\}\_\{i,k\}=p\(z=k\\mid q\_\{i\}\),\\qquad\\gamma^\{K\}\_\{j,k\}=p\(z=k\\mid k\_\{j\}\)\.Thus, the query vectorqiq\_\{i\}is represented by the responsibility vector

γiQ=\(γi,1Q,…,γi,KQ\),\\gamma^\{Q\}\_\{i\}=\(\\gamma^\{Q\}\_\{i,1\},\\ldots,\\gamma^\{Q\}\_\{i,K\}\),and the key vectorkjk\_\{j\}is represented by

γjK=\(γj,1K,…,γj,KK\)\.\\gamma^\{K\}\_\{j\}=\(\\gamma^\{K\}\_\{j,1\},\\ldots,\\gamma^\{K\}\_\{j,K\}\)\.GMA compares query and key positions through theoverlapof these responsibility vectors:

A~i​jGMA=∑k=1Kγi,kQ​γj,kK=⟨γiQ,γjK⟩,\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{ij\}=\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}=\\langle\\gamma^\{Q\}\_\{i\},\\gamma^\{K\}\_\{j\}\\rangle,which defines the un\-normalized attention weight matrix \(as compared to its normalised version Eq\.[9](https://arxiv.org/html/2606.18283#S3.E9)\)

A~GMA=ΓQ​\(ΓK\)⊤\.\\widetilde\{A\}^\{\\mathrm\{GMA\}\}=\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\.\(11\)Therefore, two tokens have high GMA affinity if their query and key representations assign high posterior probability mass to similar Gaussian components\.

The second part of the workflow is latent\-memory routing, which avoids materializing the full pairwise affinity matrix in the efficient implementation\. Rather than explicitly materializing the fullN×NN\\times Naffinity matrixA~GMA\\widetilde\{A\}^\{\\mathrm\{GMA\}\}, the key responsibilities first write values into aKK\-slot latent memory \(Eq\.[5](https://arxiv.org/html/2606.18283#S3.E5)\):

V~=\(ΓK\)⊤​VX∈ℝK×dv,Z=\(ΓK\)⊤​𝟏N∈ℝK\.\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\},\\qquad Z=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\}\\in\\mathbb\{R\}^\{K\}\.Component\-wise \(Eq\.[6](https://arxiv.org/html/2606.18283#S3.E6)\),

V~k=∑j=1Nγj,kK​VX,j,Zk=∑j=1Nγj,kK\.\\tilde\{V\}\_\{k\}=\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\},\\qquad Z\_\{k\}=\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\.Thus,V~k\\tilde\{V\}\_\{k\}is thememory slotassociated with Gaussian componentkk: it stores a responsibility\-weighted sum of the value vectors whose keys are assigned to that component\.

Finally, the query responsibilities read from this key\-induced latent memory \(Eq\.[7](https://arxiv.org/html/2606.18283#S3.E7)\):

O=ΓQ​V~ΓQ​Z\+ϵ\.O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}\.For a single output tokenii, this becomes \(Eq\.[8](https://arxiv.org/html/2606.18283#S3.E8)\)

Oi=∑k=1Kγi,kQ​V~k∑k=1Kγi,kQ​Zk\+ϵ=∑k=1Kγi,kQ​∑j=1Nγj,kK​VX,j∑k=1Kγi,kQ​∑j=1Nγj,kK\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z\_\{k\}\+\\epsilon\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\+\\epsilon\}\.Equivalently, GMA inducesnormalizedattention\-style weights \(Eq\.[9](https://arxiv.org/html/2606.18283#S3.E9)\)

Ai​jGMA=∑k=1Kγi,kQ​γj,kK∑ℓ=1N∑k=1Kγi,kQ​γℓ,kK\+ϵ,A^\{\\mathrm\{GMA\}\}\_\{ij\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\}\{\\sum\_\{\\ell=1\}^\{N\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{\\ell,k\}\+\\epsilon\},so that the output for query positioniiis \(Eq\.[10](https://arxiv.org/html/2606.18283#S3.E10)\)

Oi=∑j=1NAi​jGMA​VX,j\.O\_\{i\}=\\sum\_\{j=1\}^\{N\}A^\{\\mathrm\{GMA\}\}\_\{ij\}V\_\{X,j\}\.
In short, GMA first compares query and key positions inresponsibility space, throughA~i​jGMA=⟨γiQ,γjK⟩,\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{ij\}=\\langle\\gamma^\{Q\}\_\{i\},\\gamma^\{K\}\_\{j\}\\rangle,but it avoids forming all such pairwise affinities explicitly\. Instead, it computesV~=\(ΓK\)⊤​VX\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}first, and then reads this key\-induced memory usingO=ΓQ​V~ΓQ​Z\+ϵ\.O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}\.So GMA is best understood asresponsibility\-space affinity followed by latent\-memory routing, rather than as dot\-product matching against memory components\. Appendix[C](https://arxiv.org/html/2606.18283#A3)gives a more dimension\-explicit walkthrough of the same computation\.

### 3\.4Memory and computational costs

For one sequence and one attention head, assume thediagonalcovariance parameterization above and an associative implementation that does not materializeAGMAA^\{\\mathrm\{GMA\}\}\. Excluding the standard linear projections used to formQX,KX,VXQ\_\{X\},K\_\{X\},V\_\{X\}, the main GMA routing operations have the following costs\. Computing Gaussian\-mixture responsibilities for both query and key projections costs approximately

2​N​K​dr\+2​N​K,2NKd\_\{r\}\+2NK,where the first term accounts for evaluating diagonal Gaussian log\-densities overKKcomponents for bothQXQ\_\{X\}andKXK\_\{X\}, and the second term accounts for the row\-wise normalization over components\. The write stepV~=\(ΓK\)⊤​VX\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}costsN​K​dvNKd\_\{v\}, and computingZ=\(ΓK\)⊤​𝟏NZ=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\}costsN​KNK\. The read stepΓQ​V~\\Gamma^\{Q\}\\tilde\{V\}costsN​K​dvNKd\_\{v\}, computing the denominatorΓQ​Z\+ϵ\\Gamma^\{Q\}Z\+\\epsiloncostsN​KNK, and the final row\-wise normalization of the output costsN​dvNd\_\{v\}\. Thus, the total routing cost is

𝒞GMA≈2​N​K​dr\+2​N​K​dv\+4​N​K\+N​dv\.\\mathcal\{C\}\_\{\\mathrm\{GMA\}\}\\approx 2NKd\_\{r\}\+2NKd\_\{v\}\+4NK\+Nd\_\{v\}\.\(12\)Equivalently,

𝒞GMA=𝒪​\(N​K​dr\+N​K​dv\)\.\\mathcal\{C\}\_\{\\mathrm\{GMA\}\}=\\mathcal\{O\}\(NKd\_\{r\}\+NKd\_\{v\}\)\.Including the projection matrices adds the usual cost

𝒞proj=𝒪​\(N​dmodel​\(2​dr\+dv\)\),\\mathcal\{C\}\_\{\\mathrm\{proj\}\}=\\mathcal\{O\}\\\!\\left\(Nd\_\{\\mathrm\{model\}\}\(2d\_\{r\}\+d\_\{v\}\)\\right\),for formingQX=X​WQQ\_\{X\}=XW\_\{Q\},KX=X​WKK\_\{X\}=XW\_\{K\}, andVX=X​WVV\_\{X\}=XW\_\{V\}\. Hence, for fixeddr,dvd\_\{r\},d\_\{v\}, number of heads, and fixed number of mixture componentsKK, GMA scales linearly with sequence lengthNNwhile avoiding explicit𝒪​\(N2\)\\mathcal\{O\}\(N^\{2\}\)attention storage555The matrix\-multiplication associativity underlying this cost saving is discussed in Appendix[B](https://arxiv.org/html/2606.18283#A2)\.\.

The attention\-specific activationstorageis similarly linear inNN\. The two responsibility matrices require2​N​K2NKentries, the latent memory and normalizer requireK​dv\+KKd\_\{v\}\+Kentries, and the output requiresN​dvNd\_\{v\}entries\. Thus, aside from standard projection activations, the routing\-specific storage is

2​N​K\+K​dv\+K\+N​dv=𝒪​\(N​K\+K​dv\+N​dv\),2NK\+Kd\_\{v\}\+K\+Nd\_\{v\}=\\mathcal\{O\}\(NK\+Kd\_\{v\}\+Nd\_\{v\}\),which reduces to𝒪​\(N​K\)\\mathcal\{O\}\(NK\)when focusing on the dominant responsibility storage for fixeddvd\_\{v\}andKK\.

### 3\.5Causal GMA for Autoregressive Modelling

The formulation above is bidirectional: each output position may aggregate values from all key positions\. For autoregressive language modelling, however, the output at positioniimust depend only on prefix positionsj≤ij\\leq i\. Therefore, the non\-causal aggregation in Eq\. \([5](https://arxiv.org/html/2606.18283#S3.E5)\) is replaced by prefix versions of the latent memory and normalizer\. For each positioniiand componentkk, define

V~k\(i\)=∑j≤iγj,kK​VX,j∈ℝdv,Zk\(i\)=∑j≤iγj,kK∈ℝ\.\\tilde\{V\}^\{\(i\)\}\_\{k\}=\\sum\_\{j\\leq i\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\\in\\mathbb\{R\}^\{d\_\{v\}\},\\qquad Z^\{\(i\)\}\_\{k\}=\\sum\_\{j\\leq i\}\\gamma^\{K\}\_\{j,k\}\\in\\mathbb\{R\}\.\([6](https://arxiv.org/html/2606.18283#S3.E6)b\)HereV~k\(i\)\\tilde\{V\}^\{\(i\)\}\_\{k\}andZk\(i\)Z^\{\(i\)\}\_\{k\}are the causal prefix counterparts ofV~k\\tilde\{V\}\_\{k\}andZkZ\_\{k\}in Eq\. \([6](https://arxiv.org/html/2606.18283#S3.E6)\)\. In other words,V~k\(i\)\\tilde\{V\}^\{\(i\)\}\_\{k\}stores the value information assigned to componentkkusing only tokens up to positionii, whileZk\(i\)Z^\{\(i\)\}\_\{k\}stores the corresponding prefix key\-responsibility mass\.

The causal GMA output for positioniiis then666Unlike the bidirectional case in Eq\. \([7](https://arxiv.org/html/2606.18283#S3.E7)\), there is no single full\-sequence matrix expression of the formO=ΓQ​V~/\(ΓQ​Z\+ϵ\)O=\\Gamma^\{Q\}\\tilde\{V\}/\(\\Gamma^\{Q\}Z\+\\epsilon\)unless one introduces an additional prefix\-memory tensor\. The reason is that the causal memoryV~k\(i\)\\tilde\{V\}^\{\(i\)\}\_\{k\}and normalizerZk\(i\)Z^\{\(i\)\}\_\{k\}depend on the output positionii: each query position reads from a different prefix\-restricted memory built only from tokensj≤ij\\leq i\. Therefore, we wrote the causal formulation row\-wise as in Eq\. \([8b](https://arxiv.org/html/2606.18283#S3.Ex22)\)\.

Oi=∑k=1Kγi,kQ​V~k\(i\)∑k=1Kγi,kQ​Zk\(i\)\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}^\{\(i\)\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z^\{\(i\)\}\_\{k\}\+\\epsilon\}\.\([8](https://arxiv.org/html/2606.18283#S3.E8)b\)Thus, query positioniireads only from a prefix\-restricted latent memory, rather than from the full\-sequence memory used in the bidirectional case\.

Equivalently, causal GMA induces attention weights supported only on the prefix:

Ai​jcGMA=𝟏​\{j≤i\}​∑k=1Kγi,kQ​γj,kK∑ℓ≤i∑k=1Kγi,kQ​γℓ,kK\+ϵ\.A^\{\\mathrm\{cGMA\}\}\_\{ij\}=\\mathbf\{1\}\\\{j\\leq i\\\}\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\}\{\\sum\_\{\\ell\\leq i\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{\\ell,k\}\+\\epsilon\}\.\([9](https://arxiv.org/html/2606.18283#S3.E9)b\)Therefore,

Oi=∑j≤iAi​jcGMA​VX,jO\_\{i\}=\\sum\_\{j\\leq i\}A^\{\\mathrm\{cGMA\}\}\_\{ij\}V\_\{X,j\}\([10](https://arxiv.org/html/2606.18283#S3.E10)b\)up to the numerical stabilizer\. In implementation, the prefix memoriesV~k\(i\)\\tilde\{V\}^\{\(i\)\}\_\{k\}and prefix normalizersZk\(i\)Z^\{\(i\)\}\_\{k\}are computed by cumulative sums along the sequence dimension, so the causal variant preserves the same fixed\-KKlinear scaling in sequence length while enforcing autoregressive causality\.

### 3\.6End\-to\-End Parameter Learning

Traditional Gaussian mixture models \(GMMs\) are typically fit to static datasets by maximum\-likelihood estimation using the Expectation\-Maximization \(EM\) algorithmDempsteret al\.\([1977](https://arxiv.org/html/2606.18283#bib.bib27)\); McLachlan and Peel \([2000](https://arxiv.org/html/2606.18283#bib.bib28)\); Bishop \([2006](https://arxiv.org/html/2606.18283#bib.bib29)\); Huang \([2025](https://arxiv.org/html/2606.18283#bib.bib23)\)\. However, deploying an iterative EM loop inside each attention layer during every forward pass would be computationally prohibitive and would disrupt the continuous gradient flow required for end\-to\-end neural network training\. Instead, we treat the GMM parameters in each GMA layer, i\.e\. the meansμk\\mu\_\{k\}, diagonal covariancesΣk\\Sigma\_\{k\}, and mixture priorsπk\\pi\_\{k\}, as fully differentiable learnable parameters optimized jointly with the rest of the model by standard backpropagation\.

This treatment is related in spirit toMixture Density Networks\(MDNs\)Bishop \([1994](https://arxiv.org/html/2606.18283#bib.bib30)\); Huang \([2026a](https://arxiv.org/html/2606.18283#bib.bib1)\), which combine neural networks with mixture density models and train mixture parameters by gradient\-based optimization\. The role of the mixture model, however, is different: MDNs parameterize an output conditional density, whereas GMA uses shared mixture components to define latent responsibility\-based routing inside an attention layer\. Thus, the GMA components are not fitted by standalone maximum\-likelihood density estimation on a fixed dataset; rather, they are task\-adapted latent routing components shaped by the predictive objective\.

For one GMA layer, the trainable routing parameters are

ΘGMA=\{WQ,WK,WV,μ,ω,α\},\\Theta\_\{\\mathrm\{GMA\}\}=\\\{W\_\{Q\},W\_\{K\},W\_\{V\},\\mu,\\omega,\\alpha\\\},\(13\)whereWQ,WK,WVW\_\{Q\},W\_\{K\},W\_\{V\}are the query, key, and value projection matrices,μ\\mucontains the Gaussian component means,ω\\omegaparameterizes the diagonal covariance entries, andα\\alphaparameterizes the mixture\-prior logits\. These parameters are updated by minimizing the downstream task lossℒtask\\mathcal\{L\}\_\{\\mathrm\{task\}\}, such as cross\-entropy for language modelling or classification\.

To ensure that the GMM parameters remain valid during gradient\-based training, we use the following reparameterizations\.

- \-Means \(μk\\mu\_\{k\}\)\.The latent routing centers are represented by unconstrained learnable parametersμk∈ℝdr\\mu\_\{k\}\\in\\mathbb\{R\}^\{d\_\{r\}\}, updated directly by the optimizer\.
- \-Covariances \(Σk\\Sigma\_\{k\}\)\.To guarantee that each diagonal covariance matrix remains strictly positive\-definite, we maintain unconstrained learnable parametersω∈ℝK×dr\\omega\\in\\mathbb\{R\}^\{K\\times d\_\{r\}\}and set σk,j2=log⁡\(1\+exp⁡\(ωk,j\)\)\+ϵσ,j=1,…,dr,\\sigma\_\{k,j\}^\{2\}=\\log\(1\+\\exp\(\\omega\_\{k,j\}\)\)\+\\epsilon\_\{\\sigma\},\\qquad j=1,\\ldots,d\_\{r\},\(14\)whereϵσ\>0\\epsilon\_\{\\sigma\}\>0is a small numerical constant\. The covariance of componentkkis thenΣk=diag⁡\(σk,12,…,σk,dr2\)\\Sigma\_\{k\}=\\operatorname\{diag\}\(\\sigma\_\{k,1\}^\{2\},\\ldots,\\sigma\_\{k,d\_\{r\}\}^\{2\}\)\.
- \-Mixture priors \(πk\\pi\_\{k\}\)\.To ensure that the mixture priors lie on the probability simplex, we maintain unconstrained logitsα∈ℝK\\alpha\\in\\mathbb\{R\}^\{K\}and apply a softmax transformation: πk=exp⁡\(αk\)∑ℓ=1Kexp⁡\(αℓ\),k=1,…,K\.\\pi\_\{k\}=\\frac\{\\exp\(\\alpha\_\{k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\exp\(\\alpha\_\{\\ell\}\)\},\\qquad k=1,\\ldots,K\.\(15\)

The training process is therefore fully differentiable777For details, Appendix[D](https://arxiv.org/html/2606.18283#A4)provides the corresponding gradient derivations for the reparameterized means, diagonal covariances, and mixture\-prior logits, showing explicitly how gradients pass through the Gaussian responsibility computation during backpropagation\.\. During the forward pass, the input sequenceXXis projected intoQXQ\_\{X\},KXK\_\{X\}, andVXV\_\{X\}via Eq\. \([4](https://arxiv.org/html/2606.18283#S3.E4)\); the reparameterized GMM parameters define the query and key responsibility matricesΓQ\\Gamma^\{Q\}andΓK\\Gamma^\{K\}through the Gaussian responsibility computation in Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\); and the outputOOis computed by the normalized latent\-memory operation in Eq\. \([7](https://arxiv.org/html/2606.18283#S3.E7)\)\. During the backward pass, gradients ofℒtask\\mathcal\{L\}\_\{\\mathrm\{task\}\}propagate through the Gaussian responsibility computation in Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\), the softmax mixture\-prior parameterization in Eq\. \([15](https://arxiv.org/html/2606.18283#S3.E15)\), the softplus covariance parameterization in Eq\. \([14](https://arxiv.org/html/2606.18283#S3.E14)\), and the latent aggregation and broadcasting operations in Eqs\. \([5](https://arxiv.org/html/2606.18283#S3.E5)\) and \([7](https://arxiv.org/html/2606.18283#S3.E7)\)\. Consequently, no EM loop or auxiliary clustering loss is required inside the attention layer\. The resulting components are interpreted as task\-adapted latent routing components rather than as classical GMM clusters obtained by standalone maximum\-likelihood estimation\.

## 4Theoretical Analysis

To understand the learning dynamics and representational capacity of GMA, we analyze its gradient flow, its induced non\-negative low\-rank affinity structure, and the local stability of the Gaussian responsibility map\. Throughout this section,xi∈ℝdrx\_\{i\}\\in\\mathbb\{R\}^\{d\_\{r\}\}denotes a generic routing vector, which may be either a query routing vectorqi=\(QX\)i,:q\_\{i\}=\(Q\_\{X\}\)\_\{i,:\}or a key routing vectorki=\(KX\)i,:k\_\{i\}=\(K\_\{X\}\)\_\{i,:\}after the projection Eq\.\([4](https://arxiv.org/html/2606.18283#S3.E4)\)\.

### 4\.1Gradient Flow and Optimization Stability

A primary concern when replacing dot\-product attention with Gaussian\-mixture routing is the stability of the responsibility gradients, since Gaussian log\-densities contain Mahalanobis\-distance terms888A related pitfall is the “Mahalanobis Trace Trap”Huang \([2026a](https://arxiv.org/html/2606.18283#bib.bib1)\): if the same mini\-batch is used both to estimate an empirical covariance matrix and to evaluate Mahalanobis distances against its inverse, the quadratic term can algebraically collapse to a trace constant and lose its gradient signal\. The same issue can arise component\-wise in a GMM if eachΣk\\Sigma\_\{k\}is recomputed as the responsibility\-weighted empirical covariance of the same batch being scored\. GMA avoids this failure mode becauseμk\\mu\_\{k\},Σk\\Sigma\_\{k\}, andπk\\pi\_\{k\}are learned layer parameters, reparameterized as in Eqs\. \([14](https://arxiv.org/html/2606.18283#S3.E14)\) and \([15](https://arxiv.org/html/2606.18283#S3.E15)\), rather than empirical covariance estimates recomputed inside the forward pass\.and exponentials\. We show that the softmax\-normalized responsibility map yields a bounded and responsibility\-modulated gradient structure\.

For a routing vectorxi∈ℝdrx\_\{i\}\\in\\mathbb\{R\}^\{d\_\{r\}\}and componentkk, define the pre\-normalized log\-density score

si,k=log⁡πk−dr2​log⁡\(2​π\)−12​log⁡\|Σk\|−12​\(xi−μk\)⊤​Σk−1​\(xi−μk\)\.s\_\{i,k\}=\\log\\pi\_\{k\}\-\\frac\{d\_\{r\}\}\{2\}\\log\(2\\pi\)\-\\frac\{1\}\{2\}\\log\|\\Sigma\_\{k\}\|\-\\frac\{1\}\{2\}\(x\_\{i\}\-\\mu\_\{k\}\)^\{\\top\}\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\)\.\(16\)Then the responsibility in Eq\. \([3](https://arxiv.org/html/2606.18283#S3.E3)\) can equivalently be written as a softmax over these scores:

γi,k=exp⁡\(si,k\)∑ℓ=1Kexp⁡\(si,ℓ\)\.\\gamma\_\{i,k\}=\\frac\{\\exp\(s\_\{i,k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\exp\(s\_\{i,\\ell\}\)\}\.\([3](https://arxiv.org/html/2606.18283#S3.E3)b\)
Letℒ\\mathcal\{L\}be the downstream loss and letgi,j=∂ℒ/∂γi,jg\_\{i,j\}=\\partial\\mathcal\{L\}/\\partial\\gamma\_\{i,j\}denote the upstream gradient arriving at the responsibility vector for tokenii\. By the softmax Jacobian999For the binary sigmoid, the derivative isσ​\(x\)​\(1−σ​\(x\)\)\\sigma\(x\)\(1\-\\sigma\(x\)\)\. The softmax generalizes this: its Jacobian has diagonal termsγk​\(1−γk\)\\gamma\_\{k\}\(1\-\\gamma\_\{k\}\)and off\-diagonal terms−γj​γk\-\\gamma\_\{j\}\\gamma\_\{k\}, which account for competition between components\.,

∂γi,j∂si,k=γi,j​\(𝟏​\{j=k\}−γi,k\)\.\\frac\{\\partial\\gamma\_\{i,j\}\}\{\\partial s\_\{i,k\}\}=\\gamma\_\{i,j\}\\left\(\\mathbf\{1\}\\\{j=k\\\}\-\\gamma\_\{i,k\}\\right\)\.\(17\)Therefore, the gradient entering the scoresi,ks\_\{i,k\}is101010Using∂γi,j/∂si,k=γi,j​\(𝟏​\{j=k\}−γi,k\)\\partial\\gamma\_\{i,j\}/\\partial s\_\{i,k\}=\\gamma\_\{i,j\}\(\\mathbf\{1\}\\\{j=k\\\}\-\\gamma\_\{i,k\}\), we have∑jgi,j​γi,j​\(𝟏​\{j=k\}−γi,k\)=gi,k​γi,k−γi,k​∑jγi,j​gi,j=γi,k​\(gi,k−∑jγi,j​gi,j\)\\sum\_\{j\}g\_\{i,j\}\\gamma\_\{i,j\}\(\\mathbf\{1\}\\\{j=k\\\}\-\\gamma\_\{i,k\}\)=g\_\{i,k\}\\gamma\_\{i,k\}\-\\gamma\_\{i,k\}\\sum\_\{j\}\\gamma\_\{i,j\}g\_\{i,j\}=\\gamma\_\{i,k\}\\bigl\(g\_\{i,k\}\-\\sum\_\{j\}\\gamma\_\{i,j\}g\_\{i,j\}\\bigr\)\.

∂ℒ∂si,k=∑j=1K∂ℒ∂γi,j​∂γi,j∂si,k=γi,k​\(gi,k−∑j=1Kγi,j​gi,j\)\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial s\_\{i,k\}\}=\\sum\_\{j=1\}^\{K\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\gamma\_\{i,j\}\}\\frac\{\\partial\\gamma\_\{i,j\}\}\{\\partial s\_\{i,k\}\}=\\gamma\_\{i,k\}\\left\(g\_\{i,k\}\-\\sum\_\{j=1\}^\{K\}\\gamma\_\{i,j\}g\_\{i,j\}\\right\)\.\(18\)Thus, the gradient passed into each component score is a responsibility\-weighted and mean\-centered upstream gradient\. For any learnable scalar parameterθ\\thetathat enters the score, such as a component meanμk,m\\mu\_\{k,m\}, a covariance parameterωk,m\\omega\_\{k,m\}\(Eq\. \([14](https://arxiv.org/html/2606.18283#S3.E14)\)\), a mixture\-prior logitαk\\alpha\_\{k\}\(Eq\. \([15](https://arxiv.org/html/2606.18283#S3.E15)\)\), or an entry of the projection matricesWQW\_\{Q\}andWKW\_\{K\}through the routing vectorsqiq\_\{i\}andkik\_\{i\}\(Eq\. \([4](https://arxiv.org/html/2606.18283#S3.E4)\)\), the chain rule gives

∂ℒ∂θ=∑i∑k=1K∂ℒ∂si,k​∂si,k∂θ,\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\theta\}=\\sum\_\{i\}\\sum\_\{k=1\}^\{K\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial s\_\{i,k\}\}\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\theta\},\(19\)where∂ℒ/∂si,k\\partial\\mathcal\{L\}/\\partial s\_\{i,k\}is given by Eq\. \([18](https://arxiv.org/html/2606.18283#S4.E18)\)\. Equivalently, expanding the softmax Jacobian explicitly gives

∂ℒ∂θ=∑i∑j=1K∑k=1K∂ℒ∂γi,j​∂γi,j∂si,k​∂si,k∂θ\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\theta\}=\\sum\_\{i\}\\sum\_\{j=1\}^\{K\}\\sum\_\{k=1\}^\{K\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\gamma\_\{i,j\}\}\\frac\{\\partial\\gamma\_\{i,j\}\}\{\\partial s\_\{i,k\}\}\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\theta\}\.\([19](https://arxiv.org/html/2606.18283#S4.E19)b\)This expression is the appropriate backpropagation form for the responsibility\-based GMA layer\.

The familiar factorγi,k​\(1−γi,k\)\\gamma\_\{i,k\}\(1\-\\gamma\_\{i,k\}\)appears, from Eq\.\([17](https://arxiv.org/html/2606.18283#S4.E17)\) withj=kj=k, as the direct sensitivity of a component’s own responsibility to its own score:

∂γi,k∂si,k=γi,k​\(1−γi,k\)\.\\frac\{\\partial\\gamma\_\{i,k\}\}\{\\partial s\_\{i,k\}\}=\\gamma\_\{i,k\}\(1\-\\gamma\_\{i,k\}\)\.\([17](https://arxiv.org/html/2606.18283#S4.E17)b\)For example, since

∇μksi,k=Σk−1​\(xi−μk\),\\nabla\_\{\\mu\_\{k\}\}s\_\{i,k\}=\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\),\(20\)the direct derivative ofγi,k\\gamma\_\{i,k\}with respect to the component mean is

∇μkγi,k=γi,k​\(1−γi,k\)​Σk−1​\(xi−μk\)\.\\nabla\_\{\\mu\_\{k\}\}\\gamma\_\{i,k\}=\\gamma\_\{i,k\}\(1\-\\gamma\_\{i,k\}\)\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\)\.\(21\)Forj≠kj\\neq k, the cross\-component derivative is

∇μkγi,j=−γi,j​γi,k​Σk−1​\(xi−μk\)\.\\nabla\_\{\\mu\_\{k\}\}\\gamma\_\{i,j\}=\-\\gamma\_\{i,j\}\\gamma\_\{i,k\}\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\)\.\(22\)
These identities show that Gaussian routing gradients are naturally modulated by the uncertainty of the responsibility assignment\. The direct self\-sensitivity in Eq\. \([17b](https://arxiv.org/html/2606.18283#S4.Ex27)\) is largest when the component assignment is uncertain and vanishes asγi,k\\gamma\_\{i,k\}approaches0or11\. In the full loss gradient, Eq\. \([18](https://arxiv.org/html/2606.18283#S4.E18)\) shows that the update is further centered by the responsibility\-weighted average upstream gradient\. This gives a structured gradient\-modulation mechanism, although it does not by itself rule out component under\-use, component collapse, or poor conditioning; these properties must still be monitored empirically through component\-usage diagnostics\.

### 4\.2Non\-Negative Low\-Rank Affinity Interpretation

We have the implicit unnormalized GMA affinity matrix \(Eq\.[11](https://arxiv.org/html/2606.18283#S3.E11)\)

A~GMA=ΓQ​\(ΓK\)⊤,\\widetilde\{A\}^\{\\mathrm\{GMA\}\}=\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\},whereΓQ,ΓK∈ℝN×K\\Gamma^\{Q\},\\Gamma^\{K\}\\in\\mathbb\{R\}^\{N\\times K\}are non\-negative responsibility matrices whose rows lie on the probability simplex\. Therefore,A~GMA\\widetilde\{A\}^\{\\mathrm\{GMA\}\}is non\-negative and satisfies

rank⁡\(A~GMA\)≤K\.\\operatorname\{rank\}\(\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\)\\leq K\.\(23\)
The normalized affinity used by GMA, i\.e\. Eq\. \([9](https://arxiv.org/html/2606.18283#S3.E9)\), is obtained by row\-wise normalization:

Ai​jGMA=A~i​jGMA∑ℓ=1NA~i​ℓGMA\+ϵ\.A^\{\\mathrm\{GMA\}\}\_\{ij\}=\\frac\{\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{ij\}\}\{\\sum\_\{\\ell=1\}^\{N\}\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{i\\ell\}\+\\epsilon\}\.\([9](https://arxiv.org/html/2606.18283#S3.E9)b\)Equivalently, ifdi=∑ℓ=1NA~i​ℓGMA\+ϵd\_\{i\}=\\sum\_\{\\ell=1\}^\{N\}\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{i\\ell\}\+\\epsilon, then

AGMA=diag\(d\)−1A~GMA\.A^\{\\mathrm\{GMA\}\}=\\operatorname\{diag\}\(d\)^\{\-1\}\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\.Sincediag\(d\)−1\\operatorname\{diag\}\(d\)^\{\-1\}is a diagonal left multiplication with positive entries, the normalized affinity also has rank at mostKK\.

Thus, GMA induces a constrained non\-negative low\-rank affinity structure\. Here “affinity” refers to the induced token\-to\-token relatednessA~i​jGMA=∑k=1Kγi,kQ​γj,kK\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{ij\}=\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}, while “constrained” refers to the fact that the factor matricesΓQ\\Gamma^\{Q\}andΓK\\Gamma^\{K\}are not arbitrary non\-negative matrices: each row is a posterior probability distribution over latent Gaussian components\. This is related to the intuition behind non\-negative matrix factorization \(NMF\), in which non\-negativity encourages additive rather than subtractive combinationsLee and Seung \([1999](https://arxiv.org/html/2606.18283#bib.bib2)\)\. However, GMA is not an NMF algorithm in the classical optimization sense: it does not solve a reconstruction or approximation problem for a given matrix\. Instead, it learns two probabilistic responsibility factors as part of a task\-trained neural sequence model\.

This distinction is important\. The productΓQ​\(ΓK\)⊤\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}gives a rank\-at\-most\-KKnon\-negative affinity, but the effective rank may be smaller if some components are unused or if different responsibility columns become dependent\. Consequently, the interpretation is a*constrained non\-negative low\-rank affinity*interpretation rather than a claim that GMA always produces an exact rank\-KKdecomposition\. Similarly, component specialization and parts\-based structure are possible consequences of the learned responsibility factors, but they are not guaranteed by the factorization alone\. They should therefore be evaluated empirically using the responsibility matrices, for example through component\-usage statistics, responsibility heatmaps, top\-token analysis, or alignment with token\-category labels\.

### 4\.3Local Lipschitz Continuity and Routing Stability

Standard dot\-product attention can be sensitive to large embedding norms, which motivates the scaling factor1/dk1/\\sqrt\{d\_\{k\}\}in Eq\. \([1](https://arxiv.org/html/2606.18283#S1.E1)\)\. In GMA, routing is instead determined by Gaussian responsibilities\. We now show that, for fixed learned GMM parameters, under variance lower bounds and bounded routing vectors, the responsibility map is locally Lipschitz\.

###### Theorem 1\(Local Lipschitz Continuity of GMA Responsibilities\)\.

Assume that the diagonal variances satisfyσk,m2≥ϵσ\>0\\sigma^\{2\}\_\{k,m\}\\geq\\epsilon\_\{\\sigma\}\>0for all componentskkand coordinatesm=1,…,drm=1,\\ldots,d\_\{r\}, and that routing vectors and component means lie in a bounded set such that

‖x−μk‖2≤Rfor all​k\.\\\|x\-\\mu\_\{k\}\\\|\_\{2\}\\leq R\\qquad\\text\{for all \}k\.Then, for fixed GMM parameters, the responsibility mapx↦γ​\(x\)∈ΔK−1x\\mapsto\\gamma\(x\)\\in\\Delta^\{K\-1\}is locally Lipschitz continuous\. In particular, it has a finite Lipschitz constant depending onRR,ϵσ\\epsilon\_\{\\sigma\}, andKK\.

###### Proof\.

For each component, the score from Eq\. \([16](https://arxiv.org/html/2606.18283#S4.E16)\) can be written as

sk​\(x\)=log⁡πk−dr2​log⁡\(2​π\)−12​log⁡\|Σk\|−12​\(x−μk\)⊤​Σk−1​\(x−μk\)\.s\_\{k\}\(x\)=\\log\\pi\_\{k\}\-\\frac\{d\_\{r\}\}\{2\}\\log\(2\\pi\)\-\\frac\{1\}\{2\}\\log\|\\Sigma\_\{k\}\|\-\\frac\{1\}\{2\}\(x\-\\mu\_\{k\}\)^\{\\top\}\\Sigma\_\{k\}^\{\-1\}\(x\-\\mu\_\{k\}\)\.Its gradient with respect toxxis111111This quantity is the score of the component log\-density with respect to the input variable, i\.e\.∇xlog⁡pk​\(x\)\\nabla\_\{x\}\\log p\_\{k\}\(x\)\. In classical statistics, the term “score” often refers to the gradient of the log\-likelihood with respect to model parameters; in score\-based modelling convention, the gradient is taken with respect to the data or representation variableHuang \([2022](https://arxiv.org/html/2606.18283#bib.bib31)\)\.

∇xsk​\(x\)=−Σk−1​\(x−μk\)\.\\nabla\_\{x\}s\_\{k\}\(x\)=\-\\Sigma\_\{k\}^\{\-1\}\(x\-\\mu\_\{k\}\)\.SinceΣk\\Sigma\_\{k\}is diagonal121212Here we use the diagonal covariance parameterization from Eq\. \([14](https://arxiv.org/html/2606.18283#S3.E14)\), namelyΣk=diagσk,12,…,σk,dr2\)\\Sigma\_\{k\}=\\operatorname\{diag\}\\sigma^\{2\}\_\{k,1\},\\ldots,\\sigma^\{2\}\_\{k,d\_\{r\}\}\)\. The notation‖Σk−1‖2\\\|\\Sigma\_\{k\}^\{\-1\}\\\|\_\{2\}denotes the spectral/operator norm of the inverse covariance, not the determinant\. For a diagonal positive\-definite matrix, this norm is the largest diagonal entry ofΣk−1\\Sigma\_\{k\}^\{\-1\}\.andσk,m2≥ϵσ\\sigma^\{2\}\_\{k,m\}\\geq\\epsilon\_\{\\sigma\}for allm=1,…,drm=1,\\ldots,d\_\{r\}, we have

‖Σk−1‖2=max1≤m≤dr⁡1σk,m2≤1ϵσ\.\\\|\\Sigma\_\{k\}^\{\-1\}\\\|\_\{2\}=\\max\_\{1\\leq m\\leq d\_\{r\}\}\\frac\{1\}\{\\sigma^\{2\}\_\{k,m\}\}\\leq\\frac\{1\}\{\\epsilon\_\{\\sigma\}\}\.Using the boundedness assumption‖x−μk‖2≤R\\\|x\-\\mu\_\{k\}\\\|\_\{2\}\\leq Rand the operator\-norm inequality, we obtain

‖∇xsk​\(x\)‖2=‖−Σk−1​\(x−μk\)‖2≤‖Σk−1‖2​‖x−μk‖2≤1ϵσ​R=Rϵσ\.\\\|\\nabla\_\{x\}s\_\{k\}\(x\)\\\|\_\{2\}=\\\|\-\\Sigma\_\{k\}^\{\-1\}\(x\-\\mu\_\{k\}\)\\\|\_\{2\}\\leq\\\|\\Sigma\_\{k\}^\{\-1\}\\\|\_\{2\}\\\|x\-\\mu\_\{k\}\\\|\_\{2\}\\leq\\frac\{1\}\{\\epsilon\_\{\\sigma\}\}R=\\frac\{R\}\{\\epsilon\_\{\\sigma\}\}\.Recall that a differentiable mapffis Lipschitz on a region𝒳\\mathcal\{X\}if there exists a finite constantL<∞L<\\inftysuch that

‖f​\(x\)−f​\(x′\)‖2≤L​‖x−x′‖2,x,x′∈𝒳\.\\\|f\(x\)\-f\(x^\{\\prime\}\)\\\|\_\{2\}\\leq L\\\|x\-x^\{\\prime\}\\\|\_\{2\},\\qquad x,x^\{\\prime\}\\in\\mathcal\{X\}\.For differentiable functions, a sufficient condition is that the Jacobian norm is uniformly bounded on𝒳\\mathcal\{X\}:

supx∈𝒳‖Jf​\(x\)‖2<∞\.\\sup\_\{x\\in\\mathcal\{X\}\}\\\|J\_\{f\}\(x\)\\\|\_\{2\}<\\infty\.Applying this to the score maps​\(x\)=\(s1​\(x\),…,sK​\(x\)\)s\(x\)=\(s\_\{1\}\(x\),\\ldots,s\_\{K\}\(x\)\), its Jacobian has rows∇xsk​\(x\)⊤\\nabla\_\{x\}s\_\{k\}\(x\)^\{\\top\}\. Since each row satisfies‖∇xsk​\(x\)‖2≤R/ϵσ\\\|\\nabla\_\{x\}s\_\{k\}\(x\)\\\|\_\{2\}\\leq R/\\epsilon\_\{\\sigma\}, the Frobenius norm gives

‖Js​\(x\)‖2≤‖Js​\(x\)‖F=\(∑k=1K‖∇xsk​\(x\)‖22\)1/2≤K​Rϵσ\.\\\|J\_\{s\}\(x\)\\\|\_\{2\}\\leq\\\|J\_\{s\}\(x\)\\\|\_\{F\}=\\left\(\\sum\_\{k=1\}^\{K\}\\\|\\nabla\_\{x\}s\_\{k\}\(x\)\\\|\_\{2\}^\{2\}\\right\)^\{1/2\}\\leq\\frac\{\\sqrt\{K\}R\}\{\\epsilon\_\{\\sigma\}\}\.Thus, the score map has a finite local Lipschitz constant, with one simple upper bound proportional toK​R/ϵσ\\sqrt\{K\}R/\\epsilon\_\{\\sigma\}\.

The softmax map from scores to responsibilities has a uniformly bounded Jacobian onℝK\\mathbb\{R\}^\{K\}, and is therefore globally Lipschitz\. Hence the composition

x↦s​\(x\)↦softmax⁡\(s​\(x\)\)=γ​\(x\)x\\mapsto s\(x\)\\mapsto\\operatorname\{softmax\}\(s\(x\)\)=\\gamma\(x\)is locally Lipschitz\. The resulting Lipschitz constant is finite and depends onRR,ϵσ\\epsilon\_\{\\sigma\}, andKK\. ∎

This result shows that, under the stated boundedness and variance\-lower\-bound assumptions, small perturbations of a routing vector cannot cause unbounded changes in its responsibility vector\. The statement should be interpreted as a local routing\-stability guarantee, not as a global robustness guarantee for the entire Transformer architecture\.

For fixed model parameters and fixed sequence lengthNN, the normalized GMA output is also locally stable on bounded domains\. The latent aggregationV~=\(ΓK\)⊤​VX\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}and the normalizerZ=\(ΓK\)⊤​𝟏NZ=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\}are finite sums of products of locally Lipschitz responsibility functions and bounded value vectors, while the broadcasting step divides byΓQ​Z\+ϵ\\Gamma^\{Q\}Z\+\\epsilon\. Sinceϵ\>0\\epsilon\>0and all responsibilities are non\-negative, the denominator is bounded below byϵ\\epsilon\. Therefore, on bounded sets of routing vectors and values, the full GMA mapping defined by Eq\. \([7](https://arxiv.org/html/2606.18283#S3.E7)\) is a composition of locally Lipschitz operations and is itself locally Lipschitz\. The corresponding Lipschitz constant may depend onNN,KK, the parameter bounds, the variance lower bound, and the value\-vector bounds\.

## 5Experiments131313Part of the experimental code was developed with kind assistance from ChatGPT\-5\.5\(OpenAI,[2026](https://arxiv.org/html/2606.18283#bib.bib39)\), which the authors gratefully acknowledge\.

To empirically validate GMA, we evaluate four aspects: \(1\) controlled systems profiling of memory and throughput, \(2\) long\-context classification accuracy on LRA, \(3\) autoregressive language modelling on WikiText\-103, and \(4\) the interpretable latent structures induced by learned responsibility matrices\. All experiments in this paper evaluate GMA inself\-attention\-style settings rather than cross\-attention settings\. Specifically, the systems profiling and LRA experiments use bidirectional sequence mixing, the WikiText\-103 experiment uses causal self\-attention\-style sequence mixing, and the interpretability analysis examines the learned query responsibilities of the trained causal GMA model\.

### Experimental Setup and Baselines

We compare GMA against standard dot\-product Multi\-Head Attention \(MHA\)\(Vaswaniet al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib11)\)and a suite of efficient sequence\-modelling baselines chosen to represent distinct architectural trade\-offs:

- 1Low\-Rank & Static Projections:Linformer\(Wanget al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib5)\)projects the sequence dimensionNNinto a smaller fixed dimensionkk\. This baseline tests a static low\-rank compression approach against GMA’s input\-dependent responsibility routing\.
- 2Kernel Approximations:Performer\(Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4)\)uses random orthogonal features, andLinear Transformer\(Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7)\)uses an ELU\-based \(exponential linear unit\) feature map\. These baselines compare GMA against fixed or random feature\-map approaches to linear attention\.
- 3Modern State\-Space Models:Mamba\(Gu and Dao,[2024](https://arxiv.org/html/2606.18283#bib.bib8)\)is a strong selective state\-space model with linear scaling in sequence length\. We include it to situate GMA against contemporary sequence\-modelling alternatives that do not use attention\-style token\-to\-token routing\.

Across all experiments, GMA uses the multi\-head normalized formulation described in Section[3](https://arxiv.org/html/2606.18283#S3), with mixture priorsπk\\pi\_\{k\}initialized uniformly\. The number of Gaussian componentsKKis chosen according to the purpose of each experiment\. In the controlled systems profiling experiment, we report a component ablation overK∈\{64,128,256,512\}K\\in\\\{64,128,256,512\\\}to characterize the efficiency–capacity trade\-off\. In the LRA accuracy experiments, we useK=128K=128as a moderate setting that keeps responsibility\-memory cost manageable while retaining sufficient routing capacity\. In the WikiText\-103 language\-modelling experiments, we evaluate bothK=128K=128andK=256K=256causal GMA variants to measure how increasing the number of components affects perplexity and throughput\. All models are implemented inPyTorch\(Paszkeet al\.,[2019](https://arxiv.org/html/2606.18283#bib.bib46)\)and trained using theAdamWoptimizer\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2606.18283#bib.bib3)\)\.

### 5\.1Controlled Systems Profiling

To isolate the computational behavior of each sequence\-mixing mechanism, we first conduct141414Unless otherwise stated, all mechanisms are evaluated on a single NVIDIA A100\-SXM4\-40GB GPU using PyTorch 2\.10\.0 with CUDA 12\.8, bfloat16 mixed precision, batch size 8, 12 heads, and model hidden dimensiondmodel=768d\_\{\\mathrm\{model\}\}=768\.controlled systems profiling using synthetic sequences\. This experiment does not measure task accuracy; instead, it empirically tests the memory and throughput behavior predicted by the GMA cost analysis in Section[3\.4](https://arxiv.org/html/2606.18283#S3.SS4)\. Specifically, we measure incremental peak GPU memory allocation and forward–backward throughput for asingle attention/mixing blockunder a shared input shape\. This separates the computational scaling claims from the downstream representation\-learning claims\.

Table 1:Single\-block systems profiling of computational efficiency\. Memory and throughput are empirical measurements for a forward\-backward pass\. Parameter counts, shown in parentheses, are reported in millions for each sequence\-mixing layer\.Underlinedvalues denote the best value in each metric, i\.e\. lowest memory or highest throughput\.Note:“Standard MHA \(Eager\)” denotes a conventional explicit scaled dot\-product MHA implementation that materializes the attention score/probability tensors in the standard eager PyTorch execution path\. “Standard MHA \(SDPA\)” denotes PyTorch’s optimized scaled\-dot\-product\-attention backend, which uses fused memory\-efficient kernels when available\. Both compute standard dot\-product attention; they differ in implementation backend rather than in the mathematical attention definition\.

As shown in Table[1](https://arxiv.org/html/2606.18283#S5.T1), eager standard MHA exhibits the expected quadratic memory growth, reaching 23\.07 GB atN=4000N=4000, together with a marked throughput drop\. In contrast, GMA shows approximately linear empirical memory scaling inNNfor fixedKK\. For example, withK=256K=256, memory increases from 0\.85 GB atN=1000N=1000to 1\.70 GB atN=2000N=2000and 3\.39 GB atN=4000N=4000, while throughput remains nearly constant\. This behavior is consistent with the fixed\-KKactivation scaling analyzed in Section[3\.4](https://arxiv.org/html/2606.18283#S3.SS4)\.

TheKK\-ablation exposes the expected efficiency–capacity trade\-off\. GMA’s parameter count grows withKKbut is independent ofNN, whereas its activation memory scales through the responsibility tensorsΓQ,ΓK∈ℝB×H×N×K\\Gamma^\{Q\},\\Gamma^\{K\}\\in\\mathbb\{R\}^\{B\\times H\\times N\\times K\}, whereBB,HH,NN, andKKdenote batch size, number of heads, sequence length, and number of Gaussian mixture components, respectively\. Thus, for fixedKK, memory scales linearly with sequence length, while for fixedNN, both memory and arithmetic cost scale linearly with the number of components\. AtN=4000N=4000, memory rises from 1\.37 GB forK=64K=64to 6\.32 GB forK=512K=512, while the parameter count increases from 2\.5M to 3\.2M\. By comparison, Linformer’s static projection matrices introduce an explicit sequence\-length\-dependent parameter cost, increasing from 2\.9M parameters atN=1000N=1000to 4\.4M atN=4000N=4000\.

Overall, the profiling results empirically support the intended linear\-in\-NNscaling of GMA, but also show that this high\-level PyTorch implementation, currently without any optimisation, has larger constant factors than optimized SDPA, Linear Transformer, Mamba, and Performer baselines\. This overhead comes from computing Gaussian log densities, Mahalanobis terms, and responsibility normalizers\. The dominant responsibility activation storage is𝒪​\(B​H​N​K\)\\mathcal\{O\}\(BHNK\)\. Summing over heads, the arithmetic cost is approximately𝒪​\(B​N​K​D\)\\mathcal\{O\}\(BNKD\)when the per\-head routing and value dimensions sum to the model hidden dimensionDD\. Therefore, these results should be read as evidence for GMA’s empirical scaling behavior, not as evidence that the present implementation is the fastest raw sequence mixer\. The downstream experiments test whether this probabilistic routing cost yields competitive accuracy and interpretable responsibility structure\.

### 5\.2Long\-Context Task Accuracy: LRA

Having isolated the single\-block computational behavior, we next evaluate151515All models are evaluated under the same PyTorch pipeline, using 4 sequence\-mixing layers, model hidden dimensiondmodel=256d\_\{\\mathrm\{model\}\}=256, eight heads, the AdamW optimizer\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2606.18283#bib.bib3)\), bfloat16 mixed precision, and a warmup\-cosine learning\-rate schedule\.representational capability on the Long Range Arena \(LRA\) benchmark\(Tayet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib6)\)\. We focus on ListOps, a hierarchical reasoning task with sequence lengthN=2​KN=2\\text\{K\}, and byte\-level IMDb/Text classification with sequence lengthN=4​KN=4\\text\{K\}\. We train each model for 20 epochs over 3 random seeds and report test accuracy at the epoch with the highest validation accuracy\. For GMA, we useK=128K=128components in this experiment, matching the efficiency–capacity trade\-off identified in the controlled systems profiling experiment\.

Table 2:Task accuracy on the LRA ListOps and byte\-level IMDb/Text tasks\. Results aretest accuraciesat the best validation epoch, reported as mean±\\pmstandard deviation over 3 random seeds\. The parameter column reports millions of parameters for ListOps/Text, respectively; differences arise from task\-specific components and, for some baselines, sequence\-length dependent modules\.Underlinedvalues denote the best overall result in each metric, while boldface highlights the GMA row\.Table[2](https://arxiv.org/html/2606.18283#S5.T2)shows that GMA is competitive with standard and efficient attention baselines under the same PyTorch LRA pipeline\. On ListOps, GMA reaches41\.17%41\.17\\%, closely matching SDPA attention \(41\.33%41\.33\\%\) and outperforming Linformer, Linear Transformer, and Performer\. On Text, GMA obtains64\.97%64\.97\\%, the strongest result among the attention\-style baselines, slightly above Linear Transformer \(64\.88%64\.88\\%\) and SDPA attention \(64\.03%64\.03\\%\)\. Averaged over the two tasks, GMA achieves53\.07%53\.07\\%, giving the best average performance among the attention\-based baselines evaluated here\.

Mamba performs best overall, driven mainly by its much stronger Text result \(82\.12%82\.12\\%\), indicating that selective state\-space models provide a different and highly effective inductive bias for byte\-level classification\. Nevertheless, GMA outperforms Linformer, Linear Transformer, and Performer on both ListOps and Text in this experiment, and gives the strongest average accuracy among the attention\-style baselines\. These results support the hypothesis that learned probabilistic routing can be competitive with static or randomized efficient\-attention mechanisms under a shared training pipeline, while retaining the fixed\-KK, linear\-in\-NNscaling and analyzable responsibility structure introduced in Sections[3\.3](https://arxiv.org/html/2606.18283#S3.SS3)and[3\.4](https://arxiv.org/html/2606.18283#S3.SS4)\.

### 5\.3Language Modelling on WikiText\-103

Efficient attention mechanisms improve the scalability of Transformers on long sequences, but many do so by introducing sparsity patterns, low\-rank projections, kernel feature maps, or other approximations to full softmax attention\(Tayet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib17); Katharopouloset al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib7); Choromanskiet al\.,[2021](https://arxiv.org/html/2606.18283#bib.bib4)\)\. Prior work has observed that such approximations can involve a model\-quality trade\-off relative to exact softmax attention, especially in settings requiring dense token\-level prediction\(Daoet al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib47); Yanget al\.,[2024](https://arxiv.org/html/2606.18283#bib.bib48)\)\. We therefore evaluate autoregressive language modelling on WikiText\-103\(Merityet al\.,[2016](https://arxiv.org/html/2606.18283#bib.bib10)\)\.

Unlike the bidirectional LRA setting, this experiment requires strictly causal sequence mixing: the representation at positioniimay depend only on tokens at positionsj≤ij\\leq i, and is then used for next\-token prediction\. For GMA, we use the causal prefix\-sum formulation described in Section[3\.5](https://arxiv.org/html/2606.18283#S3.SS5), which accumulates key\-weighted latent memories and normalizers only over prefix positions\. This prevents future\-token leakage while preserving linear\-in\-NNcomplexity for fixedKK\.

All models are trained161616We use a fixed\-budget decoder\-only setup with context length 1024, model hidden dimensiondmodel=512d\_\{\\mathrm\{model\}\}=512, 6 layers, 8 heads, bfloat16 mixed precision, AdamW optimization, and an effective batch size of 16,384 tokens\. Each model is trained for 20,000 optimization steps over 3 random seeds on an NVIDIA A100\-SXM4\-40GB GPU\.from scratch using a GPT\-2 BPE tokenizer, with no pretrained model weights\. We report validation perplexity at the best validation checkpoint171717Since the best validation checkpoint occurred at the final training step for all runs, these results should be interpreted as a fixed\-budget accuracy\-efficiency comparison rather than as full\-convergence performance\.and training throughput in tokens per second\. All models passed a prefix\-causality test, with maximum prefix error equal to zero\.

Table 3:Autoregressive language modelling on WikiText\-103 under a fixed training budget\. We report validation perplexity at the best validation checkpoint; lower is better\. Throughput is measured in training tokens per second; higher is better\.Underlinedvalues denote the best overall result in each metric, while boldface highlights the stronger GMA variant\.Table[3](https://arxiv.org/html/2606.18283#S5.T3)shows that causal GMA improves perplexity over the tested linear\-attention baselines on dense autoregressive prediction\. GMA withK=256K=256obtains a validation perplexity of41\.7241\.72, outperforming Linear Transformer \(44\.8144\.81\) and Performer \(47\.4947\.49\) under the same fixed training budget\. This suggests that learned Gaussian responsibilities can preserve token\-level predictive information more effectively than the fixed ELU feature map of Linear Transformer or the random feature approximation used by Performer in this setup\.

However, causal GMA does not close the gap to optimized causal softmax attention or Mamba in this fixed\-budget implementation\. Standard MHA with causal SDPA achieves the best perplexity \(35\.4035\.40\) and the highest throughput, while Mamba reaches a similar perplexity \(35\.4435\.44\) with lower but still strong throughput\. This does not contradict the linear\-in\-NNcomplexity analysis of GMA\. The cost estimate in Eq\. \([12](https://arxiv.org/html/2606.18283#S3.E12)\) describes asymptotic scaling for fixedKK, whereas wall\-clock throughput also depends on constant factors, kernel fusion, memory access patterns, and implementation maturity\. Moreover, at the context length used here \(N=1024N=1024\), the difference betweenN​KNKandN2N^\{2\}is still modest forK=128K=128orK=256K=256before constant factors are taken into account\. The scaling advantage of GMA is expected to become more relevant at longer sequence lengths whenK≪NK\\ll N, although realizing this advantage in wall\-clock time would require optimized kernels\. In the present high\-level, non\-optimized PyTorch implementation, GMA incurs larger constant factors from Gaussian log\-density computation, Mahalanobis terms, causal prefix accumulation, and responsibility normalization\. Increasing the number of components fromK=128K=128toK=256K=256improves perplexity from42\.6242\.62to41\.7241\.72, but reduces throughput from22\.922\.9K to13\.113\.1K tokens/s—an expected capacity\-efficiency trade\-off\.

Overall, the WikiText\-103 experiment supports GMA as a probabilistic and interpretable linear\-time attention\-style alternative that improves over the tested linear and random\-feature attention variants, while remaining less efficient and less accurate than highly optimized causal SDPA and Mamba under the current implementation\.

### 5\.4Interpretability: Latent Responsibility Structure

Motivated by the non\-negative factorization structure of the GMA responsibility matrices, we analyze the learned query responsibilities on WikiText\-103 validation sequences\. We use the previously trained causal GMA model withK=128K=128and extract the last\-layer query responsibility tensorΓQ∈ℝS×H×L×K\\Gamma^\{Q\}\\in\\mathbb\{R\}^\{S\\times H\\times L\\times K\}, whereS=64S=64is the number of validation sequences,H=8H=8is the number of attention heads,L=1024L=1024is the sequence length, andK=128K=128is the number of Gaussian components\. For each token occurrence at sequencessand positionℓ\\ell, we average the responsibility vectors over heads,

γ¯s,ℓ,kQ=1H​∑h=1Hγs,h,ℓ,kQ,k=1,…,K\.\\bar\{\\gamma\}^\{Q\}\_\{s,\\ell,k\}=\\frac\{1\}\{H\}\\sum\_\{h=1\}^\{H\}\\gamma^\{Q\}\_\{s,h,\\ell,k\},\\qquad k=1,\\ldots,K\.\(24\)This produces one head\-averaged responsibility vectorγ¯s,ℓ,:Q∈ΔK−1\\bar\{\\gamma\}^\{Q\}\_\{s,\\ell,:\}\\in\\Delta^\{K\-1\}for each token occurrence\. We then flatten the sequence\-position pair\(s,ℓ\)\(s,\\ell\)into a token indext=1,…,Tt=1,\\ldots,T, whereT=S​L=65,536T=SL=65\{,\}536, and write the resulting vector asγ¯tQ\\bar\{\\gamma\}^\{Q\}\_\{t\}\. Thus, the averaging is performed across heads for the same token occurrence, not across different sequence positions or across repeated token types\. Because this diagnostic averages responsibilities across heads, the resulting component indices should be interpreted as aggregate, head\-averaged routing channels rather than isolated per\-head components\.181818If each head has its own independently learned Gaussian components, then componentkkin one head is not necessarily aligned with componentkkin another head\. The head\-averaged analysis is therefore a coarse diagnostic of aggregate routing structure; per\-head component specialization is left for more detailed future analysis\.

We study both soft and hard responsibility structure\. The marginal component usage is defined as

pk=1T​∑t=1Tγ¯t,kQ\.p\_\{k\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\.\(25\)We use the normalized usage entropy

Husage=−∑k=1Kpk​log⁡pklog⁡KH\_\{\\mathrm\{usage\}\}=\\frac\{\-\\sum\_\{k=1\}^\{K\}p\_\{k\}\\log p\_\{k\}\}\{\\log K\}\(26\)to measure whether the latent routing channels are broadly used or collapse to a small subset\. A value close to 1 indicates broad component usage, while a value close to 0 indicates concentration on a few components\.

To measure routing sharpness, we compute the normalized mean token entropy

Htoken=1T​∑t=1T−∑k=1Kγ¯t,kQ​log⁡γ¯t,kQlog⁡K\.H\_\{\\mathrm\{token\}\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\frac\{\-\\sum\_\{k=1\}^\{K\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\\log\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\}\{\\log K\}\.\(27\)Lower values indicate sharper, more confident routing\. We also report the mean maximum responsibility

Rmax=1T​∑t=1Tmaxk⁡γ¯t,kQ,R\_\{\\max\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\max\_\{k\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\},\(28\)which gives a complementary measure of assignment sharpness\.

For hard\-assignment diagnostics, we define

zt=arg⁡maxk⁡γ¯t,kQ\.z\_\{t\}=\\arg\\max\_\{k\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\.\(29\)Letctc\_\{t\}denote the surface\-form category of tokentt, such as punctuation, numeric, capitalized, lower\-case alphabetic, subword alphabetic, newline/space, or function word\. For componentkk, let

nk,c=∑t=1T𝟏​\{zt=k,ct=c\},nk=∑cnk,c\.n\_\{k,c\}=\\sum\_\{t=1\}^\{T\}\\mathbf\{1\}\\\{z\_\{t\}=k,c\_\{t\}=c\\\},\\qquad n\_\{k\}=\\sum\_\{c\}n\_\{k,c\}\.\(30\)For components withnk\>0n\_\{k\}\>0, the component purity is

Purity⁡\(k\)=maxc⁡nk,cnk,\\operatorname\{Purity\}\(k\)=\\max\_\{c\}\\frac\{n\_\{k,c\}\}\{n\_\{k\}\},\(31\)and the weighted category purity is

Purityw=∑k:nk\>0nkT​Purity⁡\(k\)\.\\operatorname\{Purity\}\_\{\\mathrm\{w\}\}=\\sum\_\{k:n\_\{k\}\>0\}\\frac\{n\_\{k\}\}\{T\}\\operatorname\{Purity\}\(k\)\.\(32\)This measures how strongly hard GMA assignments align with surface\-form token categories\.

Finally, we compute the mutual information between hard component assignmentsZZand token categoriesCC:

I​\(Z;C\)=∑k,cp​\(k,c\)​log⁡p​\(k,c\)p​\(k\)​p​\(c\),I\(Z;C\)=\\sum\_\{k,c\}p\(k,c\)\\log\\frac\{p\(k,c\)\}\{p\(k\)p\(c\)\},\(33\)wherep​\(k,c\)=nk,c/Tp\(k,c\)=n\_\{k,c\}/T,p​\(k\)=∑cp​\(k,c\)p\(k\)=\\sum\_\{c\}p\(k,c\), andp​\(c\)=∑kp​\(k,c\)p\(c\)=\\sum\_\{k\}p\(k,c\)\. We report normalized mutual information as

NMI⁡\(Z,C\)=I​\(Z;C\)min⁡\{H​\(Z\),H​\(C\)\},\\operatorname\{NMI\}\(Z,C\)=\\frac\{I\(Z;C\)\}\{\\min\\\{H\(Z\),H\(C\)\\\}\},\(34\)whereH​\(Z\)H\(Z\)andH​\(C\)H\(C\)are the marginal entropies of component assignments and token categories, respectively\. Higher NMI indicates stronger alignment between latent routing and token surface\-form categories\.

Table 4:Interpretability statistics for last\-layer, head\-averaged GMA query responsibilities on WikiText\-103 validation sequences\. Usage entropy and token entropy are normalized bylog⁡K\\log K\. Purity, mutual information \(MI\), and normalized mutual information \(NMI\) are computed from hard assignmentszt=arg⁡maxk⁡γ¯t,kQz\_\{t\}=\\arg\\max\_\{k\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}and surface\-form token categories\.Table[4](https://arxiv.org/html/2606.18283#S5.T4)shows that the learned responsibilities use most of the available latent routing channels\. In total, 104 out of 128 components receive at least one hard assignment, and the normalized usage entropy is 0\.933, indicating that the soft usage distribution is broad rather than collapsed\. At the same time, the normalized per\-token entropy is 0\.790 and the mean maximum responsibility is 0\.103, showing that head\-averaged routing remains soft rather than nearly one\-hot\. This softness is expected because responsibilities are averaged across heads\.

The hard component assignments nevertheless carry measurable surface\-form information\. The weighted category purity is0\.4830\.483, substantially above a permutation baseline of0\.293±0\.0010\.293\\pm 0\.001that preserves the component and category marginals while randomizing their association191919The weighted category purity is computed asPurityw=∑k:nk\>0\(nk/T\)​maxc⁡nk,c/nk\\operatorname\{Purity\}\_\{\\mathrm\{w\}\}=\\sum\_\{k:n\_\{k\}\>0\}\(n\_\{k\}/T\)\\max\_\{c\}n\_\{k,c\}/n\_\{k\}, giving0\.4830\.483for the observed hard assignments\. The permutation baseline is obtained by randomly permuting the surface\-form category labels\{ct\}t=1T\\\{c\_\{t\}\\\}\_\{t=1\}^\{T\}relative to the hard component assignments\{zt\}t=1T\\\{z\_\{t\}\\\}\_\{t=1\}^\{T\}, recomputingPurityw\\operatorname\{Purity\}\_\{\\mathrm\{w\}\}for each permutation, and reporting the mean and standard deviation across permutations\. This preserves the empirical component counts and category counts while destroying their association\. The global majority\-category baseline ismaxc⁡nc/T=18,757/65,536≈0\.286\\max\_\{c\}n\_\{c\}/T=18\{,\}757/65\{,\}536\\approx 0\.286, corresponding to lower\-case alphabetic tokens\.\. For reference, the largest global token category, lower\-case alphabetic tokens, accounts for0\.2860\.286of the analyzed token set\. The mutual information between component assignments and token categories is 0\.427 nats, corresponding to a normalized mutual information of 0\.244\. Several high\-usage components show preferences for recognizable token roles, including function words, punctuation, lower\-case alphabetic tokens, numeric tokens, capitalized tokens, and subword pieces\.

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig1_gammaq_heatmap_top_components.png)Figure 2:Last\-layer, head\-averaged GMA query responsibilities on a WikiText\-103 validation sequence\. Rows show the most active latent routing channels in the sequence and columns show token positions\. Concentrated bands indicate that different token positions route through distinct latent components\.Figure[2](https://arxiv.org/html/2606.18283#S5.F2)visualizes the learned responsibility structure at the token level\. The heatmap shows that different token positions activate different subsets of latent routing channels, rather than routing uniformly through all components\. This provides qualitative evidence that the GMA responsibility matrix is an analyzable latent routing object\.

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig4_component_category_heatmap.png)Figure 3:Category composition of the most\-used GMA components under hard assignmentszt=arg⁡maxk⁡γ¯t,kQz\_\{t\}=\\arg\\max\_\{k\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\. Rows correspond to latent routing channels and columns correspond to surface\-form token categories\. Several components specialize toward categories such as function words, punctuation, capitalized tokens, lower\-case alphabetic tokens, and subword fragments\.Figure[3](https://arxiv.org/html/2606.18283#S5.F3)further shows that high\-usage components are not purely random with respect to surface\-form categories: several components have visible preferences for categories such as function words, punctuation, capitalized tokens, lower\-case alphabetic tokens, and subword fragments\. This provides evidence that the learned responsibility matrices are analyzable latent routing objects\. Additional diagnostics, including the token\-category distribution, representative component examples, permutation baselines, component\-usage histograms, token\-level assignment strips, and PCA projections, are provided in Appendix[F](https://arxiv.org/html/2606.18283#A6)\.

However, the components should not be interpreted as clean semantic classes\. Many components are mixed, and the categories used here are surface\-form categories rather than full syntactic or semantic annotations202020Additional diagnostic plots in Appendix[F](https://arxiv.org/html/2606.18283#A6), including component\-usage histograms, purity bar plots, token\-level assignment strips, and PCA projections of responsibility vectors, show partial structure with substantial overlap rather than fully separated clusters\.\. We therefore conclude that GMA learns a non\-collapsed and analyzable latent responsibility structure with moderate alignment to token surface roles, rather than a set of fully disentangled semantic concepts\.

## 6Discussion

We position GMA as a new attention\-style mechanism, not as a universally superior replacement for existing attention or state\-space models\. Its central contribution is to change both the*geometry*and the*computational ordering*of attention\. At its most abstract level, an attention score measures theaffinityorcompatibilitybetween two representations\. This affinity need not be a dot product: it may be defined by cosine similarity, learned additive scores, kernels, distances, divergences, or other task\-adapted similarity functions\. From this perspective, GMA can be seen as a responsibility\-space attention mechanism\. Rather than comparing projected queries and keys directly in the original routing representation space, GMA first maps them into posterior responsibility vectors over a shared set of learned Gaussian mixture components, and then defines their unnormalized affinity by the overlap \(Eq\. \([11](https://arxiv.org/html/2606.18283#S3.E11)\)\)

A~i​jGMA=⟨γiQ,γjK⟩=∑k=1Kγi,kQ​γj,kK\.\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{ij\}=\\langle\\gamma^\{Q\}\_\{i\},\\gamma^\{K\}\_\{j\}\\rangle=\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\.The expressiveness of this affinity depends not only on the number of mixture componentsKK, but also on the routing dimension, covariance parameterization, projection matrices, number of heads, depth, and training objective\. We discuss this broader view of attention design in Appendix[G](https://arxiv.org/html/2606.18283#A7)\.

This responsibility\-space affinity alone would not necessarily reduce the computational cost: explicitly materializingA~GMA=ΓQ​\(ΓK\)⊤\\widetilde\{A\}^\{\\mathrm\{GMA\}\}=\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}would still produce anN×NN\\times Nmatrix\. The efficiency of GMA comes from using the associativity of matrix multiplication to change the order of computation\. Instead of computing

\(ΓQ​\(ΓK\)⊤\)​VX,\\bigl\(\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\\bigr\)V\_\{X\},GMA first writes values into aKK\-slot latent memory,

V~=\(ΓK\)⊤​VX,\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\},and then reads from this memory using query responsibilities,

O=ΓQ​V~ΓQ​Z\+ϵ\.O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}\.Thus, GMA is best understood as*responsibility\-space affinity plus latent\-memory routing*: the induced token\-to\-token affinity exists algebraically, but it is not materialized in the efficient implementation\. This is the mechanism that gives GMA its fixed\-KK, linear\-in\-NNactivation scaling while retaining probabilistic responsibilities that can be inspected and analyzed\.

#### Scaling behavior and implementation constants\.

The controlled systems profiling experiment confirms this intended scaling behavior\. For fixedKK, the routing\-specific responsibility tensors scale as𝒪​\(B​H​N​K\)\\mathcal\{O\}\(BHNK\)rather than𝒪​\(B​H​N2\)\\mathcal\{O\}\(BHN^\{2\}\), and the observed memory growth is approximately linear in sequence length\. This supports the cost analysis in Section[3\.4](https://arxiv.org/html/2606.18283#S3.SS4)\. At the same time, the profiling results also show that asymptotic scaling does not directly imply superior wall\-clock efficiency at the sequence lengths tested here\. The current high\-level PyTorch implementation has larger constant factors than optimized SDPA, Linear Transformer, Performer, and Mamba implementations\. This is not surprising: GMA explicitly evaluates Gaussian log densities, Mahalanobis terms, posterior responsibilities, and normalizers, whereas several baselines benefit from simpler arithmetic, fused kernels, or highly optimized recurrent or state\-space operators\. Thus, the empirical message is twofold: GMA has the intended fixed\-KKlinear memory structure, but realizing its full speed potential will likely require specialized kernels and implementation\-level optimization\.

#### Task accuracy and inductive bias\.

The LRA results suggest that the additional cost of probabilistic routing can come with useful representational capacity\. On ListOps and byte\-level Text classification, GMA is competitive with standard SDPA attention and gives the strongest average performance among the attention\-style baselines evaluated in our PyTorch pipeline\. This indicates that theN×KN\\times Klatent bottleneck does not merely compress away useful information; learned responsibilities can preserve task\-relevant sequence structure while avoiding an explicitN×NN\\times Nattention matrix\. However, the comparison with Mamba is also important\. Mamba is substantially stronger on the byte\-level Text task, suggesting that selective state\-space dynamics provide a different and highly effective inductive bias for this setting\. Therefore, GMA offers a competitive probabilistic attention\-style alternative with a distinct routing mechanism and diagnostic structure\.

#### Autoregressive modelling\.

The WikiText\-103 experiment gives a more demanding test of causal GMA\. Causal GMA improves over the Linear Transformer and Performer baselines in validation perplexity under the same training budget, which suggests that learned Gaussian responsibilities can preserve more predictive information than the tested fixed or randomized feature\-map approximations in this setup\. Nevertheless, GMA remains behind optimized causal SDPA and Mamba in both perplexity and throughput\. Increasing the number of mixture components fromK=128K=128toK=256K=256improves perplexity but reduces throughput, exposing the expected capacity–efficiency trade\-off\. These results clarify the present role of GMA in language modelling: its immediate strength is not raw speed or state\-of\-the\-art perplexity, but the combination of linear scaling, probabilistic routing, and analyzable latent responsibilities\. Longer contexts, optimized kernels, and hybrid local–global designs are natural settings in which this trade\-off should be tested more thoroughly\.

#### Interpretability and responsibility diagnostics\.

The interpretability experiment directly examines one of the main advantages of a responsibility\-based mechanism\. The last\-layer, head\-averaged query responsibilities on WikiText\-103 use most of the available latent routing channels: 104 out of 128 components receive at least one hard assignment, and the normalized usage entropy is 0\.933\. This indicates broad component usage rather than collapse to a small subset of mixture components\. The hard assignments also carry measurable surface\-form information, with weighted category purity 0\.483 and normalized mutual information 0\.244 relative to the token categories used in the analysis\. These numbers should be interpreted carefully\. The categories are surface\-form categories derived from GPT\-2 BPE tokens, not full syntactic or semantic annotations, and many components remain mixed\. Moreover, because the diagnostic averages responsibilities over heads, the resulting components should be read as head\-averaged routing channels rather than isolated per\-head experts\. Even with these qualifications, the analysis shows that GMA exposes a latent routing structure that can be quantified, visualized, and compared against external token annotations\. This kind of diagnostic access is less direct in many random\-feature, low\-rank, or state\-space alternatives\.

#### Limitations\.

Several limitations remain\. First, the present implementation is not optimized at the kernel level, so its wall\-clock throughput underestimates what may be possible with fused Gaussian\-responsibility computation and more efficient prefix accumulation\. Second, all experiments in this paper use finite\-KKGMA; the number of components is selected manually rather than adapted automatically to the data or task\. Third, the experiments focus on self\-attention\-style sequence mixing\. Cross\-attention follows naturally from the same formulation, as discussed in Appendix[E](https://arxiv.org/html/2606.18283#A5), but it is not empirically evaluated here\. Fourth, the interpretability analysis is intentionally limited: it uses surface\-form token categories and head\-averaged responsibilities, so it does not establish that GMA components correspond to clean semantic, syntactic, or discourse\-level concepts\. Finally, although GMA induces a constrained non\-negative low\-rank affinity structure, this does not guarantee component specialization or disentanglement; these properties must be measured empirically\.

### GMA Extensions: Adaptive Mixtures and Probabilistic Expert Routing

The experiments in this paper evaluate GMA as a finite\-KKlinear\-time routing mechanism with learned mixture priors, means, and diagonal covariances\. The results suggest several directions for future work: implementation\-level optimization, adaptive mixture capacity, hybrid local–global attention, cross\-attention and multimodal applications, and probabilistic expert routing\.

#### Optimized and hybrid GMA\.

A first practical direction is to develop fused kernels for Gaussian log\-density evaluation, responsibility normalization, and causal prefix accumulation\. The current implementation is useful for validating the method and analyzing its behavior, but it does not exploit the same level of hardware specialization as SDPA or state\-space implementations\. A second direction is to combine exact local attention with long\-range GMA routing\. Such a hybrid architecture could use standard attention to capture short\-range syntax, local byte\-level patterns, or fine\-grained visual details, while using mixture routing to summarize and retrieve long\-range information through a fixed number of latent memory channels\. This may be particularly relevant for long\-context language modelling, document understanding, genomics, and multimodal settings where local detail and global context are both important\.

#### Bayesian and nonparametric GMA\.

In the finite GMA formulation used in this paper, the number of latent Gaussian componentsKKis fixed, and the mixture priorπ=\(π1,…,πK\)\\pi=\(\\pi\_\{1\},\\ldots,\\pi\_\{K\}\)is learned through unconstrained logits followed by a softmax, as in Eq\. \([15](https://arxiv.org/html/2606.18283#S3.E15)\)\. A Bayesian extension would place aDirichletprior over the mixture weights,

π∼Dirichlet⁡\(α0K,…,α0K\),\\pi\\sim\\operatorname\{Dirichlet\}\\left\(\\frac\{\\alpha\_\{0\}\}\{K\},\\ldots,\\frac\{\\alpha\_\{0\}\}\{K\}\\right\),\(35\)whereα0\>0\\alpha\_\{0\}\>0controls how concentrated or diffuse the component usage is\. Small values ofα0\\alpha\_\{0\}encourage sparse component usage, while larger values encourage more uniform use of the available mixture components\.

For a generic routing vectorxix\_\{i\}, such as a query vectorqi=\(QX\)i,:∈ℝdrq\_\{i\}=\(Q\_\{X\}\)\_\{i,:\}\\in\\mathbb\{R\}^\{d\_\{r\}\}or a key vectorki=\(KX\)i,:∈ℝdrk\_\{i\}=\(K\_\{X\}\)\_\{i,:\}\\in\\mathbb\{R\}^\{d\_\{r\}\}, the responsibility retains the same form \(Eq\.\([3](https://arxiv.org/html/2606.18283#S3.E3)\)\),

γi,k=πk​𝒩​\(xi∣μk,Σk\)∑ℓ=1Kπℓ​𝒩​\(xi∣μℓ,Σℓ\),\\gamma\_\{i,k\}=\\frac\{\\pi\_\{k\}\\mathcal\{N\}\(x\_\{i\}\\mid\\mu\_\{k\},\\Sigma\_\{k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\pi\_\{\\ell\}\\mathcal\{N\}\(x\_\{i\}\\mid\\mu\_\{\\ell\},\\Sigma\_\{\\ell\}\)\},\(36\)butπ\\piis treated as a random variable or regularized latent parameter rather than only as a point estimate\. A finite Bayesian variant could be implemented by learning a variational Dirichlet posteriorq​\(π\)=Dirichlet⁡\(a\)q\(\\pi\)=\\operatorname\{Dirichlet\}\(a\)and adding the KL penalty

KL⁡\[q​\(π\)∥Dirichlet⁡\(α0K,…,α0K\)\]\\operatorname\{KL\}\\left\[q\(\\pi\)\\,\\\|\\,\\operatorname\{Dirichlet\}\\left\(\\frac\{\\alpha\_\{0\}\}\{K\},\\ldots,\\frac\{\\alpha\_\{0\}\}\{K\}\\right\)\\right\]\(37\)to the task lossℒtask\\mathcal\{L\}\_\{\\mathrm\{task\}\}\. The forward pass could use samples fromq​\(π\)q\(\\pi\)to propagate uncertainty over mixture weights, or use the posterior mean𝔼q​\[πk\]=ak/∑ℓaℓ\\mathbb\{E\}\_\{q\}\[\\pi\_\{k\}\]=a\_\{k\}/\\sum\_\{\\ell\}a\_\{\\ell\}as a deterministic approximation\. This would provide a route toward representing uncertainty over mixture weights, rather than only learning point\-valued priors\.

A further extension is to replace the finite Dirichlet prior with aDirichlet processprior, yielding a nonparametric or truncated infinite\-mixture variant of GMA\(Ferguson,[1973](https://arxiv.org/html/2606.18283#bib.bib19); Sethuraman,[1994](https://arxiv.org/html/2606.18283#bib.bib20); Rasmussen,[1999](https://arxiv.org/html/2606.18283#bib.bib21); Blei and Jordan,[2006](https://arxiv.org/html/2606.18283#bib.bib22)\)\. In stick\-breaking form,

vk∼Beta⁡\(1,α\),πk=vk​∏ℓ<k\(1−vℓ\),v\_\{k\}\\sim\\operatorname\{Beta\}\(1,\\alpha\),\\qquad\\pi\_\{k\}=v\_\{k\}\\prod\_\{\\ell<k\}\(1\-v\_\{\\ell\}\),\(38\)whereα\\alphais the concentration parameter\. In practice, such a model would be implemented using a finite truncation at a maximum sizeKmaxK\_\{\\max\}\. The nominal routing cost would then scale as𝒪​\(N​Kmax\)\\mathcal\{O\}\(NK\_\{\\max\}\)\. The effective number of active components could be smaller if many learned stick weights become negligible, although actual wall\-clock savings would require pruning, sparse evaluation, or another mechanism that skips negligible components\. This would allow GMA to adapt its effective routing capacity to the data and task, rather than relying only on manual selection ofKK\.

#### Cross\-attention and multimodal routing\.

Although the experiments in this work focus on self\-attention\-style sequence mixing, the same responsibility\-space mechanism extends naturally to cross\-attention\. Queries can be formed from a query\-side sequenceXX, while keys and values are formed from a separate context sequenceYY, producingΓXQ\\Gamma\_\{X\}^\{Q\}andΓYK\\Gamma\_\{Y\}^\{K\}as described in Appendix[E](https://arxiv.org/html/2606.18283#A5)\. This suggests possible applications in encoder–decoder architectures, retrieval\-augmented models, vision–language models, and multimodal systems in which the query and context streams have different statistical structure\. In such settings, the mixture components may serve as shared latent routing channels between modalities or between source and target representations\. A systematic study of GMA cross\-attention is therefore a natural next step\.

#### Probabilistic expert routing\.

A second direction is to use Gaussian responsibilities as probabilistic gates over attention heads, sequence\-mixing modules, or expert networks\. Letri∈ℝdgr\_\{i\}\\in\\mathbb\{R\}^\{d\_\{g\}\}denote a gating representation of tokenii, and suppose that each expert or headh=1,…,Hh=1,\\ldots,His associated with a Gaussian routing component with parameters\(ρh,νh,Λh\)\(\\rho\_\{h\},\\nu\_\{h\},\\Lambda\_\{h\}\), whereρh≥0\\rho\_\{h\}\\geq 0,∑h=1Hρh=1\\sum\_\{h=1\}^\{H\}\\rho\_\{h\}=1,νh∈ℝdg\\nu\_\{h\}\\in\\mathbb\{R\}^\{d\_\{g\}\}, andΛh∈𝕊\+\+dg\\Lambda\_\{h\}\\in\\mathbb\{S\}^\{d\_\{g\}\}\_\{\+\+\}\. A GMA\-style gate could define

γi,hMoE=ρh​𝒩​\(ri∣νh,Λh\)∑ℓ=1Hρℓ​𝒩​\(ri∣νℓ,Λℓ\)\.\\gamma\_\{i,h\}^\{\\mathrm\{MoE\}\}=\\frac\{\\rho\_\{h\}\\mathcal\{N\}\(r\_\{i\}\\mid\\nu\_\{h\},\\Lambda\_\{h\}\)\}\{\\sum\_\{\\ell=1\}^\{H\}\\rho\_\{\\ell\}\\mathcal\{N\}\(r\_\{i\}\\mid\\nu\_\{\\ell\},\\Lambda\_\{\\ell\}\)\}\.\(39\)Hereγi,hMoE\\gamma\_\{i,h\}^\{\\mathrm\{MoE\}\}is the posterior responsibility of experthhfor tokeniiunder the learned Gaussian gating model\.

If all experts are evaluated, the output can be written as the dense mixture

yi=∑h=1Hγi,hMoEExperth\(X\)i,y\_\{i\}=\\sum\_\{h=1\}^\{H\}\\gamma\_\{i,h\}^\{\\mathrm\{MoE\}\}\\operatorname\{Expert\}\_\{h\}\(X\)\_\{i\},\(40\)whereExperth\(X\)i\\operatorname\{Expert\}\_\{h\}\(X\)\_\{i\}denotes the output of expert, head, or sequence\-mixing modulehhat token positionii\. This dense form is fully differentiable, but it is not computationally sparse as all experts are executed\. To obtain conditional\-computation savings, one would need an additional sparse selection rule, such as top\-rrrouting, stochastic sampling, or a differentiable relaxation\.

For example, let𝒮i⊂\{1,…,H\}\\mathcal\{S\}\_\{i\}\\subset\\\{1,\\ldots,H\\\}denote the selected expert set for tokenii, obtained by top\-rrselection or sampling fromγiMoE\\gamma\_\{i\}^\{\\mathrm\{MoE\}\}\. A sparse GMA\-MoE output can then be written as

yi=∑h∈𝒮iγ¯i,hMoEExperth\(X\)i,γ¯i,hMoE=γi,hMoE∑ℓ∈𝒮iγi,ℓMoE\.y\_\{i\}=\\sum\_\{h\\in\\mathcal\{S\}\_\{i\}\}\\bar\{\\gamma\}\_\{i,h\}^\{\\mathrm\{MoE\}\}\\operatorname\{Expert\}\_\{h\}\(X\)\_\{i\},\\qquad\\bar\{\\gamma\}\_\{i,h\}^\{\\mathrm\{MoE\}\}=\\frac\{\\gamma\_\{i,h\}^\{\\mathrm\{MoE\}\}\}\{\\sum\_\{\\ell\\in\\mathcal\{S\}\_\{i\}\}\\gamma\_\{i,\\ell\}^\{\\mathrm\{MoE\}\}\}\.\(41\)The discrete expert\-selection step could be implemented using anargmax\\operatorname\{argmax\}or top\-rroperator, stochastic sampling, the Gumbel–Softmax or Concrete relaxation\(Janget al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib40); Maddisonet al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib41)\), or a straight\-through variant\(Bengioet al\.,[2013](https://arxiv.org/html/2606.18283#bib.bib42); Janget al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib40); Huang,[2026b](https://arxiv.org/html/2606.18283#bib.bib32)\)\. This formulation is related to conventional MoE routing\(Shazeeret al\.,[2017](https://arxiv.org/html/2606.18283#bib.bib43); Lepikhinet al\.,[2020](https://arxiv.org/html/2606.18283#bib.bib45); Feduset al\.,[2022](https://arxiv.org/html/2606.18283#bib.bib44)\), but differs in its probabilistic interpretation\. Standard MoE gates are often implemented as learned softmax or noisy top\-kkclassifiers over experts\. A GMA\-style gate instead defines expert probabilities as posterior responsibilities under a learned latent\-density model\. This could provide a principled way to combine conditional computation with uncertainty\-aware and interpretable expert selection, while preserving the possibility of sparse execution through top\-rror sampled routing\.

Taken together, these directions suggest that GMA is best viewed not only as a single efficient\-attention module, but also as an instance of a broader design principle: attention mechanisms can be designed by choosing the space in which compatibility is measured, the normalization used to turn compatibility into routing weights, and the computational ordering used to retrieve values\. The present paper studies one concrete realization of that principle using finite Gaussian mixtures and responsibility\-space routing\. In its current form, GMA is therefore*not all you need*: it is a probabilistic, interpretable, linear\-time attention\-style alternative whose value is clearest when latent routing structure, analyzability, and fixed\-KKscaling matter alongside task accuracy and implementation efficiency\.

## 7Conclusion

We introducedGaussian Mixture Attention\(GMA\), a probabilistic attention\-style sequence mixer that rethinks attention asrouting through a learned latent responsibility space\. Instead of comparing every query directly with every key, GMA maps queries and keys to posterior responsibility vectors overKKlearned Gaussian mixture components\. Their overlap defines an implicit responsibility\-spaceaffinity, and an efficient implementation avoids materializing the correspondingN×NN\\times Nmatrix by first writing values into aKK\-slot latent memory and then reading from that memory with query responsibilities\. This combination of responsibility\-space affinity and associative latent\-memory routing gives GMA fixed\-KKlinear\-in\-NNactivation scaling while preserving a normalized attention\-style interpretation\.

Methodologically, GMA opens a complementary route for attention design: rather than modifying only sparsity patterns, low\-rank projections, kernel feature maps, or implementation kernels, it changes the space in whichcompatibilityis computed\. The resulting responsibility matrices are not merely computational intermediates; they are probabilistic routing objects that can be analyzed, visualized, and compared with token\-level annotations\. We developed bidirectional and causal variants of GMA, gave an end\-to\-end differentiable parameterization of the Gaussian mixture components, and analyzed its gradient structure, constrained non\-negative low\-rank affinity interpretation, and local routing stability\.

Empirically, the results support a balanced view of GMA\. Controlled systems profiling confirms the intended approximatelylinear memory growth with sequence lengthfor fixedKK, but also shows that the present high\-level PyTorch implementation has larger constant factors than optimized SDPA, Linear Transformer, Performer, and Mamba implementations\. On LRA ListOps and byte\-level Text classification tasks, GMA is competitive with standard attention and achieves the strongest average performance among the attention\-style baselines evaluated in our pipeline\. On WikiText\-103, causal GMA improves over Linear Transformer and Performer in validation perplexity, but remains behind optimized causal SDPA and Mamba in both perplexity and throughput\. The interpretability analysis further shows broad, non\-collapsed component usage and moderate alignment between hard responsibility assignments and surface\-form token categories\.

These findings position GMAnotas a universal replacement for optimized softmax attention or state\-space sequence models, but as a probabilistic, interpretable, fixed\-KKlinear\-time attention\-style alternative\. Its advantages are clearest when fixed\-KKscaling, normalized latent routing, and access to analyzable responsibility matrices are valuable\. Its present limitations are also clear: the implementation is not yet kernel\-optimized, the number of components is manually selected, the experiments focus on self\-attention\-style settings, and the current interpretability analysis does not establish clean semantic or syntactic disentanglement\. Future work may pursue fused and hardware\-aware GMA kernels, hybrid local–global architectures, adaptive Bayesian or nonparametric component selection, cross\-attention and multimodal extensions, probabilistic expert\-routing variants, and richer linguistic or task\-specific analyses of learned mixture components\. In this sense, GMA isnot“all you need”, but it opens a complementary route for designing attention mechanisms through latent probabilistic routing spaces\.

## Author Contributions

Dr\. Yongchao Huanginitiated the idea of Gaussian Mixture Attention, formulated the methodological and theoretical framework, developed the mathematical derivations, co\-designed and ran the experimental codebase, and drafted the manuscript\.Hassan Razacontributed to the empirical evaluation by co\-designing and running the Long Range Arena \(LRA\) and WikiText\-103 experiments\. Both authors reviewed and approved the final manuscript\. Despite careful review, some errors or inaccuracies may remain; readers are therefore advised to exercise due caution when consulting the material\.

Code availability:[GitHub](https://github.com/YongchaoHuang/Gaussian_mixture_attention)

## References

- \[1\]Ž\. Avsec, V\. Agarwal, D\. Visentin, J\. R\. Ledsam, A\. Grabska\-Barwinska, K\. R\. Taylor, Y\. Assael, J\. Jumper, P\. Kohli, and D\. R\. Kelley\(2021\)Effective gene expression prediction from sequence by integrating long\-range interactions\.Nature Methods18,pp\. 1196–1203\.Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p1.12)\.
- \[2\]D\. Bahdanau, K\. Cho, and Y\. Bengio\(2016\)Neural machine translation by jointly learning to align and translate\.External Links:1409\.0473,[Link](https://arxiv.org/abs/1409.0473)Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px1.p1.2)\.
- \[3\]I\. Beltagy, M\. E\. Peters, and A\. Cohan\(2020\)Longformer: the long\-document transformer\.External Links:2004\.05150,[Link](https://arxiv.org/abs/2004.05150)Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p1.12),[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px2.p1.1)\.
- \[4\]Y\. Bengio, N\. Léonard, and A\. Courville\(2013\)Estimating or propagating gradients through stochastic neurons for conditional computation\.External Links:1308\.3432,[Link](https://arxiv.org/abs/1308.3432)Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[5\]C\. M\. Bishop\(1994\)Mixture density networks\.Technical reportTechnical ReportNCRG/94/004,Neural Computing Research Group, Aston University\.Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px8.p1.1),[§3\.6](https://arxiv.org/html/2606.18283#S3.SS6.p2.1)\.
- \[6\]C\. M\. Bishop\(2006\)Pattern recognition and machine learning\.Springer,New York\.Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px8.p1.1),[§3\.6](https://arxiv.org/html/2606.18283#S3.SS6.p1.3)\.
- \[7\]D\. M\. Blei and M\. I\. Jordan\(2006\)Variational inference for dirichlet process mixtures\.Bayesian Analysis1\(1\),pp\. 121–143\.Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px2.p3.5)\.
- \[8\]S\. Boyd and L\. Vandenberghe\(2018\)Introduction to applied linear algebra – vectors, matrices, and least squares\.Cambridge University Press\.Cited by:[Appendix B](https://arxiv.org/html/2606.18283#A2.p2.2),[footnote 22](https://arxiv.org/html/2606.18283#footnote22)\.
- \[9\]R\. Child, S\. Gray, A\. Radford, and I\. Sutskever\(2019\)Generating long sequences with sparse transformers\.External Links:1904\.10509,[Link](https://arxiv.org/abs/1904.10509)Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px2.p1.1)\.
- \[10\]K\. M\. Choromanski, V\. Likhosherstov, D\. Dohan, X\. Song, A\. Gane, T\. Sarlos, P\. Hawkins, J\. Q\. Davis, A\. Mohiuddin, L\. Kaiser, D\. B\. Belanger, L\. J\. Colwell, and A\. Weller\(2021\)Rethinking attention with performers\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=Ua6zuk0WRH)Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.SS0.SSS0.Px3.p2.2),[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[item 4](https://arxiv.org/html/2606.18283#S1.I1.i4.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p1.12),[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px4.p1.1),[2](https://arxiv.org/html/2606.18283#S5.I1.i2.p1.1),[§5\.3](https://arxiv.org/html/2606.18283#S5.SS3.p1.1)\.
- \[11\]T\. H\. Cormen, C\. E\. Leiserson, R\. L\. Rivest, and C\. Stein\(2009\)Introduction to algorithms\.3 edition,MIT Press\.Cited by:[Appendix B](https://arxiv.org/html/2606.18283#A2.p1.2)\.
- \[12\]R\. Csordás, P\. Piekos, K\. Irie, and J\. Schmidhuber\(2024\)SwitchHead: accelerating transformers with mixture\-of\-experts attention\.InProceedings of the 38th International Conference on Neural Information Processing Systems,NIPS ’24,Red Hook, NY, USA\.External Links:ISBN 9798331314385Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px9.p1.1)\.
- \[13\]T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. Ré\(2022\)FLASHATTENTION: fast and memory\-efficient exact attention with io\-awareness\.InProceedings of the 36th International Conference on Neural Information Processing Systems,NIPS ’22,Red Hook, NY, USA\.External Links:ISBN 9781713871088Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px5.p1.1),[§5\.3](https://arxiv.org/html/2606.18283#S5.SS3.p1.1)\.
- \[14\]A\. P\. Dempster, N\. M\. Laird, and D\. B\. Rubin\(1977\)Maximum likelihood from incomplete data via the em algorithm\.Journal of the Royal Statistical Society\. Series B \(Methodological\)39\(1\),pp\. 1–38\.External Links:ISSN 00359246,[Link](http://www.jstor.org/stable/2984875)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px8.p1.1),[§3\.6](https://arxiv.org/html/2606.18283#S3.SS6.p1.3)\.
- \[15\]A\. Dosovitskiy, L\. Beyer, A\. Kolesnikov, D\. Weissenborn, X\. Zhai, T\. Unterthiner, M\. Dehghani, M\. Minderer, G\. Heigold, S\. Gelly, J\. Uszkoreit, and N\. Houlsby\(2021\)An image is worth 16x16 words: transformers for image recognition at scale\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=YicbFdNTTy)Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p1.6)\.
- \[16\]W\. Fedus, B\. Zoph, and N\. Shazeer\(2022\-01\)Switch transformers: scaling to trillion parameter models with simple and efficient sparsity\.J\. Mach\. Learn\. Res\.23\(1\)\.External Links:ISSN 1532\-4435Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px9.p1.1),[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[17\]T\. S\. Ferguson\(1973\)A bayesian analysis of some nonparametric problems\.The Annals of Statistics1\(2\),pp\. 209–230\.External Links:ISSN 00905364, 21688966,[Link](http://www.jstor.org/stable/2958008)Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px2.p3.5)\.
- \[18\]A\. Gu and T\. Dao\(2024\)Mamba: linear\-time sequence modeling with selective state spaces\.External Links:2312\.00752,[Link](https://arxiv.org/abs/2312.00752)Cited by:[item 4](https://arxiv.org/html/2606.18283#S1.I1.i4.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px7.p1.1),[3](https://arxiv.org/html/2606.18283#S5.I1.i3.p1.1)\.
- \[19\]A\. Gu, K\. Goel, and C\. Ré\(2022\)Efficiently modeling long sequences with structured state spaces\.InInternational Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px7.p1.1)\.
- \[20\]Y\. Huang\(2022\)Classification via score\-based generative modelling\.External Links:2207\.11091,[Link](https://arxiv.org/abs/2207.11091)Cited by:[footnote 11](https://arxiv.org/html/2606.18283#footnote11)\.
- \[21\]Y\. Huang\(2025\)Sampling via gaussian mixture approximations\.External Links:2509\.25232,[Link](https://arxiv.org/abs/2509.25232)Cited by:[§3\.2](https://arxiv.org/html/2606.18283#S3.SS2.p1.9),[§3\.6](https://arxiv.org/html/2606.18283#S3.SS6.p1.3)\.
- \[22\]Y\. Huang\(2026\)Gaussian joint embeddings for self\-supervised representation learning\.External Links:2603\.26799,[Link](https://arxiv.org/abs/2603.26799)Cited by:[§3\.2](https://arxiv.org/html/2606.18283#S3.SS2.p1.9),[§3\.6](https://arxiv.org/html/2606.18283#S3.SS6.p2.1),[footnote 8](https://arxiv.org/html/2606.18283#footnote8)\.
- \[23\]Y\. Huang\(2026\)Neural bayesian sequential routing\.External Links:2605\.26147,[Link](https://arxiv.org/abs/2605.26147)Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[24\]R\. A\. Jacobs, M\. I\. Jordan, S\. J\. Nowlan, and G\. E\. Hinton\(1991\)Adaptive mixtures of local experts\.Neural Computation3\(1\),pp\. 79–87\.External Links:[Document](https://dx.doi.org/10.1162/neco.1991.3.1.79)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px9.p1.1)\.
- \[25\]A\. Jaegle, S\. Borgeaud, J\. Alayrac, C\. Doersch, C\. Ionescu, D\. Ding, S\. Koppula, D\. Zoran, A\. Brock, E\. Shelhamer, O\. J\. Henaff, M\. Botvinick, A\. Zisserman, O\. Vinyals, and J\. Carreira\(2022\)Perceiver IO: a general architecture for structured inputs & outputs\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=fILj7WpI-g)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px6.p1.1)\.
- \[26\]A\. Jaegle, F\. Gimeno, A\. Brock, A\. Zisserman, O\. Vinyals, and J\. Carreira\(2021\)Perceiver: general perception with iterative attention\.InInternational Conference on Machine Learning,pp\. 4651–4664\.Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px6.p1.1)\.
- \[27\]E\. Jang, S\. Gu, and B\. Poole\(2017\)Categorical reparameterization with gumbel\-softmax\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=rkE3y85ee)Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[28\]A\. Katharopoulos, A\. Vyas, N\. Pappas, and F\. Fleuret\(2020\-13–18 Jul\)Transformers are RNNs: fast autoregressive transformers with linear attention\.InProceedings of the 37th International Conference on Machine Learning,H\. D\. III and A\. Singh \(Eds\.\),Proceedings of Machine Learning Research, Vol\.119,pp\. 5156–5165\.External Links:[Link](https://proceedings.mlr.press/v119/katharopoulos20a.html)Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.SS0.SSS0.Px3.p2.2),[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[item 4](https://arxiv.org/html/2606.18283#S1.I1.i4.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p1.12),[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px4.p1.1),[2](https://arxiv.org/html/2606.18283#S5.I1.i2.p1.1),[§5\.3](https://arxiv.org/html/2606.18283#S5.SS3.p1.1)\.
- \[29\]N\. Kitaev, L\. Kaiser, and A\. Levskaya\(2020\)Reformer: the efficient transformer\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=rkgNKkHtvB)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px2.p1.1)\.
- \[30\]D\. D\. Lee and H\. S\. Seung\(1999\-10\)Learning the parts of objects by non\-negative matrix factorization\.Nature401\(6755\),pp\. 788–791\.External Links:[Document](https://dx.doi.org/10.1038/44565),[Link](https://doi.org/10.1038/44565)Cited by:[§4\.2](https://arxiv.org/html/2606.18283#S4.SS2.p3.3)\.
- \[31\]D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. Chen\(2020\)GShard: scaling giant models with conditional computation and automatic sharding\.External Links:2006\.16668,[Link](https://arxiv.org/abs/2006.16668)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px9.p1.1),[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[32\]I\. Loshchilov and F\. Hutter\(2019\)Decoupled weight decay regularization\.External Links:1711\.05101,[Link](https://arxiv.org/abs/1711.05101)Cited by:[§5](https://arxiv.org/html/2606.18283#S5.SSx1.p3.6),[footnote 15](https://arxiv.org/html/2606.18283#footnote15)\.
- \[33\]T\. Luong, H\. Pham, and C\. D\. Manning\(2015\-09\)Effective approaches to attention\-based neural machine translation\.InProceedings of the 2015 Conference on Empirical Methods in Natural Language Processing,L\. Màrquez, C\. Callison\-Burch, and J\. Su \(Eds\.\),Lisbon, Portugal,pp\. 1412–1421\.External Links:[Link](https://aclanthology.org/D15-1166/),[Document](https://dx.doi.org/10.18653/v1/D15-1166)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px1.p1.2)\.
- \[34\]C\. J\. Maddison, A\. Mnih, and Y\. W\. Teh\(2017\)The concrete distribution: a continuous relaxation of discrete random variables\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=S1jE5L5gl)Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[35\]A\. F\. T\. Martins and R\. F\. Astudillo\(2016\)From softmax to sparsemax: a sparse model of attention and multi\-label classification\.InProceedings of the 33rd International Conference on International Conference on Machine Learning \- Volume 48,ICML’16,pp\. 1614–1623\.Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.SS0.SSS0.Px5.p1.1)\.
- \[36\]G\. McLachlan and D\. Peel\(2000\)Finite mixture models\.Wiley Series in Probability and Statistics,John Wiley & Sons\.External Links:ISBN 9780471006268,[Document](https://dx.doi.org/10.1002/0471721182),[Link](https://doi.org/10.1002/0471721182)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px8.p1.1),[§3\.6](https://arxiv.org/html/2606.18283#S3.SS6.p1.3)\.
- \[37\]S\. Merity, C\. Xiong, J\. Bradbury, and R\. Socher\(2016\)Pointer sentinel mixture models\.External Links:1609\.07843Cited by:[item 4](https://arxiv.org/html/2606.18283#S1.I1.i4.p1.1),[§5\.3](https://arxiv.org/html/2606.18283#S5.SS3.p1.1)\.
- \[38\]OpenAI\(2026\)ChatGPT\-5\.5\.Note:Large language modelUsed for assistance with experimental code developmentExternal Links:[Link](https://chatgpt.com/)Cited by:[footnote 13](https://arxiv.org/html/2606.18283#footnote13)\.
- \[39\]A\. Paszke, S\. Gross, F\. Massa, A\. Lerer, J\. Bradbury, G\. Chanan, T\. Killeen, Z\. Lin, N\. Gimelshein, L\. Antiga, A\. Desmaison, A\. Köpf, E\. Yang, Z\. DeVito, M\. Raison, A\. Tejani, S\. Chilamkurthy, B\. Steiner, L\. Fang, J\. Bai, and S\. Chintala\(2019\)PyTorch: an imperative style, high\-performance deep learning library\.InProceedings of the 33rd International Conference on Neural Information Processing Systems,Cited by:[§5](https://arxiv.org/html/2606.18283#S5.SSx1.p3.6)\.
- \[40\]B\. Peters, V\. Niculae, and A\. F\. Martins\(2019\)Sparse sequence\-to\-sequence models\.InProc\. ACL,External Links:[Link](https://www.aclweb.org/anthology/P19-1146)Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.SS0.SSS0.Px5.p1.1)\.
- \[41\]A\. Radford, J\. W\. Kim, C\. Hallacy, A\. Ramesh, G\. Goh, S\. Agarwal, G\. Sastry, A\. Askell, P\. Mishkin, J\. Clark, G\. Krueger, and I\. Sutskever\(2021\-18–24 Jul\)Learning transferable visual models from natural language supervision\.InProceedings of the 38th International Conference on Machine Learning,M\. Meila and T\. Zhang \(Eds\.\),Proceedings of Machine Learning Research, Vol\.139,pp\. 8748–8763\.External Links:[Link](https://proceedings.mlr.press/v139/radford21a.html)Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p1.6)\.
- \[42\]C\. E\. Rasmussen\(1999\)The infinite gaussian mixture model\.InProceedings of the 13th International Conference on Neural Information Processing Systems,NIPS’99,Cambridge, MA, USA,pp\. 554–560\.Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px2.p3.5)\.
- \[43\]J\. Sethuraman\(1994\)A constructive definition of dirichlet priors\.Statistica Sinica4\(2\),pp\. 639–650\.External Links:ISSN 10170405, 19968507,[Link](http://www.jstor.org/stable/24305538)Cited by:[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px2.p3.5)\.
- \[44\]J\. Shah, G\. Bikshandi, Y\. Zhang, V\. Thakkar, P\. Ramani, and T\. Dao\(2024\)FlashAttention\-3: fast and accurate attention with asynchrony and low\-precision\.InProceedings of the 38th International Conference on Neural Information Processing Systems,NIPS ’24,Red Hook, NY, USA\.External Links:ISBN 9798331314385Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px5.p1.1)\.
- \[45\]N\. Shazeer, \*\. Mirhoseini, \*\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. Dean\(2017\)Outrageously large neural networks: the sparsely\-gated mixture\-of\-experts layer\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=B1ckMDqlg)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px9.p1.1),[§6](https://arxiv.org/html/2606.18283#S6.SSx1.SSS0.Px4.p3.8)\.
- \[46\]G\. Strang\(2016\)Introduction to linear algebra\.Wellesley\-Cambridge Press\.Cited by:[Appendix B](https://arxiv.org/html/2606.18283#A2.p1.2)\.
- \[47\]Y\. Tay, M\. Dehghani, S\. Abnar, Y\. Shen, D\. Bahri, P\. Pham, J\. Rao, L\. Yang, S\. Ruder, and D\. Metzler\(2020\)Long range arena: a benchmark for efficient transformers\.External Links:2011\.04006,[Link](https://arxiv.org/abs/2011.04006)Cited by:[item 4](https://arxiv.org/html/2606.18283#S1.I1.i4.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p1.12),[§5\.2](https://arxiv.org/html/2606.18283#S5.SS2.p1.3)\.
- \[48\]Y\. Tay, M\. Dehghani, D\. Bahri, and D\. Metzler\(2022\-12\)Efficient transformers: a survey\.ACM Comput\. Surv\.55\(6\)\.External Links:ISSN 0360\-0300,[Link](https://doi.org/10.1145/3530811),[Document](https://dx.doi.org/10.1145/3530811)Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§5\.3](https://arxiv.org/html/2606.18283#S5.SS3.p1.1)\.
- \[49\]A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, Ł\. Kaiser, and I\. Polosukhin\(2017\)Attention is all you need\.InProceedings of the 31st International Conference on Neural Information Processing Systems,NIPS’17,Red Hook, NY, USA,pp\. 6000–6010\.External Links:ISBN 9781510860964Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.SS0.SSS0.Px2.p1.6),[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[item 4](https://arxiv.org/html/2606.18283#S1.I1.i4.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p1.6),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px1.p1.2),[§5](https://arxiv.org/html/2606.18283#S5.SSx1.p1.1)\.
- \[50\]S\. Wang, B\. Z\. Li, M\. Khabsa, H\. Fang, and H\. Ma\(2020\)Linformer: self\-attention with linear complexity\.External Links:2006\.04768,[Link](https://arxiv.org/abs/2006.04768)Cited by:[§1](https://arxiv.org/html/2606.18283#S1.p1.12),[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px3.p1.2),[1](https://arxiv.org/html/2606.18283#S5.I1.i1.p1.2)\.
- \[51\]Y\. Xiong, Z\. Zeng, R\. Chakraborty, M\. Tan, G\. Fung, Y\. Li, and V\. Singh\(2021\)Nyströmformer: a nyström\-based algorithm for approximating self\-attention\.InProceedings of the AAAI Conference on Artificial Intelligence,Vol\.35,pp\. 14138–14148\.Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px3.p1.2)\.
- \[52\]S\. Yang, B\. Wang, Y\. Shen, R\. Panda, and Y\. Kim\(2024\)Gated linear attention transformers with hardware\-efficient training\.InProceedings of the 41st International Conference on Machine Learning,ICML’24\.Cited by:[§5\.3](https://arxiv.org/html/2606.18283#S5.SS3.p1.1)\.
- \[53\]T\. yen, T\. M\. Nguyen, D\. D\. Le, D\. K\. Nguyen, V\. Tran, R\. G\. Baraniuk, N\. Ho, and S\. J\. Osher\(2022\)Improving transformers with probabilistic attention keys\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px8.p1.1)\.
- \[54\]M\. Zaheer, G\. Guruganesh, A\. Dubey, J\. Ainslie, C\. Alberti, S\. Ontanon, P\. Pham, A\. Ravula, Q\. Wang, L\. Yang, and A\. Ahmed\(2020\)Big bird: transformers for longer sequences\.InProceedings of the 34th International Conference on Neural Information Processing Systems,NIPS ’20,Red Hook, NY, USA\.External Links:ISBN 9781713829546Cited by:[Appendix G](https://arxiv.org/html/2606.18283#A7.p1.1),[§1](https://arxiv.org/html/2606.18283#S1.p1.12),[§1](https://arxiv.org/html/2606.18283#S1.p2.1),[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px2.p1.1)\.
- \[55\]S\. Zhang and Y\. Feng\(2021\)Modeling concentrated cross\-attention for neural machine translation with gaussian mixture model\.InFindings of the Association for Computational Linguistics: EMNLP 2021,pp\. 1401–1411\.External Links:[Document](https://dx.doi.org/10.18653/v1/2021.findings-emnlp.121),[Link](https://aclanthology.org/2021.findings-emnlp.121.pdf)Cited by:[§2](https://arxiv.org/html/2606.18283#S2.SS0.SSS0.Px8.p1.1)\.

## Appendix ANotation for Gaussian Mixture Attention

Table[5](https://arxiv.org/html/2606.18283#A1.T5)summarizes the notation used in the GMA methodology, theoretical analysis, and experiments\.

Table 5:Notation used in the GMA methodology and analysis\.
## Appendix BMotivation: the Matrix Multiplication Association Rule

The linear\-time implementation of GMA relies on a standard principle from matrix algebra and matrix\-chain multiplication\. Matrix multiplication isassociative: whenever the dimensions are compatible,

\(A​B\)​C=A​\(B​C\)\.\(AB\)C=A\(BC\)\.\(42\)Thus, different parenthesizations give the same final matrix\. However, they can have very different computational costs\. This is the classical matrix\-chain multiplication observation: although all valid parenthesizations of a matrix product are algebraically equivalent, the choice of parenthesization can have a dramatic impact on the number of scalar multiplications required\[[11](https://arxiv.org/html/2606.18283#bib.bib24)\]\. The associative law itself is a basic linear algebra identity\[[46](https://arxiv.org/html/2606.18283#bib.bib25)\]\.

To see this concretely, let

A∈ℝm×n,B∈ℝn×p,C∈ℝp×q\.A\\in\\mathbb\{R\}^\{m\\times n\},\\qquad B\\in\\mathbb\{R\}^\{n\\times p\},\\qquad C\\in\\mathbb\{R\}^\{p\\times q\}\.Computing\(A​B\)​C\(AB\)Cfirst formsA​B∈ℝm×pAB\\in\\mathbb\{R\}^\{m\\times p\}, so the approximate multiplication cost is\[[8](https://arxiv.org/html/2606.18283#bib.bib26)\]

𝒞\(A​B\)​C=m​n​p\+m​p​q\.\\mathcal\{C\}\_\{\(AB\)C\}=mnp\+mpq\.\(43\)In contrast, computingA​\(B​C\)A\(BC\)first formsB​C∈ℝn×qBC\\in\\mathbb\{R\}^\{n\\times q\}, giving cost

𝒞A​\(B​C\)=n​p​q\+m​n​q\.\\mathcal\{C\}\_\{A\(BC\)\}=npq\+mnq\.\(44\)These two costs can differ substantially, even though the final result is identical\. For example, take

A∈ℝ1000×10,B∈ℝ10×1000,C∈ℝ1000×10\.A\\in\\mathbb\{R\}^\{1000\\times 10\},\\qquad B\\in\\mathbb\{R\}^\{10\\times 1000\},\\qquad C\\in\\mathbb\{R\}^\{1000\\times 10\}\.Then computing\(A​B\)​C\(AB\)Cfirst forms the large intermediate matrixA​B∈ℝ1000×1000AB\\in\\mathbb\{R\}^\{1000\\times 1000\}\. The cost is

1000⋅10⋅1000\+1000⋅1000⋅10=2×107\.1000\\cdot 10\\cdot 1000\+1000\\cdot 1000\\cdot 10=2\\times 10^\{7\}\.By contrast, computingA​\(B​C\)A\(BC\)first forms the much smaller intermediate matrixB​C∈ℝ10×10BC\\in\\mathbb\{R\}^\{10\\times 10\}\. The cost is

10⋅1000⋅10\+1000⋅10⋅10=2×105\.10\\cdot 1000\\cdot 10\+1000\\cdot 10\\cdot 10=2\\times 10^\{5\}\.Thus, the two parenthesizations produce the same output, but the second is cheaper by a factor of100100\.

This principle directly explains the efficient implementation of GMA\. Ignoring the row\-wise normalizer for the moment, the numerator of the GMA output in Eq\. \([7](https://arxiv.org/html/2606.18283#S3.E7)\) has the algebraic form

ΓQ​\(ΓK\)⊤​VX\.\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\.If we identify

A=ΓQ∈ℝN×K,B=\(ΓK\)⊤∈ℝK×N,C=VX∈ℝN×dv,A=\\Gamma^\{Q\}\\in\\mathbb\{R\}^\{N\\times K\},\\qquad B=\(\\Gamma^\{K\}\)^\{\\top\}\\in\\mathbb\{R\}^\{K\\times N\},\\qquad C=V\_\{X\}\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\},then the explicit\-affinity computation corresponds to

\(ΓQ​\(ΓK\)⊤\)​VX\.\\bigl\(\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\\bigr\)V\_\{X\}\.\(45\)This first forms the implicit token\-to\-token affinity matrix

ΓQ​\(ΓK\)⊤∈ℝN×N,\\Gamma^\{Q\}\(\\Gamma^\{K\}\)^\{\\top\}\\in\\mathbb\{R\}^\{N\\times N\},which costs𝒪​\(N2​K\)\\mathcal\{O\}\(N^\{2\}K\)arithmetic and requires𝒪​\(N2\)\\mathcal\{O\}\(N^\{2\}\)intermediate storage before multiplying byVXV\_\{X\}\.

GMA instead uses the associative parenthesization

ΓQ​\(\(ΓK\)⊤​VX\)\.\\Gamma^\{Q\}\\bigl\(\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\\bigr\)\.\(46\)The inner product

V~=\(ΓK\)⊤​VX∈ℝK×dv\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\}\(47\)constructs a compact latent memory with onlyKKslots\. The final multiplicationΓQ​V~\\Gamma^\{Q\}\\tilde\{V\}then reads from this memory\. The cost of these two matrix multiplications is

N​K​dv\+N​K​dv=2​N​K​dv,NKd\_\{v\}\+NKd\_\{v\}=2NKd\_\{v\},\(48\)and the intermediate memory has size onlyK×dvK\\times d\_\{v\}\. Therefore, GMA obtains the same implicit affinity interpretation without materializing theN×NN\\times Naffinity matrix\.

The full normalized GMA computation also includes the denominator

Z=\(ΓK\)⊤​𝟏N,ΓQ​Z\+ϵ,Z=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\},\\qquad\\Gamma^\{Q\}Z\+\\epsilon,\(49\)which adds only𝒪​\(N​K\)\\mathcal\{O\}\(NK\)arithmetic\. In addition, computing the two responsibility matricesΓQ\\Gamma^\{Q\}andΓK\\Gamma^\{K\}under diagonal Gaussian components costs𝒪​\(2​N​K​dr\)\\mathcal\{O\}\(2NKd\_\{r\}\), up to lower\-order normalization terms\. Thus, the overall routing computation remains linear inNNfor fixedKK,drd\_\{r\}, anddvd\_\{v\}, while retaining an implicit normalized attention matrix\.

## Appendix CUnderstanding GMA in Detail212121In this section, the author would like to express sincere gratitude to those whose textbooks, lectures, and scholarly discussions have shaped his understanding of numerical linear algebra, including Gilbert Strang, Stephen Boyd, Steven Roman, Nick Higham, and many others\.

We have briefly touched on the intuition behind GMA as responsibility\-space affinity followed by latent\-memory routing at the end of the core methodology in Section[3\.3](https://arxiv.org/html/2606.18283#S3.SS3)\. This appendix gives a more dimension\-explicit walkthrough of the same computation\. The goal is to clarify how key responsibilities write values into a latent memory, how query responsibilities read from this memory, how the implicit token\-to\-token affinity arises, and why the computation avoids materializing an explicitN×NN\\times Nattention matrix\.

### C\.1Objects and Dimensions

Consider a sequence of lengthNN\. LetVX∈ℝN×dvV\_\{X\}\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\}denote the value matrix, whereVX,j∈ℝdvV\_\{X,j\}\\in\\mathbb\{R\}^\{d\_\{v\}\}is the value vector at positionjj\. GMA computes two responsibility matrices:

ΓK∈ℝN×K,ΓQ∈ℝN×K\.\\Gamma^\{K\}\\in\\mathbb\{R\}^\{N\\times K\},\\qquad\\Gamma^\{Q\}\\in\\mathbb\{R\}^\{N\\times K\}\.The rowγjK=\(γj,1K,…,γj,KK\)\\gamma^\{K\}\_\{j\}=\(\\gamma^\{K\}\_\{j,1\},\\ldots,\\gamma^\{K\}\_\{j,K\}\)gives the posterior responsibilities of theKKlatent Gaussian components for key positionjj\. Similarly,γiQ=\(γi,1Q,…,γi,KQ\)\\gamma^\{Q\}\_\{i\}=\(\\gamma^\{Q\}\_\{i,1\},\\ldots,\\gamma^\{Q\}\_\{i,K\}\)gives the responsibilities of the same latent components for query positionii\. Thus,γj,kK\\gamma^\{K\}\_\{j,k\}measures how much key positionjjbelongs to latent componentkk, whileγi,kQ\\gamma^\{Q\}\_\{i,k\}measures how much query positioniireads from componentkk\.

Each row of a responsibility matrix lies on the probability simplex:

∑k=1Kγj,kK=1,∑k=1Kγi,kQ=1\.\\sum\_\{k=1\}^\{K\}\\gamma^\{K\}\_\{j,k\}=1,\\qquad\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}=1\.Therefore, GMA routes tokens through normalized probability vectors over latent components rather than through unconstrained similarity scores\.

### C\.2The Write Step: Keys Write Values into Latent Memory

The key responsibilities first aggregate the value vectors intoKKlatent memory slots:

V~=\(ΓK\)⊤​VX∈ℝK×dv,Z=\(ΓK\)⊤​𝟏N∈ℝK,\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\},\\qquad Z=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\}\\in\\mathbb\{R\}^\{K\},where𝟏N\\mathbf\{1\}\_\{N\}is the all\-ones vector\. Component\-wise, this means

V~k=∑j=1Nγj,kK​VX,j∈ℝdv,Zk=∑j=1Nγj,kK\.\\tilde\{V\}\_\{k\}=\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\\in\\mathbb\{R\}^\{d\_\{v\}\},\\qquad Z\_\{k\}=\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\.Thus,V~k\\tilde\{V\}\_\{k\}is the responsibility\-weighted sum of all value vectors assigned to latent componentkk, andZkZ\_\{k\}records the total key\-responsibility mass written into that component\.

An intuitive normalized memory slot can be written as

Mk=V~kZk\+ϵ=∑j=1Nγj,kK​VX,j∑j=1Nγj,kK\+ϵ∈ℝdv\.M\_\{k\}=\\frac\{\\tilde\{V\}\_\{k\}\}\{Z\_\{k\}\+\\epsilon\}=\\frac\{\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\}\{\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\+\\epsilon\}\\in\\mathbb\{R\}^\{d\_\{v\}\}\.This shows thatMkM\_\{k\}is the average value stored in latent componentkk, up to the numerical stabilizerϵ\\epsilon\. In practice, the implementation does not need to explicitly materializeMkM\_\{k\}; it is enough to keep the unnormalized memoryV~k\\tilde\{V\}\_\{k\}and the normalizerZkZ\_\{k\}\.

### C\.3The Read Step: Queries Read from Latent Memory

After the write step, query responsibilities read from the latent memory\. In matrix form, the GMA output is

O=ΓQ​V~ΓQ​Z\+ϵ∈ℝN×dv,O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\},where the denominator is broadcast across the value dimension and the division is applied row\-wise\. For query positionii, this is

Oi=∑k=1Kγi,kQ​V~k∑k=1Kγi,kQ​Zk\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z\_\{k\}\+\\epsilon\}\.Substituting the write step into this expression gives

Oi=∑k=1Kγi,kQ​∑j=1Nγj,kK​VX,j∑k=1Kγi,kQ​∑j=1Nγj,kK\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\sum\_\{j=1\}^\{N\}\\gamma^\{K\}\_\{j,k\}\+\\epsilon\}\.This equation makes the full workflow explicit: key responsibilities distribute values into latent components, and query responsibilities select a normalized mixture of those latent components\.

Ignoring the numerical stabilizer for clarity, we can also express the read operation as attention over normalized memory slots\. IfZk\>0Z\_\{k\}\>0, defineMk=V~k/ZkM\_\{k\}=\\tilde\{V\}\_\{k\}/Z\_\{k\}and

αi,k=γi,kQ​Zk∑ℓ=1Kγi,ℓQ​Zℓ\.\\alpha\_\{i,k\}=\\frac\{\\gamma^\{Q\}\_\{i,k\}Z\_\{k\}\}\{\\sum\_\{\\ell=1\}^\{K\}\\gamma^\{Q\}\_\{i,\\ell\}Z\_\{\\ell\}\}\.Then∑kαi,k=1\\sum\_\{k\}\\alpha\_\{i,k\}=1, and

Oi=∑k=1Kαi,k​Mk\.O\_\{i\}=\\sum\_\{k=1\}^\{K\}\\alpha\_\{i,k\}M\_\{k\}\.This form is useful for intuition\. The query does not simply average memory slots usingγi,kQ\\gamma^\{Q\}\_\{i,k\}alone\. Instead, the effective read weightαi,k\\alpha\_\{i,k\}also depends onZkZ\_\{k\}, the amount of key mass stored in componentkk\. For fixed query responsibilityγi,kQ\\gamma^\{Q\}\_\{i,k\}, a component with larger stored key mass receives a larger effective read weight than an almost\-empty component\.

### C\.4Implicit Token\-to\-Token Attention

Although GMA does not explicitly construct anN×NN\\times Nattention matrix, it still induces an attention\-style weighting over value tokens\. Combining the write and read equations gives

Oi=∑j=1NAi​jGMA​VX,j,O\_\{i\}=\\sum\_\{j=1\}^\{N\}A^\{\\mathrm\{GMA\}\}\_\{ij\}V\_\{X,j\},where

Ai​jGMA=∑k=1Kγi,kQ​γj,kK∑ℓ=1N∑k=1Kγi,kQ​γℓ,kK\+ϵ\.A^\{\\mathrm\{GMA\}\}\_\{ij\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\}\{\\sum\_\{\\ell=1\}^\{N\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{\\ell,k\}\+\\epsilon\}\.Thus, two tokens interact strongly when their query and key responsibilities overlap in the latent mixture space\. In standard attention, tokeniiattends to tokenjjthrough a direct dot productqi⊤​kjq\_\{i\}^\{\\top\}k\_\{j\}\. In GMA, tokeniiattends to tokenjjthrough their shared responsibility mass∑kγi,kQ​γj,kK\\sum\_\{k\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\.

Whenϵ=0\\epsilon=0, the induced weights are row\-normalized:

∑j=1NAi​jGMA=1\.\\sum\_\{j=1\}^\{N\}A^\{\\mathrm\{GMA\}\}\_\{ij\}=1\.Withϵ\>0\\epsilon\>0, the row sum is slightly below one, since the stabilizer adds extra positive mass to the denominator\. The stabilizer is included only for numerical safety\.

### C\.5Workflow Intuition

The full GMA computation can be understood as a three\-stage routing process:

- •First, the Gaussian mixture maps each token representation to a soft assignment overKKlatent components\. This producesΓK\\Gamma^\{K\}for keys andΓQ\\Gamma^\{Q\}for queries\. These matrices are probabilistic: each row is a distribution over latent components\.
- •Second, key responsibilities write the sequence values into a compact latent memory\. Instead of storingNNseparate value vectors for direct pairwise comparison, GMA storesKKresponsibility\-weighted memory slots\. Each slot summarizes the values assigned to one latent component, together with a normalizer recording how much total key mass was assigned there\.
- •Third, query responsibilities read from this latent memory\. A query token selects a mixture of memory slots according to its own responsibilities and the amount of key mass stored in those slots\. The output is therefore a normalized mixture of component\-level summaries\.

In short, the workflow can be summarised as:

keys write values into​K​latent memory slots,\\text\{keys write values into \}K\\text\{ latent memory slots,\}and

queries read mixtures of those slots to produce token outputs\.

### C\.6Why This Avoids the Quadratic Bottleneck

Standard dot\-product attention forms or implicitly represents all pairwise token\-to\-token scores, producing an intermediateN×NN\\times Nattention matrix\. GMA instead uses the twoN×KN\\times Kresponsibility matricesΓK\\Gamma^\{K\}andΓQ\\Gamma^\{Q\}, together with the latent memoryV~∈ℝK×dv\\tilde\{V\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\}and normalizerZ∈ℝKZ\\in\\mathbb\{R\}^\{K\}\. For fixedKK, the dominant activation storage is therefore linear inNN\.

The main costs are:

responsibility computation:​𝒪​\(N​K​dr\),\\text\{responsibility computation: \}\\mathcal\{O\}\(NKd\_\{r\}\),wheredrd\_\{r\}is the routing representation dimension, and

latent write/read operations:​𝒪​\(N​K​dv\)\.\\text\{latent write/read operations: \}\\mathcal\{O\}\(NKd\_\{v\}\)\.Thus, for fixedKK, GMA scales linearly with sequence length while retaining a normalized attention\-style interpretation\.

### C\.7Causal GMA

For autoregressive language modelling, positioniimust not read information from future positionsj\>ij\>i\. Causal GMA replaces the global write statistics by prefix\-restricted statistics:

V~k\(i\)=∑j≤iγj,kK​VX,j,Zk\(i\)=∑j≤iγj,kK\.\\tilde\{V\}^\{\(i\)\}\_\{k\}=\\sum\_\{j\\leq i\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\},\\qquad Z^\{\(i\)\}\_\{k\}=\\sum\_\{j\\leq i\}\\gamma^\{K\}\_\{j,k\}\.The output at positioniiis then

Oi=∑k=1Kγi,kQ​V~k\(i\)∑k=1Kγi,kQ​Zk\(i\)\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}^\{\(i\)\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z^\{\(i\)\}\_\{k\}\+\\epsilon\}\.This is the same write\-read idea, but the memory available to positioniiis restricted to its prefix\. In implementation, the prefix memories are computed by cumulative sums along the sequence dimension\. Therefore, causal GMA preserves the same fixed\-KKlinear scaling while enforcing autoregressive causality\.

## Appendix DDetailed Gradient Derivations for Reparameterized GMM

In Section[3\.6](https://arxiv.org/html/2606.18283#S3.SS6), we defined the trainable parameters of a GMA layer asΘGMA=\{WQ,WK,WV,μ,ω,α\}\\Theta\_\{\\mathrm\{GMA\}\}=\\\{W\_\{Q\},W\_\{K\},W\_\{V\},\\mu,\\omega,\\alpha\\\}in Eq\. \([13](https://arxiv.org/html/2606.18283#S3.E13)\)\. Among these,WQ,WK,WVW\_\{Q\},W\_\{K\},W\_\{V\}are the standard query, key, and value projection matrices \(Eq\. \([4](https://arxiv.org/html/2606.18283#S3.E4)\)\), whileμ,ω,α\\mu,\\omega,\\alphaparameterize the Gaussian mixture routing mechanism:μ\\mucontains the component means,ω\\omegaparameterizes the diagonal covariance entries \(Eq\. \([14](https://arxiv.org/html/2606.18283#S3.E14)\)\), andα\\alphaparameterizes the mixture\-prior logits \(Eq\. \([15](https://arxiv.org/html/2606.18283#S3.E15)\)\)\. This appendix derives the gradients for the reparameterized GMM routing parametersμ,ω,α\\mu,\\omega,\\alpha\. The projection matricesWQW\_\{Q\}andWKW\_\{K\}receive gradients through the routing vectorsqi=\(QX\)i,:q\_\{i\}=\(Q\_\{X\}\)\_\{i,:\}andki=\(KX\)i,:k\_\{i\}=\(K\_\{X\}\)\_\{i,:\}, whileWVW\_\{V\}receives gradients through the latent aggregation and output operations\.

Throughout,xi∈ℝdrx\_\{i\}\\in\\mathbb\{R\}^\{d\_\{r\}\}denotes a generic routing vector, which may be either a query routing vectorqiq\_\{i\}or a key routing vectorkik\_\{i\}\. We useθ\\thetato denote a generic scalar GMM routing parameter, i\.e\.θ∈\{μk,m,ωk,m,αm\}\\theta\\in\\\{\\mu\_\{k,m\},\\omega\_\{k,m\},\\alpha\_\{m\}\\\}for appropriate component and coordinate indices\. For componentkk, the pre\-normalized log\-density score under thediagonalcovariance assumption is a coordinate\-wise version of Eq\. \([16](https://arxiv.org/html/2606.18283#S4.E16)\):

si,k=log⁡πk−dr2​log⁡\(2​π\)−12​∑m=1drlog⁡\(σk,m2\)−12​∑m=1dr\(xi,m−μk,m\)2σk,m2\.s\_\{i,k\}=\\log\\pi\_\{k\}\-\\frac\{d\_\{r\}\}\{2\}\\log\(2\\pi\)\-\\frac\{1\}\{2\}\\sum\_\{m=1\}^\{d\_\{r\}\}\\log\(\\sigma\_\{k,m\}^\{2\}\)\-\\frac\{1\}\{2\}\\sum\_\{m=1\}^\{d\_\{r\}\}\\frac\{\(x\_\{i,m\}\-\\mu\_\{k,m\}\)^\{2\}\}\{\\sigma\_\{k,m\}^\{2\}\}\.\(50\)The responsibility is the softmax\-normalized score \(Eq\.\([3b](https://arxiv.org/html/2606.18283#S4.Ex25)\)\)

γi,k=exp⁡\(si,k\)∑ℓ=1Kexp⁡\(si,ℓ\)\.\\gamma\_\{i,k\}=\\frac\{\\exp\(s\_\{i,k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\exp\(s\_\{i,\\ell\}\)\}\.\(51\)Let

gi,j=∂ℒ∂γi,jg\_\{i,j\}=\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\gamma\_\{i,j\}\}denote the upstream gradient arriving at the responsibility vector for tokenii\. By the softmax Jacobian \(Eq\.\([17](https://arxiv.org/html/2606.18283#S4.E17)\)\),

∂γi,j∂si,k=γi,j​\(𝟏​\{j=k\}−γi,k\)\.\\frac\{\\partial\\gamma\_\{i,j\}\}\{\\partial s\_\{i,k\}\}=\\gamma\_\{i,j\}\\left\(\\mathbf\{1\}\\\{j=k\\\}\-\\gamma\_\{i,k\}\\right\)\.Therefore, the gradient entering the scoresi,ks\_\{i,k\}is \(Eq\.\([18](https://arxiv.org/html/2606.18283#S4.E18)\)\)

δi,k≡∂ℒ∂si,k=∑j=1K∂ℒ∂γi,j​∂γi,j∂si,k=γi,k​\(gi,k−∑j=1Kγi,j​gi,j\)\.\\delta\_\{i,k\}\\equiv\\frac\{\\partial\\mathcal\{L\}\}\{\\partial s\_\{i,k\}\}=\\sum\_\{j=1\}^\{K\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\gamma\_\{i,j\}\}\\frac\{\\partial\\gamma\_\{i,j\}\}\{\\partial s\_\{i,k\}\}=\\gamma\_\{i,k\}\\left\(g\_\{i,k\}\-\\sum\_\{j=1\}^\{K\}\\gamma\_\{i,j\}g\_\{i,j\}\\right\)\.This is the quantity that multiplies each parameter derivative∂si,k/∂θ\\partial s\_\{i,k\}/\\partial\\thetain backpropagation\. The familiar factorγi,k​\(1−γi,k\)\\gamma\_\{i,k\}\(1\-\\gamma\_\{i,k\}\)appears as the self\-sensitivity∂γi,k/∂si,k\\partial\\gamma\_\{i,k\}/\\partial s\_\{i,k\}\(Eq\.\([17b](https://arxiv.org/html/2606.18283#S4.Ex27)\)\), while the full gradient also includes cross\-component softmax terms \(Eq\.\([22](https://arxiv.org/html/2606.18283#S4.E22)\)\)\.

For any learnable scalar parameterθ\\thetathat enters the scores, such as a component ofμ\\mu,ω\\omega, orα\\alpha, or a projection parameter entering the routing vectors throughWQW\_\{Q\}orWKW\_\{K\}, the chain rule gives

∂ℒ∂θ=∑i∑k=1K∂ℒ∂si,k​∂si,k∂θ,\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\theta\}=\\sum\_\{i\}\\sum\_\{k=1\}^\{K\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial s\_\{i,k\}\}\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\theta\},\(cc\.Eq\.[19](https://arxiv.org/html/2606.18283#S4.E19)\)where∂ℒ/∂si,k\\partial\\mathcal\{L\}/\\partial s\_\{i,k\}is given by Eq\. \([18](https://arxiv.org/html/2606.18283#S4.E18)\)\. Equivalently, expanding the softmax Jacobian explicitly gives

∂ℒ∂θ=∑i∑j=1K∑k=1K∂ℒ∂γi,j​∂γi,j∂si,k​∂si,k∂θ\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\theta\}=\\sum\_\{i\}\\sum\_\{j=1\}^\{K\}\\sum\_\{k=1\}^\{K\}\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\gamma\_\{i,j\}\}\\frac\{\\partial\\gamma\_\{i,j\}\}\{\\partial s\_\{i,k\}\}\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\theta\}\.\(cc\.Eq\.[19b](https://arxiv.org/html/2606.18283#S4.Ex26)\)This expression is the appropriate backpropagation form for the responsibility\-based GMA layer\. We now derive the score derivatives for the reparameterized GMM parameters\.

### D\.1Derivation for Latent Means

The means are unconstrained parametersμk∈ℝdr\\mu\_\{k\}\\in\\mathbb\{R\}^\{d\_\{r\}\}\. Differentiating Eq\. \([50](https://arxiv.org/html/2606.18283#A4.E50)\) with respect to themm\-th coordinate ofμk\\mu\_\{k\}gives

∂si,k∂μk,m=xi,m−μk,mσk,m2\.\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\mu\_\{k,m\}\}=\\frac\{x\_\{i,m\}\-\\mu\_\{k,m\}\}\{\\sigma\_\{k,m\}^\{2\}\}\.\(52\)In vector form,

∇μksi,k=Σk−1​\(xi−μk\)\.\\nabla\_\{\\mu\_\{k\}\}s\_\{i,k\}=\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\)\.\(53\)Thus, the full loss gradient with respect to the mean of componentkkis

∂ℒ∂μk=∑iδi,k​Σk−1​\(xi−μk\)\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\mu\_\{k\}\}=\\sum\_\{i\}\\delta\_\{i,k\}\\Sigma\_\{k\}^\{\-1\}\(x\_\{i\}\-\\mu\_\{k\}\)\.\(54\)Hence, the mean update is driven by precision\-scaled residuals, weighted by the score\-level gradientδi,k\\delta\_\{i,k\}induced by the responsibility softmax\.

### D\.2Derivation for Reparameterized Diagonal Covariances

To ensure positive diagonal covariance entries, GMA maintains unconstrained parametersω∈ℝK×dr\\omega\\in\\mathbb\{R\}^\{K\\times d\_\{r\}\}and sets \(Eq\.\([14](https://arxiv.org/html/2606.18283#S3.E14)\)\)

σk,m2=softplus⁡\(ωk,m\)\+ϵσ=log⁡\(1\+exp⁡\(ωk,m\)\)\+ϵσ,m=1,…,dr\.\\sigma\_\{k,m\}^\{2\}=\\operatorname\{softplus\}\(\\omega\_\{k,m\}\)\+\\epsilon\_\{\\sigma\}=\\log\(1\+\\exp\(\\omega\_\{k,m\}\)\)\+\\epsilon\_\{\\sigma\},\\qquad m=1,\\ldots,d\_\{r\}\.\(55\)By the chain rule,

∂si,k∂ωk,m=∂si,k∂σk,m2​∂σk,m2∂ωk,m\.\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\omega\_\{k,m\}\}=\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\sigma\_\{k,m\}^\{2\}\}\\frac\{\\partial\\sigma\_\{k,m\}^\{2\}\}\{\\partial\\omega\_\{k,m\}\}\.\(56\)First, differentiating Eq\. \([50](https://arxiv.org/html/2606.18283#A4.E50)\) with respect to the varianceσk,m2\\sigma\_\{k,m\}^\{2\}gives

∂si,k∂σk,m2=−12​σk,m2\+\(xi,m−μk,m\)22​\(σk,m2\)2=12​σk,m2​\(\(xi,m−μk,m\)2σk,m2−1\)\.\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\sigma\_\{k,m\}^\{2\}\}=\-\\frac\{1\}\{2\\sigma\_\{k,m\}^\{2\}\}\+\\frac\{\(x\_\{i,m\}\-\\mu\_\{k,m\}\)^\{2\}\}\{2\(\\sigma\_\{k,m\}^\{2\}\)^\{2\}\}=\\frac\{1\}\{2\\sigma\_\{k,m\}^\{2\}\}\\left\(\\frac\{\(x\_\{i,m\}\-\\mu\_\{k,m\}\)^\{2\}\}\{\\sigma\_\{k,m\}^\{2\}\}\-1\\right\)\.\(57\)Second, the derivative of the softplus reparameterization is

∂σk,m2∂ωk,m=exp⁡\(ωk,m\)1\+exp⁡\(ωk,m\)=σsig​\(ωk,m\),\\frac\{\\partial\\sigma\_\{k,m\}^\{2\}\}\{\\partial\\omega\_\{k,m\}\}=\\frac\{\\exp\(\\omega\_\{k,m\}\)\}\{1\+\\exp\(\\omega\_\{k,m\}\)\}=\\sigma\_\{\\mathrm\{sig\}\}\(\\omega\_\{k,m\}\),\(58\)whereσsig​\(⋅\)\\sigma\_\{\\mathrm\{sig\}\}\(\\cdot\)denotes the logisticsigmoid function\. Combining Eqs\. \([57](https://arxiv.org/html/2606.18283#A4.E57)\) and \([58](https://arxiv.org/html/2606.18283#A4.E58)\), we obtain

∂si,k∂ωk,m=12​σk,m2​\(\(xi,m−μk,m\)2σk,m2−1\)​σsig​\(ωk,m\)\.\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\omega\_\{k,m\}\}=\\frac\{1\}\{2\\sigma\_\{k,m\}^\{2\}\}\\left\(\\frac\{\(x\_\{i,m\}\-\\mu\_\{k,m\}\)^\{2\}\}\{\\sigma\_\{k,m\}^\{2\}\}\-1\\right\)\\sigma\_\{\\mathrm\{sig\}\}\(\\omega\_\{k,m\}\)\.\(59\)Therefore, the full loss gradient for the covariance parameter is

∂ℒ∂ωk,m=∑iδi,k​12​σk,m2​\(\(xi,m−μk,m\)2σk,m2−1\)​σsig​\(ωk,m\)\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\omega\_\{k,m\}\}=\\sum\_\{i\}\\delta\_\{i,k\}\\frac\{1\}\{2\\sigma\_\{k,m\}^\{2\}\}\\left\(\\frac\{\(x\_\{i,m\}\-\\mu\_\{k,m\}\)^\{2\}\}\{\\sigma\_\{k,m\}^\{2\}\}\-1\\right\)\\sigma\_\{\\mathrm\{sig\}\}\(\\omega\_\{k,m\}\)\.\(60\)The sigmoid factor attenuates the gradient passed through the softplus reparameterization, while the additive lower boundϵσ\\epsilon\_\{\\sigma\}prevents the diagonal variances from becoming singular\. This improves numerical stability, although it does not by itself eliminate the need for appropriate initialization, learning\-rate control, or other optimization safeguards\.

### D\.3Derivation for Reparameterized Mixture Priors

To satisfy the simplex constraint, the mixture priors are parameterized by unconstrained logitsα∈ℝK\\alpha\\in\\mathbb\{R\}^\{K\}\(Eq\.\([15](https://arxiv.org/html/2606.18283#S3.E15)\)\):

πk=exp⁡\(αk\)∑ℓ=1Kexp⁡\(αℓ\)\.\\pi\_\{k\}=\\frac\{\\exp\(\\alpha\_\{k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\exp\(\\alpha\_\{\\ell\}\)\}\.\(cc\.Eq\.[15](https://arxiv.org/html/2606.18283#S3.E15)\)Becausesi,ks\_\{i,k\}contains the termlog⁡πk\\log\\pi\_\{k\}\(Eq\.\([50](https://arxiv.org/html/2606.18283#A4.E50)\)\), we need the derivative oflog⁡πk\\log\\pi\_\{k\}with respect to an arbitrary logitαm\\alpha\_\{m\}\. Using the softmax Jacobian,

∂πk∂αm=πk​\(𝟏​\{k=m\}−πm\)\.\\frac\{\\partial\\pi\_\{k\}\}\{\\partial\\alpha\_\{m\}\}=\\pi\_\{k\}\(\\mathbf\{1\}\\\{k=m\\\}\-\\pi\_\{m\}\)\.\(61\)Therefore,

∂si,k∂αm=∂log⁡πk∂αm=1πk​∂πk∂αm=𝟏​\{k=m\}−πm\.\\frac\{\\partial s\_\{i,k\}\}\{\\partial\\alpha\_\{m\}\}=\\frac\{\\partial\\log\\pi\_\{k\}\}\{\\partial\\alpha\_\{m\}\}=\\frac\{1\}\{\\pi\_\{k\}\}\\frac\{\\partial\\pi\_\{k\}\}\{\\partial\\alpha\_\{m\}\}=\\mathbf\{1\}\\\{k=m\\\}\-\\pi\_\{m\}\.\(62\)The full loss gradient with respect to prior logitαm\\alpha\_\{m\}is then

∂ℒ∂αm=∑i∑k=1Kδi,k​\(𝟏​\{k=m\}−πm\)\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\alpha\_\{m\}\}=\\sum\_\{i\}\\sum\_\{k=1\}^\{K\}\\delta\_\{i,k\}\\left\(\\mathbf\{1\}\\\{k=m\\\}\-\\pi\_\{m\}\\right\)\.\(63\)Since the score\-level gradientsδi,k\\delta\_\{i,k\}arise from a softmax Jacobian, they satisfy∑k=1Kδi,k=0\\sum\_\{k=1\}^\{K\}\\delta\_\{i,k\}=0for each tokenii\. Hence Eq\. \([63](https://arxiv.org/html/2606.18283#A4.E63)\) can also be written as

∂ℒ∂αm=∑iδi,m\.\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\alpha\_\{m\}\}=\\sum\_\{i\}\\delta\_\{i,m\}\.\(64\)At the score level, increasingαm\\alpha\_\{m\}increases the log\-prior contribution of componentmmby1−πm1\-\\pi\_\{m\}and decreases the log\-prior contribution of every other component byπm\\pi\_\{m\}\. The actual parameter update is then determined by the downstream loss through the score\-level gradientsδi,k\\delta\_\{i,k\}\.

## Appendix ESelf\-Attention, Cross\-Attention, and GMA

This appendix clarifies the distinction betweenself\-attentionandcross\-attention, and shows how the GMA routing mechanism extends naturally to both settings\. To avoid notational ambiguity, we writeKattK\_\{\\mathrm\{att\}\}for the standard\-attention key matrix, whileKKdenotes the number of Gaussian mixture components in GMA\.

### E\.1Standard Self\-Attention

In self\-attention, the queries, keys, and values are all projected from the same input sequence\. Let

X∈ℝN×dmodel,X\\in\\mathbb\{R\}^\{N\\times d\_\{\\mathrm\{model\}\}\},whereNNis the sequence length\. Standard self\-attention forms

Q=X​WQ,Katt=X​WK,V=X​WV,Q=XW\_\{Q\},\\qquad K\_\{\\mathrm\{att\}\}=XW\_\{K\},\\qquad V=XW\_\{V\},with

Q,Katt∈ℝN×dk,V∈ℝN×dv\.Q,K\_\{\\mathrm\{att\}\}\\in\\mathbb\{R\}^\{N\\times d\_\{k\}\},\\qquad V\\in\\mathbb\{R\}^\{N\\times d\_\{v\}\}\.The attention matrix and output are

A=softmax⁡\(Q​Katt⊤dk\),O=A​V\.A=\\operatorname\{softmax\}\\left\(\\frac\{QK\_\{\\mathrm\{att\}\}^\{\\top\}\}\{\\sqrt\{d\_\{k\}\}\}\\right\),\\qquad O=AV\.Elementwise,

Ai​j=exp⁡\(qi⊤​kj/dk\)∑ℓ=1Nexp⁡\(qi⊤​kℓ/dk\),Oi=∑j=1NAi​j​vj\.A\_\{ij\}=\\frac\{\\exp\(q\_\{i\}^\{\\top\}k\_\{j\}/\\sqrt\{d\_\{k\}\}\)\}\{\\sum\_\{\\ell=1\}^\{N\}\\exp\(q\_\{i\}^\{\\top\}k\_\{\\ell\}/\\sqrt\{d\_\{k\}\}\)\},\\qquad O\_\{i\}=\\sum\_\{j=1\}^\{N\}A\_\{ij\}v\_\{j\}\.Thus, tokeniiattends to tokensjjwithin the same sequence\.

For causal self\-attention, as used in autoregressive language modelling, future positions are masked:

Ai​j=0for​j\>i\.A\_\{ij\}=0\\qquad\\text\{for \}j\>i\.Equivalently,

Ai​j=𝟏​\{j≤i\}​exp⁡\(qi⊤​kj/dk\)∑ℓ≤iexp⁡\(qi⊤​kℓ/dk\)\.A\_\{ij\}=\\frac\{\\mathbf\{1\}\\\{j\\leq i\\\}\\exp\(q\_\{i\}^\{\\top\}k\_\{j\}/\\sqrt\{d\_\{k\}\}\)\}\{\\sum\_\{\\ell\\leq i\}\\exp\(q\_\{i\}^\{\\top\}k\_\{\\ell\}/\\sqrt\{d\_\{k\}\}\)\}\.The representation at positioniimay depend only on tokens at positionsj≤ij\\leq i, and is then used for next\-token prediction\.

### E\.2Standard Cross\-Attention

In cross\-attention, the queries come from one sequence, while keys and values come from another sequence\. Let

X∈ℝNq×dmodelX\\in\\mathbb\{R\}^\{N\_\{q\}\\times d\_\{\\mathrm\{model\}\}\}denote a query\-side sequence, and let

Y∈ℝNk×dmodelY\\in\\mathbb\{R\}^\{N\_\{k\}\\times d\_\{\\mathrm\{model\}\}\}denote a context or source sequence\. Cross\-attention forms

QX=X​WQ,KY=Y​WK,VY=Y​WV,Q\_\{X\}=XW\_\{Q\},\\qquad K\_\{Y\}=YW\_\{K\},\\qquad V\_\{Y\}=YW\_\{V\},where

QX∈ℝNq×dk,KY∈ℝNk×dk,VY∈ℝNk×dv\.Q\_\{X\}\\in\\mathbb\{R\}^\{N\_\{q\}\\times d\_\{k\}\},\\qquad K\_\{Y\}\\in\\mathbb\{R\}^\{N\_\{k\}\\times d\_\{k\}\},\\qquad V\_\{Y\}\\in\\mathbb\{R\}^\{N\_\{k\}\\times d\_\{v\}\}\.The attention matrix is rectangular:

A=softmax⁡\(QX​KY⊤dk\)∈ℝNq×Nk,A=\\operatorname\{softmax\}\\left\(\\frac\{Q\_\{X\}K\_\{Y\}^\{\\top\}\}\{\\sqrt\{d\_\{k\}\}\}\\right\)\\in\\mathbb\{R\}^\{N\_\{q\}\\times N\_\{k\}\},and the output is

OX=A​VY∈ℝNq×dv\.O\_\{X\}=AV\_\{Y\}\\in\\mathbb\{R\}^\{N\_\{q\}\\times d\_\{v\}\}\.Elementwise,

Ai​j=exp⁡\(qX,i⊤​kY,j/dk\)∑ℓ=1Nkexp⁡\(qX,i⊤​kY,ℓ/dk\),OX,i=∑j=1NkAi​j​vY,j\.A\_\{ij\}=\\frac\{\\exp\(q\_\{X,i\}^\{\\top\}k\_\{Y,j\}/\\sqrt\{d\_\{k\}\}\)\}\{\\sum\_\{\\ell=1\}^\{N\_\{k\}\}\\exp\(q\_\{X,i\}^\{\\top\}k\_\{Y,\\ell\}/\\sqrt\{d\_\{k\}\}\)\},\\qquad O\_\{X,i\}=\\sum\_\{j=1\}^\{N\_\{k\}\}A\_\{ij\}v\_\{Y,j\}\.Thus, query tokeniiin sequenceXXattends to context tokenjjin sequenceYY\. Self\-attention is the special case in which the query\-side and context\-side sequences coincide\.

### E\.3GMA Self\-Attention

The GMA mechanism developed in Section[3](https://arxiv.org/html/2606.18283#S3)is presented primarily in theself\-attentionsetting\. Given

X∈ℝN×dmodel,X\\in\\mathbb\{R\}^\{N\\times d\_\{\\mathrm\{model\}\}\},we form

QX=X​WQ,KX=X​WK,VX=X​WV\.Q\_\{X\}=XW\_\{Q\},\\qquad K\_\{X\}=XW\_\{K\},\\qquad V\_\{X\}=XW\_\{V\}\.Instead of computing all pairwise dot productsqi⊤​kjq\_\{i\}^\{\\top\}k\_\{j\}, GMA maps query and key projections to responsibility vectors overKKlearned Gaussian components:

ΓQ=Γ​\(QX\),ΓK=Γ​\(KX\),\\Gamma^\{Q\}=\\Gamma\(Q\_\{X\}\),\\qquad\\Gamma^\{K\}=\\Gamma\(K\_\{X\}\),where

ΓQ,ΓK∈ℝN×K\.\\Gamma^\{Q\},\\Gamma^\{K\}\\in\\mathbb\{R\}^\{N\\times K\}\.For a generic routing vectorxx, the responsibility of componentkkis

γk​\(x\)=πk​𝒩​\(x∣μk,Σk\)∑ℓ=1Kπℓ​𝒩​\(x∣μℓ,Σℓ\)\.\\gamma\_\{k\}\(x\)=\\frac\{\\pi\_\{k\}\\mathcal\{N\}\(x\\mid\\mu\_\{k\},\\Sigma\_\{k\}\)\}\{\\sum\_\{\\ell=1\}^\{K\}\\pi\_\{\\ell\}\\mathcal\{N\}\(x\\mid\\mu\_\{\\ell\},\\Sigma\_\{\\ell\}\)\}\.The key\-side responsibilities write values into a latent memory:

V~=\(ΓK\)⊤​VX,Z=\(ΓK\)⊤​𝟏N,\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\},\\qquad Z=\(\\Gamma^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\},where

V~∈ℝK×dv,Z∈ℝK\.\\tilde\{V\}\\in\\mathbb\{R\}^\{K\\times d\_\{v\}\},\\qquad Z\\in\\mathbb\{R\}^\{K\}\.The query\-side responsibilities then read from this latent memory:

O=ΓQ​V~ΓQ​Z\+ϵ,O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\},where the denominator is broadcast row\-wise across the value dimension\. Equivalently, for tokenii,

Oi=∑k=1Kγi,kQ​V~k∑k=1Kγi,kQ​Zk\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z\_\{k\}\+\\epsilon\}\.The induced GMA self\-attention weight from tokeniito tokenjjis

Ai​jGMA=∑k=1Kγi,kQ​γj,kK∑ℓ=1N∑k=1Kγi,kQ​γℓ,kK\+ϵ\.A^\{\\mathrm\{GMA\}\}\_\{ij\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\}\{\\sum\_\{\\ell=1\}^\{N\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{\\ell,k\}\+\\epsilon\}\.Thus, GMA replaces direct pairwise dot\-product comparison with responsibility\-space affinity followed by latent\-memory routing\.

For causal GMA self\-attention, the key\-side latent memory is accumulated only over prefix positions:

V~k\(i\)=∑j≤iγj,kK​VX,j,Zk\(i\)=∑j≤iγj,kK\.\\tilde\{V\}^\{\(i\)\}\_\{k\}=\\sum\_\{j\\leq i\}\\gamma^\{K\}\_\{j,k\}V\_\{X,j\},\\qquad Z^\{\(i\)\}\_\{k\}=\\sum\_\{j\\leq i\}\\gamma^\{K\}\_\{j,k\}\.The causal output is

Oi=∑k=1Kγi,kQ​V~k\(i\)∑k=1Kγi,kQ​Zk\(i\)\+ϵ\.O\_\{i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\tilde\{V\}^\{\(i\)\}\_\{k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}Z^\{\(i\)\}\_\{k\}\+\\epsilon\}\.The corresponding induced causal GMA weight is

Ai​jcGMA=𝟏​\{j≤i\}​∑k=1Kγi,kQ​γj,kK∑ℓ≤i∑k=1Kγi,kQ​γℓ,kK\+ϵ\.A^\{\\mathrm\{cGMA\}\}\_\{ij\}=\\mathbf\{1\}\\\{j\\leq i\\\}\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\}\{\\sum\_\{\\ell\\leq i\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{\\ell,k\}\+\\epsilon\}\.This is the form used in the WikiText\-103 language\-modelling experiment\.

### E\.4GMA Cross\-Attention

GMA can also be applied tocross\-attention\. Let

X∈ℝNq×dmodelX\\in\\mathbb\{R\}^\{N\_\{q\}\\times d\_\{\\mathrm\{model\}\}\}be the query\-side sequence and

Y∈ℝNk×dmodelY\\in\\mathbb\{R\}^\{N\_\{k\}\\times d\_\{\\mathrm\{model\}\}\}be the context\-side sequence\. We form

QX=X​WQ,KY=Y​WK,VY=Y​WV\.Q\_\{X\}=XW\_\{Q\},\\qquad K\_\{Y\}=YW\_\{K\},\\qquad V\_\{Y\}=YW\_\{V\}\.The query\-side and context\-side responsibilities are

ΓXQ=Γ​\(QX\)∈ℝNq×K,ΓYK=Γ​\(KY\)∈ℝNk×K\.\\Gamma\_\{X\}^\{Q\}=\\Gamma\(Q\_\{X\}\)\\in\\mathbb\{R\}^\{N\_\{q\}\\times K\},\\qquad\\Gamma\_\{Y\}^\{K\}=\\Gamma\(K\_\{Y\}\)\\in\\mathbb\{R\}^\{N\_\{k\}\\times K\}\.The context sequence writes values into a latent memory:

V~Y=\(ΓYK\)⊤​VY,ZY=\(ΓYK\)⊤​𝟏Nk\.\\tilde\{V\}\_\{Y\}=\(\\Gamma\_\{Y\}^\{K\}\)^\{\\top\}V\_\{Y\},\\qquad Z\_\{Y\}=\(\\Gamma\_\{Y\}^\{K\}\)^\{\\top\}\\mathbf\{1\}\_\{N\_\{k\}\}\.The query sequence reads from this context\-side memory:

OX=ΓXQ​V~YΓXQ​ZY\+ϵ∈ℝNq×dv\.O\_\{X\}=\\frac\{\\Gamma\_\{X\}^\{Q\}\\tilde\{V\}\_\{Y\}\}\{\\Gamma\_\{X\}^\{Q\}Z\_\{Y\}\+\\epsilon\}\\in\\mathbb\{R\}^\{N\_\{q\}\\times d\_\{v\}\}\.Elementwise,

OX,i=∑k=1KγX,i,kQ​V~Y,k∑k=1KγX,i,kQ​ZY,k\+ϵ\.O\_\{X,i\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{X,i,k\}\\tilde\{V\}\_\{Y,k\}\}\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{X,i,k\}Z\_\{Y,k\}\+\\epsilon\}\.The induced GMA cross\-attention weight from query tokeniiinXXto context tokenjjinYYis

Ai​jGMA​\-​cross=∑k=1KγX,i,kQ​γY,j,kK∑ℓ=1Nk∑k=1KγX,i,kQ​γY,ℓ,kK\+ϵ\.A^\{\\mathrm\{GMA\\text\{\-\}cross\}\}\_\{ij\}=\\frac\{\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{X,i,k\}\\gamma^\{K\}\_\{Y,j,k\}\}\{\\sum\_\{\\ell=1\}^\{N\_\{k\}\}\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{X,i,k\}\\gamma^\{K\}\_\{Y,\\ell,k\}\+\\epsilon\}\.Thus, GMA cross\-attention has the same structure as GMA self\-attention, except that query responsibilities are computed from the query\-side sequence and key responsibilities/value memories are computed from the context\-side sequence\.

The experiments in this paper focus on self\-attention\-style sequence mixing: bidirectional self\-attention\-style mixing for LRA and causal self\-attention\-style mixing for WikiText\-103\. A systematic empirical study of GMA cross\-attention, for example in encoder–decoder or multimodal architectures, is left to future work\.

## Appendix FAdditional Interpretability Diagnostics

This appendix reports additional diagnostics for the interpretability analysis in Section[5\.4](https://arxiv.org/html/2606.18283#S5.SS4)\. All results use the trained WikiText\-103 causal GMA model withK=128K=128\. We extract the final\-layer query responsibility tensor, average it over heads for each token occurrence, and analyzeS=64S=64validation sequences of lengthL=1024L=1024, givingT=S​L=65,536T=SL=65\{,\}536token\-level responsibility vectors\. Following Section[5\.4](https://arxiv.org/html/2606.18283#S5.SS4), we write the resulting head\-averaged responsibility vector asγ¯tQ∈ΔK−1\\bar\{\\gamma\}^\{Q\}\_\{t\}\\in\\Delta^\{K\-1\}fort=1,…,Tt=1,\\ldots,T, and define the hard assignmentzt=arg⁡maxk⁡γ¯t,kQz\_\{t\}=\\arg\\max\_\{k\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\.

### F\.1Token Category Distribution

The token categories used in the purity and mutual\-information analysis are surface\-form categories derived from GPT\-2 BPE tokens\. Table[6](https://arxiv.org/html/2606.18283#A6.T6)shows the category frequencies in the analyzed validation subset\.

Table 6:Surface\-form token category distribution in the analyzed WikiText\-103 validation subset\.The largest single category is lower\-case alphabetic tokens, with fraction18,757/65,536≈0\.28618\{,\}757/65\{,\}536\\approx 0\.286\. This gives a simple majority\-category baseline for interpreting the weighted category purity\. The observed weighted purity of0\.4830\.483is substantially above this baseline\.

### F\.2Representative Component Specialization

Table[7](https://arxiv.org/html/2606.18283#A6.T7)shows representative routing channels with their hard\-assignment frequency, dominant category, purity, and representative tokens\. These examples illustrate that some channels are broad mixed routers, whereas others show clearer surface\-form specialization\.

Table 7:Representative GMA routing channels under hard assignmentszt=arg⁡maxk⁡γ¯t,kQz\_\{t\}=\\arg\\max\_\{k\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}\. Purity is the fraction of tokens assigned to the channel that belong to its dominant surface\-form category\.The component examples show a mixture of broad and specialized behavior\. High\-usage channels such as component 22 have low purity but absorb common mixed tokens\. Other channels are more interpretable: component 50 is strongly associated with lower\-case alphabetic words, components 5 and 116 are function\-word\-heavy, component 17 is punctuation\-like, component 73 captures capitalized fragments, and component 2 captures subword alphabetic fragments\. These examples should be interpreted as surface\-form tendencies of head\-averaged routing channels, not as clean syntactic or semantic categories\.

### F\.3Permutation Baseline

To check whether the observed component\-category association could arise simply from imbalanced category frequencies, we compare against a permutation baseline\. The category labels\{ct\}t=1T\\\{c\_\{t\}\\\}\_\{t=1\}^\{T\}are randomly shuffled while keeping the learned hard component assignments\{zt\}t=1T\\\{z\_\{t\}\\\}\_\{t=1\}^\{T\}fixed\. This preserves both the marginal component distribution and the marginal category distribution, but destroys any systematic alignment between them\.

Table 8:Permutation sanity check for component\-category alignment\. The permutation baseline shuffles token categories while preserving hard component assignments and category frequencies\.The permutation baseline confirms that the observed purity and mutual information are not explained by the marginal token\-category frequencies alone\. The learned component assignments contain substantially more surface\-form category information than randomized assignments with the same marginals\.

### F\.4Additional Visualizations

Figure[4](https://arxiv.org/html/2606.18283#A6.F4)shows the marginal component usage and per\-component category purity\. The component usage distribution has normalized entropy0\.9330\.933, indicating broad usage of the latent routing channels\. The purity plot shows that several high\-usage channels have a clear dominant surface\-form category, although many remain mixed\.

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig2_component_usage_histogram.png)

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig3_component_purity_top_components.png)

Figure 4:Additional component\-level diagnostics\. Top: marginal component usage distributionpk=T−1​∑t=1Tγ¯t,kQp\_\{k\}=T^\{\-1\}\\sum\_\{t=1\}^\{T\}\\bar\{\\gamma\}^\{Q\}\_\{t,k\}, showing broad component usage and no severe collapse\. Bottom: dominant\-category purity for the most\-used components\.Figure[5](https://arxiv.org/html/2606.18283#A6.F5)gives a qualitative visualization of hard component assignments on one WikiText\-103 validation segment\. Tokens are colored by their dominant component\. Repeated colors over related local token roles provide a visual indication of latent routing specialization\.

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig5_token_component_strip.png)Figure 5:Qualitative token strip in which each token is colored by its dominant GMA component\. This visualization illustrates local routing patterns across a validation sequence\.Figure[6](https://arxiv.org/html/2606.18283#A6.F6)projects token responsibility vectors to two dimensions using PCA\. The first two principal components explain approximately12\.0%12\.0\\%and7\.9%7\.9\\%of the variance, respectively\. The category\-colored projection shows partial but overlapping separation between surface\-form categories, while the component\-colored projection shows that hard component assignments occupy structured regions of the responsibility space\.

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig6_pca_responsibilities_by_category.png)

![Refer to caption](https://arxiv.org/html/2606.18283v1/figures/fig6b_pca_responsibilities_by_component.png)

Figure 6:PCA projections of token responsibility vectors\. Top: points colored by surface\-form token category\. Bottom: points colored by hard GMA component assignment\. The projections show structured but overlapping responsibility geometry\.

## Appendix GAttention, Similarity, Distance, Divergence, and Correlation

This appendix provides a broaderconceptualperspective on attention mechanisms\. At an abstract level, an attention layer assigns weights to value vectors by first computing a compatibility score between a query representation and a set of key representations\. Different attention mechanisms can therefore be viewed as different choices of compatibility function\. Dot\-product attention uses an inner product\[[49](https://arxiv.org/html/2606.18283#bib.bib11)\]; additive attention uses a learned neural alignment/scoring function\[[2](https://arxiv.org/html/2606.18283#bib.bib35)\]; kernelized attention uses feature\-map similarities and associative matrix products\[[28](https://arxiv.org/html/2606.18283#bib.bib7),[10](https://arxiv.org/html/2606.18283#bib.bib4)\]; sparse attention modifies which query\-key pairs are compared\[[9](https://arxiv.org/html/2606.18283#bib.bib36),[3](https://arxiv.org/html/2606.18283#bib.bib14),[54](https://arxiv.org/html/2606.18283#bib.bib15)\]; and GMA replaces direct pairwise comparison with routing through latent Gaussian mixture responsibilities, as introduced in this work\.

#### Similarity and distance\.

A similarity function assigns large values to objects that are considered compatible or related\. A distance function assigns small values to objects that are close\. These two notions are related, but not identical\. A distance can be converted into a similarity through a monotone decreasing transformation, for example

κ​\(x,y\)=exp⁡\(−β​d​\(x,y\)\),κ​\(x,y\)=11\+d​\(x,y\),s​\(x,y\)=−d​\(x,y\),\\kappa\(x,y\)=\\exp\(\-\\beta d\(x,y\)\),\\qquad\\kappa\(x,y\)=\\frac\{1\}\{1\+d\(x,y\)\},\\qquad s\(x,y\)=\-d\(x,y\),whereβ\>0\\beta\>0\. However, a similarity need not be the reciprocal of a distance, and not every similarity corresponds to a valid metric\.

A metric distanced​\(x,y\)d\(x,y\)satisfies non\-negativity, identity of indiscernibles, symmetry, and the triangle inequality:

d​\(x,y\)≥0,d​\(x,y\)=0⇔x=y,d​\(x,y\)=d​\(y,x\),d​\(x,z\)≤d​\(x,y\)\+d​\(y,z\)\.d\(x,y\)\\geq 0,\\qquad d\(x,y\)=0\\Leftrightarrow x=y,\\qquad d\(x,y\)=d\(y,x\),\\qquad d\(x,z\)\\leq d\(x,y\)\+d\(y,z\)\.Common pointwise distances include theℓ1\\ell\_\{1\}distance, Euclidean distance, squared Euclidean distance, Mahalanobis distance, cosine distance, Hamming distance, and edit distance\. In representation learning, these distances can be used directly as negative attention scores or indirectly through kernels such as the radial basis function \(RBF\) kernel

κ​\(x,y\)=exp⁡\(−‖x−y‖222​σ2\)\.\\kappa\(x,y\)=\\exp\\left\(\-\\frac\{\\\|x\-y\\\|\_\{2\}^\{2\}\}\{2\\sigma^\{2\}\}\\right\)\.The RBF kernel is a similarity derived from Euclidean distance: nearby points receive high similarity, while distant points receive exponentially small similarity\.

#### Dot products, cosine similarity, and bilinear scores\.

Standard scaled dot\-product attention uses

s​\(qi,kj\)=qi⊤​kjdk,s\(q\_\{i\},k\_\{j\}\)=\\frac\{q\_\{i\}^\{\\top\}k\_\{j\}\}\{\\sqrt\{d\_\{k\}\}\},whereqiq\_\{i\}is a query vector andkjk\_\{j\}is a key vector\[[49](https://arxiv.org/html/2606.18283#bib.bib11)\]\. This is a bilinear compatibility score: it is linear inqiq\_\{i\}whenkjk\_\{j\}is fixed, and linear inkjk\_\{j\}whenqiq\_\{i\}is fixed\. Cosine similarity normalizes the dot product by the vector norms,

cos⁡\(qi,kj\)=qi⊤​kj‖qi‖2​‖kj‖2,\\operatorname\{cos\}\(q\_\{i\},k\_\{j\}\)=\\frac\{q\_\{i\}^\{\\top\}k\_\{j\}\}\{\\\|q\_\{i\}\\\|\_\{2\}\\\|k\_\{j\}\\\|\_\{2\}\},so that compatibility depends on direction rather than magnitude\. More general bilinear attention scores take the form

s​\(qi,kj\)=qi⊤​W​kj,s\(q\_\{i\},k\_\{j\}\)=q\_\{i\}^\{\\top\}Wk\_\{j\},whereWWis a learned matrix\. Dot\-product attention is the special caseW=IW=I, up to projection matrices and scaling\.

From this viewpoint, standard attention can be interpreted as a learned similarity\-based retrieval mechanism\. Given a query, the model asks which keys are most compatible with it, normalizes the resulting scores, and uses the normalized weights to aggregate values:

Ai​j=exp⁡\(s​\(qi,kj\)\)∑ℓ=1Nexp⁡\(s​\(qi,kℓ\)\),Oi=∑j=1NAi​j​Vj\.A\_\{ij\}=\\frac\{\\exp\(s\(q\_\{i\},k\_\{j\}\)\)\}\{\\sum\_\{\\ell=1\}^\{N\}\\exp\(s\(q\_\{i\},k\_\{\\ell\}\)\)\},\\qquad O\_\{i\}=\\sum\_\{j=1\}^\{N\}A\_\{ij\}V\_\{j\}\.Thus, attention is not merely a similarity function; it is a normalized similarity\-weighted aggregation mechanism\.

#### Correlation and kernel similarity\.

Correlation is another form of relatedness\. The Pearson correlation coefficient between scalar random variablesXXandYYis

ρ​\(X,Y\)=Cov⁡\(X,Y\)Var⁡\(X\)​Var⁡\(Y\)\.\\rho\(X,Y\)=\\frac\{\\operatorname\{Cov\}\(X,Y\)\}\{\\sqrt\{\\operatorname\{Var\}\(X\)\\operatorname\{Var\}\(Y\)\}\}\.It measures normalizedlineardependence\. Cosine similarity can be viewed as a closely related geometric quantity: it is the normalized dot product between two vectors\. If vectors are centered, cosine similarity and Pearson correlation are closely aligned222222See e\.g\. Eq\.\(3\.7\) in\[[8](https://arxiv.org/html/2606.18283#bib.bib26)\]\.\.

Kernel functions generalize inner\-product similarity by implicitly mapping inputs into a feature space:

κ​\(x,y\)=⟨ϕ​\(x\),ϕ​\(y\)⟩ℋ\.\\kappa\(x,y\)=\\langle\\phi\(x\),\\phi\(y\)\\rangle\_\{\\mathcal\{H\}\}\.Examples include the linear kernel, polynomial kernel, RBF kernel, Laplacian kernel, Matérn kernels and string kernels\. Kernelized attention mechanisms build on this idea by replacing the softmax attention kernel with feature\-map approximations or positive kernel decompositions\. Linear Transformer\[[28](https://arxiv.org/html/2606.18283#bib.bib7)\]and Performer\[[10](https://arxiv.org/html/2606.18283#bib.bib4)\]are examples: they approximate or reformulate attention through feature maps so that the attention computation can be reassociated and evaluated in linear time\[[28](https://arxiv.org/html/2606.18283#bib.bib7),[10](https://arxiv.org/html/2606.18283#bib.bib4)\]\. In this sense, designing efficient attention can be viewed as designing a similarity function whose algebraic structure permits efficient aggregation\.

#### Divergences between distributions\.

When the objects being compared are probability distributions rather than vectors, one often uses divergences or distributional distances\. The Kullback\-Leibler divergence is

DKL​\(p∥q\)=∫p​\(x\)​log⁡p​\(x\)q​\(x\)​d​x\.D\_\{\\mathrm\{KL\}\}\(p\\\|q\)=\\int p\(x\)\\log\\frac\{p\(x\)\}\{q\(x\)\}\\,dx\.It is not a metric because it is generally asymmetric and does not satisfy the triangle inequality\. The Jensen\-Shannon divergence symmetrizes and smooths KL:

DJS​\(p∥q\)=12​DKL​\(p∥m\)\+12​DKL​\(q∥m\),m=12​\(p\+q\)\.D\_\{\\mathrm\{JS\}\}\(p\\\|q\)=\\frac\{1\}\{2\}D\_\{\\mathrm\{KL\}\}\(p\\\|m\)\+\\frac\{1\}\{2\}D\_\{\\mathrm\{KL\}\}\(q\\\|m\),\\qquad m=\\frac\{1\}\{2\}\(p\+q\)\.Other distributional comparisons include total variation distance, Hellinger distance, Wasserstein distance, Fisher\-Rao distance, maximum mean discrepancy \(MMD\), and Bregman divergences\. These quantities suggest possible distribution\-aware attention mechanisms in which queries and keys are not treated as points, but as distributions with uncertainty\. In such a setting, compatibility could be based on negative divergence, transport cost, kernel mean embedding distance, or overlap between probabilistic representations\.

#### Attention as normalized compatibility\.

The common structure behind many attention mechanisms is

Ai​j=Normalizej⁡\[Compat⁡\(qi,kj\)\],Oi=∑jAi​j​Vj\.A\_\{ij\}=\\operatorname\{Normalize\}\_\{j\}\\left\[\\operatorname\{Compat\}\(q\_\{i\},k\_\{j\}\)\\right\],\\qquad O\_\{i\}=\\sum\_\{j\}A\_\{ij\}V\_\{j\}\.The compatibility function may be a dot product, a bilinear score, a neural network, a negative distance, a kernel value, a negative divergence, or a probabilistic responsibility\. The normalization may be a softmax, sparsemax\[[35](https://arxiv.org/html/2606.18283#bib.bib37)\], entmax\[[40](https://arxiv.org/html/2606.18283#bib.bib38)\], top\-kknormalization, kernel normalization, or another stochastic normalizer\. This perspective separates two design choices: first, how compatibility is measured; second, how compatibility is converted into aggregation weights\.

This also clarifies why the relationship between similarity and attention is subtle\. A similarity score is usually symmetric if the same representation space and symmetric function are used\. Attention, however, is generally asymmetric: queries and keys are produced by different projections, and the normalization is row\-wise\. Even ifs​\(qi,kj\)=s​\(kj,qi\)s\(q\_\{i\},k\_\{j\}\)=s\(k\_\{j\},q\_\{i\}\)were symmetric, the normalized attention matrix need not be symmetric because

Ai​j=exp⁡\(s​\(qi,kj\)\)∑ℓexp⁡\(s​\(qi,kℓ\)\)A\_\{ij\}=\\frac\{\\exp\(s\(q\_\{i\},k\_\{j\}\)\)\}\{\\sum\_\{\\ell\}\\exp\(s\(q\_\{i\},k\_\{\\ell\}\)\)\}normalizes separately for each query positionii\.

#### GMA as responsibility\-based compatibility\.

GMA follows a different route from direct pairwise similarity\. Instead of first forming allN×NN\\times Nquery\-key scores, it maps queries and keys into posterior responsibility vectors overKKlearned Gaussian mixture components\. For a query vectorqiq\_\{i\}and key vectorkjk\_\{j\}, the induced unnormalized affinity is \(Eq\.[11](https://arxiv.org/html/2606.18283#S3.E11)\)

A~i​jGMA=∑k=1Kγi,kQ​γj,kK\.\\widetilde\{A\}^\{\\mathrm\{GMA\}\}\_\{ij\}=\\sum\_\{k=1\}^\{K\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}\.Thus,GMA compares tokens through their distributions over latent routing components, rather than through a dot product in the projected routing representation space\. Equivalently, the affinity is an inner product in the probability simplex of component responsibilities\. This makes GMA a responsibility\-space similarity mechanism\.

The computational philosophy is also different\. Standard attention first forms pairwise query\-key compatibilities and then aggregates values\. GMA first uses key responsibilities to write values into a compact latent memory \(Eq\.[5](https://arxiv.org/html/2606.18283#S3.E5)\),

V~=\(ΓK\)⊤​VX,\\tilde\{V\}=\(\\Gamma^\{K\}\)^\{\\top\}V\_\{X\},and then uses query responsibilities to read from this memory \(Eq\.[7](https://arxiv.org/html/2606.18283#S3.E7)\):

O=ΓQ​V~ΓQ​Z\+ϵ\.O=\\frac\{\\Gamma^\{Q\}\\tilde\{V\}\}\{\\Gamma^\{Q\}Z\+\\epsilon\}\.In this sense, GMA replaces direct pairwise similarity with a latent mixture routing mechanism\. The learned Gaussian components define a multi\-modal routing space, the key responsibilities construct component\-indexed memory slots, and the query responsibilities determine how each token reads from those slots\.

#### Implications for attention design\.

This perspective suggests that new attention mechanisms can be designed by choosing different notions of compatibility or routing geometry\. Possible choices include:

Dot\-product attentionq⊤​kinner\-product similarity,Cosine attentionq⊤​k‖q‖​‖k‖angular similarity,Mahalanobis attention−\(q−k\)⊤​Σ−1​\(q−k\)anisotropic distance,RBF\-kernel attentionexp⁡\(−‖q−k‖2/2​σ2\)distance\-induced similarity,Polynomial\-kernel attention\(q⊤​k\+c\)phigher\-order feature similarity,Divergence\-based attention−D​\(pq∥pk\)distributional compatibility,Transport\-based attention−Wp​\(pq,pk\)geometry\-aware distributional distance,GMA∑kγi,kQ​γj,kKresponsibility\-space affinity\.\\begin\{array\}\[\]\{lll\}\\text\{Dot\-product attention\}&q^\{\\top\}k&\\text\{inner\-product similarity\},\\\\\[5\.69054pt\] \\text\{Cosine attention\}&\\frac\{q^\{\\top\}k\}\{\\\|q\\\|\\\|k\\\|\}&\\text\{angular similarity\},\\\\\[8\.53581pt\] \\text\{Mahalanobis attention\}&\-\(q\-k\)^\{\\top\}\\Sigma^\{\-1\}\(q\-k\)&\\text\{anisotropic distance\},\\\\\[5\.69054pt\] \\text\{RBF\-kernel attention\}&\\exp\(\-\\\|q\-k\\\|^\{2\}/2\\sigma^\{2\}\)&\\text\{distance\-induced similarity\},\\\\\[5\.69054pt\] \\text\{Polynomial\-kernel attention\}&\(q^\{\\top\}k\+c\)^\{p\}&\\text\{higher\-order feature similarity\},\\\\\[5\.69054pt\] \\text\{Divergence\-based attention\}&\-D\(p\_\{q\}\\\|p\_\{k\}\)&\\text\{distributional compatibility\},\\\\\[5\.69054pt\] \\text\{Transport\-based attention\}&\-W\_\{p\}\(p\_\{q\},p\_\{k\}\)&\\text\{geometry\-aware distributional distance\},\\\\\[5\.69054pt\] \\text\{GMA\}&\\sum\_\{k\}\\gamma^\{Q\}\_\{i,k\}\\gamma^\{K\}\_\{j,k\}&\\text\{responsibility\-space affinity\}\.\\end\{array\}The main challenge is not only to define a meaningful compatibility score, but also to make it computationally tractable\. Pairwise similarities usually require anN×NN\\times Nscore matrix\. Efficient attention mechanisms therefore need additional structure: sparsity, low rank, kernel feature maps, recurrence, state\-space dynamics, or latent routing\. GMA belongs to the latent\-routing family: it replaces direct token\-to\-token comparison with interaction through a fixed number of learned probabilistic components\.

Overall, attention can be understood as normalized compatibility\-based aggregation\. Dot\-product attention is one important instance of this principle, but not the only one\. Distances, similarities, correlations, kernels, divergences, and probabilistic responsibilities all provide possible foundations for attention design\. GMA contributes to this broader view by showing how Gaussian\-mixture responsibilities can define a linear\-time, interpretable attention\-style routing mechanism\.