SMat-Attention: Structured Long-Context Sequence Modeling

# SMat-Attention: Structured Long-ContextSequence Modeling
Source: [https://arxiv.org/html/2609.36062](https://arxiv.org/html/2609.36062)
## SMat\-Attention: Structured Long\-Context Sequence Modeling

Abdullah Ateyeh\*Affiliation:University of California, BerkeleyEmail:[abdullah\_ateyeh@berkeley\.edu](mailto:)Archer Wang\*Affiliation:Massachusetts Institute of TechnologyEmail:[archerw@mit\.edu](mailto:)Marin SoljačićAffiliation:Massachusetts Institute of TechnologyEmail:[soljacic@mit\.edu](mailto:)

###### Abstract

Long\-context sequence models face a fundamental tradeoff: softmax attention uses flexible token\-level interactions at quadratic cost, whereas linear attention obtains linear\-time training and constant\-time decoding by compressing history into a fixed\-size state\. In this work, we ask whether we can connect these regimes through a tunable notion of structure\. To this end, we introduce Structured Matrix Attention \(SMat\-Attention\) via a family of causal masks with structured long\-range routing whose row supports have VC\-dimensiondd\. In our construction,d=1d=1recovers the standard causal mask, and increasingddpermits richer subset\-routing patterns\. We give chunkwise forward and backward algorithms to enable hardware\-efficiency\. For sequences of lengthTT, the hard\-routing construction takesO⁡\(T2−3/d\+T\)O\(T^\{2\-3/d\}\+T\)work, despite the mask being dense, for our prescribed family\. In fixed\-horizon streaming, decoding after the distant prefix takes constant time per token usingO⁡\(T1−1/d\)O\(T^\{1\-1/d\}\)cached states\. SMat\-Attention therefore makes VC\-dimension an explicit knob governing access\-pattern complexity, prefill cost, and decoding memory\. Empirically, subset\-routing and rule\-assisted multi\-key retrieval experiments illustrate the masks’ routing expressiveness\. Extensions to Mamba\-2 and Gated DeltaNet using learned routing with top\-kkquery reads retain subquadratic prefill, improve recall accuracy over the backbones in several settings, and achieve comparable small\-scale language\-modeling performance\.

## 1Introduction

Attention is a foundational building block of modern deep learning\([Bahdanau et al\., 2014](https://arxiv.org/html/2609.36062#bib.bib17)\)and serves as the core mechanism for modeling token interactions in Transformer architectures\([Vaswani et al\., 2017](https://arxiv.org/html/2609.36062#bib.bib18)\)\. Given key, query, and value matrices𝐐,𝐊\{\\mathbf\{Q\}\},\{\\mathbf\{K\}\}, and𝐕\{\\mathbf\{V\}\}, softmax attention computes

Attention⁡\(𝐐,𝐊,𝐕\)=softmax⁡\(𝐐𝐊⊤dk\)​𝐕\.\\mathrm\{Attention\}\(\\mathbf\{Q\},\\mathbf\{K\},\\mathbf\{V\}\)=\\mathrm\{softmax\}\\left\(\\frac\{\\mathbf\{Q\}\\mathbf\{K\}^\{\\top\}\}\{\\sqrt\{d\_\{k\}\}\}\\right\)\\mathbf\{V\}\.\(1\)This operation gives each query direct access to token\-level information, but its prefill computation grows quadratically with sequence length and its decoding cache grows linearly\([Vaswani et al\., 2017](https://arxiv.org/html/2609.36062#bib.bib18)\)\. However, fundamentally, long\-context sequence modeling requires retaining useful information and selecting which parts of the past should influence each query\. Hardware\-aware kernels improve execution efficiency\([Dao, 2023](https://arxiv.org/html/2609.36062#bib.bib24);[Shah et al\., 2024](https://arxiv.org/html/2609.36062#bib.bib10);[Liu et al\., 2024](https://arxiv.org/html/2609.36062#bib.bib13);[Kwon et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib12)\), while sparse methods such as Native Sparse Attention and MoBA reduce the interactions evaluated for each query\([Yuan et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib37);[Lu et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib9)\)\.

Recurrent alternatives such as linear attention address these costs by compressing the history into a fixed\-size recurrent state\([Katharopoulos et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib5)\)\. Modern variants improve how the model maintains this state\. Structured state\-space models \(SSMs\)\([Fu et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib27);[Gu et al\., 2022](https://arxiv.org/html/2609.36062#bib.bib32)\)compress history with time\-invariant recurrences; Mamba and Mamba\-2 make this recurrence input\-dependent through selective gating\([Gu and Dao, 2023](https://arxiv.org/html/2609.36062#bib.bib26);[Dao and Gu, 2024](https://arxiv.org/html/2609.36062#bib.bib28)\), while DeltaNet and Gated DeltaNet use structured transition matrices\([Schlag et al\., 2021](https://arxiv.org/html/2609.36062#bib.bib23);[Yang et al\., 2024a](https://arxiv.org/html/2609.36062#bib.bib6);[Yang et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib11)\)that update via the delta rule\([Schmidhuber, 1992](https://arxiv.org/html/2609.36062#bib.bib31);[Widrow and Hoff, 1960](https://arxiv.org/html/2609.36062#bib.bib47)\)\. These mechanisms improve retention and retrieval; however, their fixed\-size hidden state still constrains associative recall over long contexts\([Arora et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib1)\)\.

These advances highlight the role of structure in efficient sequence modeling\. For instance, linear attention exploits its causal prefix structure to reuse accumulated key–value summaries, yieldingO⁡\(T\)O\(T\)computation\. Gated variants, in turn, extend this approach through semiseparable structure\([Dao and Gu, 2024](https://arxiv.org/html/2609.36062#bib.bib28)\), while long\-convolution models exploit Toeplitz structure to compute their outputs inO⁡\(T​log⁡T\)O\(T\\log T\)time using FFT\([Poli et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib15);[Qin et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib51)\)\. Log\-Linear Attention\([Guo et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib4)\)further expands this design space by changing the organization of memory: it organizes recurrent summaries through a Fenwick\-tree hierarchy, achievingO⁡\(T​log⁡T\)O\(T\\log T\)computation andO⁡\(log⁡T\)O\(\\log T\)decoding memory\. Other recent approaches route tokens among multiple recurrent states\([Du et al\., 2026](https://arxiv.org/html/2609.36062#bib.bib7)\)or let the compressed memory grow with context length\([Behrouz et al\., 2026](https://arxiv.org/html/2609.36062#bib.bib46);[Goldstein et al\., 2026](https://arxiv.org/html/2609.36062#bib.bib52)\), or as a latent vector\([Anand et al\., 2026](https://arxiv.org/html/2609.36062#bib.bib22)\)\. These approaches motivate studying not only how much memory a model retains, but also which subsets of stored information each query can access\. Therefore, we investigate the following question:*can the combinatorial richness of long\-range access patterns be made an explicit architectural parameter, with corresponding guarantees on computation and memory?*

We study structured long\-range access as an intermediate regime and how its complexity governs computation and memory\. To make this complexity explicit, we use the VC dimension of a causal mask’s row supports\([Kearns and Vazirani, 1994](https://arxiv.org/html/2609.36062#bib.bib42);[Vapnik and Chervonenkis, 1971](https://arxiv.org/html/2609.36062#bib.bib19)\)\. Each row specifies the keys available to a query; under this lens, the VC dimension measures the largest number of keys on which the rows realize every possible subset\. Recent connections between VC dimension and matrix multiplication\([Anand et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib3)\)motivate constructing attention mechanisms that couple this combinatorial parameter to computational guarantees\.

We introduce Structured Matrix Attention \(SMat\-Attention\), a family of causal attention masks built from point\-hyperplane incidences over finite fields, whose row\-support VC dimensionddexplicitly controls long\-range access complexity\. Our framework recovers ordinary causal masking whend=1d=1, and exploits additional structure to support richer access patterns with subquadratic attention\. Building on structured masked\-attention formulations\([Choromanski et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib45)\), we establish the following results:

1. 1\.We exploit the resulting incidence structure to derive chunkwise forward and backward algorithms\. For sequences of lengthTT, we show that this mechanism takesO⁡\(T2−3/d\+T\)O\(T^\{2\-3/d\}\+T\)work, yieldingO⁡\(T\)O\(T\)\-attention ford≤3d\\leq 3andO⁡\(T2−3/d\)O\(T^\{2\-3/d\}\)\-attention ford\>3d\>3, despite the full causal mask havingΘ⁡\(T2\)\\Theta\(T^\{2\}\)nonzero entries\.
2. 2\.We show that after processing a fixed distant prefix, SMat\-Attention supportsTT\-independent per\-token decoding usingO⁡\(T1−1/d\)O\(T^\{1\-1/d\}\)cached states, revealing an explicit tradeoff between long\-range access complexity and memory\.
3. 3\.We extend the construction to Mamba\-2 and Gated DeltaNet with learned content hashing and a learned four\-read selector\. The extension inherits the tabulation, cache and VC bounds of the above; its selector addsΘ⁡\(T2−1/d\)\\Theta\(T^\{2\-1/d\}\)prefill work andO⁡\(T1−1/d\)O\(T^\{1\-1/d\}\)per decoded token\. Controlled tasks show benefits consistent with increased routing expressiveness; learned SMat extensions improve mean recall accuracy over native backbones in several tested settings and remain competitive on small\-scale PG\-19 language modeling\.

## 2Preliminaries

LetTTbe the length of the input sequence\. Following[Vaswani et al\. \(2017\)](https://arxiv.org/html/2609.36062#bib.bib18), attention linearly projects the input tokens into𝐐∈ℝT×dQ​K\\mathbf\{Q\}\\in\\mathbb\{R\}^\{T\\times d\_\{QK\}\},𝐊∈ℝT×dQ​K\\mathbf\{K\}\\in\\mathbb\{R\}^\{T\\times d\_\{QK\}\}and𝐕∈ℝT×dv\\mathbf\{V\}\\in\\mathbb\{R\}^\{T\\times d\_\{v\}\}, the queries, keys and values\. Following[Choromanski et al\. \(2023\)](https://arxiv.org/html/2609.36062#bib.bib45), the*general masked kernel attention*is

𝖠𝗍𝗍K​\(𝐐,𝐊,𝐕,𝐌\)=𝐃−1​𝐀𝐕,𝐀=𝐌⊙𝒦⁡\(𝐐,𝐊\),𝐃=𝖽𝗂𝖺𝗀⁡\(𝐀𝟏T\),\\mathsf\{Att\}\_\{K\}\(\{\\mathbf\{Q\}\},\{\\mathbf\{K\}\},\{\\mathbf\{V\}\},\{\\mathbf\{M\}\}\)=\{\\mathbf\{D\}\}^\{\-1\}\{\\mathbf\{A\}\}\{\\mathbf\{V\}\},\\qquad\{\\mathbf\{A\}\}=\{\\mathbf\{M\}\}\\odot\\mathcal\{K\}\(\{\\mathbf\{Q\}\},\{\\mathbf\{K\}\}\),\\qquad\{\\mathbf\{D\}\}=\\mathsf\{diag\}\(\{\\mathbf\{A\}\}\\mathbf\{1\}\_\{T\}\),where⊙\\odotis the entrywise product,𝒦:ℝdQ​K×ℝdQ​K→ℝ\\mathcal\{K\}:\\mathbb\{R\}^\{d\_\{QK\}\}\\times\\mathbb\{R\}^\{d\_\{QK\}\}\\to\\mathbb\{R\}is a kernel,𝒦​\(𝐐,𝐊\)i​j=K⁡\(𝐪i,𝐤j\)\\mathcal\{K\}\(\{\\mathbf\{Q\}\},\{\\mathbf\{K\}\}\)\_\{ij\}=K\(\\mathbf\{q\}\_\{i\},\\mathbf\{k\}\_\{j\}\)for𝐪i\\mathbf\{q\}\_\{i\}theiith row of𝐐\{\\mathbf\{Q\}\}and𝐤j\\mathbf\{k\}\_\{j\}thejjth row of𝐊\{\\mathbf\{K\}\}, and𝟏T\\mathbf\{1\}\_\{T\}is the all\-ones vector\. Softmax attention is the special caseK⁡\(x,y\)=exp⁡\(x⊤​ydQ​K\)K\(x,y\)=\\exp\(\\frac\{x^\{\\top\}y\}\{\\sqrt\{d\_\{QK\}\}\}\)where𝐌=exp⁡\(𝐍\)\{\\mathbf\{M\}\}=\\exp\(\{\\mathbf\{N\}\}\)entrywise and𝐍\{\\mathbf\{N\}\}is the logits mask\.

Finite\-feature attention\.Assuming the kernel has a nonnegative feature factorization of dimensionrr,𝒦⁡\(q,k\)=ϕ𝐐​\(q\)⊤​ϕ𝐊​\(k\)\\mathcal\{K\}\(q,k\)=\\phi\_\{\\mathbf\{Q\}\}\(q\)^\{\\top\}\\phi\_\{\\mathbf\{K\}\}\(k\), whereϕ𝐐​\(q\),ϕ𝐊​\(k\)∈ℝ≥0r\\phi\_\{\\mathbf\{Q\}\}\(q\),\\phi\_\{\\mathbf\{K\}\}\(k\)\\in\\mathbb\{R\}^\{r\}\_\{\\geq 0\}\. Letϕi:=ϕ𝐐​\(qi\)\\phi\_\{i\}:=\\phi\_\{\\mathbf\{Q\}\}\(q\_\{i\}\)andψj:=hj​ϕ𝐊​\(kj\)\\psi\_\{j\}:=h\_\{j\}\\phi\_\{\\mathbf\{K\}\}\(k\_\{j\}\), wherehj≥0h\_\{j\}\\geq 0is an optional key gate \(sethj≡1h\_\{j\}\\equiv 1for ungated attention\)\. The normalizer is carried along with the values by appending a constant coordinate: withp:=dv\+1p:=d\_\{v\}\+1,

v¯j:=\[vj1\]∈ℝp,𝐙j:=ψj​v¯j⊤∈ℝr×p\.\\bar\{v\}\_\{j\}:=\\begin\{bmatrix\}v\_\{j\}\\\\ 1\\end\{bmatrix\}\\in\\mathbb\{R\}^\{p\},\\qquad\{\\mathbf\{Z\}\}\_\{j\}:=\\psi\_\{j\}\\bar\{v\}\_\{j\}^\{\\top\}\\in\\mathbb\{R\}^\{r\\times p\}\.\(2\)For any nonnegative mask𝐌\{\\mathbf\{M\}\}the*augmented output*and the attention output are

y¯i=ϕi⊤∑j=1T𝐌i​j𝐙j∈ℝp,oi=y¯i,1:dvy¯i,p\.\\bar\{y\}\_\{i\}=\\phi\_\{i\}^\{\\top\}\\sum\_\{j=1\}^\{T\}\{\\mathbf\{M\}\}\_\{ij\}\{\\mathbf\{Z\}\}\_\{j\}\\in\\mathbb\{R\}^\{p\},\\qquad o\_\{i\}=\\frac\{\\bar\{y\}\_\{i,1:d\_\{v\}\}\}\{\\bar\{y\}\_\{i,p\}\}\.\(3\)Through the augmentation,y¯i\\bar\{y\}\_\{i\}is*linear*in𝐌\{\\mathbf\{M\}\}: the last coordinate accumulates the denominator along with the numerator, and the single nonlinearity is the final division\. So if𝐌=∑ℓ𝐌\[ℓ\]\\smash\{\{\\mathbf\{M\}\}=\\sum\_\{\\ell\}\{\\mathbf\{M\}\}^\{\[\\ell\]\}\}, the contributionsy¯i\[ℓ\]\\smash\{\\bar\{y\}\_\{i\}^\{\[\\ell\]\}\}can be computed independently, in different orders and with different computational kernels, as long as they are summed before the division\. Section[3\.1](https://arxiv.org/html/2609.36062#S3.SS1)does this with two summands\. Importantly, we require the feature map to be finite and nonnegative, so softmax attention is covered only through a kernel approximation such as[Choromanski et al\. \(2022\)](https://arxiv.org/html/2609.36062#bib.bib36)\.

VC dimension of a mask\.A binary mask𝐌∈\{0,1\}T×T\{\\mathbf\{M\}\}\\in\\\{0,1\\\}^\{T\\times T\}defines a set system on the key indices: rowiiis the setSi​\(𝐌\)=\{j∈\[T\]:𝐌i​j=1\}S\_\{i\}\(\{\\mathbf\{M\}\}\)=\\\{j\\in\[T\]:\{\\mathbf\{M\}\}\_\{ij\}=1\\\}of keys visible to queryii\. If we let𝒮⁡\(𝐌\)=\{S1​\(𝐌\),…,ST​\(𝐌\)\}\\mathcal\{S\}\(\{\\mathbf\{M\}\}\)=\\\{S\_\{1\}\(\{\\mathbf\{M\}\}\),\\dots,S\_\{T\}\(\{\\mathbf\{M\}\}\)\\\}, thenVC⁡\(𝐌\)\\mathrm\{VC\}\(\{\\mathbf\{M\}\}\)is the VC dimension of𝒮⁡\(𝐌\)\\mathcal\{S\}\(\{\\mathbf\{M\}\}\), i\.e\. the largestkkfor which some set ofkkkeys is shattered by the rows\. The causal mask𝐋T\\mathbf\{L\}\_\{T\}hasVC=1\\mathrm\{VC\}=1: its rows are the prefixes\{1,…,i\}\\\{1,\\dots,i\\\}, which are totally ordered, so no two keys can be shattered, i\.e\. no query sees a later key without also seeing every earlier one\. On the other hand, an unconstrainedTT\-row mask can have VC\-dimension as large as⌊log⁡T⌋\\lfloor\\log T\\rfloor\. The parameterddinterpolates between these, and Theorems[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2)and[3\.3](https://arxiv.org/html/2609.36062#S3.Thmthm3)price the interpolation\.

We construct a family of causal masks𝐌\(1\),𝐌\(2\),…\\mathbf\{M\}^\{\(1\)\},\\mathbf\{M\}^\{\(2\)\},\\dotsindexed by their VC dimension\. We show that attention under𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}can be computed inO⁡\(T2−3/d\+T\)O\(T^\{2\-3/d\}\+T\)work\. Each mask is built from incidences between points and hyperplanes over a finite field, which gives it a computationally favorable structure\. Ford≤3d\\leq 3the forward pass is linear in the sequence length, and decoding runs from a cache ofO⁡\(T1−1/d\)O\(T^\{1\-1/d\}\)states\.

𝐋n\\mathbf\{L\}\_\{n\}𝟎\\mathbf\{0\}𝐆=𝐑​𝐂​𝐄⊤\\mathbf\{G\}=\\mathbf\{R\}\\,\\mathbf\{C\}\\,\\mathbf\{E\}^\{\\top\}𝐋m\\mathbf\{L\}\_\{m\}distant\[n\]\[n\]recent\(n,T\]\(n,T\]distant queriesrecent queriesFigure 1:Block layout of𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}; shaded regions are nonzero\. Both diagonal blocks are ordinary causal masks, and all long\-range structure lives in𝐆\\mathbf\{G\}\.Block form\.Letnnbe the number of*distant*tokens, andm=T−nm=T\-nbe the remaining*recent*tokens\. Let𝐋s\\mathbf\{L\}\_\{s\}be the inclusive lower\-triangular all\-ones matrix of orderss\. Every mask in the family has the form

𝐌\(d\)=\(𝐋n𝟎𝐆𝐋m\),𝐆∈\{0,1\}m×n,\\mathbf\{M\}^\{\(d\)\}=\\begin\{pmatrix\}\\mathbf\{L\}\_\{n\}&\\mathbf\{0\}\\\\\[2\.0pt\] \\mathbf\{G\}&\\mathbf\{L\}\_\{m\}\\end\{pmatrix\},\\qquad\\mathbf\{G\}\\in\\\{0,1\\\}^\{m\\times n\},\(4\)so distant and recent tokens are each causally masked, and𝐆\\mathbf\{G\}encodes all of the long\-range interaction\.

The construction of𝐆\\mathbf\{G\}\.Letd≥2d\\geq 2andqqbe a prime\. The geometry is the ambient space of\(d−1\)\(d\-1\)\-dimensional vectors over the finite field𝔽q\\mathbb\{F\}\_\{q\}, given by𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}, with its affine hyperplanesHa,b=\{x:a⊤​x=b\}H\_\{a,b\}=\\\{x:a^\{\\top\}x=b\\\}, one per normalized directiona≠0a\\neq 0and offsetbb, of which there are

B=q⁡\(qd−1−1\)q−1=Θ⁡\(qd−1\)\.B=\\frac\{q\(q^\{d\-1\}\-1\)\}\{q\-1\}=\\Theta\\big\(q^\{d\-1\}\\big\)\.\(5\)Each distant key is assigned a*profile*, a point of the ambient space, and each recent query a*type*, one of the hyperplanes,

prof:\[n\]→𝔽qd−1,type:\[m\]→\{0,…,B−1\},\\operatorname\{prof\}:\[n\]\\to\\mathbb\{F\}\_\{q\}^\{d\-1\},\\qquad\\operatorname\{type\}:\[m\]\\to\\\{0,\\dots,B\-1\\\},and a recent query attends to a distant key when that key’s profile lies on the query’s hyperplane:

𝐆i​j=𝟙\{prof\(j\)∈Htype⁡\(i\)\},\\mathbf\{G\}\_\{ij\}=\\mathbbm\{1\}\\\{\\operatorname\{prof\}\(j\)\\in H\_\{\\operatorname\{type\}\(i\)\}\\\},\(6\)i\.e\. recent query of typehhattends to its own prefix among the recent tokens and every distant token whose profile lies onHhH\_\{h\}\.𝐆\\mathbf\{G\}is the only component of the mask which isdd\-dependent and its structure dictates the VC\-dimension \(Theorem[3\.1](https://arxiv.org/html/2609.36062#S3.Thmthm1)\)\. In order to establish a VC lower bound for Theorem[3\.1](https://arxiv.org/html/2609.36062#S3.Thmthm1)\(iii\), we further characterizeprof\\operatorname\{prof\}andtype\\operatorname\{type\}\. Lete1,…,ed−1e\_\{1\},\\dots,e\_\{d\-1\}be the standard basis of𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}and, forR⊆\[d−1\]R\\subseteq\[d\-1\], letHRH\_\{R\}be the witness hyperplane\. We impose the condition that there are distinct distant positionsj1,…,jd−1j\_\{1\},\\dots,j\_\{d\-1\}withprof⁡\(jℓ\)=eℓ\\mathrm\{prof\}\(j\_\{\\ell\}\)=e\_\{\\ell\}, and a recent indexτ\\tausuch that, for everyR⊆\[d−1\]R\\subseteq\[d\-1\],HRH\_\{R\}occurs at indicesiR−<τ≤iR\+i\_\{R\}^\{\-\}<\\tau\\leq i\_\{R\}^\{\+\}\. We give two concrete examples:

1. 1\.The*positional*assignment takesprof⁡\(j\)\\operatorname\{prof\}\(j\)to be the base\-qqdigits of\(j−1\)modqd−1\(j\-1\)\\bmod q^\{d\-1\}andtype⁡\(i\)=imodB\\operatorname\{type\}\(i\)=i\\bmod B\.
2. 2\.The*content\-based*assignment fixes a hashxxwhich maps each token to a point of𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}\. This mapping can be fixed or learned\. Letutu\_\{t\}be the hashed vector at positiontt\. A distant key hasprof⁡\(j\)=uj\\operatorname\{prof\}\(j\)=u\_\{j\}and for a hashed directionaa, a recent query hastype⁡\(i\)=Ha,a⊤​un\+i\\operatorname\{type\}\(i\)=H\_\{a,a^\{\\top\}u\_\{n\+i\}\}\(the hyperplane through its own cell\)\. A query sees every distant key whose token repeats its own\.

Scaling with sequence length\.We define the family by prescribing the field size as a function of sequence length and VC\-dimension\. For each fixedd≥2d\\geq 2, we choose a primeq=Θ⁡\(n1/d\)q=\\Theta\(n^\{1/d\}\), wheren,m=Θ⁡\(T\)n,m=\\Theta\(T\)\. WritingP=qd−1P=q^\{d\-1\}for the number of profile cells, this givesP=Θ⁡\(T1−1/d\)P=\\Theta\(T^\{1\-1/d\}\)andB=Θ⁡\(T1−1/d\)B=\\Theta\(T^\{1\-1/d\}\)\. For each fixeddd, a sufficiently largeTTensuresP≤nP\\leq nand2​B≤m2B\\leq m, as required by our positional construction\. Moreover, ford=1d=1, when𝔽q0\\mathbb\{F\}\_\{q\}^\{0\}is a single point, we useB=1B=1andH0=𝔽q0H\_\{0\}=\\mathbb\{F\}\_\{q\}^\{0\}\. Then𝐆\\mathbf\{G\}is all ones and the block form becomes𝐌\(1\)=𝐋T\\mathbf\{M\}^\{\(1\)\}=\\mathbf\{L\}\_\{T\}, ordinary causal masking\.

###### Theorem 3\.1\(Properties of hard\-routing SMat masks𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}\)\.

LetT≥2T\\geq 2and1≤d<⌊log2⁡T⌋1\\leq d<\\lfloor\\log\_\{2\}T\\rfloor\. Then*\(i\)*𝐌\(d\)∈\{0,1\}T×T\\mathbf\{M\}^\{\(d\)\}\\in\\\{0,1\\\}^\{T\\times T\}is causal with𝐌t​t\(d\)=1\\mathbf\{M\}^\{\(d\)\}\_\{tt\}=1for alltt;*\(ii\)*𝐌\(1\)=𝐋T\\mathbf\{M\}^\{\(1\)\}=\\mathbf\{L\}\_\{T\},*\(iii\)*VC⁡\(𝐌\(d\)\)=d\\mathrm\{VC\}\(\\mathbf\{M\}^\{\(d\)\}\)=d; and*\(iv\)*the number of nonzero entries of𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}, given bynnz⁡\(𝐌\(d\)\)\\mathrm\{nnz\}\(\\mathbf\{M\}^\{\(d\)\}\), satisfiesnnz⁡\(𝐌\(d\)\)=Θ⁡\(T2\)\\mathrm\{nnz\}\(\\mathbf\{M\}^\{\(d\)\}\)=\\Theta\(T^\{2\}\)\.

### 3\.1Chunkwise SMat\-attention

Splitting the mask\.Split the block form into

𝐌\(d\)=\(𝐋n𝟎𝟎𝐋m\)⏟two independent causal masks\+\(𝟎𝟎𝐆𝟎\)⏟long range\.\\mathbf\{M\}^\{\(d\)\}=\\underbrace\{\\begin\{pmatrix\}\\mathbf\{L\}\_\{n\}&\\mathbf\{0\}\\\\ \\mathbf\{0\}&\\mathbf\{L\}\_\{m\}\\end\{pmatrix\}\}\_\{\\text\{two independent causal masks\}\}\+\\underbrace\{\\begin\{pmatrix\}\\mathbf\{0\}&\\mathbf\{0\}\\\\ \\mathbf\{G\}&\\mathbf\{0\}\\end\{pmatrix\}\}\_\{\\text\{long range\}\}\.\(7\)Sincey¯i\\bar\{y\}\_\{i\}is linear in the mask, the two attention branches can be computed separately and added before the final division\. By construction𝐆i​j\\mathbf\{G\}\_\{ij\}depends oniionly throughtype⁡\(i\)\\operatorname\{type\}\(i\)and onjjonly throughprof⁡\(j\)\\operatorname\{prof\}\(j\)\. It therefore factors through the point–hyperplane incidence matrix

𝐂∈\{0,1\}B×qd−1,𝐂h​x=𝟙\{x∈Hh\},𝐆i​j=𝐂type⁡\(i\),prof⁡\(j\)\.\\mathbf\{C\}\\in\\\{0,1\\\}^\{B\\times q^\{d\-1\}\},\\qquad\\mathbf\{C\}\_\{hx\}=\\mathbbm\{1\}\\\{x\\in H\_\{h\}\\\},\\qquad\\mathbf\{G\}\_\{ij\}=\\mathbf\{C\}\_\{\\operatorname\{type\}\(i\),\\operatorname\{prof\}\(j\)\}\.\(8\)Writing𝐄∈\{0,1\}n×qd−1\{\\mathbf\{E\}\}\\in\\\{0,1\\\}^\{n\\times q^\{d\-1\}\}and𝐑∈\{0,1\}m×B\{\\mathbf\{R\}\}\\in\\\{0,1\\\}^\{m\\times B\}for the one\-hot matrices ofprof\\operatorname\{prof\}andtype\\operatorname\{type\},𝐆=𝐑​𝐂​𝐄⊤\\mathbf\{G\}=\{\\mathbf\{R\}\}\\,\\mathbf\{C\}\\,\{\\mathbf\{E\}\}^\{\\top\}\. Distant keys sharing a profile can be pooled once and reused by every query that sees that profile, and recent queries sharing a type can be answered from one aggregated state\.

The long\-range branch\.We apply the factorization𝐆=𝐑​𝐂​𝐄⊤\\mathbf\{G\}=\{\\mathbf\{R\}\}\\,\\mathbf\{C\}\\,\{\\mathbf\{E\}\}^\{\\top\}one factor at a time\.𝐄⊤\{\\mathbf\{E\}\}^\{\\top\}pools the distant states into a*profile table*\{Fx\}\\\{F\_\{x\}\\\}, one state per profile;𝐂\\mathbf\{C\}converts this into a*type table*\{Uh\}\\\{U\_\{h\}\\\}, one state per query type;𝐑\{\\mathbf\{R\}\}contracts each query against the entry for its own type:

Fx:=∑j≤n:prof⁡\(j\)=x𝐙j,Uh:=∑x𝐂h​xFx,F\_\{x\}:=\\\!\\\!\\\!\\sum\_\{j\\leq n:\\,\\operatorname\{prof\}\(j\)=x\}\\\!\\\!\\\!\{\\mathbf\{Z\}\}\_\{j\},\\qquad U\_\{h\}:=\\sum\_\{x\}\\mathbf\{C\}\_\{hx\}F\_\{x\},\(9\)wherey¯n\+ilr=ϕn\+i⊤​Utype⁡\(i\)\\bar\{y\}^\{\\mathrm\{lr\}\}\_\{n\+i\}=\\phi\_\{n\+i\}^\{\\top\}U\_\{\\operatorname\{type\}\(i\)\}fori∈\[m\]i\\in\[m\], andy¯ilr=0\\bar\{y\}^\{\\mathrm\{lr\}\}\_\{i\}=0fori≤ni\\leq n\. Queries of a common type are contracted together in one matrix product and their rows scattered back into chronological order\.

![Refer to caption](https://arxiv.org/html/2609.36062v1/figures/hyperplane_aggregation.png)Figure 2:Hyperplane aggregation and shared\-state reads in SMat\-Attention\.\(a\)A profile table is indexed by𝔽33\\mathbb\{F\}\_\{3\}^\{3\}\. Three coordinate hyperplanes each aggregate their incident profile states into a type state \(A\-C\)\. The central profile belongs to all hyperplanes, so its state contributes to all the sums\.\(b\)Each type state supplies every query of that type\. The horizontal lines carry fixed shared states throughout these reads\. The output scatters its results back to this order\. The figure shows augmented long\-range contributions, before causal addition and final normalization\.The causal branch\.Following standard chunkwise formulations of linear attention\([Hua et al\., 2022](https://arxiv.org/html/2609.36062#bib.bib25);[Beck et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib53);[Yang et al\., 2024a](https://arxiv.org/html/2609.36062#bib.bib6)\), partition\[n\]\[n\]and\(n,T\]\(n,T\]separately intoNcN\_\{c\}chronological chunksI1,…,INcI\_\{1\},\\dots,I\_\{N\_\{c\}\}of width at mostcc; Setn=c​⌊T/2​c⌋n=c\\lfloor T/2c\\rfloor, so no chunk crosses block boundaries\. Each chunk forms its summaryDb=∑j∈Ib𝐙jD\_\{b\}=\\sum\_\{j\\in I\_\{b\}\}\{\\mathbf\{Z\}\}\_\{j\}, an exclusive prefix scan\([Blelloch, 1990](https://arxiv.org/html/2609.36062#bib.bib35);[Yau et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib21)\)over chunks gives the incoming stateSbin=∑a<bDaS^\{\\mathrm\{in\}\}\_\{b\}=\\sum\_\{a<b\}D\_\{a\}, and all chunks then run in parallel:

y¯iloc=ϕi⊤​Sbin⏟earlier chunks\+ϕi⊤​∑j∈Ib,j≤i𝐙j⏟within chunk​b,y¯i=y¯ilr\+y¯iloc,i∈Ib,\\bar\{y\}^\{\\,\\mathrm\{loc\}\}\_\{i\}=\\underbrace\{\\phi\_\{i\}^\{\\top\}S^\{\\mathrm\{in\}\}\_\{b\}\}\_\{\\text\{earlier chunks\}\}\+\\underbrace\{\\phi\_\{i\}^\{\\top\}\\\!\\\!\\sum\_\{j\\in I\_\{b\},\\,j\\leq i\}\\\!\\\!\{\\mathbf\{Z\}\}\_\{j\}\}\_\{\\text\{within chunk \}b\},\\qquad\\bar\{y\}\_\{i\}=\\bar\{y\}^\{\\,\\mathrm\{lr\}\}\_\{i\}\+\\bar\{y\}^\{\\,\\mathrm\{loc\}\}\_\{i\},\\qquad i\\in I\_\{b\},\(10\)followed by a row\-wise division\.

Algorithm 1SMat forward pass for one attention head0:

𝐐,𝐊,𝐕\{\\mathbf\{Q\}\},\{\\mathbf\{K\}\},\{\\mathbf\{V\}\}; maps

prof,type\\operatorname\{prof\},\\operatorname\{type\}; incidence

𝐂\\mathbf\{C\}; chunk width

cc;

n=c​⌊T/2​c⌋n=c\\lfloor T/2c\\rfloor
1:Compute

ϕi\\phi\_\{i\}and

𝐙j\{\\mathbf\{Z\}\}\_\{j\}as in Equation[2](https://arxiv.org/html/2609.36062#S2.E2)\.

2:Pool

Fx=∑prof⁡\(j\)=x𝐙jF\_\{x\}=\\sum\_\{\\operatorname\{prof\}\(j\)=x\}\{\\mathbf\{Z\}\}\_\{j\}over

j≤nj\\leq nby grouped reduction\.

3:Tabulate

Uh=∑x𝐂h​x​FxU\_\{h\}=\\sum\_\{x\}\\mathbf\{C\}\_\{hx\}F\_\{x\}for every type

h∈\{0,…,B−1\}h\\in\\\{0,\\dots,B\-1\\\}\.

4:for alltypes

hhin parallel: contract all queries of type

hhagainst

UhU\_\{h\}and scatter the rows into chronological order\.

5:for allchunks

bbin parallel:

Db=∑j∈Ib𝐙jD\_\{b\}=\\sum\_\{j\\in I\_\{b\}\}\{\\mathbf\{Z\}\}\_\{j\}\.

6:Exclusive prefix scan over chunks, separately within

\[n\]\[n\]and within

\(n,T\]\(n,T\]\.

7:for allchunks

bbin parallel: evaluate Equation[10](https://arxiv.org/html/2609.36062#S3.E10), add the long\-range term, and normalize once\.

Implicit incidence multiplication\.We compute the type tableU=𝐂​FU=\{\\mathbf\{C\}\}Fwithout materializing𝐂\{\\mathbf\{C\}\}\. Let𝒜\\mathcal\{A\}contain the normalized nonzero directions in𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}, with the first nonzero coordinate of eachaaequal to one\. Index each type by its defining pair\(a,b\)\(a,b\), using the same ordering astype\\operatorname\{type\}\. For a fixed directionaa, every pointxxlies in exactly one hyperplaneHa,bH\_\{a,b\}, namely the one withb=a⊤​xb=a^\{\\top\}x\. Algorithm[2](https://arxiv.org/html/2609.36062#alg2)therefore groups the profile states by this offset and accumulatesUa,b=∑x:a⊤​x=bFxU\_\{a,b\}=\\sum\_\{x:\\,a^\{\\top\}x=b\}F\_\{x\}, where the offsets use finite\-field arithmetic and the state additions use ordinary arithmetic inℝr×p\\mathbb\{R\}^\{r\\times p\}\.

Algorithm 2Implicit incidence tabulation0:Profile states

\{Fx\}\\\{F\_\{x\}\\\}, field size

qq, dimension

dd
1:if

d=1d=1then

U0←F0U\_\{0\}\\leftarrow F\_\{0\}, andreturn

UU
2:for

a∈𝒜a\\in\\mathcal\{A\}in paralleldo

3:

Ua,b←0r×pU\_\{a,b\}\\leftarrow 0\_\{r\\times p\}for all

b∈𝔽qb\\in\\mathbb\{F\}\_\{q\}
4:for

x∈𝔽qd−1x\\in\\mathbb\{F\}\_\{q\}^\{d\-1\}do

5:

b←a⊤​xb\\leftarrow a^\{\\top\}x/⋆\\starFinite\-field arithmetic

6:

Ua,b←Ua,b\+FxU\_\{a,b\}\\leftarrow U\_\{a,b\}\+F\_\{x\}/⋆\\starOrdinary state addition

7:return

UU

Ford≥2d\\geq 2, the algorithm performs

\|𝒜\|​qd−1=qd−1−1q−1​qd−1=B​qd−2=Θ⁡\(q2​d−3\)\|\\mathcal\{A\}\|q^\{d\-1\}=\\frac\{q^\{d\-1\}\-1\}\{q\-1\}\\,q^\{d\-1\}=Bq^\{d\-2\}=\\Theta\(q^\{2d\-3\}\)\(11\)state additions: one for each nonzero of𝐂\{\\mathbf\{C\}\}\. We enumerate points with a base\-qqcounter and maintaina⊤​xa^\{\\top\}xas its entries change\. Over a full traversal, the counter changesO⁡\(qd−1\)O\(q^\{d\-1\}\)entries, so generating the offsets takesO⁡\(\|𝒜\|​qd−1\)O\(\|\\mathcal\{A\}\|q^\{d\-1\}\)field operations\. So, the total work isO⁡\(q2​d−3​r​p\)O\(q^\{2d\-3\}rp\), and the profile and type tables occupyO⁡\(\(qd−1\+B\)​r​p\)=O⁡\(qd−1​r​p\)O\(\(q^\{d\-1\}\+B\)rp\)=O\(q^\{d\-1\}rp\)words\. Forq=Θ⁡\(T1/d\)q=\\Theta\(T^\{1/d\}\), these bounds becomeO⁡\(T2−3/d​r​p\)O\(T^\{2\-3/d\}rp\)work andO⁡\(T1−1/d​r​p\)O\(T^\{1\-1/d\}rp\)words\. Ifd=1d=1, we directly use the single profile state\.

![Refer to caption](https://arxiv.org/html/2609.36062v1/figures/prefill_cost.png)Figure 3:Prefill cost against fused causal attention \(bf16,88heads,r=64r=64, chunk128128; median of1515timed runs\)\.###### Theorem 3\.2\(Chunkwise SMat\-attention\)\.

Fixd≥1d\\geq 1, and let𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}be the mask of Section[3](https://arxiv.org/html/2609.36062#S3)under the prescribed sequence\-length scaling\. Assume the kernel has a nonnegative feature factorization of dimensionrr, and every normalizery¯i,p\\bar\{y\}\_\{i,p\}is positive\. Then, for largeTT, Algorithm[1](https://arxiv.org/html/2609.36062#alg1)computes the attention outputso1,…,oTo\_\{1\},\\dots,o\_\{T\}, in

O⁡\(\(T2−3/d\+T\)​r​p\+T​c​\(r\+p\)\)O\\big\(\(T^\{2\-3/d\}\+T\)\\,rp\\;\+\\;T\\,c\\,\(r\+p\)\\big\)\(12\)work, excluding evaluation of the feature maps, usingO⁡\(\(T1−1/d\+T/c\)​r​p\+c2\)O\\big\(\(T^\{1\-1/d\}\+T/c\)rp\+c^\{2\}\\big\)words of working memory\. The backward pass has the same asymptotic cost\.

The proof of Theorem[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2)is in Appendix[E](https://arxiv.org/html/2609.36062#A5)\.

### 3\.2Gating

Recent variants of linear attention, such as gated linear attention\([Yang et al\., 2024a](https://arxiv.org/html/2609.36062#bib.bib6)\)and Mamba\-2\([Dao and Gu, 2024](https://arxiv.org/html/2609.36062#bib.bib28)\), yield increased performance by weighting each edge according to the product of the gates between its source and target\. To integrate our binary mask𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}into these continuous layers, we introduce a gating mechanism\. Leta1,…,aT∈\(0,1\]a\_\{1\},\\ldots,a\_\{T\}\\in\(0,1\]be per\-token scalar decay gates, andλi=∑k≤ilog⁡ak\\lambda\_\{i\}=\\sum\_\{k\\leq i\}\\log a\_\{k\}\. We replace the binary mask𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}with the continuous mask𝐌~\(d\)\\widetilde\{\\mathbf\{M\}\}^\{\(d\)\}, given by:

𝐌~i​j\(d\)=𝐌i​j\(d\)​exp⁡\(λi−λj\),j≤i\\widetilde\{\\mathbf\{M\}\}^\{\(d\)\}\_\{ij\}=\\mathbf\{M\}^\{\(d\)\}\_\{ij\}\\exp\(\\lambda\_\{i\}\-\\lambda\_\{j\}\),\\quad j\\leq i\(13\)Note that the write scale of these layers is the key gatehjh\_\{j\}from equation[2](https://arxiv.org/html/2609.36062#S2.E2), which is already subsumed within𝐙j\{\\mathbf\{Z\}\}\_\{j\}, and that the entries of𝐌~\(d\)\\widetilde\{\\mathbf\{M\}\}^\{\(d\)\}are non\-negative with the same support as𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}\. We show this additional gating does not affect training or decoding complexity in Lemma[4](https://arxiv.org/html/2609.36062#Thmlemma4)\.

### 3\.3Memory\-efficient decoding

For SMat\-Attention with hard\-routing, we compress the distant block once and reuse its summaries throughout the decoding process\. For this, we begin by fixing the distance/recent boundarynnin advance\. After processing the distant tokens, we tabulate the type states\{Uh\}h<B\\\{U\_\{h\}\\\}\_\{h<B\}and discard the intermediate profile states\{Fx\}\\\{F\_\{x\}\\\}\. At each subsequent stept\>nt\>n, we update a single running stateSt=St−1\+ZtS\_\{t\}=S\_\{t\-1\}\+Z\_\{t\}, initialized withSn=0S\_\{n\}=0, and contract the query featureϕt\\phi\_\{t\}againstUtype​\(t−n\)\+StU\_\{\\mathrm\{type\}\}\(t\-n\)\+S\_\{t\}before normalization\. Figure[4](https://arxiv.org/html/2609.36062#S3.F4)illustrates this separation between fixed distant memory and evolving recent memory\. The resulting decoder retainsB\+1B\+1states and performsO⁡\(r​p\)O\(rp\)state\-update and readout work per token after tabulation\. Appendix[F](https://arxiv.org/html/2609.36062#A6)relaxes the fixed boundary, extending the construction to decoding at any length without knowingTTin advance\.

![Refer to caption](https://arxiv.org/html/2609.36062v1/figures/decoding.png)Figure 4:Cached decoding for additive SMat\-Attention with hard\-routing\.Left:distant key\-value contributions pool into profile states, which are added via the incidence structure to form the type cache\. Colors identify the contributing profiles within each cached summary\.Right:successive decoding steps access the fixed cache while maintaining one running state for recent tokens\. The input vectors at each step are the query, key, and value\. The key and value update the running state before answering the query\. Repeated cache banks depict successive views of the shared memory\.###### Theorem 3\.3\(Streaming decoding\)\.

Assume the hypotheses of Theorem[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2), and suppose the total sequence lengthTTis fixed in advance so that the distant/recent boundarynnis known\. Once Algorithm[1](https://arxiv.org/html/2609.36062#alg1)has been run over the distant block, each subsequent token can be decoded inO⁡\(r​p\)O\(rp\)time, independent ofTT, from a cache ofO⁡\(T1−1/d​r​p\)O\(T^\{1\-1/d\}rp\)words\.

We provide the proof of Theorem[3\.3](https://arxiv.org/html/2609.36062#S3.Thmthm3)in Appendix[F](https://arxiv.org/html/2609.36062#A6)\.

### 3\.4Recurrent architectures and learned routing

We extend SMat’s profile organization to recurrent sequence models\. The local branch processes the distant and recent blocks separately, resetting its recurrent state and short convolution at the boundary\. Distant tokens update profile memoriesFxF\_\{x\}\(specified in Appendix[G](https://arxiv.org/html/2609.36062#A7)\), which are aggregated into type summariesUh=∑xCh​x​FxU\_\{h\}=\\sum\_\{x\}C\_\{hx\}F\_\{x\}\. For a recent query, the memory contribution is combined with the local output before backbone normalization and output projectionzi=zilocal\+λi​q~i⊤​∑hRi​h​Uhz\_\{i\}=z\_\{i\}^\{\\mathrm\{local\}\}\+\\lambda\_\{i\}\\widetilde\{q\}\_\{i\}^\{\\top\}\\sum\_\{h\}R\_\{ih\}U\_\{h\}, whereRi​hR\_\{ih\}assigns queries to types andλi\\lambda\_\{i\}controls the memory contribution\. To combine information across subsets, we introduce an independent query\-dependent scorer over types\. Each query selects the four highest\-scoring types per head and applies a softmax over their scores to obtain the read weightsRi​hR\_\{ih\}, retaining a fixed number of summary reads\.

The learned selector remains subquadratic because each query scores a sublinear number of cached types,B=Θ⁡\(T1−1/d\)B=\\Theta\(T^\{1\-1/d\}\), rather than all tokens\. Across the sequence, scoring costsΘ⁡\(T​B\)=Θ⁡\(T2−1/d\)\\Theta\(TB\)=\\Theta\(T^\{2\-1/d\}\), a factor ofT1/dT^\{1/d\}fewer scores than token\-level all\-pairs scoring\. Type\-summary construction is also subquadratic: Algorithm[2](https://arxiv.org/html/2609.36062#alg2)costsO⁡\(T2−3/d​r​p\)O\(T^\{2\-3/d\}rp\), while a dense contraction costsO⁡\(T2−2/d​r​p\)O\(T^\{2\-2/d\}rp\)\. These routing and aggregation costs are subquadratic for every fixed finited≥2d\\geq 2and fixed model dimensions\. The linear\-time bound ford≤3d\\leq 3applies to the hard\-routing algorithm; learned selection instead incurs the subquadratic scoring cost above\.

For learned content routing, a token’s write address is computed from its current and preceding hidden states, while each recent query selects summaries using only its own causal representation\. A learned projection quantizes the write representation into a discrete profile in𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}\. We train this hash jointly with the backbone using a straight\-through gradient estimator and an auxiliary load\-balancing loss\. The four\-read selector is learned separately through the softmax weights of its selected types\. Thus, the model learns both where to store information and which summaries to retrieve; Appendix[G\.1](https://arxiv.org/html/2609.36062#A7.SS1)gives the details\. We show in Appendix[D](https://arxiv.org/html/2609.36062#A4)that the VC\-dimension of the support of the learned mask continues to beO⁡\(d\)O\(d\)\.

## 4Experiments

We evaluate SMat on controlled synthetic tasks designed to probe routing and retrieval, followed by long\-context language modeling\.

Implementation and training details\.We implement our models in PyTorch with custom Triton kernels\([Tillet et al\., 2019](https://arxiv.org/html/2609.36062#bib.bib54)\), and our experiments are run on single NVIDIA A100 GPUs\. We use the Zoology training pipeline for MQAR\([Arora et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib1)\), and AdamW with cosine learning\-rate decay\. MQAR, joint context–key recall, and the 750M\-token PG\-19 experiments use the learned four\-read routing variant of Section[3\.4](https://arxiv.org/html/2609.36062#S3.SS4)\. Subset routing and multi\-key retrieval use the mask\-based constructions\. All SMat masks were gated in the recurrent section of the mask\. We provide our code111[https://anonymous\.4open\.science/r/smat\_attention/README\.md](https://anonymous.4open.science/r/smat_attention/README.md)and defer details to Appendix[G](https://arxiv.org/html/2609.36062#A7)\.

Subset Routing\.We first test whether models using the proposed mask can realize the subset\-selection patterns predicted by its VC dimension\. LetJ=\{j1,…,jk\}J=\\\{j\_\{1\},\\ldots,j\_\{k\}\\\}denotekklandmark positions in the context\. Each landmarkjℓj\_\{\\ell\}stores an independent random payloadxℓx\_\{\\ell\}in a distinct output channel,vjℓ=xℓ​eℓv\_\{j\_\{\\ell\}\}=x\_\{\\ell\}e\_\{\\ell\}, while all non\-landmark values are zero\. A query specifies a subsetA⊆JA\\subseteq J, and the target isyA=∑jℓ∈Axℓ​eℓy\_\{A\}=\\sum\_\{j\_\{\\ell\}\\in A\}x\_\{\\ell\}e\_\{\\ell\}\. Payloads and requested subsets are resampled across examples, preventing the model from memorizing fixed input\-output mappings\.

![Refer to caption](https://arxiv.org/html/2609.36062v1/figures/d_subset_routing.png)Figure 5:Routing\-pattern match against the requested dimension k, after training with the mask fixed, T = 1024, mean of three seeds\.This task tests the access patterns characterized by VC dimension\. Ifkklandmarks are shattered by the row supports of𝐌\{\\mathbf\{M\}\}with VC\-dimensiondd, then every subset of those landmarks can be selected by some query and𝐌\{\\mathbf\{M\}\}can realize all subset\-selection patterns over someddlandmarks, but not over anyd\+1d\+1landmarks\. So, we expect performance to degrade once the requested routing dimension exceedsdd\. The task captures a basic requirement of long\-context retrieval: selecting several relevant pieces of information while ignoring other nearby or similar context\.

Multi\-key retrieval\.We next test content\-based retrieval when a single query must recover multiple items\. We randomly placeNNkey\-payload pairs throughout the distant context\. Each query specifieskktarget keys and must retrieve the corresponding payloads\. Since the locations of the pairs vary across examples, the model cannot solve the task using fixed positional routing\. For SMat, any set of at mostd−1d\-1target profiles in𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}lies on a common affine hyperplane\. This guarantees that the target profiles lie on a common hyperplane, but does not guarantee exact selection: distractors or profile collisions may also lie on that hyperplane\. We report exact\-support accuracy \(Appendix[G\.2](https://arxiv.org/html/2609.36062#A7.SS2)\), where a query is ‘correct’ only if every requested item is present in the output and every unrequested item is absent\.

Table 1:Exact routing\-pattern match \(%\) on multi\-key retrieval, by the numberkkof marked positions\. Payloads are resampled each batch and evaluation uses fresh ones\. Mean \(std\) over 3 seeds\. SMat performs better at higherkkasddincreases\.Multi\-query associative recall\.We train on the Zoology MQAR\([Arora et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib1)\)mixture with 4\-64 key\-value pairs at sequence lengths 64–256, using two\-layer models with head and state dimension 16\. Evaluation uses 1000 held\-out examples per load at sequence lengths 64\-256\. Accuracy is averaged over query tokens within each example, then equally over examples, giving equal weight to each load\. Log\-Linear Attention\([Guo et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib4)\)uses the same training harness and budget as the corresponding backbone\. Table[2](https://arxiv.org/html/2609.36062#S4.T2)reports retrieval accuracy for Mamba\-2 and Gated DeltaNet with and without SMat\. At every width, for both backbones, the SMat variants perform well\.

Table 2:Final MQAR accuracy \(%\), 32 epochs\. Mean \(std\) over 3 seeds\.Joint Context–Key RecallWe adapt the multi\-query joint recall task of[Zhan et al\. \(2025\)](https://arxiv.org/html/2609.36062#bib.bib20), which requires retrieving values using both a context and a key\. We represent each binding as an explicit\(context,key,value\)\(\\text\{context\},\\text\{key\},\\text\{value\}\)record and independently shuffle the records and queries\. Keys repeat across contexts, while values are sampled independently, requiring joint identification of the requested record\. We train on 180K examples spanning 4, 16, 128, 256, and 512 bindings, using two\-layer, width\-64 models for 32 epochs\. Table[3](https://arxiv.org/html/2609.36062#S4.T3)reports validation accuracy averaged across the five memory loads\. Among the tested SMat settings,d=3d=3achieves the highest mean accuracy and lowest sample standard deviation on both backbones\. It outperforms the native backbones in mean accuracy, while scoring above Log\-Linear on GDN and below it on Mamba\-2\. These results do not establish the cause of the differences acrossdd\.

Table 3:Average validation accuracy \(%\) on joint context–key recall, averaged over five memory loads\. All models use width 64, 32 epochs, learning rate 0\.003\. Mean \(std\) over 3 seeds\.Natural Language Modeling\.We use language modeling to evaluate whether the additional routing structure from SMat preserves the modeling capabilities of the underlying architectures\. We evaluate language modeling on PG\-19\([Rae et al\., 2019](https://arxiv.org/html/2609.36062#bib.bib2)\)using the GPT\-2 tokenizer\. We train separate models at context lengths of 16K and 32K, each on 300M tokens, and report per\-token loss on held\-out tests at the corresponding training context length\. All models use 8 layers, a hidden width of 384, and the same training data order\. SMat, Mamba\-2, and the Transformer have 27\.8M, 26\.9M, and 33\.5M parameters, respectively\. Across both context lengths, SMat variants achieve slightly lower negative log\-likelihood \(NLL\) than the similarly\-sized Mamba\-2 baseline, with small differences amongd=2,3,4d=2,3,4, as seen in Table[16](https://arxiv.org/html/2609.36062#A7.T16)in Appendix[G](https://arxiv.org/html/2609.36062#A7)\.

We further evaluate SMat augmentation of GDN and Mamba\-2 at a larger scale, using 16\-layer models with hidden width 768, a 16K context length, and 750M training tokens\. As shown in[16](https://arxiv.org/html/2609.36062#A7.T16), all SMat variants achieve roughly the same perplexity as its baselines, suggesting that SMat’s performance in language\-modeling is comparable across both backbone architectures in this setting\.

## 5Conclusion and Future Work

We introduced SMat\-Attention, a framework for explicitly trading off the flexibility of long\-range access against computation and memory\. Rather than compressing the entire past into a fixed\-size state or allowing unrestricted token\-level interactions, SMat\-Attention provides an intermediate regime in which the richness of long\-range routing is controlled by a single parameter, the VC dimensiondd\. This structure yields provably subquadratic training and constant\-time decoding for the hard\-routing masks, and the learned extensions inherit those bounds up to a subquadratic selector term\. Controlled routing experiments show benefits consistent with this expressiveness\. Learned SMat extensions improve mean associative and joint context–key recall accuracy in several tested settings, while remaining comparable in performance at small\-scale language modeling\.

We focus on a structured finite\-feature setting and a fixed\-horizon decoding formulation, while some empirical variants introduce additional learned routing and memory updates\. Extending the framework to more adaptive routing schemes, dynamic contexts, and larger\-scale language models is a natural direction for future work\. Our results suggest that explicitly controlling the complexity of long\-range access patterns is a useful way to navigate tradeoffs between expressivity and computation\.

## AI use statement

In this work, we used generative AI tools for literature searches, coding implementation, and to aid in the presentation of our experimental results\. We also used AI assistance to edit the manuscript, help design scientific figures, and validate our mathematical claims\. The authors take full responsibility for verifying the correctness and originality of all the material in the manuscript, including the theoretical claims, experimental results, code, and figures\.

## Ethics Statement

In this work, we investigate the computational properties of the attention mechanism used in the transformer, and study tradeoffs between the memory, expressiveness, and computation across various representations\. Our evaluation uses synthetic tasks and Google DeepMind’s PG\-19 corpus for measuring the ability of language models to process long\-range contexts\. In general, efficiency gains in this line of work may broaden the access to long\-context modeling, while also lowering the cost of processing potentially sensitive pieces of textual information\. Moreover, models using the proposed mechanism are subject to the standard bias, privacy, and misuse risks associated with language models, and the attention mechanism on its own does not provide safeguards against these risks\. We discuss broader impacts in Appendix[A](https://arxiv.org/html/2609.36062#A1)\.

## Reproducibility Statement

Section[3](https://arxiv.org/html/2609.36062#S3)specifies the mask construction and attention algorithms, including pseudocode for the forward pass and implicit incidence tabulation\. We provide proofs of our theoretical claims in Appendices[C](https://arxiv.org/html/2609.36062#A3)\-[F](https://arxiv.org/html/2609.36062#A6)\. Section[G](https://arxiv.org/html/2609.36062#A7)describes the experimental tasks, datasets and synthetic\-data generation procedures, model configurations, optimization settings, training budgets, evaluation protocols, and computing hardware\. Finally, we provide an anonymous link to a faithful implementation of our code in the main body of our paper, along with documentation on how to run the experiments\.

## 6Acknowledgements

We are deeply grateful to Jan van den Brand, Jacob Abernethy, Peter Bartlett, Sarah Liaw, Avi Feller, and Ali Behrouz for sharing their helpful ideas and insightful discussions\. Emile Anand is supported by NSF Grant CCF 2338816\. Abdullah Ateyeh and Archer Wang are supported by the NSF graduate research fellowship\. This research was also sponsored by the Department of the Air Force Artificial Intelligence Accelerator and was accomplished under Cooperative Agreement Number FA8750\-19\-2\-1000\. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Department of the Air Force or the U\.S\. Government\. The U\.S\. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein\. In addition, this work is supported by the National Science Foundation under Cooperative Agreement PHY\-2019786 \(The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, http://iaifi\.org/\)\.

## References

- Anandet al\.\(2026\)E\. Anand, A\. Ateyeh, X\. Cao, and M\. DabagiaContinuous latent contexts enable efficient online learning in transformers\.External Links:2605\.09867,[Link](https://arxiv.org/abs/2605.09867)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p3.1)\.
- Anandet al\.\(2025\)E\. Anand, J\. van den Brand, and R\. McCartyThe structural complexity of matrix\-vector multiplication\.External Links:2502\.21240,[Link](https://arxiv.org/abs/2502.21240)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p4.1),[§1](https://arxiv.org/html/2609.36062#S1.p4.1)\.
- Aroraet al\.\(2023\)S\. Arora, S\. Eyuboglu, A\. Timalsina, I\. Johnson, M\. Poli, J\. Zou, A\. Rudra, and C\. RéZoology: measuring and improving recall in efficient language models\.External Links:2312\.04927,[Link](https://arxiv.org/abs/2312.04927)Cited by:[§G\.2](https://arxiv.org/html/2609.36062#A7.SS2.SSS0.Px1.p4.1),[§G\.2](https://arxiv.org/html/2609.36062#A7.SS2.SSS0.Px1.p9.1),[§1](https://arxiv.org/html/2609.36062#S1.p2.1),[§4](https://arxiv.org/html/2609.36062#S4.p2.1),[§4](https://arxiv.org/html/2609.36062#S4.p6.1)\.
- Aroraet al\.\(2024\)S\. Arora, S\. Eyuboglu, M\. Zhang, A\. Timalsina, S\. Alberti, D\. Zinsley, J\. Zou, A\. Rudra, and C\. RéSimple linear attention language models balance the recall\-throughput tradeoff\.InProceedings of the 41st International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.235\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Bahdanauet al\.\(2014\)D\. Bahdanau, K\. Cho, and Y\. BengioNeural machine translation by jointly learning to align and translate\.External Links:1409\.0473,[Link](https://arxiv.org/abs/1409.0473)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.1)\.
- Becket al\.\(2025\)M\. Beck, K\. Pöppel, P\. Lippe, and S\. HochreiterTiled flash linear attention: more efficient linear rnn and xlstm kernels\.Advances in Neural Information Processing Systems38,pp\. 75093–75148\.Cited by:[§3\.1](https://arxiv.org/html/2609.36062#S3.SS1.p3.1)\.
- Behrouzet al\.\(2026\)A\. Behrouz, Z\. Li, Y\. Deng, P\. Zhong, M\. Razaviyayn, and V\. MirrokniMemory caching: RNNs with growing memory\.InProceedings of the 43rd International Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/2602.24281),2602\.24281Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p3.1)\.
- Beltagyet al\.\(2020\)I\. Beltagy, M\. E\. Peters, and A\. CohanLongformer: the long\-document transformer\.External Links:2004\.05150,[Link](https://arxiv.org/abs/2004.05150)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p5.1)\.
- Blelloch \(1990\)G\. E\. BlellochPrefix sums and their applications\.Technical reportTechnical ReportCMU\-CS\-90\-190,Carnegie Mellon University, Department of Computer Science\.External Links:[Link](https://www.cs.cmu.edu/~scandal/papers/CMU-CS-90-190.html)Cited by:[§3\.1](https://arxiv.org/html/2609.36062#S3.SS1.p3.1)\.
- Choromanskiet al\.\(2022\)K\. Choromanski, V\. Likhosherstov, D\. Dohan, X\. Song, A\. Gane, T\. Sarlos, P\. Hawkins, J\. Davis, A\. Mohiuddin, L\. Kaiser, D\. Belanger, L\. Colwell, and A\. WellerRethinking attention with performers\.External Links:2009\.14794,[Link](https://arxiv.org/abs/2009.14794)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1),[§2](https://arxiv.org/html/2609.36062#S2.p2.3)\.
- Choromanskiet al\.\(2023\)K\. Choromanski, H\. Lin, H\. Chen, T\. Zhang, A\. Sehanobish, V\. Likhosherstov, J\. Parker\-Holder, T\. Sarlos, A\. Weller, and T\. WeingartenFrom block\-toeplitz matrices to differential equations on graphs: towards a general theory for scalable masked transformers\.External Links:2107\.07999,[Link](https://arxiv.org/abs/2107.07999)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1),[§1](https://arxiv.org/html/2609.36062#S1.p5.1),[§2](https://arxiv.org/html/2609.36062#S2.p1.1)\.
- Daoet al\.\(2022\)T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. RéFlashAttention: fast and memory\-efficient exact attention with io\-awareness\.External Links:2205\.14135,[Link](https://arxiv.org/abs/2205.14135)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p5.1)\.
- Dao and Gu \(2024\)T\. Dao and A\. GuTransformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality\.InProceedings of the 41st International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.235,pp\. 10041–10071\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1),[§1](https://arxiv.org/html/2609.36062#S1.p2.1),[§1](https://arxiv.org/html/2609.36062#S1.p3.1),[§3\.2](https://arxiv.org/html/2609.36062#S3.SS2.p1.1)\.
- Dao \(2023\)T\. DaoFlashAttention\-2: faster attention with better parallelism and work partitioning\.External Links:2307\.08691,[Link](https://arxiv.org/abs/2307.08691)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.2)\.
- Duet al\.\(2026\)J\. Du, W\. Sun, D\. Lan, J\. Hu, T\. Zhang, and Y\. ChengMoM: linear sequence modeling with mixture\-of\-memories\.InThe Fourteenth International Conference on Learning Representations,External Links:2502\.13685Cited by:[§G\.2](https://arxiv.org/html/2609.36062#A7.SS2.SSS0.Px1.p11.1),[§1](https://arxiv.org/html/2609.36062#S1.p3.1)\.
- Fein\-Ashleyet al\.\(2025\)J\. Fein\-Ashley, N\. Gupta, R\. Kannan, and V\. PrasannaSPECTRE: an fft\-based efficient drop\-in replacement to self\-attention for long contexts\.External Links:2502\.18394,[Link](https://arxiv.org/abs/2502.18394)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1)\.
- Fuet al\.\(2023\)D\. Y\. Fu, T\. Dao, K\. K\. Saab, A\. W\. Thomas, A\. Rudra, and C\. RéHungry hungry hippos: towards language modeling with state space models\.External Links:2212\.14052,[Link](https://arxiv.org/abs/2212.14052)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Goldsteinet al\.\(2026\)D\. Goldstein, N\. Singhal, and E\. CheahKey\-value means: transformers with expandable block\-recurrent compressed memory\.Note:arXiv preprintExternal Links:2605\.09877,[Link](https://arxiv.org/abs/2605.09877)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p3.1)\.
- Gu and Dao \(2023\)A\. Gu and T\. DaoMamba: Linear\-Time Sequence Modeling with Selective State Spaces\.arXiv preprint arXiv:2312\.00752\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1),[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Guet al\.\(2022\)A\. Gu, K\. Goel, and C\. RéEfficiently modeling long sequences with structured state spaces\.External Links:2111\.00396,[Link](https://arxiv.org/abs/2111.00396)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Guoet al\.\(2025\)H\. Guo, S\. Yang, T\. Goel, E\. P\. Xing, T\. Dao, and Y\. KimLog\-linear attention\.External Links:2506\.04761,[Link](https://arxiv.org/abs/2506.04761)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p3.1),[§G\.2](https://arxiv.org/html/2609.36062#A7.SS2.SSS0.Px1.p9.1),[§1](https://arxiv.org/html/2609.36062#S1.p3.1),[§4](https://arxiv.org/html/2609.36062#S4.p6.1)\.
- Huaet al\.\(2022\)W\. Hua, Z\. Dai, H\. Liu, and Q\. V\. LeTransformer quality in linear time\.External Links:2202\.10447,[Link](https://arxiv.org/abs/2202.10447)Cited by:[§3\.1](https://arxiv.org/html/2609.36062#S3.SS1.p3.1)\.
- Jianget al\.\(2018\)Q\. Jiang, X\. Cui, and W\. LiDeep discrete supervised hashing\.IEEE Transactions on Image Processing27\(12\),pp\. 5996–6009\.External Links:ISSN 1941\-0042,[Document](https://dx.doi.org/10.1109/tip.2018.2864894)Cited by:[§G\.1](https://arxiv.org/html/2609.36062#A7.SS1.p3.1)\.
- Kachamet al\.\(2023\)P\. Kacham, V\. Mirrokni, and P\. ZhongPolySketchFormer: fast transformers via sketching polynomial kernels\.arXiv preprint arXiv:2310\.01655\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Katharopouloset al\.\(2020\)A\. Katharopoulos, A\. Vyas, N\. Pappas, and F\. FleuretTransformers are rnns: fast autoregressive transformers with linear attention\.External Links:2006\.16236,[Link](https://arxiv.org/abs/2006.16236)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1),[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1),[Appendix B](https://arxiv.org/html/2609.36062#A2.p6.1),[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Katsch \(2024\)T\. KatschGateLoop: fully data\-controlled linear recurrence for sequence modeling\.External Links:2311\.01927,[Link](https://arxiv.org/abs/2311.01927)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Kearns and Vazirani \(1994\)M\. J\. Kearns and U\. V\. VaziraniAn introduction to computational learning theory\.MIT Press,Cambridge, MA, USA\.External Links:ISBN 978\-0\-262\-11193\-5Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p4.1),[§1](https://arxiv.org/html/2609.36062#S1.p4.1)\.
- Kitaevet al\.\(2020\)N\. Kitaev, Ł\. Kaiser, and A\. LevskayaReformer: the efficient transformer\.External Links:2001\.04451,[Link](https://arxiv.org/abs/2001.04451)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p5.1),[§G\.2](https://arxiv.org/html/2609.36062#A7.SS2.SSS0.Px1.p4.1)\.
- Kwonet al\.\(2023\)W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. H\. Yu, J\. Gonzalez, H\. Zhang, and I\. StoicaEfficient memory management for large language model serving with pagedattention\.InProceedings of the 29th symposium on operating systems principles,pp\. 611–626\.Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.2)\.
- Liet al\.\(2019\)S\. Li, X\. Jin, Y\. Xuan, X\. Zhou, W\. Chen, Y\. Wang, and X\. YanEnhancing the locality and breaking the memory bottleneck of transformer on time series forecasting\.Advances in neural information processing systems32\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p5.1)\.
- Liuet al\.\(2024\)H\. Liu, M\. Zaharia, and P\. AbbeelRing attention with blockwise transformers for near\-infinite context\.InInternational Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.2)\.
- Luet al\.\(2025\)E\. Lu, Z\. Jiang, J\. Liu, Y\. Du, T\. Jiang, C\. Hong, S\. Liu, W\. He, E\. Yuan, Y\. Wang, Z\. Huang, H\. Yuan, S\. Xu, X\. Xu, G\. Lai, Y\. Chen, H\. Zheng, J\. Yan, J\. Su, Y\. Wu, N\. Y\. Zhang, Z\. Yang, X\. Zhou, M\. Zhang, and J\. QiuMoBA: mixture of block attention for long\-context LLMs\.External Links:2502\.13189,[Link](https://arxiv.org/abs/2502.13189)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.2)\.
- Luoet al\.\(2021\)S\. Luo, S\. Li, T\. Cai, D\. He, D\. Peng, S\. Zheng, G\. Ke, L\. Wang, and T\. LiuStable, fast and accurate: kernelized attention with relative positional encoding\.Advances in Neural Information Processing Systems34,pp\. 22795–22807\.Cited by:[Appendix C](https://arxiv.org/html/2609.36062#A3.p1.2.1)\.
- Massaroliet al\.\(2023\)S\. Massaroli, M\. Poli, D\. Y\. Fu, H\. Kumbong, R\. N\. Parnichkun, A\. Timalsina, D\. W\. Romero, Q\. McIntyre, B\. Chen, A\. Rudra, C\. Zhang, C\. Re, S\. Ermon, and Y\. BengioLaughing hyena distillery: extracting compact recurrences from convolutions\.External Links:2310\.18780,[Link](https://arxiv.org/abs/2310.18780)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1)\.
- Penget al\.\(2024\)B\. Peng, D\. Goldstein, Q\. Anthony, A\. Albalak, E\. Alcaide, S\. Biderman, E\. Cheah, X\. Du, T\. Ferdinan, H\. Hou, P\. Kazienko, K\. K\. GV, J\. Kocoń, B\. Koptyra, S\. Krishna, R\. M\. Jr\., J\. Lin, N\. Muennighoff, F\. Obeid, A\. Saito, G\. Song, H\. Tu, C\. Wirawan, S\. Woźniak, R\. Zhang, B\. Zhao, Q\. Zhao, P\. Zhou, J\. Zhu, and R\. ZhuEagle and finch: rwkv with matrix\-valued states and dynamic recurrence\.External Links:2404\.05892,[Link](https://arxiv.org/abs/2404.05892)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Penget al\.\(2021\)H\. Peng, N\. Pappas, D\. Yogatama, R\. Schwartz, N\. A\. Smith, and L\. KongRandom feature attention\.External Links:2103\.02143,[Link](https://arxiv.org/abs/2103.02143)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Poliet al\.\(2023\)M\. Poli, S\. Massaroli, E\. Nguyen, D\. Y\. Fu, T\. Dao, S\. Baccus, Y\. Bengio, S\. Ermon, and C\. ReHyena hierarchy: towards larger convolutional language models\.InProceedings of the 40th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.202,pp\. 28043–28078\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1),[§1](https://arxiv.org/html/2609.36062#S1.p3.1)\.
- Qinet al\.\(2023\)Z\. Qin, X\. Han, W\. Sun, B\. He, D\. Li, D\. Li, Y\. Dai, L\. Kong, and Y\. ZhongToeplitz neural network for sequence modeling\.External Links:2305\.04749,[Link](https://arxiv.org/abs/2305.04749)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p2.1),[§1](https://arxiv.org/html/2609.36062#S1.p3.1)\.
- Qinet al\.\(2024\)Z\. Qin, S\. Yang, W\. Sun, X\. Shen, D\. Li, W\. Sun, and Y\. ZhongHGRN2: gated linear rnns with state expansion\.External Links:2404\.07904,[Link](https://arxiv.org/abs/2404.07904)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Raeet al\.\(2019\)J\. W\. Rae, A\. Potapenko, S\. M\. Jayakumar, and T\. P\. LillicrapCompressive transformers for long\-range sequence modelling\.External Links:1911\.05507,[Link](https://arxiv.org/abs/1911.05507)Cited by:[§4](https://arxiv.org/html/2609.36062#S4.p8.1)\.
- Royet al\.\(2020\)A\. Roy, M\. Saffar, A\. Vaswani, and D\. GrangierEfficient content\-based sparse attention with routing transformers\.External Links:2003\.05997,[Link](https://arxiv.org/abs/2003.05997)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p5.1)\.
- Schlaget al\.\(2021\)I\. Schlag, K\. Irie, and J\. SchmidhuberLinear transformers are secretly fast weight programmers\.InProceedings of the 38th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.139,pp\. 9355–9366\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1),[Appendix B](https://arxiv.org/html/2609.36062#A2.p6.1),[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Schmidhuber \(1992\)J\. SchmidhuberLearning to control fast\-weight memories: an alternative to dynamic recurrent networks\.Neural Computation4\(1\),pp\. 131–139\.External Links:ISSN 0899\-7667,[Document](https://dx.doi.org/10.1162/neco.1992.4.1.131),[Link](https://doi.org/10.1162/neco.1992.4.1.131),https://direct\.mit\.edu/neco/article\-pdf/4/1/131/812242/neco\.1992\.4\.1\.131\.pdfCited by:[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Shahet al\.\(2024\)J\. Shah, G\. Bikshandi, Y\. Zhang, V\. Thakkar, P\. Ramani, and T\. DaoFlashAttention\-3: fast and accurate attention with asynchrony and low\-precision\.InAdvances in Neural Information Processing Systems,Vol\.37\.External Links:[Document](https://dx.doi.org/10.52202/079017-2193),2407\.08608Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.2)\.
- Sunet al\.\(2023\)Y\. Sun, L\. Dong, S\. Huang, S\. Ma, Y\. Xia, J\. Xue, J\. Wang, and F\. WeiRetentive network: a successor to transformer for large language models\.External Links:2307\.08621,[Link](https://arxiv.org/abs/2307.08621)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1)\.
- Tilletet al\.\(2019\)P\. Tillet, H\. Kung, and D\. CoxTriton: an intermediate language and compiler for tiled neural network computations\.InProceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages,pp\. 10–19\.Cited by:[§4](https://arxiv.org/html/2609.36062#S4.p2.1)\.
- Vapnik and Chervonenkis \(1971\)V\. N\. Vapnik and A\. Y\. ChervonenkisOn the uniform convergence of relative frequencies of events to their probabilities\.Theory of Probability and Its Applications16\(2\),pp\. 264–280\.External Links:[Document](https://dx.doi.org/10.1137/1116025)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p4.1)\.
- Vaswaniet al\.\(2017\)A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. PolosukhinAttention is all you need\.External Links:1706\.03762,[Link](https://arxiv.org/abs/1706.03762)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p1.1),[§1](https://arxiv.org/html/2609.36062#S1.p1.2),[§2](https://arxiv.org/html/2609.36062#S2.p1.1)\.
- Widrow and Hoff \(1960\)B\. Widrow and M\. E\. HoffAdaptive switching circuits\.In1960 IRE WESCON Convention Record,Vol\.4,pp\. 96–104\.Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Yanget al\.\(2025\)S\. Yang, J\. Kautz, and A\. HatamizadehGated delta networks: improving mamba2 with delta rule\.External Links:2412\.06464,[Link](https://arxiv.org/abs/2412.06464)Cited by:[§1](https://arxiv.org/html/2609.36062#S1.p2.1)\.
- Yanget al\.\(2024a\)S\. Yang, B\. Wang, Y\. Shen, R\. Panda, and Y\. KimGated linear attention transformers with hardware\-efficient training\.InProceedings of the 41st International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.235,pp\. 56501–56523\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1),[Appendix B](https://arxiv.org/html/2609.36062#A2.p6.1),[§1](https://arxiv.org/html/2609.36062#S1.p2.1),[§3\.1](https://arxiv.org/html/2609.36062#S3.SS1.p3.1),[§3\.2](https://arxiv.org/html/2609.36062#S3.SS2.p1.1)\.
- Yanget al\.\(2024b\)S\. Yang, B\. Wang, Y\. Zhang, Y\. Shen, and Y\. KimParallelizing linear transformers with the delta rule over sequence length\.InAdvances in Neural Information Processing Systems,Vol\.37\.Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p1.1),[Appendix B](https://arxiv.org/html/2609.36062#A2.p6.1)\.
- Yauet al\.\(2025\)M\. Yau, S\. Gupta, V\. Engelmayer, K\. Irie, S\. Jegelka, and J\. AndreasSequential\-parallel duality in prefix scannable models\.Note:arXiv:2506\.10918Cited by:[§3\.1](https://arxiv.org/html/2609.36062#S3.SS1.p3.1)\.
- Yuanet al\.\(2025\)J\. Yuan, H\. Gao, D\. Dai, J\. Luo, L\. Zhao, Z\. Zhang, Z\. Xie, Y\. X\. Wei, L\. Wang, Z\. Xiao, Y\. Wang, C\. Ruan, M\. Zhang, W\. Liang, and W\. ZengNative sparse attention: hardware\-aligned and natively trainable sparse attention\.External Links:2502\.11089,[Link](https://arxiv.org/abs/2502.11089)Cited by:[Appendix B](https://arxiv.org/html/2609.36062#A2.p5.1),[§1](https://arxiv.org/html/2609.36062#S1.p1.2)\.
- Zaheeret al\.\(2020\)M\. Zaheer, G\. Guruganesh, A\. Dubey, J\. Ainslie, C\. Alberti, S\. Ontanon, P\. Pham, A\. Ravula, Q\. Wang, L\. Yang, and A\. AhmedBig 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 B](https://arxiv.org/html/2609.36062#A2.p5.1)\.
- Zhanet al\.\(2025\)Z\. Zhan, J\. Zhao, Z\. Zhu, and J\. TangOvercoming long context limitations of state space models via context dependent sparse attention\.InAdvances in Neural Information Processing Systems,Vol\.38\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2025/hash/3e4ea825a66bf1e1e33eca3e81fe7886-Abstract-Conference.html)Cited by:[§4](https://arxiv.org/html/2609.36062#S4.p7.1)\.

Outline of the Appendices\.

- •Section[A](https://arxiv.org/html/2609.36062#A1)explains the broader societal impacts of our work,
- •Section[B](https://arxiv.org/html/2609.36062#A2)lists the related work,
- •Section[C](https://arxiv.org/html/2609.36062#A3)states an auxiliary lemma to motivate our kernel attention mechanism,
- •Section[D](https://arxiv.org/html/2609.36062#A4)proves our theorem to characterize the properties of the masks,
- •Section[E](https://arxiv.org/html/2609.36062#A5)proves the chunkwise SMat\-attention result,
- •Section[F](https://arxiv.org/html/2609.36062#A6)discusses extension to horizon\-free decoding, and
- •Section[G](https://arxiv.org/html/2609.36062#A7)lists the training details

## Appendix ABroader Impacts

SMat\-Attention targets the computational and memory costs of long\-context sequence modeling\. Namely, more efficient prefill and decoding could lower the energy and hardware requirements of deploying long\-context models, potentially widening access to researchers and practitioners without large compute budgets\. By making long\-range access complexity an explicit architectural parameter, our framework also offers a more interpretable theory on which parts of the context a model can attend to, which may aid analysis of how long\-context models retrieve and use information\. At the same time, cheaper long\-context inference lowers the barrier to processing large volumes of personal or sensitive text, and the general risks of language models, including the generation of misleading or harmful content, apply to systems built on this mechanism\.

## Appendix BRelated Work

Table 4:SMat sits between linear and softmax attention\. It keeps theO⁡\(r​p\)O\(rp\)per\-token decoding cost of linear attention and pays for VC dimensionddin cache size rather than in decoding time\.Kernel and recurrent attention\.Kernelized attention factorizes the content kernel asK⁡\(q,k\)=ϕ𝐐​\(q\)⊤​ϕ𝐊​\(k\)K\(q,k\)=\\phi\_\{\\mathbf\{Q\}\}\(q\)^\{\\top\}\\phi\_\{\\mathbf\{K\}\}\(k\), allowing causal attention to be accumulated in finite\-dimensional recurrent states\[[Kacham et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib55)\]\. For fixed feature and value dimensions, this yields linear work in the sequence length, as well as a recurrent state whose size is independent ofTT\[[Katharopoulos et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib5)\]\. Random\-feature methods such as Performer\[[Choromanski et al\., 2022](https://arxiv.org/html/2609.36062#bib.bib36)\]approximate the softmax kernel within this framework\. More recent architectures enrich the recurrent state and its update\[[Sun et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib34),[Katsch, 2024](https://arxiv.org/html/2609.36062#bib.bib33),[Qin et al\., 2024](https://arxiv.org/html/2609.36062#bib.bib16),[Peng et al\., 2024](https://arxiv.org/html/2609.36062#bib.bib49)\]: for instance, gated linear attention introduces input\-dependent retention\[[Peng et al\., 2021](https://arxiv.org/html/2609.36062#bib.bib50),[Yang et al\., 2024a](https://arxiv.org/html/2609.36062#bib.bib6)\], whereas DeltaNet uses key\-conditioned corrective updates\[[Schlag et al\., 2021](https://arxiv.org/html/2609.36062#bib.bib23),[Yang et al\., 2024b](https://arxiv.org/html/2609.36062#bib.bib14)\]\. These mechanisms improve how a fixed\-size state is maintained, but retaining a sequence\-length\-independent state creates a capacity\-recall tradeoff on tasks requiring access to many independent items from the context\[[Arora et al\., 2024](https://arxiv.org/html/2609.36062#bib.bib43)\]\. Our work differs: SMat\-Attention accepts any supplied finite nonnegative feature factorization and changes the*causal support pattern*over which the resulting features interact\.

Structured sequence mixing and hierarchical memory\.A recurring theme in efficient sequence modeling is that computational savings arise from algebraic structure in the sequence mixing matrix since many efficient sequence models can be interpreted as multiplication by a structured causal matrix\. For instance, linear attention induces a lower\-triangular structured operator\[[Katharopoulos et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib5)\], whereas long\-convolution architectures such as Hyena\[[Massaroli et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib30),[Poli et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib15)\]use Toeplitz\-like operators\[[Qin et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib51),[Fein\-Ashley et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib48)\]with FFT\-based multiplication\. Similarly, selective state\-space models \(SSMs\) such as Mamba induce input\-dependent semi\-separable mixing matrices\[[Gu and Dao, 2023](https://arxiv.org/html/2609.36062#bib.bib26)\]\. Mamba\-2 makes this matrix perspective explicit through structured state\-space duality, showing an equivalence between SSM recurrences and multiplication by semi\-separable matrices\[[Dao and Gu, 2024](https://arxiv.org/html/2609.36062#bib.bib28)\]\. More generally,[Choromanski et al\. \[2023\]](https://arxiv.org/html/2609.36062#bib.bib45)showed that efficient multiplication by a mask can be lifted to efficient finite\-feature masked attention, encompassing causal prefix masks, relative\-position operators, and a variety of graph\-derived masks\. These examples suggest treating the structure of the sequence mixing matrix itself as a design space\.

The most closely related work to ours is Log\-Linear Attention, which replaces linear attention’s single prefix state withO⁡\(log⁡T\)O\(\\log T\)summaries of disjoint dyadic buckets maintained through a Fenwick\-tree schedule\[[Guo et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib4)\]\. Its hierarchical matrix structure supportsO⁡\(T​log⁡T\)O\(T\\log T\)parallel training andO⁡\(log⁡T\)O\(\\log T\)time and memory per decoded token, while query\-dependent coefficients select among temporal scales\. We explore a different structural axis in SMat\-Attention by constructing overlapping binary access patterns from point\-hyperplane incidences, and quantifying their combinatorial richness via the VC\-dimension\. Therefore, SMat Attention’s states summarize geometric profiles rather than temporal buckets, and our resulting guarantee is different: we show that after a fixed, known distant prefix has been processed, SMat\-Attention decodes each subsequent token in time independent ofTTusingO⁡\(T1−1/d\)O\(T^\{1\-1/d\}\)cached feature states\.

VC dimension and structured matrix multiplication\.Beyond sequence modeling, a line of work studies when structured matrices result in fast matrix\-vector multiplication\. The VC\-dimension is a combinatorial complexity metric that classically measures the richness of a set system through the subsets it realizes\[[Kearns and Vazirani, 1994](https://arxiv.org/html/2609.36062#bib.bib42)\]\. Recent work connects this quantity to the complexity of matrix\-vector multiplication\. For instance, after anO~​\(T2\)\\tilde\{O\}\(T^\{2\}\)preprocessing of aT×TT\\times TBoolean matrix𝐌\\mathbf\{M\}of VC\-dimensiondd,[Anand et al\. \[2025\]](https://arxiv.org/html/2609.36062#bib.bib3)gives an algorithm for multiplying𝐌\\mathbf\{M\}by an arbitrary vector inO~​\(T2−1/d\)\\widetilde\{O\}\(T^\{2\-1/d\}\)time\. This connection motivates our use of VC dimension as an access\-complexity measure; our construction additionally exploits point–hyperplane incidence structure to obtain a sharper specialized algorithm\.

Sparse and hardware\-efficient attention\.Hardware\-aware algorithms such as FlashAttention reorganize exact softmax attention into on\-chip tiles, substantially reducing memory traffic without changing its worst\-case quadratic arithmetic complexity\[[Dao et al\., 2022](https://arxiv.org/html/2609.36062#bib.bib41)\]\. Sparse\-attention methods instead reduce the number of evaluated query\-key pairs through local and global windows, random edges, or content\-dependent selection\[[Beltagy et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib40),[Zaheer et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib39),[Yuan et al\., 2025](https://arxiv.org/html/2609.36062#bib.bib37),[Li et al\., 2019](https://arxiv.org/html/2609.36062#bib.bib56)\]\. Conversely, SMat\-Attention’s complete causal mask hasΘ⁡\(T2\)\\Theta\(T^\{2\}\)non\-zeros, and therefore cannot be evaluated efficiently by enumerating all permitted interactions\. Although its long\-range block containsΘ⁡\(T2−1/d\)\\Theta\(T^\{2\-1/d\}\)token\-level edges, sparsity alone would only yield the correspondingΘ⁡\(T2−1/d\)\\Theta\(T^\{2\-1/d\}\)computation\. Our sharper bound comes from additional reuse: distant keys with the same profile are pooled once, queries with the same type share an aggregate, and the two causal blocks are handled by prefix scans and local dense tiles\. Thus, SMat\-Attention combines hardware\-friendly intrachunk computation with algebraic reuse across chunks; its speedup is not merely a consequence of deleting attention edges\. Additionally, the Reformer’s LSH attention\[[Kitaev et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib38)\]and the Routing Transformer\[[Roy et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib29)\]restrict each query to the keys sharing its hash or cluster, and our content\-based assignment uses a similar device\. The structures differ ford≥3d\\geq 3: hash and cluster buckets partition the keys, so a query sees exactly one cell, whereas each of our hyperplanes containsqd−2q^\{d\-2\}profiles, so a query reads a structured union of cells\. Ford=2d=2the hyperplanes are single points and SMat reduces to bucketed linear attention\.

Linear attention and its variants\.Linear attention replaces the softmax kernel with a feature map that factorizes the attention matrix, allowingOt=ϕ​\(qt\)⊤​∑j≤tϕ⁡\(kj\)​vj⊤O\_\{t\}=\\phi\(q\_\{t\}\)^\{\\top\}\\sum\_\{j\\leq t\}\\phi\(k\_\{j\}\)v\_\{j\}^\{\\top\}to be computed using a fixed\-size recurrent state\. This reduces autoregressive decoding from linear to constant cost per token and allows for efficient parallel training, but compresses the entire history into a fixed\-size state\[[Katharopoulos et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib5)\]\. However, vanilla linear attention compresses history into a fixed\-size state by accumulating key–value associations, leading to memory\-capacity limitations and interference between stored associations\[[Schlag et al\., 2021](https://arxiv.org/html/2609.36062#bib.bib23)\]\. Gated variants augment this recurrence with input\-dependent retention factors that modulate the existing state, allowing the model to selectively preserve or forget past information while retaining efficient recurrent inference\. For example, gated linear attention updates the state asSt=At⊙St−1\+vt​kt⊤S\_\{t\}=A\_\{t\}\\odot S\_\{t\-1\}\+v\_\{t\}k\_\{t\}^\{\\top\}, where the gateAtA\_\{t\}determines which parts of the previous state are retained, while preserving their efficient recurrent and parallel forms\[[Yang et al\., 2024a](https://arxiv.org/html/2609.36062#bib.bib6)\]\. Gating improves memory management by controlling how much of the existing state is retained, but it does not directly account for what value is already stored at a particular key\. Delta\-rule models make the update key\-specific by using the current prediction error,St=St−1​\(I−βt​kt​kt⊤\)\+βt​vt​kt⊤S\_\{t\}=S\_\{t\-1\}\(I\-\\beta\_\{t\}k\_\{t\}k\_\{t\}^\{\\top\}\)\+\\beta\_\{t\}v\_\{t\}k\_\{t\}^\{\\top\}, so that writing atktk\_\{t\}partially removes the value currently associated with that key before insertingvtv\_\{t\}\[[Yang et al\., 2024b](https://arxiv.org/html/2609.36062#bib.bib14)\]\.

## Appendix CAuxiliary Lemmas

###### Lemma 1\.

Assume that the mask𝐌∈ℝT×T\{\\mathbf\{M\}\}\\in\\mathbb\{R\}^\{T\\times T\}supports matrix\-vector multiplication in timef𝐌​\(T\)f\_\{\\mathbf\{M\}\}\(T\)\. Then, the general masked kernel attention algorithm with mask𝐌\{\\mathbf\{M\}\}can be implemented in timeO⁡\(\(f𝐌​\(T\)\+T\)​r​dv\)O\(\(f\_\{\\mathbf\{M\}\}\(T\)\+T\)rd\_\{v\}\)\.

###### Proof\.

Note that theii’th token representation obtained from the general masked kernel attention has the form

𝖠𝗍𝗍K​\(𝐐,𝐊,𝐕,𝐌\)i=ϕ​\(𝐪i⊤\)⊤​∑j=1T𝐌i,j​ϕ​\(𝐤j⊤\)⊤​vjϕ​\(𝐪i⊤\)⊤​∑j=1T𝐌i,j​ϕ​\(𝐤j⊤\)\.\\mathsf\{Att\}\_\{K\}\(\{\\mathbf\{Q\}\},\{\\mathbf\{K\}\},\{\\mathbf\{V\}\},\{\\mathbf\{M\}\}\)\_\{i\}=\\frac\{\\phi\(\\mathbf\{q\}\_\{i\}^\{\\top\}\)^\{\\top\}\\sum\_\{j=1\}^\{T\}\{\\mathbf\{M\}\}\_\{i,j\}\\phi\(\\mathbf\{k\}\_\{j\}^\{\\top\}\)^\{\\top\}v\_\{j\}\}\{\\phi\(\\mathbf\{q\}\_\{i\}^\{\\top\}\)^\{\\top\}\\sum\_\{j=1\}^\{T\}\{\\mathbf\{M\}\}\_\{i,j\}\\phi\(\\mathbf\{k\}\_\{j\}^\{\\top\}\)\}\.Then, following[Luo et al\. \[2021\]](https://arxiv.org/html/2609.36062#bib.bib44), let

𝐇1=\(∑j=1T𝐌i,j​ϕ​\(𝐤j⊤\)​𝐯j\)i=1T∈ℝr×dv\\mathbf\{H\}^\{1\}=\\left\(\\sum\_\{j=1\}^\{T\}\{\\mathbf\{M\}\}\_\{i,j\}\\phi\(\\mathbf\{k\}\_\{j\}^\{\\top\}\)\\mathbf\{v\}\_\{j\}\\right\)\_\{i=1\}^\{T\}\\in\\mathbb\{R\}^\{r\\times d\_\{v\}\}and

𝐇2=\(∑j=1T𝐌i,j​ϕ​\(𝐤j⊤\)⊤\)i=1T∈ℝ1×r\.\\mathbf\{H\}^\{2\}=\\left\(\\sum\_\{j=1\}^\{T\}\{\\mathbf\{M\}\}\_\{i,j\}\\phi\(\\mathbf\{k\}\_\{j\}^\{\\top\}\)^\{\\top\}\\right\)\_\{i=1\}^\{T\}\\in\\mathbb\{R\}^\{1\\times r\}\.Note that if𝐃1\\mathbf\{D\}^\{1\}and𝐃2\{\\mathbf\{D\}\}^\{2\}are the vectorized forms of𝐇1\{\\mathbf\{H\}\}^\{1\}and𝐇2\{\\mathbf\{H\}\}^\{2\}\(respectively\), where each element of the sequence is vectorized and the resulting vectors are stacked into matrices, then𝐃1=𝐌𝐕1\{\\mathbf\{D\}\}^\{1\}=\{\\mathbf\{M\}\}\{\\mathbf\{V\}\}^\{1\}and𝐃2=𝐌𝐕2\{\\mathbf\{D\}\}^\{2\}=\{\\mathbf\{M\}\}\{\\mathbf\{V\}\}^\{2\}, where theii’th rows of𝐕1\{\\mathbf\{V\}\}^\{1\}and𝐕2\{\\mathbf\{V\}\}^\{2\}are given as𝐕i1=vec⁡\(ϕ​\(𝐤i\)⊤​𝐯i\)\{\\mathbf\{V\}\}^\{1\}\_\{i\}=\\mathrm\{vec\}\(\\phi\(\\mathbf\{k\}\_\{i\}\)^\{\\top\}\\mathbf\{v\}\_\{i\}\)and𝐕i2=ϕ​\(𝐤i⊤\)⊤\{\\mathbf\{V\}\}\_\{i\}^\{2\}=\\phi\(\\mathbf\{k\}\_\{i\}^\{\\top\}\)^\{\\top\}\. Therefore, computing𝐇1\{\\mathbf\{H\}\}^\{1\}and𝐇2\{\\mathbf\{H\}\}^\{2\}takes timef𝐌​\(T\)​r​dvf\_\{\\mathbf\{M\}\}\(T\)rd\_\{v\}, and so𝖠𝗍𝗍i\\mathsf\{Att\}\_\{i\}can be computed in timeO⁡\(\(f𝐌​\(T\)\+T\)​r​dv\)O\(\(f\_\{\\mathbf\{M\}\}\(T\)\+T\)rd\_\{v\}\), which completes the proof\.∎

## Appendix DVC dimension of the causal incidence masks

###### Lemma 2\(Lower triangular matrices\)\.

TheT×TT\\times Tbinary lower triangular matrix𝐋T\{\\mathbf\{L\}\}\_\{T\}has VC\-dimension00forT=1T=1and11forT≥2T\\geq 2\.

###### Proof\.

ForT=1T=1, the only row support is\{1\}\\\{1\\\}, so no singleton is shattered\.

ForT≥2T\\geq 2, note that column22is excluded by row11and included by row22, and so the singleton\{2\}\\\{2\\\}is shattered, which provesVC⁡\(𝐋T\)≥1\\operatorname\{VC\}\(\{\\mathbf\{L\}\}\_\{T\}\)\\geq 1\. Next, take any two distinct column/row indicesx,yx,y, withx<yx<y\. In𝐋T\{\\mathbf\{L\}\}\_\{T\}, the upper\-right entries are all00’s\. So, if a row indicator functions is11at a later index, we cannot have independent labelings\. Therefore, since no configuration of two points can achieve all22=42^\{2\}=4binary label combinations, the VC\-dimension is less than22, completing the proof\.∎

###### Lemma 3\(Affine hyperplanes\)\.

LetD≥1D\\geq 1and letqqbe a prime power\. Then, the set system of all proper affine hyperplanes in𝔽qD\\mathbb\{F\}\_\{q\}^\{D\}has VC dimension exactlyDD\.

###### Proof\.

We first record an affine\-dependence observation\. Suppose profilesxj∈𝔽qDx\_\{j\}\\in\\mathbb\{F\}\_\{q\}^\{D\}, indexed by a setCC, satisfy

∑j∈Cλjxj=0,∑j∈Cλj=0,λj≠0\(j∈C\)\.\\sum\_\{j\\in C\}\\lambda\_\{j\}x\_\{j\}=0,\\qquad\\sum\_\{j\\in C\}\\lambda\_\{j\}=0,\\qquad\\lambda\_\{j\}\\neq 0\\quad\(j\\in C\)\.For everya∈Ca\\in C, any affine hyperplane containing allxjx\_\{j\}withj∈C∖\{a\}j\\in C\\setminus\\\{a\\\}also containsxax\_\{a\}: indeed,

xa=−∑j∈C∖\{a\}λjλaxj,−∑j∈C∖\{a\}λjλa=1\.x\_\{a\}=\-\\sum\_\{j\\in C\\setminus\\\{a\\\}\}\\frac\{\\lambda\_\{j\}\}\{\\lambda\_\{a\}\}x\_\{j\},\\qquad\-\\sum\_\{j\\in C\\setminus\\\{a\\\}\}\\frac\{\\lambda\_\{j\}\}\{\\lambda\_\{a\}\}=1\.This observation also applies to repeated profiles at distinct indices\.

IfD\+1D\+1points were shattered, the all\-included trace would place them in a proper hyperplane of affine dimensionD−1D\-1\. They are therefore affinely dependent\. Restricting a nonzero dependence to its nonzero coefficients gives a setCCsatisfying\(∗\)\(\*\)\. The preceding observation rules out the traceC∖\{a\}C\\setminus\\\{a\\\}, contradicting shattering\. This proves the upper bound\.

For the lower bound, lete1,…,eDe\_\{1\},\\ldots,e\_\{D\}be the standard basis\. For each𝐑⊆\[D\]\{\\mathbf\{R\}\}\\subseteq\[D\], define

HR=\{\{x∈𝔽qD:∑ℓ∉Rxℓ=0\},R≠\[D\],\{x∈𝔽qD:∑ℓ=1Dxℓ=1\},R=\[D\]\.H\_\{R\}=\\begin\{cases\}\\displaystyle\\left\\\{x\\in\\mathbb\{F\}\_\{q\}^\{D\}:\\sum\_\{\\ell\\notin R\}x\_\{\\ell\}=0\\right\\\},&R\\neq\[D\],\\\\\[4\.0pt\] \\displaystyle\\left\\\{x\\in\\mathbb\{F\}\_\{q\}^\{D\}:\\sum\_\{\\ell=1\}^\{D\}x\_\{\\ell\}=1\\right\\\},&R=\[D\]\.\\end\{cases\}\(14\)Each defining normal is nonzero, so everyHRH\_\{R\}is a proper affine hyperplane\. Moreover,eℓ∈HRe\_\{\\ell\}\\in H\_\{R\}if and only ifℓ∈R\\ell\\in R\. Thus the standard basis is shattered\. The construction is valid over every finite field, including𝔽2\\mathbb\{F\}\_\{2\}, which completes the proof\.∎

See[3\.1](https://arxiv.org/html/2609.36062#S3.Thmthm1)

###### Proof\.

Let𝐌\(d\)\{\\mathbf\{M\}\}^\{\(d\)\}be the mask constructed in Section[3](https://arxiv.org/html/2609.36062#S3)\. Letqqbe a prime\. Then, for every1≤d<⌊log2⁡T⌋1\\leq d<\\lfloor\\log\_\{2\}T\\rfloorfor sufficiently largeTT, we prove the above properties\.

Recall the block form from Section[3](https://arxiv.org/html/2609.36062#S3):

𝐌\(d\)=\(𝐋n0𝐆𝐋m\),𝐆i​j=𝟏\{prof\(j\)∈Htype⁡\(i\)\}\.\{\\mathbf\{M\}\}^\{\(d\)\}=\\begin\{pmatrix\}\{\\mathbf\{L\}\}\_\{n\}&0\\\\ \{\\mathbf\{G\}\}&\{\\mathbf\{L\}\}\_\{m\}\\end\{pmatrix\},\\qquad\{\\mathbf\{G\}\}\_\{ij\}=\\mathbf\{1\}\\\{\\operatorname\{prof\}\(j\)\\in H\_\{\\operatorname\{type\}\(i\)\}\\\}\.The two diagonal blocks are inclusive causal triangles, and every entry of𝐆\{\\mathbf\{G\}\}connects a recent query to a distant key\. This proves causality, the unit diagonal, and𝐌\(d\)≤𝐋T\{\\mathbf\{M\}\}^\{\(d\)\}\\leq\{\\mathbf\{L\}\}\_\{T\}\.

We prove \(iii\) and \(iv\) for everydd\. First, ford=1d=1, the mask is𝐋T\{\\mathbf\{L\}\}\_\{T\}which clearly satisfies \(i\) and \(ii\) and \(iv\)\. Then, from[lemma2](https://arxiv.org/html/2609.36062#Thmlemma2), \(iii\) is satisfied\. Moreover, note that

nnz⁡\(𝐌\(d\)\)≥n⁡\(n\+1\)\+m⁡\(m\+1\)2≥\(n\+m\)24=T24,\\mathrm\{nnz\}\(\{\\mathbf\{M\}\}^\{\(d\)\}\)\\geq\\frac\{n\(n\+1\)\+m\(m\+1\)\}\{2\}\\geq\\frac\{\(n\+m\)^\{2\}\}\{4\}=\\frac\{T^\{2\}\}\{4\},while causality givesnnz⁡\(𝐌\(d\)\)≤T⁡\(T\+1\)/2\\operatorname\{nnz\}\(\{\\mathbf\{M\}\}^\{\(d\)\}\)\\leq T\(T\+1\)/2\.

Finally, it remains to prove the VC\-dimension claim ford≥2d\\geq 2\. For this, letD=d−1D=d\-1\.

Upper bound\.Suppose a setJJofd\+1d\+1column indices is shattered\. Split it into distant and recent indices,J=Jdist⊔JrecJ=J\_\{\\mathrm\{dist\}\}\\sqcup J\_\{\\mathrm\{rec\}\}\. The traces on the recent block are empty for distant rows and prefixes for recent rows\. In particular, ifu<vu<vare recent indices, every row containingvvalso containsuu\. Thus\|Jrec\|≤1\|J\_\{\\mathrm\{rec\}\}\|\\leq 1\. We proceed by casework:

If\|Jrec\|=1\|J\_\{\\mathrm\{rec\}\}\|=1, consider the labelings in which that recent coordinate is one\. Their realizing rows must be recent rows, whose distant supports are\{j∈\[n\]:prof⁡\(j\)∈Htype⁡\(i\)\}\\\{j\\in\[n\]:\\operatorname\{prof\}\(j\)\\in H\_\{\\operatorname\{type\}\(i\)\}\\\}\. Therefore, thed=D\+1d=D\+1distant indices would be shattered by proper affine hyperplanes in𝔽qD\\mathbb\{F\}\_\{q\}^\{D\}\. If two of these indices share a profile, they cannot be independently labeled\. Otherwise, their profiles would form a shattered set ofD\+1D\+1points, contradicting Lemma[3](https://arxiv.org/html/2609.36062#Thmlemma3)\. Hence this case is impossible\.

On the other hand, if\|Jrec\|=0\|J\_\{\\mathrm\{rec\}\}\|=0, alld\+1=D\+2d\+1=D\+2indices are distant\. Their profiles are affinely dependent, so there exist a setC⊆JC\\subseteq Jand coefficients satisfying

∑j∈Cλjprof\(j\)=0,∑j∈Cλj=0,λj≠0\(j∈C\),\\sum\_\{j\\in C\}\\lambda\_\{j\}\\operatorname\{prof\}\(j\)=0,\\qquad\\sum\_\{j\\in C\}\\lambda\_\{j\}=0,\\qquad\\lambda\_\{j\}\\neq 0\\quad\(j\\in C\),where\|C\|≥2\|C\|\\geq 2\. By the affine\-dependence observation in Lemma[3](https://arxiv.org/html/2609.36062#Thmlemma3), for everya∈Ca\\in C, any affine hyperplane containing the profiles indexed byC∖\{a\}C\\setminus\\\{a\\\}also containsprof⁡\(a\)\\operatorname\{prof\}\(a\)\. So, no recent row realizes the traceC∖\{a\}C\\setminus\\\{a\\\}onCC\.

Since shatteringJJimplies shatteringCC, all the tracesC∖\{a\}C\\setminus\\\{a\\\}would have to be realized by distant rows\. But distant row supports are nested prefixes, whereas the\|C\|≥2\|C\|\\geq 2sets\{C∖\{a\}:a∈C\}\\bigl\\\{C\\setminus\\\{a\\\}:a\\in C\\bigr\\\}are pairwise incomparable, and a chain cannot realize all of them\. Therefore, this contradiction provesVC⁡\(𝐌\(d\)\)≤d\\operatorname\{VC\}\(\{\\mathbf\{M\}\}^\{\(d\)\}\)\\leq d\.

Lower bound\.We now use the prescribed assignment conditions from Section[3](https://arxiv.org/html/2609.36062#S3)\. These provide distinct distant positionsj1,…,jDj\_\{1\},\\ldots,j\_\{D\}withprof⁡\(jℓ\)=eℓ\\operatorname\{prof\}\(j\_\{\\ell\}\)=e\_\{\\ell\}forℓ∈\[D\]\\ell\\in\[D\], and a recent indexτ∈\[m\]\\tau\\in\[m\]such that each witness hyperplaneHRH\_\{R\},𝐑⊆\[D\]\{\\mathbf\{R\}\}\\subseteq\[D\], occurs at recent query indicesiR−,iR\+i\_\{R\}^\{\-\},i\_\{R\}^\{\+\}satisfying

iR−<τ≤iR\+,Htype⁡\(iR−\)=Htype⁡\(iR\+\)=HR\.i\_\{R\}^\{\-\}<\\tau\\leq i\_\{R\}^\{\+\},\\qquad H\_\{\\operatorname\{type\}\(i\_\{R\}^\{\-\}\)\}=H\_\{\\operatorname\{type\}\(i\_\{R\}^\{\+\}\)\}=H\_\{R\}\.By the definition of these witness hyperplanes,eℓ∈HRe\_\{\\ell\}\\in H\_\{R\}if and only ifℓ∈R\\ell\\in R\.

Consider theddcolumn indicesJ∗=\{j1,…,jD,n\+τ\}J\_\{\*\}=\\\{j\_\{1\},\\ldots,j\_\{D\},n\+\\tau\\\}\. FixA⊆J∗A\\subseteq J\_\{\*\}and set𝐑=\{ℓ:jℓ∈A\}\{\\mathbf\{R\}\}=\\\{\\ell:j\_\{\\ell\}\\in A\\\}\. Ifn\+τ∉An\+\\tau\\notin A, use rown\+iR−n\+i\_\{R\}^\{\-\}; ifn\+τ∈An\+\\tau\\in A, use rown\+iR\+n\+i\_\{R\}^\{\+\}\. In both cases, the incidence condition gives exactly the required trace on the distant landmarks\. The recent causal triangle includes columnn\+τn\+\\tauprecisely when the recent query index is at leastτ\\tau, so the final coordinate also has its required label\. Thus the selected row has trace exactlyAAonJ∗J\_\{\*\}\.

Therefore, every subset ofJ∗J\_\{\*\}is realized, givingVC⁡\(𝐌\(d\)\)≥d\\operatorname\{VC\}\(\{\\mathbf\{M\}\}^\{\(d\)\}\)\\geq d\. Together with the upper bound, this provesVC⁡\(𝐌\(d\)\)=d\\operatorname\{VC\}\(\{\\mathbf\{M\}\}^\{\(d\)\}\)=d\.∎

VC dimension of learned multi\-read routing\.Fix a sequence and head, and letD=d−1≥1D=d\-1\\geq 1\. With one profile per distant key andkkpositively weighted reads, each cross\-block support is a union ofkkaffine\-hyperplane traces\. WritingBkB\_\{k\}for this binary cross\-block mask and

M^k=\(Ln0BkLm\),\\widehat\{M\}\_\{k\}=\\begin\{pmatrix\}L\_\{n\}&0\\\\ B\_\{k\}&L\_\{m\}\\end\{pmatrix\},we have, for an absolute constantCC,

VC⁡\(M^k\)≤C⁡\(d−1\)​k​log2⁡\(2​k\)\+2\.\\operatorname\{VC\}\(\\widehat\{M\}\_\{k\}\)\\leq C\(d\-1\)k\\log\_\{2\}\(2k\)\+2\.In particular,VC⁡\(M^4\)≤25​d−24=O⁡\(d\)\\operatorname\{VC\}\(\\widehat\{M\}\_\{4\}\)\\leq 25d\-24=O\(d\)\.

###### Proof\.

Affine hyperplanes in𝔽qD\\mathbb\{F\}\_\{q\}^\{D\}have VC dimensionDD\. By Sauer\-Shelah’s lemma, the number of cross\-block traces on anys≥Ds\\geq Dkeys satisfies

ΠBk​\(s\)≤\(∑r=0D\(sr\)\)k≤\(e​sD\)D​k\.\\Pi\_\{B\_\{k\}\}\(s\)\\leq\\left\(\\sum\_\{r=0\}^\{D\}\\binom\{s\}\{r\}\\right\)^\{k\}\\leq\\left\(\\frac\{es\}\{D\}\\right\)^\{Dk\}\.Thus, shattering requires2s≤\(e​s/D\)D​k2^\{s\}\\leq\(es/D\)^\{Dk\}\. Settingu=s/\(D​k\)u=s/\(Dk\)gives2u≤e​k​u2^\{u\}\\leq eku, henceu=O⁡\(log2⁡\(2​k\)\)u=O\(\\log\_\{2\}\(2k\)\)and

VC⁡\(Bk\)=O⁡\(D​k​log2⁡\(2​k\)\)\.\\operatorname\{VC\}\(B\_\{k\}\)=O\\\!\\left\(Dk\\log\_\{2\}\(2k\)\\right\)\.Fork=4k=4, settingu=s/Du=s/Dinstead yields2u≤\(e​u\)42^\{u\}\\leq\(eu\)^\{4\}, which impliesu<25u<25\. Consequently,VC⁡\(B4\)≤25​D−1\\operatorname\{VC\}\(B\_\{4\}\)\\leq 25D\-1\.

Finally, letv=VC⁡\(Bk\)v=\\operatorname\{VC\}\(B\_\{k\}\)\. A shattered column set contains at most one recent column, since recent\-column supports are nested prefixes\. If it contains one, fixing that column to one forces recent rows to shatter all selected distant columns, giving size at mostv\+1v\+1\. If all columns are distant, fixing the earliest to zero and the latest to one excludes every distant prefix row, so recent rows must shatter the remaining columns\. ThusVC⁡\(M^k\)≤v\+2\\operatorname\{VC\}\(\\widehat\{M\}\_\{k\}\)\\leq v\+2, proving both claims\. ∎

Dependence on the construction\.The bounds above characterize our family under its prescribed field\-size scaling\. Other choices ofqqcan preserve the same mask VC dimension while changing the number of profile cells, the number of tokens sharing each profile, and the computational cost\. In particular, fixedqqgives a fixed number of profiles asTTgrows\. Our scaling instead lets the profile and type tables grow with sequence length\.

## Appendix EChunking Proofs

We now provide the proof for our result in Theorem[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2)\.

See[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2)

###### Proof\.

Forward Pass\.Before the final division,y¯i\\bar\{y\}\_\{i\}is linear in the mask, and the two sub\-masks of𝐌\(d\)\{\\mathbf\{M\}\}^\{\(d\)\}have disjoint support\. Thus, it suffices to evaluate each branch separately, add the outputs, and then normalize\.

For the long\-range branch, a recent tokenn\+in\+ireceivesϕn\+i⊤​∑j≤n𝐆i​j​𝐙j\\phi\_\{n\+i\}^\{\\top\}\\sum\_\{j\\leq n\}\\mathbf\{G\}\_\{ij\}\{\\mathbf\{Z\}\}\_\{j\}\. Substituting𝐆i​j=𝐂type⁡\(i\),prof⁡\(j\)\\mathbf\{G\}\_\{ij\}=\\mathbf\{C\}\_\{\\operatorname\{type\}\(i\),\\operatorname\{prof\}\(j\)\}and grouping the sum by profile,

∑j≤n𝐆i​j𝐙j=∑x∈𝔽qd−1𝐂type⁡\(i\),x∑j≤n:prof⁡\(j\)=x𝐙j=∑x𝐂type⁡\(i\),xFx=Utype⁡\(i\),\\sum\_\{j\\leq n\}\\mathbf\{G\}\_\{ij\}\{\\mathbf\{Z\}\}\_\{j\}=\\sum\_\{x\\in\\mathbb\{F\}\_\{q\}^\{d\-1\}\}\\mathbf\{C\}\_\{\\operatorname\{type\}\(i\),x\}\\\!\\\!\\\!\\sum\_\{j\\leq n:\\,\\operatorname\{prof\}\(j\)=x\}\\\!\\\!\\\!\{\\mathbf\{Z\}\}\_\{j\}=\\sum\_\{x\}\\mathbf\{C\}\_\{\\operatorname\{type\}\(i\),x\}F\_\{x\}=U\_\{\\operatorname\{type\}\(i\)\},which is the entry of the type table computed by the algorithm\. The regrouping is valid becauseprof\\operatorname\{prof\}partitions\[n\]\[n\]into disjoint sets\. Furthermore, because this long\-range branch only applies to recent queries, the term evaluates to zero for all distant queriesy¯ilr=0\\bar\{y\}^\{\\,\\mathrm\{lr\}\}\_\{i\}=0fori≤ni\\leq n\.

To be more explicit, each augmented state flattens into a row, so𝐙∈ℝn×r​p\{\\mathbf\{Z\}\}\\in\\mathbb\{R\}^\{n\\times rp\}\. Recall the one\-hot matrices of equation[8](https://arxiv.org/html/2609.36062#S3.E8):𝐄∈\{0,1\}n×qd−1\{\\mathbf\{E\}\}\\in\\\{0,1\\\}^\{n\\times q^\{d\-1\}\}with𝐄j​x=𝟙\{prof\(j\)=x\}\{\\mathbf\{E\}\}\_\{jx\}=\\mathbbm\{1\}\\\{\\operatorname\{prof\}\(j\)=x\\\}, and𝐑∈\{0,1\}m×B\{\\mathbf\{R\}\}\\in\\\{0,1\\\}^\{m\\times B\}with𝐑i​h=𝟙\{type\(i\)=h\}\{\\mathbf\{R\}\}\_\{ih\}=\\mathbbm\{1\}\\\{\\operatorname\{type\}\(i\)=h\\\}, so that𝐆=𝐑​𝐂​𝐄⊤\\mathbf\{G\}=\{\\mathbf\{R\}\}\\,\\mathbf\{C\}\\,\{\\mathbf\{E\}\}^\{\\top\}\. The long\-range states of all recent queries are the rows of

𝐆​𝐙=𝐑​𝐂​𝐄⊤​𝐙=𝐑⁡\(𝐂⁡\(𝐄⊤​𝐙⏟𝐅∈ℝqd−1×r​p\)⏟U∈ℝB×r​p\)⏟∈ℝm×r​p,\\mathbf\{G\}\\,\{\\mathbf\{Z\}\}\\;=\\;\{\\mathbf\{R\}\}\\,\\mathbf\{C\}\\,\{\\mathbf\{E\}\}^\{\\top\}\{\\mathbf\{Z\}\}\\;=\\;\\underbrace\{\{\\mathbf\{R\}\}\\,\\big\(\\,\\underbrace\{\\mathbf\{C\}\\,\\big\(\\,\\underbrace\{\{\\mathbf\{E\}\}^\{\\top\}\{\\mathbf\{Z\}\}\}\_\{\\textstyle\{\\mathbf\{F\}\}\\in\\mathbb\{R\}^\{q^\{d\-1\}\\times rp\}\}\\,\\big\)\}\_\{\\textstyle U\\in\\mathbb\{R\}^\{B\\times rp\}\}\\,\\big\)\}\_\{\\textstyle\\in\\mathbb\{R\}^\{m\\times rp\}\},\(15\)and Algorithm[1](https://arxiv.org/html/2609.36062#alg1)evaluates equation[15](https://arxiv.org/html/2609.36062#A5.E15)from the right:𝐄⊤​Z\{\\mathbf\{E\}\}^\{\\top\}Zis the pooling step,𝐂⁡\(⋅\)\\mathbf\{C\}\(\\cdot\)the tabulation step, and the outer𝐑\{\\mathbf\{R\}\}the per\-type read, which is a row gather whose inverse permutation is the scatter back into chronological order\.

For the causal branch, fixi∈Ibi\\in I\_\{b\}\. Becausennis chunk\-aligned,IbI\_\{b\}lies entirely in\[n\]\[n\]or entirely in\(n,T\]\(n,T\]\. The incoming state collects the keys of all earlier chunks in that block and the within\-chunk term collects\{j∈Ib:j≤i\}\\\{j\\in I\_\{b\}:j\\leq i\\\}, so their union is\{j≤i\}\\\{j\\leq i\\\}wheni≤ni\\leq nand\{n<j≤i\}\\\{n<j\\leq i\\\}wheni\>ni\>n\. These are the row supports of𝐋n\\mathbf\{L\}\_\{n\}and𝐋m\\mathbf\{L\}\_\{m\}\. Summing the branches and dividing once gives the output for𝐌\(d\)\\mathbf\{M\}^\{\(d\)\}\.

For cost, the three factors of equation[15](https://arxiv.org/html/2609.36062#A5.E15)correspond to three counts, each obtained from the last by collapsing one index, each worth a factor ofΘ⁡\(T1/d\)\\Theta\(T^\{1/d\}\)\. Note,m,n=Θ⁡\(T\)m,n=\\Theta\(T\),B=Θ⁡\(qd−1\)B=\\Theta\(q^\{d\-1\}\)andq=Θ⁡\(T1/d\)q=\\Theta\(T^\{1/d\}\)\.

From𝐆\\mathbf\{G\}itself, we have a hyperplane of𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}containingqd−2q^\{d\-2\}of theqd−1q^\{d\-1\}points \(a1/q1/qof the entries are11s\), so a recent query is incident toΘ⁡\(n/q\)\\Theta\(n/q\)distant keys and

nnz⁡\(𝐆\)=Θ⁡\(m​n/q\)=Θ⁡\(T2−1/d\),\\operatorname\{nnz\}\(\\mathbf\{G\}\)=\\Theta\\big\(mn/q\\big\)=\\Theta\\big\(T^\{2\-1/d\}\\big\),\(16\)Applying𝐑\{\\mathbf\{R\}\}collapses the queries – rowiiof𝐆\\mathbf\{G\}depends oniionly throughtype⁡\(i\)\\operatorname\{type\}\(i\)– so themmrows take onlyBBdistinct values, and the\(type,token\)\(\\text\{type\},\\text\{token\}\)incidences number

Θ⁡\(B​n/q\)=Θ⁡\(n​qd−2\)=Θ⁡\(T2−2/d\),\\Theta\\big\(Bn/q\\big\)=\\Theta\\big\(nq^\{d\-2\}\\big\)=\\Theta\(T^\{2\-2/d\}\),which produces an additional saving ofm/B=Θ⁡\(T1/d\)m/B=\\Theta\(T^\{1/d\}\)\. ApplyingE⊤E^\{\\top\}collapses the keys in the same way: columnjjdepends onjjonly throughprof⁡\(j\)\\operatorname\{prof\}\(j\), so thenncolumns take onlyqd−1q^\{d\-1\}distinct values, and thus the\(type,profile\)\(\\text\{type\},\\text\{profile\}\)incidences are the nonzeros of𝐂\\mathbf\{C\},

nnz⁡\(𝐂\)=Θ⁡\(B​qd−2\)=Θ⁡\(q2​d−3\)=Θ⁡\(T2−3/d\),\\operatorname\{nnz\}\(\\mathbf\{C\}\)=\\Theta\\big\(Bq^\{d\-2\}\\big\)=\\Theta\\big\(q^\{2d\-3\}\\big\)=\\Theta\\big\(T^\{2\-3/d\}\\big\),\(17\)a further saving ofn/qd−1=Θ⁡\(T1/d\)n/q^\{d\-1\}=\\Theta\(T^\{1/d\}\)\. Thus, tabulation performs one state addition of sizer​prpper nonzero of𝐂\\mathbf\{C\}, forO⁡\(T2−3/d​r​p\)O\(T^\{2\-3/d\}rp\)\.

The remaining steps are linear inTT\. Pooling reads each distant token once and adds it into one bucket; the per\-type contractions costr​prpper recent query, for∑hmh​r​p=m​r​p\\sum\_\{h\}m\_\{h\}\\,rp=m\\,rpin total, independent ofBB; the chunk summaries and the scan touch each token and each of theT/cT/cchunk boundaries a constant number of times\. Together these areO⁡\(T​r​p\)O\(Trp\)\. Within a chunk, calculating the local attention scores costsO⁡\(c2​r\)O\(c^\{2\}r\)and applying them to the augmented values costsO⁡\(c2​p\)O\(c^\{2\}p\); summing overT/cT/cchunks givesO⁡\(T​c​\(r\+p\)\)O\(Tc\(r\+p\)\)\. Together, we have the stated bound\.

Backward Pass\.For the Long\-range branch, let the three steps of equation[9](https://arxiv.org/html/2609.36062#S3.E9)be

𝒫:𝐙↦F,\\displaystyle\\mathcal\{P\}:\{\\mathbf\{Z\}\}\\mapsto F,Fx=∑j≤n:prof⁡\(j\)=x𝐙j,\\displaystyle F\_\{x\}=\\\!\\\!\\\!\\sum\_\{j\\leq n:\\,\\operatorname\{prof\}\(j\)=x\}\\\!\\\!\\\!\{\\mathbf\{Z\}\}\_\{j\},𝒞:F↦U,\\displaystyle\\mathcal\{C\}:F\\mapsto U,Uh=∑x𝐂h​x​Fx,\\displaystyle U\_\{h\}=\\sum\_\{x\}\\mathbf\{C\}\_\{hx\}F\_\{x\},ℛΦ:U↦y¯lr,\\displaystyle\\mathcal\{R\}\_\{\\Phi\}:U\\mapsto\\bar\{y\}^\{\\mathrm\{lr\}\},y¯n\+ilr=ϕn\+i⊤​Utype⁡\(i\),\\displaystyle\\bar\{y\}^\{\\mathrm\{lr\}\}\_\{n\+i\}=\\phi\_\{n\+i\}^\{\\top\}U\_\{\\operatorname\{type\}\(i\)\},so thaty¯lr=ℛΦ​𝒞​𝒫​𝐙\\bar\{y\}^\{\\mathrm\{lr\}\}=\\mathcal\{R\}\_\{\\Phi\}\\,\\mathcal\{C\}\\,\\mathcal\{P\}\\,\\mathbf\{Z\}\.

Let

g¯i:=∂ℒ∂y¯i∈ℝp\\bar\{g\}\_\{i\}:=\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\\bar\{y\}\_\{i\}\}\\in\\mathbb\{R\}^\{p\}denote the gradient of the loss with respect to the augmented output\. Pairing withg¯\\bar\{g\}and moving one factor at a time across the inner product gives

U¯h=∑i:type⁡\(i\)=hϕn\+ig¯n\+i⊤,F¯=𝐂⊤U¯,𝐙¯j=F¯prof⁡\(j\),ϕ¯n\+i=Utype⁡\(i\)g¯n\+i\.\\bar\{U\}\_\{h\}=\\\!\\\!\\\!\\sum\_\{i:\\,\\operatorname\{type\}\(i\)=h\}\\\!\\\!\\\!\\phi\_\{n\+i\}\\,\\bar\{g\}\_\{n\+i\}^\{\\top\},\\quad\\bar\{F\}=\\mathbf\{C\}^\{\\top\}\\bar\{U\},\\quad\\bar\{\\mathbf\{Z\}\}\_\{j\}=\\bar\{F\}\_\{\\operatorname\{prof\}\(j\)\},\\quad\\bar\{\\phi\}\_\{n\+i\}=U\_\{\\operatorname\{type\}\(i\)\}\\,\\bar\{g\}\_\{n\+i\}\.\(18\)The two grouped steps switch roles: the forward pass reduces overprof\\operatorname\{prof\}and broadcasts overtype\\operatorname\{type\}, while equation[18](https://arxiv.org/html/2609.36062#A5.E18)reduces overtype\\operatorname\{type\}and broadcasts overprof\\operatorname\{prof\}\. Each is one pass over the tokens it owns atO⁡\(r​p\)O\(rp\)per token, so both costO⁡\(T​r​p\)O\(Trp\)\.

For the middle step, the incidence structure is biregular, i\.e\. every hyperplane of𝔽qd−1\\mathbb\{F\}\_\{q\}^\{d\-1\}containsqd−2q^\{d\-2\}points, and also every point lies on\(qd−1−1\)/\(q−1\)=B/q\(q^\{d\-1\}\-1\)/\(q\-1\)=B/qhyperplanes \(one per normalized directionaa\), sincex∈Ha,bx\\in H\_\{a,b\}forcesb=a⊤​xb=a^\{\\top\}x\. Hence,nnz⁡\(𝐂⊤\)=nnz⁡\(𝐂\)=Θ⁡\(T2−3/d\)\\operatorname\{nnz\}\(\\mathbf\{C\}^\{\\top\}\)=\\operatorname\{nnz\}\(\\mathbf\{C\}\)=\\Theta\(T^\{2\-3/d\}\), and𝐂⊤\\mathbf\{C\}^\{\\top\}has constant column degree just as𝐂\\mathbf\{C\}has constant row degree\. The tabulation step therefore applies to𝐂⊤\\mathbf\{C\}^\{\\top\}with the point\-to\-hyperplane table in place of the hyperplane\-to\-point one, givingO⁡\(T2−3/d​r​p\)O\(T^\{2\-3/d\}rp\)forF¯\\bar\{F\}, matching the forward calculation\.

For the causal branch, fix a chunkIbI\_\{b\}and stack its rows asΦb∈ℝc×r\\Phi\_\{b\}\\in\\mathbb\{R\}^\{c\\times r\},Ψb∈ℝc×r\\Psi\_\{b\}\\in\\mathbb\{R\}^\{c\\times r\},V¯b∈ℝc×p\\bar\{V\}\_\{b\}\\in\\mathbb\{R\}^\{c\\times p\}, and writeA\(b\):=tril⁡\(Φb​Ψb⊤\)A^\{\(b\)\}:=\\operatorname\{tril\}\(\\Phi\_\{b\}\\Psi\_\{b\}^\{\\top\}\)for the within\-chunk score tile, so that equation[10](https://arxiv.org/html/2609.36062#S3.E10)readsYb=Φb​Sbin\+A\(b\)​V¯bY\_\{b\}=\\Phi\_\{b\}S^\{\\mathrm\{in\}\}\_\{b\}\+A^\{\(b\)\}\\bar\{V\}\_\{b\}\. Differentiating this at fixedtril\\operatorname\{tril\}pattern,

S¯bin\\displaystyle\\bar\{S\}^\{\\mathrm\{in\}\}\_\{b\}=Φb⊤​Y¯b,\\displaystyle=\\Phi\_\{b\}^\{\\top\}\\bar\{Y\}\_\{b\},V¯¯b\\displaystyle\\qquad\\bar\{\\bar\{V\}\}\_\{b\}=\(A\(b\)\)⊤​Y¯b,\\displaystyle=\\big\(A^\{\(b\)\}\\big\)^\{\\top\}\\bar\{Y\}\_\{b\},\(19\)Φ¯b\\displaystyle\\bar\{\\Phi\}\_\{b\}=Y¯b​\(Sbin\)⊤\+tril⁡\(Y¯b​V¯b⊤\)​Ψb,\\displaystyle=\\bar\{Y\}\_\{b\}\\big\(S^\{\\mathrm\{in\}\}\_\{b\}\\big\)^\{\\top\}\+\\operatorname\{tril\}\\\!\\big\(\\bar\{Y\}\_\{b\}\\bar\{V\}\_\{b\}^\{\\top\}\\big\)\\Psi\_\{b\},Ψ¯b\\displaystyle\\qquad\\bar\{\\Psi\}\_\{b\}=tril⁡\(Y¯b​V¯b⊤\)⊤​Φb\.\\displaystyle=\\operatorname\{tril\}\\\!\\big\(\\bar\{Y\}\_\{b\}\\bar\{V\}\_\{b\}^\{\\top\}\\big\)^\{\\\!\\top\}\\Phi\_\{b\}\.Each is onec×cc\\times cbyc×rc\\times rorc×pc\\times pcontraction, soO​\(c2​\(r\+p\)\)O\(c^\{2\}\(r\+p\)\)per chunk andO⁡\(T​c​\(r\+p\)\)O\(Tc\(r\+p\)\)over theT/cT/cchunks, the same as the forward tile\. Note thatA\(b\)A^\{\(b\)\}andtril⁡\(Y¯b​V¯b⊤\)\\operatorname\{tril\}\(\\bar\{Y\}\_\{b\}\\bar\{V\}\_\{b\}^\{\\top\}\)are rebuilt fromΦb,Ψb,Y¯b,V¯b\\Phi\_\{b\},\\Psi\_\{b\},\\bar\{Y\}\_\{b\},\\bar\{V\}\_\{b\}when the chunk is visited, so onec2c^\{2\}tile is live at a time \(same as in the forward pass\)\.

For the scan,Sbin=∑a<bDaS^\{\\mathrm\{in\}\}\_\{b\}=\\sum\_\{a<b\}D\_\{a\}gives

∑b⟨S¯bin,∑a<bDa⟩=∑a⟨∑b\>aS¯bin,Da⟩\\sum\_\{b\}\\big\\langle\\bar\{S\}^\{\\mathrm\{in\}\}\_\{b\},\\textstyle\\sum\_\{a<b\}D\_\{a\}\\big\\rangle=\\sum\_\{a\}\\Big\\langle\\textstyle\\sum\_\{b\>a\}\\bar\{S\}^\{\\mathrm\{in\}\}\_\{b\},\\;D\_\{a\}\\Big\\ranglesoD¯a=∑b\>aS¯bin\\bar\{D\}\_\{a\}=\\sum\_\{b\>a\}\\bar\{S\}^\{\\mathrm\{in\}\}\_\{b\}\. In other words, the adjoint of an exclusive prefix scan is an exclusive suffix scan over the same chunks, run separately within\[n\]\[n\]and within\(n,T\]\(n,T\], at the sameO⁡\(T/c\)O\(T/c\)state additions\. BroadcastingD¯b\\bar\{D\}\_\{b\}back to the tokens ofIbI\_\{b\}and adding the two contributions of equation[18](https://arxiv.org/html/2609.36062#A5.E18)and equation[19](https://arxiv.org/html/2609.36062#A5.E19)accumulates𝐙¯j\\bar\{\\mathbf\{Z\}\}\_\{j\}, from whichψ¯j=𝐙¯j​v¯j\\bar\{\\psi\}\_\{j\}=\\bar\{\\mathbf\{Z\}\}\_\{j\}\\bar\{v\}\_\{j\}andv¯¯j=𝐙¯j⊤​ψj\\bar\{\\bar\{v\}\}\_\{j\}=\\bar\{\\mathbf\{Z\}\}\_\{j\}^\{\\top\}\\psi\_\{j\}follow pointwise inO⁡\(r​p\)O\(rp\)per token\.

Finally sinceoi=y¯i,1:dv/y¯i,po\_\{i\}=\\bar\{y\}\_\{i,1:d\_\{v\}\}/\\bar\{y\}\_\{i,p\}withy¯i,p\>0\\bar\{y\}\_\{i,p\}\>0by assumption, its Jacobian is row\-local: witho¯i:=∂ℒ∂oi\\bar\{o\}\_\{i\}:=\\frac\{\\partial\\mathcal\{L\}\}\{\\partial o\_\{i\}\},

g¯i,1:dv=o¯iy¯i,p,g¯i,p=−o¯i⊤y¯i,1:dvy¯i,p2=−o¯i⊤​oiy¯i,p,\\bar\{g\}\_\{i,1:d\_\{v\}\}=\\frac\{\\bar\{o\}\_\{i\}\}\{\\bar\{y\}\_\{i,p\}\},\\qquad\\bar\{g\}\_\{i,p\}=\-\\frac\{\\bar\{o\}\_\{i\}^\{\\top\}\\,\\bar\{y\}\_\{i,1:d\_\{v\}\}\}\{\\bar\{y\}\_\{i,p\}^\{2\}\}=\-\\frac\{\\bar\{o\}\_\{i\}^\{\\top\}o\_\{i\}\}\{\\bar\{y\}\_\{i,p\}\},atO⁡\(p\)O\(p\)per token andO⁡\(T​p\)O\(Tp\)overall, which is dominated\.

Every step of Algorithm[1](https://arxiv.org/html/2609.36062#alg1)is therefore matched by an adjoint of the same shape and arithmetic\. The backward pass runs inO⁡\(\(T2−3/d\+T\)​r​p\+T​c​\(r\+p\)\)O\\big\(\(T^\{2\-3/d\}\+T\)rp\+Tc\(r\+p\)\\big\)work andO⁡\(\(T1−1/d\+T/c\)​r​p\+c2\)O\\big\(\(T^\{1\-1/d\}\+T/c\)rp\+c^\{2\}\\big\)words, the bounds of equation[12](https://arxiv.org/html/2609.36062#S3.E12)\.∎

###### Lemma 4\(Gated SMat\)\.

Under equation[13](https://arxiv.org/html/2609.36062#S3.E13)with scalar per\-token gates, Algorithm[1](https://arxiv.org/html/2609.36062#alg1)computes the exact outputs within the complexity bound of equation[12](https://arxiv.org/html/2609.36062#S3.E12), usingO⁡\(T\)O\(T\)additional space\.

Lemmas[4](https://arxiv.org/html/2609.36062#Thmlemma4)–[5](https://arxiv.org/html/2609.36062#Thmlemma5)concern additive kernel attention with optional scalar decay; they do not cover the delta\-rule recurrence used in the GDN extension of Section[3\.4](https://arxiv.org/html/2609.36062#S3.SS4)\.

###### Proof\.

We first computeλ1,…,λT\\lambda\_\{1\},\\dots,\\lambda\_\{T\}with one prefix sum, inO⁡\(T\)O\(T\)work andO⁡\(T\)O\(T\)space\. Every factor used afterwards has the formexp⁡\(λi−λj\)\\exp\(\\lambda\_\{i\}\-\\lambda\_\{j\}\)withj≤ij\\leq i\. Sinceλ\\lambdais non\-increasing, each factor lies in\(0,1\]\(0,1\], so we never formexp⁡\(λi\)\\exp\(\\lambda\_\{i\}\)orexp⁡\(−λj\)\\exp\(\-\\lambda\_\{j\}\)separately and no step can overflow\.

*Causal branch\.*Lete⁡\(b\)e\(b\)be the last index of chunkIbI\_\{b\}\. Each entry of the within\-chunk tile becomes\(ϕi⊤​ψj\)​exp⁡\(λi−λj\)\(\\phi\_\{i\}^\{\\top\}\\psi\_\{j\}\)\\exp\(\\lambda\_\{i\}\-\\lambda\_\{j\}\), which is one entrywise product on a tile that is already materialized\. The chunk summary becomesDb=∑j∈Ibexp⁡\(λe⁡\(b\)−λj\)​𝐙jD\_\{b\}=\\sum\_\{j\\in I\_\{b\}\}\\exp\(\\lambda\_\{e\(b\)\}\-\\lambda\_\{j\}\)\{\\mathbf\{Z\}\}\_\{j\}, which rescales each𝐙j\{\\mathbf\{Z\}\}\_\{j\}by a scalar prior to the sum\. The carried state is read asexp⁡\(λi−λe⁡\(b−1\)\)​ϕi⊤​Sbin\\exp\(\\lambda\_\{i\}\-\\lambda\_\{e\(b\-1\)\}\)\\,\\phi\_\{i\}^\{\\top\}S^\{\\mathrm\{in\}\}\_\{b\}and updated asSb\+1in=exp⁡\(λe⁡\(b\)−λe⁡\(b−1\)\)​Sbin\+DbS^\{\\mathrm\{in\}\}\_\{b\+1\}=\\exp\(\\lambda\_\{e\(b\)\}\-\\lambda\_\{e\(b\-1\)\}\)\\,S^\{\\mathrm\{in\}\}\_\{b\}\+D\_\{b\}, adding one scalar multiplication per chunk\. The scan is otherwise unchanged\.

*Long\-range branch\.*For a distant keyj≤nj\\leq nand a recent queryn\+in\+i, the factor splits at the boundary,

exp⁡\(λn\+i−λj\)=exp⁡\(λn\+i−λn\)⋅exp⁡\(λn−λj\),\\exp\(\\lambda\_\{n\+i\}\-\\lambda\_\{j\}\)=\\exp\(\\lambda\_\{n\+i\}\-\\lambda\_\{n\}\)\\cdot\\exp\(\\lambda\_\{n\}\-\\lambda\_\{j\}\),into a query\-side and a key\-side scalar, each at\(0,1\]\(0,1\]\. The key\-side scalar is folded into𝐙j\{\\mathbf\{Z\}\}\_\{j\}before pooling, and the query\-side scalar multiplies the query’s read from the type table\. Neither modifies𝐂\\mathbf\{C\}, the profile table or the type table, so the long\-range branch is the ungated computation applied to rescaled inputs\.

The added work isO⁡\(T\)O\(T\)exponentials,O⁡\(T​c\)O\(Tc\)for the tile products andO⁡\(T​r\)O\(Tr\)to rescale features, all dominated by terms already in equation[12](https://arxiv.org/html/2609.36062#S3.E12)\. The added space is theO⁡\(T\)O\(T\)values ofλ\\lambda\. ∎

Note, the same factorization gives streaming decoding\. Letat=exp⁡\(λt−λt−1\)∈\(0,1\]a\_\{t\}=\\exp\(\\lambda\_\{t\}\-\\lambda\_\{t\-1\}\)\\in\(0,1\]\. At the boundary, cache the type\-major states

Uh\(n\)=∑j≤nCh,prof⁡\(j\)​exp⁡\(λn−λj\)​𝐙j,U\_\{h\}^\{\(n\)\}=\\sum\_\{j\\leq n\}C\_\{h,\\operatorname\{prof\}\(j\)\}\\,\\exp\(\\lambda\_\{n\}\-\\lambda\_\{j\}\)\\,\{\\mathbf\{Z\}\}\_\{j\},which is the ungated cache with each𝐙j\{\\mathbf\{Z\}\}\_\{j\}rescaled\. Fort\>nt\>n, maintain the recent stateSt=at​St−1\+𝐙tS\_\{t\}=a\_\{t\}S\_\{t\-1\}\+\{\\mathbf\{Z\}\}\_\{t\}and the scalargt=at​gt−1g\_\{t\}=a\_\{t\}g\_\{t\-1\}, starting fromSn=0S\_\{n\}=0andgn=1g\_\{n\}=1\. By induction,St=∑n<j≤texp⁡\(λt−λj\)​𝐙jS\_\{t\}=\\sum\_\{n<j\\leq t\}\\exp\(\\lambda\_\{t\}\-\\lambda\_\{j\}\)\{\\mathbf\{Z\}\}\_\{j\}andgt=exp⁡\(λt−λn\)g\_\{t\}=\\exp\(\\lambda\_\{t\}\-\\lambda\_\{n\}\), so the augmented output is

y¯t=ϕt⊤​\[gt​Utype⁡\(t−n\)\(n\)\+St\]\.\\bar\{y\}\_\{t\}=\\phi\_\{t\}^\{\\top\}\\\!\\left\[\\,g\_\{t\}\\,U^\{\(n\)\}\_\{\\operatorname\{type\}\(t\-n\)\}\+S\_\{t\}\\right\]\.Each token costs the same as as the ungated computation plus two scalar multiplications\.

## Appendix FDecoding and Horizon\-Free Decoding

We first restate and prove Theorem[3\.3](https://arxiv.org/html/2609.36062#S3.Thmthm3)\. See[3\.3](https://arxiv.org/html/2609.36062#S3.Thmthm3)

###### Proof\.

The type table\{Uh\}h<B\\\{U\_\{h\}\\\}\_\{h<B\}depends only on the distant tokens\. Since every distant position precedes every recent one, it is fixed once the distant block has been consumed, and hence, the profile states\{Fx\}\\\{F\_\{x\}\\\}can be discarded\. By the decomposition in the proof of Theorem[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2), the augmented output at recent positionn\+in\+iis

y¯n\+i=ϕn\+i⊤​\(Utype⁡\(i\)\+∑n<j≤n\+i𝐙j\),\\bar\{y\}\_\{n\+i\}=\\phi\_\{n\+i\}^\{\\top\}\\Big\(U\_\{\\operatorname\{type\}\(i\)\}\+\\\!\\\!\\sum\_\{n<j\\leq n\+i\}\\\!\\\!\{\\mathbf\{Z\}\}\_\{j\}\\Big\),\(20\)and the second term is a single running stateSSmaintained by the in\-place updateS←S\+𝐙n\+iS\\leftarrow S\+\{\\mathbf\{Z\}\}\_\{n\+i\}\. The indextype⁡\(i\)=imodB\\operatorname\{type\}\(i\)=i\\bmod Bis one arithmetic operation\. Forming𝐙n\+i\{\\mathbf\{Z\}\}\_\{n\+i\}, updatingSS, and contractingϕn\+i⊤​\(Utype⁡\(i\)\+S\)\\phi\_\{n\+i\}^\{\\top\}\(U\_\{\\operatorname\{type\}\(i\)\}\+S\)each costO⁡\(r​p\)O\(rp\)\. The retained state is theBBtabulated states together withSS\. For the gated variant from Equation[13](https://arxiv.org/html/2609.36062#S3.E13), the running state update becomesS←an\+i​S\+𝐙n\+iS\\leftarrow a\_\{n\+i\}S\+\{\\mathbf\{Z\}\}\_\{n\+i\}and the read ofUtype⁡\(i\)U\_\{\\operatorname\{type\}\(i\)\}carries the scalarexp⁡\(λn\+i−λn\)\\exp\(\\lambda\_\{n\+i\}\-\\lambda\_\{n\}\), maintained by one addition per token\.∎

The streaming\-decoding theorem \(Theorem[3\.3](https://arxiv.org/html/2609.36062#S3.Thmthm3)\) assumes that the total lengthTTof the token sequence is known in advance, since the distant/recent boundarynnand the field sizeqqare both functions ofTTand together determine the structure of the mask\. Hence, if more tokens were to be added beyondTT, the layer degrades to linear attention\. In this section we replace that assumption with a fixed training horizonTmaxT\_\{\\max\}and a recent window that advances in steps ofcd​cc\_\{dc\}\.

###### Definition 1\(Stepped\-window mask\)\.

Fix a training horizonTmaxT\_\{\\max\}, a block lengthcd​cc\_\{dc\}dividingTmaxT\_\{\\max\}and a multiple of the chunk widthcc, and the geometry\(q,𝐂,B\)\(q,\{\\mathbf\{C\}\},B\)from the SMat construction of Section[3](https://arxiv.org/html/2609.36062#S3)forT=TmaxT=T\_\{\\max\}\. Fort≥1t\\geq 1letb⁡\(t\)=⌊\(t−1\)/cd​c⌋b\(t\)=\\lfloor\(t\-1\)/c\_\{dc\}\\rfloorbe the block of positionttandn⁡\(t\)=max⁡\{0,\(b⁡\(t\)−1\)​cd​c\}n\(t\)=\\max\\\{0,\\,\(b\(t\)\-1\)c\_\{dc\}\\\}be its number of distant positions\. For anyTT, includingT\>TmaxT\>T\_\{\\max\},

𝐌t​j\(step\)=\{1n⁡\(t\)<j≤t\(recent window\),𝐂type⁡\(t\),prof⁡\(j\)j≤n⁡\(t\)\(distant, through​𝐆\),0j\>t\.\{\\mathbf\{M\}\}^\{\(\\mathrm\{step\}\)\}\_\{tj\}=\\begin\{cases\}1&n\(t\)<j\\leq t\\quad\(\\text\{recent window\}\),\\\\ \{\\mathbf\{C\}\}\_\{\\operatorname\{type\}\(t\),\\,\\operatorname\{prof\}\(j\)\}&j\\leq n\(t\)\\quad\(\\text\{distant, through \}\{\\mathbf\{G\}\}\),\\\\ 0&j\>t\.\\end\{cases\}

In the block form of Section[3](https://arxiv.org/html/2609.36062#S3)the boundaryn=c​⌊T/2​c⌋n=c\\lfloor T/2c\\rflooris a constant determined byTT\. Here, the boundary is the functionn⁡\(⋅\)n\(\\cdot\), of which the constantnnis the special casen⁡\(t\)≡nn\(t\)\\equiv n\. Sincen⁡\(t\)n\(t\)depends only ontt, the mask for a length\-ttsequence is the leadingt×tt\\times tblock of the mask for any longer sequence\. Thus, training atTmaxT\_\{\\max\}and decoding at any length uses one mask\.

In Algorithm[3](https://arxiv.org/html/2609.36062#alg3), each fresh recurrent state is initialized to zero, with updateR←R\+ZtR\\leftarrow R\+Z\_\{t\}and readoutread⁡\(R,ϕt\)=ϕt⊤​R\\operatorname\{read\}\(R,\\phi\_\{t\}\)=\\phi\_\{t\}^\{\\top\}R\. For the scalar\-gated variant, the update becomesR←at​R\+ZtR\\leftarrow a\_\{t\}R\+Z\_\{t\}\.

current𝐌\(d\)\{\\mathbf\{M\}\}^\{\(d\)\}, boundary kept atn=Tmax/2n=T\_\{\\max\}/2TmaxT\_\{\\max\}recent \(causal\) region grows withttstepped window,cd​c=Tmax/4c\_\{dc\}=T\_\{\\max\}/4TmaxT\_\{\\max\}recent region stayscd​cc\_\{dc\}to2​cd​c2c\_\{dc\}causal recurrenceincidence block𝐆\{\\mathbf\{G\}\}Figure 6:Both masks drawn at the same lengthT=1\.5​TmaxT=1\.5\\,T\_\{\\max\}, in blocks ofcd​c=Tmax/4c\_\{dc\}=T\_\{\\max\}/4tokens\.Left:the fixed\-TTmask once its boundary is kept atn=Tmax/2n=T\_\{\\max\}/2\. The incidence block𝐆\{\\mathbf\{G\}\}stops growing at two blocks and everything after the boundary falls to the recurrence\.Right:the stepped window, whose recent region stays betweencd​cc\_\{dc\}and2​cd​c2c\_\{dc\}tokens at every length while𝐆\{\\mathbf\{G\}\}keeps absorbing the older blocks\.Algorithm 3Stepped\-window decoding for one attention head0:geometry

\(q,𝐂,B\)\(q,\{\\mathbf\{C\}\},B\)fixed from

TmaxT\_\{\\max\}; block length

cd​cc\_\{dc\}; maps

prof,type\\operatorname\{prof\},\\operatorname\{type\}
1:

U←0U\\leftarrow 0;

Δold,Δcur←0\\Delta\_\{\\mathrm\{old\}\},\\Delta\_\{\\mathrm\{cur\}\}\\leftarrow 0;

Rold,Rcur←R\_\{\\mathrm\{old\}\},R\_\{\\mathrm\{cur\}\}\\leftarrowfresh recurrent states

2:for

t=1,2,…t=1,2,\\dotsdo

3:if

t\>1t\>1and

b⁡\(t\)≠b⁡\(t−1\)b\(t\)\\neq b\(t\-1\)then

4:if

b⁡\(t\)≥2b\(t\)\\geq 2then

U←U\+𝐂​ΔoldU\\leftarrow U\+\{\\mathbf\{C\}\}\\,\\Delta\_\{\\mathrm\{old\}\}
5:

Δold←Δcur\\Delta\_\{\\mathrm\{old\}\}\\leftarrow\\Delta\_\{\\mathrm\{cur\}\};

Δcur←0\\Delta\_\{\\mathrm\{cur\}\}\\leftarrow 0;

Rold←RcurR\_\{\\mathrm\{old\}\}\\leftarrow R\_\{\\mathrm\{cur\}\};

Rcur←R\_\{\\mathrm\{cur\}\}\\leftarrowfresh

6:Compute

ϕt\\phi\_\{t\}and

𝐙t=ψt​v¯t⊤\{\\mathbf\{Z\}\}\_\{t\}=\\psi\_\{t\}\\bar\{v\}\_\{t\}^\{\\top\}; update

Rold,RcurR\_\{\\mathrm\{old\}\},R\_\{\\mathrm\{cur\}\};

Δcur​\[prof⁡\(t\)\]\+=𝐙t\\Delta\_\{\\mathrm\{cur\}\}\[\\operatorname\{prof\}\(t\)\]\\mathrel\{\+\}=\{\\mathbf\{Z\}\}\_\{t\}\.

7:

y¯t←ϕt⊤​Utype⁡\(t\)\+read⁡\(Rold,ϕt\)\\bar\{y\}\_\{t\}\\leftarrow\\phi\_\{t\}^\{\\top\}U\_\{\\operatorname\{type\}\(t\)\}\+\\operatorname\{read\}\(R\_\{\\mathrm\{old\}\},\\phi\_\{t\}\);

ot←y¯t,1:dv/y¯t,po\_\{t\}\\leftarrow\\bar\{y\}\_\{t,1:d\_\{v\}\}/\\bar\{y\}\_\{t,p\}\.

###### Lemma 5\(Horizon\-free streaming decoding\)\.

Assume the hypotheses of Theorem[3\.2](https://arxiv.org/html/2609.36062#S3.Thmthm2), with the geometry fixed fromTmaxT\_\{\\max\}andcd​c=Tmax/4c\_\{dc\}=T\_\{\\max\}/4\. Then Algorithm[3](https://arxiv.org/html/2609.36062#alg3)computes the outputs of the stepped\-window mask of Definition[1](https://arxiv.org/html/2609.36062#Thmdefinition1), for every lengthTTatO⁡\(\(1\+Tmax1−3/d\)​r​p\)O\\big\(\(1\+T\_\{\\max\}^\{1\-3/d\}\)rp\\big\)work per token from a cache ofO⁡\(Tmax1−1/d\)O\(T\_\{\\max\}^\{1\-1/d\}\)states of sizer×pr\\times p\.

###### Proof\.

Index blocks so that blockβ\\betaoccupies positionsβ​cd​c\+1,…,\(β\+1\)​cd​c\\beta c\_\{dc\}\+1,\\dots,\(\\beta\+1\)c\_\{dc\}, and writeβ⁡\(t\)=b⁡\(t\)\\beta\(t\)=b\(t\)\. For0≤n≤T0\\leq n\\leq Tlet

Fx\(n\):=∑j≤n,prof⁡\(j\)=x𝐙j,x∈𝔽qd−1,F^\{\(n\)\}\_\{x\}:=\\sum\_\{j\\leq n,\\ \\operatorname\{prof\}\(j\)=x\}\{\\mathbf\{Z\}\}\_\{j\},\\qquad x\\in\\mathbb\{F\}\_\{q\}^\{d\-1\},be the profile table of the firstnnpositions, viewed as aqd−1×r​pq^\{d\-1\}\\times rpmatrix\. Sincen⁡\(t\)=\(β⁡\(t\)−1\)​cd​cn\(t\)=\(\\beta\(t\)\-1\)c\_\{dc\}forβ⁡\(t\)≥1\\beta\(t\)\\geq 1, the distant region of a query in blockβ\\betais the union of blocks0,…,β−20,\\dots,\\beta\-2\.

We claim that for everyt≥1t\\geq 1, by line[7](https://arxiv.org/html/2609.36062#alg3.l7)in the algorithm,

1. \(i\)U=𝐂​F\(n⁡\(t\)\)U=\{\\mathbf\{C\}\}F^\{\(n\(t\)\)\};
2. \(ii\)Δold\\Delta\_\{\\mathrm\{old\}\}holds the per\-profile sums of blockβ⁡\(t\)−1\\beta\(t\)\-1, andΔcur\\Delta\_\{\\mathrm\{cur\}\}those of blockβ⁡\(t\)\\beta\(t\)restricted to positions≤t\\leq t;
3. \(iii\)RoldR\_\{\\mathrm\{old\}\}is the state of the recurrence initialized at positionn⁡\(t\)\+1n\(t\)\+1and advanced throughtt\.

We prove this by induction ontt\. Fort=1t=1we haveβ⁡\(1\)=0\\beta\(1\)=0andn⁡\(1\)=0n\(1\)=0, the conditional does not fire, andU=0=𝐂​F\(0\)U=0=\{\\mathbf\{C\}\}F^\{\(0\)\}, giving \(i\)–\(iii\)\. Assume the invariant att−1t\-1and considertt\.

Ifβ⁡\(t\)=β⁡\(t−1\)\\beta\(t\)=\\beta\(t\-1\)thenn⁡\(t\)=n⁡\(t−1\)n\(t\)=n\(t\-1\)and lines 4–5 are skipped, soUUis unchanged and \(i\) persists; line[6](https://arxiv.org/html/2609.36062#alg3.l6)adds𝐙t\{\\mathbf\{Z\}\}\_\{t\}toΔcur​\[prof⁡\(t\)\]\\Delta\_\{\\mathrm\{cur\}\}\[\\operatorname\{prof\}\(t\)\]and advances both recurrences, which preserves \(ii\) and \(iii\)\.

Ifβ⁡\(t\)≠β⁡\(t−1\)\\beta\(t\)\\neq\\beta\(t\-1\)thent=β​cd​c\+1t=\\beta c\_\{dc\}\+1andt−1t\-1is the last position of blockβ−1\\beta\-1\. Forβ≤1\\beta\\leq 1we haven⁡\(t\)=0n\(t\)=0, the guard on line[4](https://arxiv.org/html/2609.36062#alg3.l4)suppresses the update, and \(i\) holds withU=0U=0\. Forβ≥2\\beta\\geq 2, the inductive hypothesis att−1t\-1givesU=𝐂​F\(n⁡\(t−1\)\)=𝐂​F\(\(β−2\)​cd​c\)U=\{\\mathbf\{C\}\}F^\{\(n\(t\-1\)\)\}=\{\\mathbf\{C\}\}F^\{\(\(\\beta\-2\)c\_\{dc\}\)\}and, by \(ii\),Δold\\Delta\_\{\\mathrm\{old\}\}equal to the per\-profile sums of blockβ−2\\beta\-2\. Those sums are preciselyF\(\(β−1\)​cd​c\)−F\(\(β−2\)​cd​c\)F^\{\(\(\\beta\-1\)c\_\{dc\}\)\}\-F^\{\(\(\\beta\-2\)c\_\{dc\}\)\}, so line[4](https://arxiv.org/html/2609.36062#alg3.l4)yields

U=𝐂​F\(\(β−2\)​cd​c\)\+𝐂⁡\(F\(\(β−1\)​cd​c\)−F\(\(β−2\)​cd​c\)\)=𝐂​F\(\(β−1\)​cd​c\)=𝐂​F\(n⁡\(t\)\),U\\;=\\;\{\\mathbf\{C\}\}F^\{\(\(\\beta\-2\)c\_\{dc\}\)\}\+\{\\mathbf\{C\}\}\\Big\(F^\{\(\(\\beta\-1\)c\_\{dc\}\)\}\-F^\{\(\(\\beta\-2\)c\_\{dc\}\)\}\\Big\)\\;=\\;\{\\mathbf\{C\}\}F^\{\(\(\\beta\-1\)c\_\{dc\}\)\}\\;=\\;\{\\mathbf\{C\}\}F^\{\(n\(t\)\)\},which shows \(i\)\. Line[5](https://arxiv.org/html/2609.36062#alg3.l5)then setsΔold\\Delta\_\{\\mathrm\{old\}\}to the sums of blockβ−1\\beta\-1and clearsΔcur\\Delta\_\{\\mathrm\{cur\}\}, and line[6](https://arxiv.org/html/2609.36062#alg3.l6)adds𝐙t\{\\mathbf\{Z\}\}\_\{t\}, giving \(ii\); and it reassignsRold←RcurR\_\{\\mathrm\{old\}\}\\leftarrow R\_\{\\mathrm\{cur\}\}, whereRcurR\_\{\\mathrm\{cur\}\}was initialized at the first position of blockβ−1\\beta\-1, namely\(β−1\)​cd​c\+1=n⁡\(t\)\+1\(\\beta\-1\)c\_\{dc\}\+1=n\(t\)\+1, giving \(iii\)\.

*Exactness\.*The augmented output is linear in the mask row, so by Definition[1](https://arxiv.org/html/2609.36062#Thmdefinition1),

y¯t=∑j𝐌t​j\(step\)​ϕt⊤​𝐙j=ϕt⊤​\(∑j≤n⁡\(t\)𝐂type⁡\(t\),prof⁡\(j\)​𝐙j⏟distant\+∑n⁡\(t\)<j≤t𝐙j⏟recent\)\.\\bar\{y\}\_\{t\}=\\sum\_\{j\}\{\\mathbf\{M\}\}^\{\(\\mathrm\{step\}\)\}\_\{tj\}\\,\\phi\_\{t\}^\{\\top\}\{\\mathbf\{Z\}\}\_\{j\}=\\phi\_\{t\}^\{\\top\}\\Big\(\\underbrace\{\\sum\_\{j\\leq n\(t\)\}\{\\mathbf\{C\}\}\_\{\\operatorname\{type\}\(t\),\\operatorname\{prof\}\(j\)\}\{\\mathbf\{Z\}\}\_\{j\}\}\_\{\\text\{distant\}\}\+\\underbrace\{\\sum\_\{n\(t\)<j\\leq t\}\{\\mathbf\{Z\}\}\_\{j\}\}\_\{\\text\{recent\}\}\\Big\)\.\(21\)Grouping the distant keys by profile gives∑j≤n𝐂h,prof⁡\(j\)​𝐙j=\(𝐂​F\(n\)\)h\\sum\_\{j\\leq n\}\{\\mathbf\{C\}\}\_\{h,\\operatorname\{prof\}\(j\)\}\{\\mathbf\{Z\}\}\_\{j\}=\(\{\\mathbf\{C\}\}F^\{\(n\)\}\)\_\{h\}, so by invariant \(i\) the first term of equation[21](https://arxiv.org/html/2609.36062#A6.E21)equalsϕt⊤​Utype⁡\(t\)\\phi\_\{t\}^\{\\top\}U\_\{\\operatorname\{type\}\(t\)\}\. The recurrence is linear and its state is initialized to zero atn⁡\(t\)\+1n\(t\)\+1, so by invariant \(iii\) it holds∑n⁡\(t\)<j≤t𝐙j\\sum\_\{n\(t\)<j\\leq t\}\{\\mathbf\{Z\}\}\_\{j\}and the second term equalsread⁡\(Rold,ϕt\)\\operatorname\{read\}\(R\_\{\\mathrm\{old\}\},\\phi\_\{t\}\)\. Both are accumulated intoy¯t\\bar\{y\}\_\{t\}before the single division on line[7](https://arxiv.org/html/2609.36062#alg3.l7), so the normalized output is exact\.

*Gating\.*For the gated mask𝐌~t​j=𝐌t​j\(step\)​exp⁡\(λt−λj\)\\widetilde\{\{\\mathbf\{M\}\}\}\_\{tj\}=\{\\mathbf\{M\}\}^\{\(\\mathrm\{step\}\)\}\_\{tj\}\\exp\(\\lambda\_\{t\}\-\\lambda\_\{j\}\), replaceF\(n\)F^\{\(n\)\}by the table referred to the last distant position,F~x\(n\):=∑j≤n,prof⁡\(j\)=xexp⁡\(λn−λj\)​𝐙j\\widetilde\{F\}^\{\(n\)\}\_\{x\}:=\\sum\_\{j\\leq n,\\ \\operatorname\{prof\}\(j\)=x\}\\exp\(\\lambda\_\{n\}\-\\lambda\_\{j\}\)\{\\mathbf\{Z\}\}\_\{j\}, and readUUasexp⁡\(λt−λn⁡\(t\)\)​ϕt⊤​Utype⁡\(t\)\\exp\(\\lambda\_\{t\}\-\\lambda\_\{n\(t\)\}\)\\,\\phi\_\{t\}^\{\\top\}U\_\{\\operatorname\{type\}\(t\)\}\. WritingN=\(β−2\)​cd​cN=\(\\beta\-2\)c\_\{dc\}andN′=\(β−1\)​cd​cN^\{\\prime\}=\(\\beta\-1\)c\_\{dc\}, the two tables satisfy

F~\(N′\)=exp⁡\(λN′−λN\)​F~\(N\)\+\(per\-profile sums of block​β−2​referred to​N′\),\\widetilde\{F\}^\{\(N^\{\\prime\}\)\}=\\exp\(\\lambda\_\{N^\{\\prime\}\}\-\\lambda\_\{N\}\)\\,\\widetilde\{F\}^\{\(N\)\}\+\\Big\(\\text\{per\-profile sums of block \}\\beta\-2\\text\{ referred to \}N^\{\\prime\}\\Big\),so invariant \(i\) is restored with line[4](https://arxiv.org/html/2609.36062#alg3.l4)replaced byU←exp⁡\(λN′−λN\)​U\+𝐂​ΔoldU\\leftarrow\\exp\(\\lambda\_\{N^\{\\prime\}\}\-\\lambda\_\{N\}\)U\+\{\\mathbf\{C\}\}\\,\\Delta\_\{\\mathrm\{old\}\}, one scalar multiply per block\. The induction and the exactness argument are otherwise unchanged\.

*Cost\.*Per token, formingϕt\\phi\_\{t\}and𝐙t\{\\mathbf\{Z\}\}\_\{t\}, the two recurrent updates, the scatter\-add intoΔcur\\Delta\_\{\\mathrm\{cur\}\}, and the read and normalization on line[7](https://arxiv.org/html/2609.36062#alg3.l7)are eachO⁡\(r​p\)O\(rp\)\. The only other work is line[4](https://arxiv.org/html/2609.36062#alg3.l4), one sparse product costingnnz⁡\(𝐂\)​r​p\\operatorname\{nnz\}\(\{\\mathbf\{C\}\}\)\\,rpscalar operations once everycd​cc\_\{dc\}tokens\. As written, line[4](https://arxiv.org/html/2609.36062#alg3.l4)performs the sparse product at the first token of each block, so the algorithm meets the bound below in amortized form\. For a worst\-case bound, the next type tableU\+𝐂​ΔoldU\+\{\\mathbf\{C\}\}\\,\\Delta\_\{\\mathrm\{old\}\}is built in a second buffer during the preceding block: the block it absorbs is complete for allcd​cc\_\{dc\}steps of that block, so copyingUU\(B​r​pB\\,rpoperations\) and the sparse product \(nnz⁡\(𝐂\)​r​p\\operatorname\{nnz\}\(\{\\mathbf\{C\}\}\)\\,rp\) can be spread over those steps, and the buffers are swapped only once the new table is complete\. SinceB≤nnz⁡\(𝐂\)B\\leq\\operatorname\{nnz\}\(\{\\mathbf\{C\}\}\), this addsO⁡\(\(1\+nnz⁡\(𝐂\)/cd​c\)​r​p\)O\\big\(\(1\+\\operatorname\{nnz\}\(\{\\mathbf\{C\}\}\)/c\_\{dc\}\)\\,rp\\big\)work to each token\. In the regular construction, the number of normalized directions is\|𝒜\|=\(qd−1−1\)/\(q−1\)=Θ⁡\(qd−2\)\|\\mathcal\{A\}\|=\(q^\{d\-1\}\-1\)/\(q\-1\)=\\Theta\(q^\{d\-2\}\)and each profile lies on one hyperplane per direction, so

nnz⁡\(𝐂\)=\|𝒜\|​qd−1=Θ⁡\(q2​d−3\)\.\\operatorname\{nnz\}\(\{\\mathbf\{C\}\}\)=\|\\mathcal\{A\}\|\\,q^\{d\-1\}=\\Theta\(q^\{2d\-3\}\)\.WithTmax=Θ⁡\(qd\)T\_\{\\max\}=\\Theta\(q^\{d\}\)this isnnz⁡\(𝐂\)=Θ⁡\(Tmax2−3/d\)\\operatorname\{nnz\}\(\{\\mathbf\{C\}\}\)=\\Theta\(T\_\{\\max\}^\{2\-3/d\}\), and sincecd​c=Θ⁡\(Tmax\)c\_\{dc\}=\\Theta\(T\_\{\\max\}\)the per\-token share isΘ⁡\(Tmax1−3/d\)\\Theta\(T\_\{\\max\}^\{1\-3/d\}\), givingO⁡\(\(1\+Tmax1−3/d\)​r​p\)O\\big\(\(1\+T\_\{\\max\}^\{1\-3/d\}\)rp\\big\), which isO⁡\(r​p\)O\(rp\)exactly whend≤3d\\leq 3\.

*Cache\.*The retained state is the type tableUU\(BBstates\), the two profile deltas \(qd−1q^\{d\-1\}states each\), the two recurrent states, and the second type table used above\. SinceB=Θ⁡\(qd−1\)=Θ⁡\(Tmax1−1/d\)B=\\Theta\(q^\{d\-1\}\)=\\Theta\(T\_\{\\max\}^\{1\-1/d\}\), the total isO⁡\(Tmax1−1/d\)O\(T\_\{\\max\}^\{1\-1/d\}\)states of sizer×pr\\times p\. ∎

#### Other dynamic boundary candidates\.

Writen⁡\(t\)n\(t\)for the boundary each schedule places at querytt, andTtrainT\_\{\\mathrm\{train\}\}for the training length\. We compare six boundary schedules\.

1. 1\.*No long\-range branch*setsn⁡\(t\)=0n\(t\)=0, so𝐆\{\\mathbf\{G\}\}is never used and the layer is the ordinary recurrence, thed=1d=1SMat member\.
2. 2\.*Boundary kept, current*is the mask of the main text withn⁡\(t\)=Ttrain/2n\(t\)=T\_\{\\mathrm\{train\}\}/2held where prefill placed it, which is what the fixed\-horizon decoder does when it is run past its horizon\.
3. 3\.*Rebuilt, current*is the same mask recomputed at the evaluation length,n⁡\(t\)=T/2n\(t\)=T/2\. It is computationally expensive and needsTTin advance\.
4. 4\.*Doubling*places the boundary atn⁡\(t\)=2⌊log2⁡t⌋−1n\(t\)=2^\{\\lfloor\\log\_\{2\}t\\rfloor\-1\}, so it doubles at each power of two\.
5. 5\.The two*window*arms are definition[1](https://arxiv.org/html/2609.36062#Thmdefinition1)withcd​c=Ttrain/4c\_\{dc\}=T\_\{\\mathrm\{train\}\}/4andcd​c=Ttrain/8c\_\{dc\}=T\_\{\\mathrm\{train\}\}/8\.

All arms share architecture, data, optimizer and step budget\. We demonstrate capabilities with thed=2d=2SMat mask: a Mamba\-2\-style decayed linear recurrence reset atn⁡\(t\)n\(t\), plus the incidence branch read by a fixed random hash of token identifiers of Section[3](https://arxiv.org/html/2609.36062#S3)\. Each result is the mean over three seeds, with the sample standard deviation as a subscript\.

#### Task and evaluation layouts\.

Each MQAR sequence contains6464key–value pairs over a noise vocabulary\. Every key occurs twice, once beside its value and once as a query, and the model must emit that value at the query\. During training the pair and query positions are drawn uniformly over the sequence\. At evaluation we use two layouts\. Under*uniform*the positions are drawn as in training, so the result measures length generalization alone\. Under*placed*all pairs are confined to a256256\-token region and all queries to a256256\-token region at the very end of the sequence, separated by a gap ofDDtokens;DDcontrols how far back a query must reach, with largeDDputting the pairs near the start of the context\. The placed layout moves the stored content across each schedule’s boundary while holding everything else fixed\.

#### Length generalization\.

Table[5](https://arxiv.org/html/2609.36062#A6.T5)evaluates each schedule out to eight times its training length under the training distribution\.

Table 5:MQAR accuracy \(%\) by evaluation length,*uniform*layout, models trained at length 1024\. The quarter\-length window is the only schedule whose accuracy rises with the evaluation length, and the only one whose spread does not\. KnowingTTin advance recovers part but not all of the fixed boundary’s loss: the rebuilt oracle ends at 81\.8 against 51\.9 for the kept boundary\. Mean \(std\) over 3 seeds\.
#### Recall relative to the boundary\.

Table[6](https://arxiv.org/html/2609.36062#A6.T6)fixes the evaluation length at81928192and sweeps the gapDD\. The final column,D=7678D=7678, places the pairs at the very start of the sequence, where every schedule routes them through𝐆\{\\mathbf\{G\}\}; it is a control rather than a hard case\.

Table 6:MQAR accuracy \(%\) by gapDDbetween the pair region and the query region,*placed*layout at evaluation length81928192\. The lower block reports how often the pairs actually reach the long\-range branch\. AtD=7678D=7678the pairs precede every boundary and the fixed schedules are the strongest arms\. For256≤D≤3072256\\leq D\\leq 3072the pairs land inside the fixed boundary’s recent region, which it can serve only from the fixed\-size recurrence, and its accuracy falls to about 34% while the quarter\-length window holds about 95%\. AtD=0D=0, the quarter\-length window obtains27\.1%27\.1\\%, below the two fixed\-boundary variants \(35\.4%35\.4\\%and36\.1%36\.1\\%\)\. The eighth\-length window instead obtains73\.4%73\.4\\%, with substantial seed variation\. Performance near the boundary therefore depends on the particular window schedule\. Mean \(std\) over 3 seeds\.
#### Language modelling\.

Table[7](https://arxiv.org/html/2609.36062#A6.T7)reports byte\-level PG\-19, trained at length 2048 and evaluated on1638416384\-byte windows\.

Table 7:PG\-19 bits per byte by position within the evaluation window, trained at length 2048 and evaluated at1638416384; lower is better\. Every schedule improves past the training length rather than degrading, and the quarter\-length window is similar to the fixed boundary\. Mean \(std\) over 3 seeds\.

## Appendix GTraining Details and additional results

All runs use an Adam\-family optimizer, gradient\-norm clipping at1\.01\.0, and a single GPU per run\. Learning\-rate schedules are linear warmup followed by cosine decay to0\.1×0\.1\\timesthe peak, except on subset routing and multi\-key subset recall, which train at a constant rate\. Weight decay is applied to matrix parameters only; biases, norms, and the recurrent parametersAlogA\_\{\\log\},Δbias\\Delta\_\{\\mathrm\{bias\}\}, andDDare excluded\. Every arm within a table shares its data order, schedule, and precision\.

Table 8:Training configuration for each experiment\. Routing trains with the mask fixed and optimizes a mean\-squared\-error readout; MKAR uses binary cross\-entropy on one sigmoid logit per pair \(multi\-label\), and MQAR, JCKR and PG\-19 use categorical \(softmax\) cross\-entropy\. PG\-19 batch sizes are per context length, chosen so that every step sees 32,768 tokens, giving 300M training tokens per arm\.### G\.1Extension Details

We describe one attention head and suppress the head index\. Lethth\_\{t\}denote the layer input at positiontt, which depends only on tokens at positions1,…,t1,\\ldots,t\. The distant/recent boundarynnand the finite geometry are fixed independently of token content\.

Causal write addresses\.The write\-side representation is a short causal convolution,

ct=∑ℓ=0L−1Kℓ​ht−ℓ,c\_\{t\}=\\sum\_\{\\ell=0\}^\{L\-1\}K\_\{\\ell\}h\_\{t\-\\ell\},with zero padding before the sequence begins; our implementation uses a depthwise convolution withL=4L=4\. This allows a write to be addressed using nearby preceding content, such as a key preceding its value\. Write addresses are computed when tokens are processed and are not revised using later queries\. LetD=d−1D=d\-1\. For each coordinatek∈\[D\]k\\in\[D\], we compute

zt,k=γk​⟨wk‖wk‖2,LN⁡\(ct\)⟩\+bk,z\_\{t,k\}=\\gamma\_\{k\}\\left\\langle\\frac\{w\_\{k\}\}\{\\\|w\_\{k\}\\\|\_\{2\}\},\\operatorname\{LN\}\(c\_\{t\}\)\\right\\rangle\+b\_\{k\},and

st,k=clip⁡\(q​Φ​\(zt,k\),0,q−ε\),s\_\{t,k\}=\\operatorname\{clip\}\\bigl\(q\\Phi\(z\_\{t,k\}\),\\,0,\\,q\-\\varepsilon\\bigr\),whereLN\\operatorname\{LN\}normalizes each token independently,Φ\\Phiis the standard normal CDF, andε\>0\\varepsilon\>0is a small numerical constant\. The discrete hash is

xθ​\(ct\)=\(⌊st,1⌋,…,⌊st,D⌋\)∈𝔽qD\.x\_\{\\theta\}\(c\_\{t\}\)=\\bigl\(\\lfloor s\_\{t,1\}\\rfloor,\\ldots,\\lfloor s\_\{t,D\}\\rfloor\\bigr\)\\in\\mathbb\{F\}\_\{q\}^\{D\}\.Thus,prof⁡\(j\)=xθ​\(cj\)\\operatorname\{prof\}\(j\)=x\_\{\\theta\}\(c\_\{j\}\)determines the profile memory updated by distant tokenjj\. The projection parameters\(wk,γk,bk\)\(w\_\{k\},\\gamma\_\{k\},b\_\{k\}\)are learned jointly with the backbone; the geometric incidence matrixCCremains fixed\.

Learning the discrete hash\.The floor operation has no useful ordinary derivative\. Therefore, following[Jiang et al\. \[2018\]](https://arxiv.org/html/2609.36062#bib.bib8), we use a straight\-through estimator: the forward pass uses a discrete address, while the backward pass uses a surrogate based on interpolation between neighboring bins\. Write

at,k=⌊st,k⌋,ft,k=st,k−at,k\.a\_\{t,k\}=\\lfloor s\_\{t,k\}\\rfloor,\\qquad f\_\{t,k\}=s\_\{t,k\}\-a\_\{t,k\}\.The corresponding interpolation assigns weights1−ft,k1\-f\_\{t,k\}andft,kf\_\{t,k\}to binsat,ka\_\{t,k\}andmin⁡\(at,k\+1,q−1\)\\min\(a\_\{t,k\}\+1,q\-1\), respectively\. In the sparse four\-read implementation, we retain only the selected write route\. Its straight\-through weight is

ωt=∏k=1D\[1\+\(1−ft,k\)−sg⁡\(1−ft,k\)\],\\omega\_\{t\}=\\prod\_\{k=1\}^\{D\}\\left\[1\+\(1\-f\_\{t,k\}\)\-\\operatorname\{sg\}\(1\-f\_\{t,k\}\)\\right\],wheresg\\operatorname\{sg\}denotes stop\-gradient\. Numerically,ωt=1\\omega\_\{t\}=1, so each token writes to exactly one profile during both training and inference\. During backpropagation, the write weight supplies a surrogate gradient to the hash parameters while the discrete address is held fixed\. This is a one\-sided, biased straight\-through estimator, rather than differentiation through the discrete bin index\.

We also encourage balanced coordinate occupancy\. Letp¯k,b\\bar\{p\}\_\{k,b\}be the average soft interpolation mass assigned to binbbof coordinatekkover distant write positions\. The auxiliary loss is

ℒbal=1D​∑k=1D∑b=0q−1p¯k,b​log⁡\(q​p¯k,b\),ℒ=ℒtask\+βbal​ℒbal\.\\mathcal\{L\}\_\{\\mathrm\{bal\}\}=\\frac\{1\}\{D\}\\sum\_\{k=1\}^\{D\}\\sum\_\{b=0\}^\{q\-1\}\\bar\{p\}\_\{k,b\}\\log\\bigl\(q\\bar\{p\}\_\{k,b\}\\bigr\),\\qquad\\mathcal\{L\}=\\mathcal\{L\}\_\{\\mathrm\{task\}\}\+\\beta\_\{\\mathrm\{bal\}\}\\mathcal\{L\}\_\{\\mathrm\{bal\}\}\.This penalizes concentration in individual coordinate bins; it does not require uniform occupancy of all joint profiles\.

Query routing and causality\.After processing the distant block, profile memories are aggregated into type summaries,

Uh=∑xCh​x​Fx\.U\_\{h\}=\\sum\_\{x\}C\_\{hx\}F\_\{x\}\.For the single\-hyperplane construction, a recent query att=n\+it=n\+ihashes its causal representation to a pointutu\_\{t\}and chooses a nonzero normalized directionata\_\{t\}using that same representation\. Its type isHat,at⊤​utH\_\{a\_\{t\},a\_\{t\}^\{\\top\}u\_\{t\}\}, with arithmetic over𝔽q\\mathbb\{F\}\_\{q\}\. Consequently, its long\-range support is

Gi​j=𝟏\{an\+i⊤prof\(j\)=an\+i⊤un\+i\},j≤n\.G\_\{ij\}=\\mathbf\{1\}\\left\\\{a\_\{n\+i\}^\{\\top\}\\operatorname\{prof\}\(j\)=a\_\{n\+i\}^\{\\top\}u\_\{n\+i\}\\right\\\},\\qquad j\\leq n\.This guarantees visibility for matching hash points\. Repeated token identities alone need not produce matching points when the hash inputs include contextual information\.

The four\-read variant instead uses an independent learned scorerℓi,h\\ell\_\{i,h\}computed fromhn\+ih\_\{n\+i\}\. Let𝒮i=Top4h⁡\(ℓi,h\)\\mathcal\{S\}\_\{i\}=\\operatorname\{Top4\}\_\{h\}\(\\ell\_\{i,h\}\)\. Its read weights are

Ri​h=\{exp⁡\(ℓi,h\)∑g∈𝒮iexp⁡\(ℓi,g\),h∈𝒮i,0,otherwise\.R\_\{ih\}=\\begin\{cases\}\\displaystyle\\frac\{\\exp\(\\ell\_\{i,h\}\)\}\{\\sum\_\{g\\in\\mathcal\{S\}\_\{i\}\}\\exp\(\\ell\_\{i,g\}\)\},&h\\in\\mathcal\{S\}\_\{i\},\\\\\[6\.0pt\] 0,&\\text\{otherwise\}\.\\end\{cases\}The selected scores receive ordinary gradients through this softmax; we do not differentiate through the top\-four indices\. Unlike the single\-hyperplane read, this selector need not choose a hyperplane through the query’s own hash point\.

Both variants are causal: every summary contains only positionsj≤n<n\+ij\\leq n<n\+i, and the query’s routing decision depends only on its causal representation\. Combined with the causal local branch, the output at positionttdepends only on tokens at positions at mosttt\.

#### Memory updates\.

The memory updates are adapted from the recurrent dynamics of the corresponding backbone, with new payloads routed to learned profile addresses\.

For Mamba\-2 extensions, we use an additive memory update:Fx\(j\)=ajGFx\(j−1\)\+𝟏\{xj=x\}ωjkjvj⊤F\_\{x\}^\{\(j\)\}=a\_\{j\}^\{G\}F\_\{x\}^\{\(j\-1\)\}\+\\mathbf\{1\}\\\{x\_\{j\}=x\\\}\\omega\_\{j\}k\_\{j\}v\_\{j\}^\{\\top\}, wherexjx\_\{j\}is the learned write profile andωj\\omega\_\{j\}is an independent sigmoid write gate\. We reuse the backbone’s input\-dependent retention factorajG=exp⁡\(−exp⁡\(θA\)​Δj\)a\_\{j\}^\{G\}=\\exp\(\-\\exp\(\\theta\_\{A\}\)\\Delta\_\{j\}\), whereθA\\theta\_\{A\}is a learned parameter andΔj\>0\\Delta\_\{j\}\>0is the input\-dependent time step\.

For GDN extensions, we use boundary\-transported additive memory, incorporating the backbone’s delta\-rule transitions\. Letkt,qtk\_\{t\},q\_\{t\}be the normalized backbone keys and queries,vtv\_\{t\}the values,βt\\beta\_\{t\}the write gate, andata\_\{t\}the scalar decay\. DefineAt=at​\(I−βt​kt​kt⊤\)A\_\{t\}=a\_\{t\}\(I\-\\beta\_\{t\}k\_\{t\}k\_\{t\}^\{\\top\}\)\. For distant tokenj≤nj\\leq nwith learned profilexjx\_\{j\}, the boundary\-transported key is

k¯j=An⋯Aj\+1βjkj,Fx=∑j≤n:xj=xk¯jvj⊤,\\bar\{k\}\_\{j\}=A\_\{n\}\\cdots A\_\{j\+1\}\\beta\_\{j\}k\_\{j\},\\qquad F\_\{x\}=\\sum\_\{j\\leq n:\\,x\_\{j\}=x\}\\bar\{k\}\_\{j\}v\_\{j\}^\{\\top\},where an empty product is the identity\. Equivalently, starting from zero,Fx\(j\)=AjFx\(j−1\)\+𝟏\{xj=x\}βjkjvj⊤F\_\{x\}^\{\(j\)\}=A\_\{j\}F\_\{x\}^\{\(j\-1\)\}\+\\mathbf\{1\}\\\{x\_\{j\}=x\\\}\\beta\_\{j\}k\_\{j\}v\_\{j\}^\{\\top\}\. Thus every profile undergoes the shared transitionAjA\_\{j\}, while only the selected profile receives the new payload\. The implementation computes transported keys and pools them additively, rather than explicitly updating every profile\.

After formingUb=∑xCb​x​FxU\_\{b\}=\\sum\_\{x\}C\_\{bx\}F\_\{x\}, a recent tokent\>nt\>nuses

q¯t=r−1/2\(At⋯An\+1\)⊤qt,zt=ztlocal\+λtq¯t⊤∑bRt​bUb\.\\bar\{q\}\_\{t\}=r^\{\-1/2\}\(A\_\{t\}\\cdots A\_\{n\+1\}\)^\{\\top\}q\_\{t\},\\qquad z\_\{t\}=z\_\{t\}^\{\\mathrm\{local\}\}\+\\lambda\_\{t\}\\bar\{q\}\_\{t\}^\{\\top\}\\sum\_\{b\}R\_\{tb\}U\_\{b\}\.HereRt​bR\_\{tb\}contains the four selected softmax read weights andλt\\lambda\_\{t\}is a learned sigmoid gate\.

### G\.2Experiments

Exact\-support match\.Both subset routing and Multi\-key subset recall tasks are scored by whether the model retrieves exactly the requested items\. A query counts as correct only if every requested item is present in the output and every unrequested item is absent\. In subset routing the output has one channel per marked position, and a channel counts as present when its magnitude exceeds1%1\\%of the largest payload in that example\. In multi\-key subset recall the output has one raw logitziz\_\{i\}per key–payload pairi∈\{1,…,8\}i\\in\\\{1,\\dots,8\\\}, with no sigmoid or softmax applied before thresholding\. A pair counts as present whenzi\>0\.5z\_\{i\}\>0\.5, i\.e\. when its sigmoid probability exceedsσ⁡\(0\.5\)≈0\.62\\sigma\(0\.5\)\\approx 0\.62, and the query is correct when\{i:zi\>0\.5\}\\\{i:z\_\{i\}\>0\.5\\\}equals the requested set\. Training uses the matching per\-pair binary cross\-entropy,−18∑i\[yilogσ\(zi\)\+\(1−yi\)log\(1−σ\(zi\)\)\]\-\\frac\{1\}\{8\}\\sum\_\{i\}\\big\[y\_\{i\}\\log\\sigma\(z\_\{i\}\)\+\(1\-y\_\{i\}\)\\log\(1\-\\sigma\(z\_\{i\}\)\)\\big\]averaged over answer rows, with akk\-hot targetyy\. Each evaluation draws2,0482\{,\}048fresh queries \(8 batches of 32 sequences, 8 answer rows each\)\.

Subset routing\.Payloads and requested subsets are resampled every batch, giving roughly10510^\{5\}examples per run against at most262^\{6\}distinct requests\. Reported values are the mean of three seeds\.

Table 9:Exact routing\-pattern match \(%\) on subset routing, by the numberkkof marked positions\.T=1024T=1024, one layer, 1500 steps, batch 64, learning rate 0\.003; payloads are resampled each batch and evaluation uses fresh ones\. Mean \(std\) over 3 seeds\. SMat performs better at higherkkasddincreases\.The causal mask has VC dimension one because its row supports form a nested family, and softmax attention on the mask can assign substantially different weights to keys within an allowed prefix\. The baseline results in Table[9](https://arxiv.org/html/2609.36062#A7.T9)therefore describe performance under the evaluated architecture, optimization, and scoring protocol\.

#### Verification of end\-to\-end cost\.

Figure[3](https://arxiv.org/html/2609.36062#S3.F3)times the prefill of a single SMat layer\. Here we measure the full cost of the models that produce these subset routing results\. Namely, we analyze training steps \(forward pass, backward pass and optimizer update\), prefill, and per\-token decoding, together with peak memory and decoding cache\. All measurements use one NVIDIA H200 GPU with PyTorch 2\.11 and Triton 3\.7\.1, random weights and random inputs \(step time depends on shapes, not on trained values\)\. We report the median of at least ten timed steps after warm\-up\.

We note some critical testing implementation details\. The SMat prefill is reported with custom Triton kernels, but SMat training was done through its Pytorch path \(we did not develop the kernels for backwards pass\)\. Softmax attention uses PyTorch SDPA, which in fp32 runs the memory\-efficient kernel\. The recurrent baselines appear twice\. The upper block runs them on optimized kernels: Mamba\-2 on the SSD kernel ofmamba\_ssm, and DeltaNet and Gated DeltaNet on the chunked kernels offlash\-linear\-attention, which agree with the reference forms to within0\.4%0\.4\\%relative error\. The DeltaNet kernel does not accept fp32 inputs, so we ran it on bf16\. Log\-Linear uses the upstream kernels\. The lower block runs the reference PyTorch implementations, as trained for Table[9](https://arxiv.org/html/2609.36062#A7.T9)\. SMat decodes with the cached decoder of Section[3\.3](https://arxiv.org/html/2609.36062#S3.SS3); softmax decodes from a preallocated KV cache, and the recurrences decode one step at a time\. Before timing, every decoder is checked against the full forward pass\. The largest relative error in fp32 is5\.3×10−45\.3\\times 10^\{\-4\}, and the SMat decoder reproduces the prefill rows to10−1610^\{\-16\}in fp64\.

Train \(ms/step\)Prefill \(ms\)DecodeCacheModel16K64K64K256K\(ms/token\)\(MB\)SMat \(d=1d=1\)12\.948\.111\.545\.60\.250\.5SMat \(d=2d=2\)14\.151\.611\.344\.50\.2599SMat \(d=3d=3\)25\.1107\.018\.363\.80\.25774SMat \(d=4d=4\)192\.81717\.9168\.6725\.40\.263439Gated SMat \(d=2d=2\)92\.7988\.943\.1211\.30\.3499Gated SMat \(d=3d=3\)104\.01044\.370\.9316\.90\.34774Softmax \(SDPA\)204\.43147\.9836\.513410\.212\.904295Mamba\-2 \(SSD kernel\)8\.629\.17\.2–a0\.220\.3DeltaNet \(FLA, bf16\)8\.931\.711\.746\.50\.240\.3Gated DeltaNet \(FLA\)11\.341\.913\.854\.70\.310\.3Log\-Linear \(upstream\)21\.997\.822\.8OOMb––*Reference implementations, as trained for Table[9](https://arxiv.org/html/2609.36062#A7.T9)*Mamba\-266\.0683\.036\.3175\.20\.220\.3DeltaNet250\.73230\.485\.4396\.40\.240\.3Gated DeltaNet279\.03377\.7121\.7534\.40\.310\.3Log\-LinearOOMc–340\.7d–––Table 10:End\-to\-end cost of the subset\-routing models: batch 64, fp32\. SMat prefill uses the Triton kernels\. Decode latency and cache size are at context length 256K\. Optimized and reference recurrences share the same one\-step decoder\.aThe SSD kernel exceeds a CUDA launch limit at batch 64 and length 256K\.bThe upstream port builds a denseT×TT\\times Tlevel table \(512 GiB at 256K\)\.cThe reference Log\-Linear builds the same dense table and runs out of memory in tr aining at 16K\.dValue at 16K\.Table[10](https://arxiv.org/html/2609.36062#A7.T10)sheds light on SMat’s improvements over other attention variants\. Up tod=3d=3, the hard route SMat is linear in time from end to end\. From 16K to 64K its training step grows by3\.73\.7–4\.3×4\.3\\times,while softmax attention’s grows by15\.4×15\.4\\times\. Prefill grows by3\.53\.5–4\.0×4\.0\\timesfrom 64K to 256K, against16×16\\timesfor softmax\. Per\-token decoding is flat at0\.250\.25ms from 1K to 256K, and0\.220\.22ms at batch 1\. Against softmax attention this means2929–65×65\\timesfaster training steps at 64K,210210–300×300\\timesfaster prefill at 256K,51×51\\timesfaster decoding at 256K, and a cache5\.5×5\.5\\times\(d=3d=3\) to43×43\\times\(d=2d=2\) smaller than the KV cache\. Against optimized linear\-time kernels, SMat’s Triton prefill is comparable: at 256K it matches DeltaNet and is faster than Gated DeltaNet\. SMat training runs on unfused PyTorch operations and is1\.21\.2–1\.8×1\.8\\timesslower than those kernels at 64K ford≤2d\\leq 2, and2\.62\.6–3\.7×3\.7\\timesslower ford=3d=3\. It remains faster than the upstream Log\-Linear kernels\. SMat’s decoding cache holdsB\+1B\+1summaries, so it exceeds the constant\-size recurrent state\. At short lengths it can also exceed the KV cache: at 1K,d=3d=3uses 36 MB against 17 MB for softmax, and the order reverses by 4K\. Atd=4d=4theO⁡\(T2−3/d\)O\(T^\{2\-3/d\}\)long\-range term dominates at these lengths\. SMat is then slower than the linear\-time kernels, but still1\.8×1\.8\\timesfaster than softmax in training at 64K and18×18\\timesfaster in prefill at 256K\. Gating adds only linear work \(Lemma[4](https://arxiv.org/html/2609.36062#Thmlemma4)\)\. Our gated causal scan, however, is an unfused PyTorch loop over chunks, so gated training is about20×20\\timesslower than ungated at 64K, while gated prefill and decoding stay within2\.52\.5–5×5\\timesand1\.4×1\.4\\timesof ungated\.

Multi\-key subset recall\.Keys and payloads are placed at random positions in the distant block, so the task cannot be solved positionally\. The hyperplane read is used throughout, and the reported metric is exact\-support match: a query counts as correct only when the retrieved support equals the requested key set\.*LSH bucketing*\[[Kitaev et al\., 2020](https://arxiv.org/html/2609.36062#bib.bib38)\]hashes keys into the same cells, but a query reads only its own cell rather than a hyperplane of cells\.*Memory\-matched linear attention*tests whether SMat’s advantage is simply the larger state it keeps, since associative recall in efficient models is known to be limited by state size\[[Arora et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib1)\]\.

Table 11:Per\-seed exact\-support accuracy \(%\) behind Table[1](https://arxiv.org/html/2609.36062#S4.T1), after 12K steps on 2,048 fresh queries per run\. SMat rows use rule\-chosen hyperplane directions \(see text\)\.To further examine these results, a network that ignores the query does best by predicting the base ratek/8k/8for every pair\. That constant predictor has lossH⁡\(k/8\)=0\.377H\(k/8\)=0\.377,0\.5620\.562and0\.6620\.662nats fork=1,2,3k=1,2,3, and its logitlog⁡k8−k<0\.5\\log\\frac\{k\}\{8\-k\}<0\.5, so it marks every pair absent\. More generally, a constant output is all\-absent or all\-present under any threshold and therefore never equals akk\-hot target, so it scores exactly00whether the threshold is placed on logits or on probabilities\. All5454linear\-baseline runs end within0\.00120\.0012nats ofH⁡\(k/8\)H\(k/8\)\(Figure[7](https://arxiv.org/html/2609.36062#A7.F7)\)\. Their zeros therefore come from models that never began to use the query\. For reference, guessingkkpairs uniformly at random scores1/\(8k\)=12\.5%1/\\binom\{8\}\{k\}=12\.5\\%,3\.6%3\.6\\%and1\.8%1\.8\\%\.

Table[11](https://arxiv.org/html/2609.36062#A7.T11)lists every run behind Table[1](https://arxiv.org/html/2609.36062#S4.T1), and Figure[7](https://arxiv.org/html/2609.36062#A7.F7)shows the training loss and exact\-support accuracy of each seed\. Softmax outcomes are bimodal\. Each run starts on the constant\-prior plateau with some runs leaving after66K and1010K steps and reaching≥99%\\geq 99\\%by the next evaluation\. Seed 1 never leaves the plateau atk=2,3k=2,3; atk=1k=1, seeds 1 and 2 had begun to leave it when training stopped\. The large standard deviations of the softmax row therefore record whether a seed escaped the plateau within the budget\. The other warmup and learning\-rate settings we tried for softmax did no better \(Table[12](https://arxiv.org/html/2609.36062#A7.T12)\)\.

It is also important to note the difference between the rule\-chosen and learned hyperplane schemas for the SMat matrix in this experiment\. In the SMat rows of Table[1](https://arxiv.org/html/2609.36062#S4.T1), the hyperplane direction at each answer row is set, during training and evaluation, by a fixed rule that picks a direction on which the hashed cells of allkkrequested keys agree whenever one exists\. The baselines have no equivalent mechanism\. This isolates what the mask can express\. When a linear head predicts the direction from the query instead \(seed 0 only\), SMat scores99\.999\.9,60\.160\.1and44\.6%44\.6\\%atd=3d=3and100\.0100\.0,74\.274\.2and59\.2%59\.2\\%atd=4d=4fork=1,2,3k=1,2,3\. Atd=2d=2the hyperplanes are single points, so there is only one hyperplane direction and the two schemas coincide\.

Table 12:Softmax on multi\-key subset recall atT=1024T=1024: exact\-support accuracy \(%\) for seeds 0 / 1 / 2 under each schedule tried\. A run that leaves the constant\-prior plateau reaches≥99%\\geq 99\\%; the runs between 1% and 25% had begun leaving it when training stopped\. The reported row is the best schedule\.Table 13:Positive control: exact\-support accuracy \(%\) fork=1k=1, seed 0, with the harness, metric, model and 12K\-step budget of Table[1](https://arxiv.org/html/2609.36062#S4.T1)and only the context length changed \(8 pairs throughout\)\.![Refer to caption](https://arxiv.org/html/2609.36062v1/figures/fig18_mkar_sanity.png)Figure 7:Multi\-key subset recall atT=1024T=1024, every seed\. Top: training binary cross\-entropy of the minibatch at each evaluation step \(log scale\); the dashed line is the constant\-prior lossH⁡\(k/8\)H\(k/8\), and all1818linear\-baseline runs per panel lie on it\. Bottom: exact\-support accuracy on2,0482\{,\}048fresh queries, evaluated every22K steps\.Table 14:Mamba\-2 SMat results across three seeds\. Mean \(std\)\.Table[14](https://arxiv.org/html/2609.36062#A7.T14)shows the results of a learned Mamba\-2–SMat on multi\-key subset recall\. The baseline Mamba\-2 achieves zero exact\-support accuracy under this training setup, whereas SMat withd=2d=2succeeds primarily atk=1k=1,d=3d=3substantially improves performance atk=2k=2, andd=4d=4achieves near\-perfect accuracy across all three settings\. This qualitatively matches the original MKAR pattern, with increasing geometric dimension enabling accurate retrieval of larger requested subsets\.

MQAR\.Run through the Zoology harness\[[Arora et al\., 2023](https://arxiv.org/html/2609.36062#bib.bib1)\]with its data pipeline unchanged: training mixtures of44,88,1616,3232, and6464key–value pairs at lengths6464–256256\(100100K examples for the44\-pair mixture,2020K for each of the others\), and10001000held\-out examples per mixture\. Evaluation batch size is3232\. The GDN backbone uses one head at width1616and two heads at widths3232and6464, with value expansion11, following[Guo et al\. \[2025\]](https://arxiv.org/html/2609.36062#bib.bib4);3232epochs is707707optimizer steps per epoch\. Log\-Linear Attention uses the authors’ released implementation with the same widths, head configuration, and value expansion as the corresponding backbone, and the same 32\-epoch budget\.

Joint context\-key recall \(JCKR\)\.Models retrieve values from shuffled context\-key\-value records, with keys shared across contexts and values sampled uniformly from 16 symbols\. Every context\-key pair is queried in random order, with answers masked and cross\-entropy applied only at query positions\. We use\(C,K\)∈\{\(1,4\),\(2,8\),\(8,16\),\(16,16\),\(32,16\)\}\(C,K\)\\in\\\{\(1,4\),\(2,8\),\(8,16\),\(16,16\),\(32,16\)\\\}, yielding sequence lengths6464–30763076, with 36000/2000/4000 train/validation/test examples per configuration\. Two\-layer models of width 64 train for 32 epochs using AdamW, learning rate3×10−33\\times 10^\{\-3\}, cosine decay, and batch size 256 \(Table[8](https://arxiv.org/html/2609.36062#A7.T8)\)\. We report final\-epoch validation accuracy averaged equally across configurations, with mean and standard deviation over three training seeds\.

Table 15:Shared\-key joint recall: final\-test accuracy averaged equally over the five evaluated binding loads\. Mean and std reported across three seeds\.Furthermore, we compare SMAT with two variants derived from Mixture\-of\-Memories \(MoM\)[Du et al\. \[2026\]](https://arxiv.org/html/2609.36062#bib.bib7), both using Gated DeltaNet updates, top\-4 routing, and no shared memory\. At each sequence length, the profile\-count variant has one independent memory matrix per SMAT profile, while the matrix\-storage variant matches the combined FP32 storage of SMAT’s profile and derived summary matrices\. All memory matrices are16×1616\\times 16\. We retain the same data splits, model width, layer and head counts, and 32\-epoch training recipe\.

\(a\) 300M training tokens

\(b\) 750M training tokens

Table 16:PG\-19 language modeling results\. Left: negative log\-likelihood \(NLL\) per token after 300M training tokens with eight layers and width 384\. Right: validation perplexity after 750M training tokens at 16K context \(seed 123\)\. Lower is better\.PG\-19\.Books are tokenized with the GPT\-2 BPE vocabulary and packed into fixed\-length windows; all arms see the same data order\. The transformer arm is attention plus a4×4\\timesMLP with eight heads, sized so that its non\-embedding parameter count is comparable to SMat’s\. An auxiliary load\-balancing term with coefficient0\.010\.01is applied to the routing hash and annealed over the first10001000steps\. Evaluation is per\-token negative log\-likelihood on the PG\-19 test split \(100100books,5×1065\\times 10^\{6\}tokens\) at the training context length\. Peak memory is roughly1515GB per arm at1616K and3030GB at3232K\.222Due to academic compute constraints, each model configuration in Table[16](https://arxiv.org/html/2609.36062#A7.T16)was trained with a single seed\.