cs.LGNov 5, 2026

Muir: Linear Attention with Sensitivity-Allocated Gating and Matched-Filter Readout

Authors: Tushaar Gangavarapu, Gil Tabak, Clemens JS Schaefer

Organizations: University of Texas at Austin · Google

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Linear attention compresses context into a fixed-size key-value state for linear-time training and constant-memory inference. Gated delta-rule models decay this state and erase its prediction for an incoming key before writing a new value; Gated DeltaNet-2 (GDN-2) generalizes this with channel-wise decay and decoupled, full-rank erase and write gates. However, we show that state updates are far less sensitive to erasure than to writes: key activation and unit normalization concentrate key energy onto a sparse subset of channels, whereas value writes scale all channels, so the erase gate can be parameterized at lower rank. Furthermore, slow-decaying state rows accumulate much larger norms than fast rows, dominating unweighted readout and forcing stronger write-time erasure to suppress interference. Guided by this analysis, we introduce Muir, which fuses all gates into a sensitivity-allocated low-rank projection, re-weights state rows with a matched-filter readout gate, and reinvests the saved parameters in a wider feed-forward network. Because standard chunkwise algorithms become numerically unstable at the large chunk sizes needed to saturate wide matrix units, we design a parallel update for stable, efficient training. Across 166M-1.34B parameters, Muir lowers perplexity by 1.8-2.0% over GDN-2 and 2.5-3.6% over Transformer, raises commonsense accuracy by up to 1.9 points, and improves extractive recall over GDN-2. At 1.34B, Muir outperforms both baselines on multi-hop and multi-needle recall, extrapolates 2,560x from its 4K training length to 10M-token streams without perplexity drift, and maintains near-100% passkey retrieval to 1M tokens with static readout. Muir achieves these gains while matching GDN-2 throughput, reaching its 100B-token validation loss 1.2x faster, and training 2.7x faster than Transformer at 32K.

Figures & tables

Figure 1: Muir outperforms GDN-2 and Transformer on language modeling and long-context throughput while improving recall over GDN-2. (a) Sensitivity-allocated low-rank gates with FFN expansion drive the bulk of the validation NLL gain over GDN-2, while dynamic matched-filter readout drives the gain on extractive recall. (b) Across 166M–1.34B166\text{M}\text{--}1.34\text{B} parameters, Muir lowers validation NLL by 18.0–20.4 millinats18.0\text{--}20.4\text{ millinats} over GDN-2 and 25.0–36.2 millinats25.0\text{--}36.2\text{ millinats} over Transformer. (c) On 88 TPU7x chips ( 0.5M0.5\text{M} tokens/step) at 1.34B1.34\text{B} with our chunkwise kernel, Muir matches GDN-2 throughput despite an additional readout gate and a 30.4%30.4\% wider FFN, and trains 1.64×1.64\times and 2.74×2.74\times faster than Transformer at 16K16\text{K} and 32K,32\text{K}, respectively.
Figure 2: Matched-filter readout and low-rank erase gates improve synthetic associative retrieval and lower optimal erase strength. (a) Full whitening worsens retrieval error relative to unit gain (GDN-2), whereas half-power scaling and the Corollary 3.8 filter reduce error and shift μ∗{\mu}^{\ast} from 0.600.60 to 0.50,0.50, with the Corollary 3.8 filter matching the static oracle. (b) On a 50/5050/50 mixture of lags ℓ∈{2,80}\ell\in\{2,80\} at μ=0.60,{\mu}=0.60, the lag-matched oracle outperforms the static oracle across all horizons, while the Proposition 3.6 filter matches the lag-matched oracle at ℓ=80\ell=80 and beats the static oracle on long and mixed horizons. (c) Low-rank erase gates rbr_{b} outperform dense gates at every bottleneck rank on long-lag recall, whereas low-rank write gates rwr_{w} degrade below dense at rw≤8r_{w}\leq 8 and recover once rw≥16.r_{w}\geq 16.
Figure 3: Muir co-designs asymmetric low-rank gate parameterization and matched-filter query readout with MXU tile alignment and iso-parameter FFN reallocation. (a) Standard residual block with a Muir token mixer. (b) Within the Muir block, xt{\bm{x}}_{t} feeds linear- Conv\mathrm{Conv} - SiLU\mathrm{SiLU} pathways for (qt,kt,vt)({\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}) alongside a single fused low-rank gate projection that produces all five gates (log⁡αt,bt,gt,wt,γt).(\log{\bm{\alpha}}_{t},{\bm{b}}_{t},{\bm{g}}_{t},{\bm{w}}_{t},{\bm{\gamma}}_{t}). (c) Asymmetric gate ranks (rf,rb,rg,rw,rq)=(64,64,128,128,128)(r_{f},r_{b},r_{g},r_{w},r_{q})=(64,64,128,128,128) sum to Rtot=512=2×256,R_{\mathrm{tot}}=512=2\times 256, which fills two 256256 -column MXU tiles in the fused down-projection and aligns (f,b)(f,b) and (g,w,q)(g,w,q) on 128128 -lane VMEM tile boundaries for batched up-projections. (d) At d=2,048,d=2{,}048, replacing the dense (b,w)(b,w) gate matrices in GDN-2 cuts gate-projection parameters by 4.5×4.5\times ; this either reduces per-layer parameters and FLOPs by 14.6%14.6\% at fixed FFN width, or frees +7.3M+7.3\text{M} parameters per layer to widen the FFN under an iso-parameter budget.
Figure 4: Boundary-pivoted chunkwise kernel eliminates 3D decay tensors, division overflow, and VPU bottlenecks. (a) Memory hierarchy of a TPU7x chip. (b) On the 256×256256\times 256 MXU, an unfused 256×128256\times 128 right-hand side V^\hat{\bm{V}} achieves only 50%50\% utilization due to zero-padded lanes, whereas the fused solve [UY]=T−1[V^K^⊙π]\begin{bmatrix}{\bm{U}}&{\bm{Y}}\end{bmatrix}={\bm{T}}^{-1}\begin{bmatrix}\hat{\bm{V}}&\hat{\bm{K}}_{\odot\pi}\end{bmatrix} achieves 100%100\% utilization. (c) At chunk size C=256,C=256, leaf size 8,8, and dk=dv=128,d_{k}=d_{v}=128, replacing naive C×C×dkC\times C\times d_{k} 3D VPU decay and sequential VPU triangular solves with boundary-pivoted MXU matrix multiplications, recursive block inversion, and a custom backward pass cuts forward-and-backward VPU compute by 24×24\times and total kernel FLOPs by 12%,12\%, routing 98%98\% of chunkwise arithmetic onto the MXU. (d) Boundary-pivoted chunkwise execution combines pre-permutation VREG fusion for q^t\hat{\bm{q}}_{t} and v^t,\hat{\bm{v}}_{t}, MXU cumulative log-decay, numerically stable VPU decay splitting, recursive block inversion of T{\bm{T}} on the MXU, and a packed MXU scan carrying S0→SC{\bm{S}}_{0}\to{\bm{S}}_{C} in VMEM (saving S0{\bm{S}}_{0} only for backward).
Figure 5: Muir consistently lowers validation NLL relative to GDN-2 and Transformer throughout 100B\bm{100\text{B}} -token pretraining. Held-out FineWeb-Edu validation token NLL difference relative to GDN-2 across 166M–1.34B166\text{M}\text{--}1.34\text{B} parameters as a function of training tokens. All three Muir variants improve over GDN-2 within the first 10B10\text{B} tokens and reach a 12–20 mnat12\text{--}20\text{ mnat} reduction by 100B100\text{B} tokens, while static and dynamic matched-filter readout provide additional reductions over Muir (no γ{\bm{\gamma}} ).
  Language ModelingZero-Shot Commonsense Reasoning5-Shot
 ModelFW-EduWiki.LMB.LMB.PIQAHella.Wino.ARC-eARC-cOBQASIQAAvg.Wino.TQA
  NLL ↓\downarrowppl ↓\downarrowppl ↓\downarrowacc ↑\uparrowacc ↑\uparrowacc n ↑\uparrowacc ↑\uparrowacc ↑\uparrowacc n ↑\uparrowacc ↑\uparrowacc ↑\uparrowacc ↑\uparrowacc ↑\uparrowacc ↑\uparrow
166MTransformer2.456931.6452.6928.4765.6735.9550.9955.3526.2820.2036.3939.9149.964.74
GDN-22.446630.8439.8930.6064.7437.0250.2856.6126.4520.4037.9240.5050.125.80
Muir (no γ{\bm{\gamma}} )2.432230.3134.8732.6865.5138.1952.0156.4827.8220.8038.9541.5552.493.77
Muir (static γ{\bm{\gamma}} )2.427329.8340.9830.5866.4337.9251.7855.3026.5420.8038.1840.9451.544.32
Table 1: Muir improves language modeling and downstream reasoning over Transformer and GDN-2 across model scale. We report FineWeb-Edu validation token NLL (FW-Edu), WikiText (Wiki.) and LAMBADA (LMB.) word-level perplexity, zero-shot accuracy across eight commonsense reasoning benchmarks and their mean (Avg.), and 5-shot WinoGrande (Wino.) and TriviaQA (TQA) accuracy after 100B100\text{B} training tokens. Best results per scale are in bold and second-best are underlined .
Figure 6: Muir improves synthetic single-needle, multi-hop, multi-needle, and real-world retrieval over GDN-2. (a) Averaged across the three single-needle tasks, all Muir variants mitigate the 4K4\text{K} degradation of GDN-2 and extrapolate to 8K,8\text{K}, where unextended RoPE Transformer degrades beyond its 4K4\text{K} training window. (b) On 55 -hop variable-binding chains, static γ{\bm{\gamma}} fails to match the varying hop lags, whereas dynamic γt{\bm{\gamma}}_{t} leads all models at every context length. (c, d) When retrieving four values bound to a single key or four distinct keys queried jointly, dynamic γt{\bm{\gamma}}_{t} matches or outperforms both GDN-2 and Transformer across 1K–4K.1\text{K}\text{--}4\text{K}. (e) Averaged across the RULER tasks (a–d) by needle depth and context length, Muir achieves the largest gains over GDN-2 on early needles ( 0–40%0\text{--}40\% depth) at 4K–8K.4\text{K}\text{--}8\text{K}. (f) On real-world document extraction at 2K,2\text{K}, Muir improves over GDN-2 on all six tasks, with dynamic γt{\bm{\gamma}}_{t} driving 69%69\% of the gain.
Figure 7: Muir extrapolates to 128K\bm{128\text{K}} -token documents and 10M\bm{10\text{M}} -token streams with the lowest perplexity at a constant 25 MB\bm{25\text{ MB}} state. (a, c) Zero-shot from 4K4\text{K} pretraining, Dynamic YaRN Transformer diverges past 8K8\text{K} and unchunked evaluation exceeds HBM at ≥ ⁣128K,\geq\!128\text{K}, while Muir (static γ{\bm{\gamma}} ) drifts upward past 32K32\text{K} ; dynamic γt{\bm{\gamma}}_{t} avoids this static-gain drift and outperforms GDN-2 across 128K128\text{K} documents and 10M10\text{M} -token streams ( 2,560×2{,}560\times training length). (b, d) After a 64K64\text{K} cooldown ( 0.3%0.3\% of pretraining), Transformer improves through 64K64\text{K} before unchunked evaluation exceeds HBM at ≥ ⁣128K,\geq\!128\text{K}, static- γ{\bm{\gamma}} drift vanishes, and all Muir variants improve monotonically below GDN-2 through 10M10\text{M} tokens ( 160×160\times beyond cooldown).
Figure 8: Static matched-filter readout extends passkey retrieval to 1M\bm{1\text{M}} tokens both zero-shot and after a 64K\bm{64\text{K}} cooldown. S-NIAH-1 accuracy across 1010 needle-depth bins and context lengths 1K–1M1\text{K}\text{--}1\text{M} ( 256×256\times pretraining length) at 1.34B1.34\text{B} scale under (top) zero-shot evaluation from 4K4\text{K} pretraining and (bottom) pure-text continual cooldown at 64K.64\text{K}. Cooldown extends Dynamic YaRN Transformer from 8K8\text{K} to 32K32\text{K} before early-depth accuracy drops at 64K64\text{K} and unchunked prefill exceeds HBM at ≥ ⁣128K,\geq\!128\text{K}, whereas recurrent horizons change little under pure-text cooldown: GDN-2 collapses past 4K–8K,4\text{K}\text{--}8\text{K}, Muir (no γ{\bm{\gamma}} ) fades past 32K,32\text{K}, Muir (dynamic γt,{\bm{\gamma}}_{t}, splitting (rw,rq)=(128,128)(r_{w},r_{q})=(128,128) for multi-lag recall) holds through 32K32\text{K} before tapering at ≥ ⁣64K,\geq\!64\text{K}, and Muir (static γ,{\bm{\gamma}}, keeping rw=256r_{w}=256 for single-lag lookup) maintains near-perfect recall to 1M1\text{M} tokens in both regimes.
Figure 9: Boundary-pivoted chunkwise execution eliminates decay overflow and maintains well-conditioned forward and backward passes. (a) Against an exact float64 sequential reference, naive unpivoted float32 decay factorization overflows and produces NaN on 25%25\% of heads, whereas the boundary-pivoted bfloat16 and float32 kernels match the reference to machine precision on both chunk outputs and final state. (b) Across synthetic stress regimes spanning κ2(T)∈[1,104],\kappa_{2}({\bm{T}})\in[1,10^{4}], the bfloat16 block-Neumann inverse stays near its precision floor throughout the well-conditioned regime ( κ2(T)<40,\kappa_{2}({\bm{T}})<40, shaded), which contains the sampled 100B100\text{B} -token checkpoints of GDN-2 and Muir (markers show medians). (c) Single-layer backward gradients ( T ⁣= ⁣128,C ⁣= ⁣64,dk ⁣= ⁣16T\!=\!128,C\!=\!64,d_{k}\!=\!16 ) match the float64 reference across all six inputs without log-decay drift.
  ConfigurationFineWeb-Edu Validation NLL ↓\downarrow1.34B PPL ↓\downarrow1.34B Accuracy ↑\uparrow
 Modelrfr_{f}rbr_{b}rgr_{g}rwr_{w}rqr_{q}dffnd_{\mathrm{ffn}}166M406M834M1.34BWiki.LMB.LM-EvalBasedTQA
BaseTransformer–––––8,1922.45692.24242.11862.033517.1710.3052.0545.9214.56
GDN-2128dd128dd–5,8882.44662.23542.10602.017716.8010.3852.4632.0816.48
OursMuir (no γ{\bm{\gamma}} )646412825607,6802.43222.22302.08762.000616.4610.2653.7932.8616.52
Muir (static γ{\bm{\gamma}} )646412825607,6802.42732.21912.08691.999416.3510.1553.4033.0217.71
Muir (ours)64641281281287,6802.42762.21742.08621.997316.4110.0053.8434.6019.32
Table 2: Sensitivity-allocated gate ranks and dynamic readout match or outperform wider output-gate and uniform-rank allocations while preserving MXU alignment and FFN width. For each configuration of gate bottleneck ranks and FFN width (shown at d=2,048d=2{,}048 ), we report FineWeb-Edu validation token NLL across all four model scales, alongside 1.34B1.34\text{B} WikiText (Wiki.) and LAMBADA (LMB.) word-level perplexity, 8-task zero-shot commonsense accuracy (LM-Eval), 6-task Based extractive recall accuracy (Based), and 5-shot TriviaQA accuracy (TQA) after 100B100\text{B} training tokens.
Figure 10: Trained models exhibit the predicted write-erase sensitivity asymmetry under both SVD rank truncation and causal gate freezing. (a) Truncating each rank- 128128 gate projection of the uniform-rank Muir checkpoint via SVD (one gate at a time across all layers) shows that the decay ( ff ) and erase ( bb ) gates saturate by rank 64,64, whereas the output ( gg ), write ( ww ), and readout gain ( qq ) gates require rank 128.128. (b) Replacing the dynamic write gate with its token-averaged channel mean ( wt←wˉ{\bm{w}}_{t}\leftarrow\bar{\bm{w}} ) across all layers degrades validation NLL 14.5×14.5\times more than freezing the erase gate ( bt←bˉ{\bm{b}}_{t}\leftarrow\bar{\bm{b}} ) in both GDN-2 and Muir.
Figure 11: Stationary state row energies obey the closed-form law and are balanced by fractional power-law readout gains. (a) Across state channels, the analytical fixed point mj∗{m}_{j}^{\ast} ( 5 ) closely tracks measured stationary row energy mj=E[∥St,j,:∥22]{m}_{j}=\mathbb{E}[\|{\bm{S}}_{t,j,:}\|_{2}^{2}] in both GDN-2 and Muir. (b) Across all heads, weighting state rows by the learned readout gain compresses the within-head p90/p10\mathrm{p90}/\mathrm{p10} energy spread by 35%35\% relative to GDN-2. (c) Within each head, learned static gains ∣γj∣|{\gamma}_{j}| follow a fractional power law γj∝ρj−p{\gamma}_{j}\propto{\rho}_{j}^{-p} ( p=0.35p=0.35 ) of stationary row RMS ρj,{\rho}_{j}, intermediate between unit readout ( p=0p=0 ) and full whitening ( p=1p=1 ).
Figure 12: Readout-exponent sweeps and erase-logit shifts confirm the optimality of fractional power-law readout and the early-vs.-late erasure trade-off. (a) Replacing the learned static gain in each head with an RMS-matched power law γj(p)∝ρj−p{\gamma}_{j}(p)\propto{\rho}_{j}^{-p} yields a U-shaped loss curve minimized near the learned exponent ( p≈0.35p\approx 0.35 ), outperforming both unit readout ( p=0p=0 ) and full whitening ( p=1p=1 ). (b) Shifting the pre-activation erase logit ( ztb←ztb+δb{\bm{z}}^{b}_{t}\leftarrow{\bm{z}}^{b}_{t}+\delta_{b} ) shows that weaker erasure degrades predictions across all positions, whereas stronger erasure leaves early tokens [0,16)[0,16) nearly unchanged while degrading late tokens [256,4K).[256,4\text{K}).
 DiagnosticStatisticScaleGDN-2MuirEffectTheory
Sens.Key participation ratioηˉeff=1/E[∥kt∥44]\bar{\eta}_{\mathrm{eff}}=1/\mathbb{E}[\|{\bm{k}}_{t}\|_{4}^{4}]166M7.435.98≪ ⁣128\ll\!128 ch.Remark 3.2
Write-erase sensitivityLayer-mean E[∥DzwSt∥F2]/E[∥DzbSt∥F2]\mathbb{E}[\|\mathrm{D}_{{\bm{z}}^{w}}{\bm{S}}_{t}\|_{F}^{2}]/\mathbb{E}[\|\mathrm{D}_{{\bm{z}}^{b}}{\bm{S}}_{t}\|_{F}^{2}]166M13.213.313×13\times ww vs. bbProposition 3.1
EnergyStationary row-law fitMedian within-head rs(mj,mj∗)r_{s}({m}_{j},{m}_{j}^{\ast})166M0.9570.960+ 0.003+\,0.003Proposition 3.4
Raw row-energy spreadMedian within-head p90/p10\mathrm{p90}/\mathrm{p10} of mj{m}_{j}1.34B14.717.2++ 17.0%Proposition 3.4
Readout energy spreadMedian within-head p90/p10\mathrm{p90}/\mathrm{p10} of γˉj2mj\bar{\gamma}_{j}^{2}{m}_{j}1.34B14.79.6−- 34.7%Corollary 3.8
EraseParallel erase rateμˉ=E[⟨kt,bt⊙kt⟩]\bar{\mu}=\mathbb{E}[\langle{\bm{k}}_{t},{\bm{b}}_{t}\odot{\bm{k}}_{t}\rangle]1.34B0.7820.652−- 16.6%Proposition 3.10
Table 3: Trained GDN-2 and Muir checkpoints confirm the predicted write-erase sensitivity asymmetry, stationary row-energy law, and readout-induced erasure reduction. Diagnostics are measured on held-out FineWeb-Edu validation sequences ( T=4,096T=4{,}096 ) at 166M166\text{M} (gate sensitivity and stationary row-law fit) and 1.34B1.34\text{B} (within-head row-energy spread and write-time erasure).
Figure 13: Muir preserves early variable bindings and routes readout to the target hop at each answer step. (a) On 4-value MV-NIAH, Transformer drops on the final write V4V_{4} and unit-readout recurrences degrade after V1,V_{1}, whereas Muir remains accurate across V1–V4.V_{1}\text{--}V_{4}. (b) On 5-hop RULER VT, Transformer recalls only the root assignment H1,H_{1}, whereas Muir maintains accuracy across all five hops H1–H5.H_{1}\text{--}H_{5}. (c) Within-model root-needle ( H1H_{1} ) retention in St{\bm{S}}_{t} from write milestone s1s_{1} ( 100%100\% ) to query milestone tqt_{\mathrm{q}} across 1616 held-out prompts (interquartile bands across prompts). (d) Within-model retention of written needles H1–H5H_{1}\text{--}H_{5} from each write ( 100%100\% ) across sequence milestones s1–s5s_{1}\text{--}s_{5} to tq.t_{\mathrm{q}}. (e) Share of readout weight q^tSt\hat{\bm{q}}_{t}{\bm{S}}_{t} (over the haystack) on each needle H1–H5H_{1}\text{--}H_{5} across the four answer-generation steps A1–A4A_{1}\text{--}A_{4} that chain H1→⋯→H5.H_{1}\to\dotsb\to H_{5}.
Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol Definition
LLN\mathbb{N}Number of residual layers
ddN\mathbb{N}Model residual-stream width
dffnd_{\mathrm{ffn}}N\mathbb{N}FFN intermediate width ( 4d4d in Transformer, 3d−2563d-256 in GDN-2, 4d−5124d-512 in Muir)
dk,dvd_{k},d_{v}N\mathbb{N}Per-head key and value dimensions ( dk=dv=128d_{k}=d_{v}=128 )
NhN_{h}N\mathbb{N}Number of attention heads ( Nh=d/dkN_{h}=d/d_{k} )
T,CT,CN\mathbb{N}Sequence length and chunk length ( C=256C=256 )
Appendix
Table 4: Summary of notation. All vectors are row vectors; hats ( ⋅^\hat{\cdot} ) denote gate-modulated vectors or decayed states, bars ( ⋅ˉ\bar{\cdot} ) denote expectations or empirical means, and asterisks ( ∗ ) denote stationary fixed points or optimal solutions.
 Per-Head Recurrence and ReadoutGate Projections
ModelDecayEraseWriteReadoutScalarLow-RankDense
Lin. Attn. [ Katharopoulos et al., 2020 ]I\mathbf{I}I\mathbf{I}vt{\bm{v}}_{t}qtSt{\bm{q}}_{t}{\bm{S}}_{t}–––
RetNet [ Sun et al., 2023 ]αI{\alpha}\mathbf{I}I\mathbf{I}vt{\bm{v}}_{t}qtSt{\bm{q}}_{t}{\bm{S}}_{t}– †–gt{\bm{g}}_{t}
Mamba-2 [ Dao and Gu, 2024 ]αtI{\alpha}_{t}\mathbf{I}I\mathbf{I}βtvt{\beta}_{t}{\bm{v}}_{t}qtSt{\bm{q}}_{t}{\bm{S}}_{t}αt,βt{\alpha}_{t},{\beta}_{t}–gt{\bm{g}}_{t}
GLA [ Yang et al., 2023 ]\Diag(αt)\Diag({\bm{\alpha}}_{t})I\mathbf{I}vt{\bm{v}}_{t}qtSt{\bm{q}}_{t}{\bm{S}}_{t}–αt{\bm{\alpha}}_{t}gt{\bm{g}}_{t}
DeltaNet [ Yang et al., 2024 ]I\mathbf{I}I−βtkt⊤kt\mathbf{I}-{\beta}_{t}{\bm{k}}_{t}^{\top}{\bm{k}}_{t}βtvt{\beta}_{t}{\bm{v}}_{t}qtSt{\bm{q}}_{t}{\bm{S}}_{t}βt{\beta}_{t}–gt{\bm{g}}_{t}
Appendix
Table 5: Comparison of per-head state recurrence, query readout, and gate projection parameterizations across linear-time sequence models. For each model, we list the per-head decay, key-erasure, and value-write factors, the query readout, and whether each learned gate projection across heads is a per-head scalar ( d→Nhd\to N_{h} ), low-rank bottleneck ( d→r→dd\to r\to d ), or dense matrix ( d→dd\to d ). ( † RetNet fixes its per-head decay scalar α{\alpha} without an input projection; ‡ RWKV-7 combines decay and removal additively, writes with a decoupled replacement key k~t⊤vt,\tilde{\bm{k}}_{t}^{\top}{\bm{v}}_{t}, and includes a low-rank cross-layer value-residual gate.)
 Row-Energy LawKey RankMax / MinRMS SpreadEnergySpearmanLog-Pearson
 Diagnostic ( μ=0.60{\mu}=0.60 )E[ηeff(kt)]\mathbb{E}[\eta_{\mathrm{eff}}({\bm{k}}_{t})]max⁡jmj/min⁡jmj\max_{j}{m}_{j}/\min_{j}{m}_{j}p90/p10(ρj)\mathrm{p90}/\mathrm{p10}({\rho}_{j})p90/p10\mathrm{p90}/\mathrm{p10}rs(mj,mj∗)r_{s}({m}_{j},{m}_{j}^{\ast})rlog⁡r_{\log}
 Stationary row energy mj{m}_{j}12.2452.4 ×\times5.69 ×\times32.3 ×\times1.00001.0000
  Normalized Retrieval MSE ↓\downarrowOptimal Erase μ∗{\mu}^{\ast}
 Geometric Horizon Readout γ{\bm{\gamma}}Horizon H ⁣= ⁣10H\!=\!10Horizon H ⁣= ⁣27H\!=\!27Horizon H ⁣= ⁣81H\!=\!81H ⁣= ⁣10H\!=\!10H ⁣= ⁣27H\!=\!27H ⁣= ⁣81H\!=\!81
GeometricUnit gain γj=1{\gamma}_{j}=1 (GDN-2)0.3599  ± \,\pm\, 0.00190.6389  ± \,\pm\, 0.00170.8607  ± \,\pm\, 0.00170.900.600.30
Full whitening γj∝ρj−1{\gamma}_{j}\propto{\rho}_{j}^{-1} ( p=1p=1 )0.4068  ± \,\pm\, 0.00250.6802  ± \,\pm\, 0.00210.8817  ± \,\pm\, 0.00160.800.500.20
Appendix
Table 6: Single-head stationary row-energy verification, geometric-horizon retrieval sweep, and bimodal two-lag readout comparison. All results are reported as mean±SD\text{mean}\pm\text{SD} across 2020 random seeds for a single recurrent head with dk=dv=64.d_{k}=d_{v}=64. (Top) Verification of the closed-form row-energy law at μ=0.60.{\mu}=0.60. (Middle) Optimal erase strength μ∗{\mu}^{\ast} and normalized retrieval MSE under geometric lag distributions of mean H∈{10,27,81}.H\in\{10,27,81\}. (Bottom) Retrieval MSE on short ( ℓ=2\ell=2 ), long ( ℓ=80\ell=80 ), and mixed ( 50/5050/50 ) queries at μ=0.60,{\mu}=0.60, along with the Spearman correlation rs(γ,α)r_{s}({\bm{\gamma}},{\bm{\alpha}}) between readout gain γ{\bm{\gamma}} and the channel decay factor α.{\bm{\alpha}}.
  Gate ConfigAssociative Recall Accuracy (%) ↑\uparrowDiagnostics
 Modelrbr_{b}rwr_{w}γt{\bm{\gamma}}_{t}OverallShortMediumLongEraseRow Ener.
     acc ↑\uparrow[8,32)[8,32)[32,128)[32,128)[128,256)[128,256)tok var ↓\downarrowmˉT\bar{m}_{T} ↓\downarrow
 Dense (b,w)(b,w) (GDN-2)64641{\bm{1}}90.40  ± \,\pm\, 3.7996.29  ± \,\pm\, 2.7777.32  ± \,\pm\, 8.4933.68  ± \,\pm\, 9.4727.9%1,593.0
Erase rbr_{b}Erase rank rb=1r_{b}=11641{\bm{1}}91.28  ± \,\pm\, 3.2697.06  ± \,\pm\, 2.1079.01  ± \,\pm\, 8.1235.77  ± \,\pm\, 9.038.4%1,451.0
Erase rank rb=2r_{b}=22641{\bm{1}}91.33  ± \,\pm\, 2.9197.04  ± \,\pm\, 1.6879.30  ± \,\pm\, 7.4635.76  ± \,\pm\, 9.453.0%1,436.1
Erase rank rb=4r_{b}=44641{\bm{1}}91.51  ± \,\pm\, 3.3797.03  ± \,\pm\, 2.3680.05  ± \,\pm\, 7.9136.91  ± \,\pm\, 9.074.0%1,446.3
Appendix
Table 7: Learned multi-query associative recall accuracy, erase-logit token variance, and state row energy across gate-rank and readout-gain configurations. All results are evaluated on a two-head recurrent model with dk=dv=32d_{k}=d_{v}=32 and residual width d=128d=128 across 2020 random seeds (mean ±\pm SD; dense gates have full head-channel rank Nhdk=64N_{h}d_{k}=64 ). We report read accuracy overall and by event lag across single-gate rank sweeps and readout-gain variants with low-rank (b,w)(b,w) at rb=rw=4r_{b}=r_{w}=4 ( rq=8r_{q}=8 for dynamic γt{\bm{\gamma}}_{t} ), along with the token-dependent fraction of projected erase-logit energy on write steps E[∥z~tb−zˉb∥22]/E[∥z~tb∥22]\mathbb{E}[\|\tilde{\bm{z}}^{b}_{t}-\bar{\bm{z}}^{b}\|_{2}^{2}]/\mathbb{E}[\|\tilde{\bm{z}}^{b}_{t}\|_{2}^{2}] ( z~tb=ztb−cb\tilde{\bm{z}}^{b}_{t}={\bm{z}}^{b}_{t}-{\bm{c}}_{b} ) and mean terminal row energy mˉT=(Nhdk)−1E[∥ST∥F2].\bar{m}_{T}=(N_{h}d_{k})^{-1}\mathbb{E}[\|{\bm{S}}_{T}\|_{F}^{2}].
 Naive ChunkwiseBoundary-Pivoted Kernel (Ours)
Kernel StageMXUVPUMXUMXUVPUMXU
 FLOPsFLOPsShareFLOPsFLOPsShare
Diag. blocks ( n0n_{0} ), Tleaf−1{\bm{T}}_{\mathrm{leaf}}^{-1} ( 12 )00O(n0dk){\mathcal{O}}(n_{0}d_{k})0.0%4n02(log⁡2n0 ⁣− ⁣1)4n_{0}^{2}(\log_{2}n_{0}\!-\!1)O(n0dk){\mathcal{O}}(n_{0}d_{k})6.3%
Off-diag. blocks ( Tcross,Mcross{\bm{T}}_{\mathrm{cross}},{\bm{M}}_{\mathrm{cross}} )007(C ⁣− ⁣n0)dk7(C\!-\!n_{0})d_{k}0.0%2(C ⁣− ⁣n0)dk2(C\!-\!n_{0})d_{k}O(rdk){\mathcal{O}}(rd_{k})96.6%
2 ⁣× ⁣22\!\times\!2 block-tri. merges ( 11 )0000–23(C2 ⁣− ⁣n02)\tfrac{2}{3}(C^{2}\!-\!n_{0}^{2})00100.0%
Tri. solve TZ=[V^   K^⊙π]{\bm{T}}{\bm{Z}}=[\hat{\bm{V}}\;\,\hat{\bm{K}}_{\odot\pi}]00C(dk ⁣+ ⁣dv)C(d_{k}\!+\!d_{v})0.0%2C(dk ⁣+ ⁣dv)2C(d_{k}\!+\!d_{v})00100.0%
Appendix
Table 8: Per-token, per-head FLOP breakdown of the boundary-pivoted chunkwise kernel and naive chunkwise execution. Stage rows report leading per-token MXU and VPU operation counts (given cumulative log-decay Λ{\bm{\Lambda}} ) as a function of chunk length C,C, head dimensions (dk,dv),(d_{k},d_{v}), leaf size n0,n_{0}, and tree depth r=log⁡2(C/n0).r=\log_{2}(C/n_{0}). MXU shares and totals are evaluated at C=256,C=256, dk=dv=128,d_{k}=d_{v}=128, and n0=8n_{0}=8 ( r=5r=5 ).
  Throughput ( K tokens/s )XLA Profile ( 4K→32K4\text{K}\to 32\text{K} )
 Model ( 1.34B1.34\text{B} )T ⁣= ⁣2KT\!=\!2\text{K}T ⁣= ⁣4KT\!=\!4\text{K}T ⁣= ⁣8KT\!=\!8\text{K}T ⁣= ⁣16KT\!=\!16\text{K}T ⁣= ⁣32KT\!=\!32\text{K}CompiledAchievedHBM Traffic
  ( B ⁣= ⁣256B\!=\!256 )( B ⁣= ⁣128B\!=\!128 )( B ⁣= ⁣64B\!=\!64 )( B ⁣= ⁣32B\!=\!32 )( B ⁣= ⁣16B\!=\!16 )TFLOP/devMFU ( %\% )(GB/dev)
p50\mathrm{p50}Transformer ( 8,1928{,}192 )523.5396.9267.9163.693.6348  → \,\to\, 65218.2  → \,\to\, 8.6613  → \,\to\, 750
GDN-2 ( 5,8885{,}888 )179.6179.3179.1264.3253.5330  → \,\to\, 3307.1  → \,\to\, 10.14,478  → \,\to\, 4,148
Muir (matched, 5,8885{,}888 )183.4183.2182.6276.8261.8285  → \,\to\, 2856.2  → \,\to\, 8.94,413  → \,\to\, 4,058
Muir (ours; iso, 7,680\bm{7{,}680} )179.5179.4175.6268.5256.1324  → \,\to\, 3247.1  → \,\to\, 10.24,445  → \,\to\, 4,089
Appendix
Table 9: End-to-end training throughput, compiled XLA FLOPs, and HBM traffic across sequence lengths at 1.34B\bm{1.34\text{B}} scale. All measurements are collected on 88 TPU7x chips over 100100 -step pure-training intervals at a fixed global batch size of B×T=524,288 tokens/step,B\times T=524{,}288\text{ tokens/step}, comparing iso-parameter Muir ( dffn=7,680d_{\mathrm{ffn}}=7{,}680 ) and matched-FFN Muir ( dffn=5,888d_{\mathrm{ffn}}=5{,}888 ) against Transformer (TPU Pallas Splash Attention) and GDN-2. (XLA profile metrics are reported per TensorCore device across the 1616 TensorCores of the 88 TPU7x chips; compiled TFLOP/dev includes activation rematerialization, whereas MFU counts non-rematerialized model FLOPs. Bottom-row XLA profile ratios are evaluated at T=32K.T=32\text{K}. )
  ConfigurationGate Ranks & FFN WidthParam Breakdown (M)
 ModelLLddNhN_{h}dk=dvd_{k}=d_{v}rfr_{f}rbr_{b}rgr_{g}rwr_{w}rqr_{q}dffnd_{\mathrm{ffn}}MixerFFNTotal
            + GatesParamsParams
166MTransformer81,0248128–––––4,09633.5567.11166.20
GDN-281,02481281281,0241281,024–2,81654.5346.14166.20
Muir (no γ{\bm{\gamma}} )81,0248128646412825603,58441.9458.72166.20
Muir (static γ{\bm{\gamma}} )81,0248128646412825603,58441.9558.72166.21
Appendix
Table 11: Architectural specifications and parameter accounting across all four iso-parameter scales ( 166M\bm{166\text{M}} to 1.34B\bm{1.34\text{B}} ). All models at a given scale share depth L,L, residual width d,d, head dimensions dk=dv=128,d_{k}=d_{v}=128, head count Nh=d/128N_{h}=d/128 (so all recurrent models are strictly iso-state with Nhdkdv=128dN_{h}d_{k}d_{v}=128d state elements per layer), vocabulary size V=32,000V=32{,}000 (untied input/output embeddings), per-layer non-embedding budget 12d2,12d^{2}, global batch size 524,288 tokens/step524{,}288\text{ tokens/step} ( B=128,B=128, T=4,096T=4{,}096 ), and 100B100\text{B} training tokens.
  PerplexityZero-Shot Accuracy (%)
 ModelFW-EduWikiTextBoolQAvg. (9)
  Token PPL ↓\downarrow( 4K4\text{K} ) PPL ↓\downarrowacc ↑\uparroww/ BoolQ ↑\uparrow
166MTransformer11.6730.1559.1142.04
GDN-211.5529.3253.4641.94
Muir (no γ{\bm{\gamma}} )11.3828.8758.9943.49
Muir (static γ{\bm{\gamma}} )11.3328.3157.1342.74
Appendix
Table 12: Supplementary language modeling perplexities, zero-shot BoolQ accuracy, and 9-task commonsense average across all four scales. Complementing Table 1 , we report FineWeb-Edu validation token perplexity (PPL), full-window ( T=4,096T=4{,}096 ) WikiText word-level perplexity, zero-shot BoolQ accuracy, and the 9-task zero-shot commonsense average including BoolQ (“Avg. (9)”).
 S-NIAH-1S-NIAH-2S-NIAH-3S-NIAH-1/2/3 Avg.
Model ( 1.34B1.34\text{B} )1K2K4K8K1K2K4K8K1K2K4K8K1K2K4K8K1–4K1–8K
Transformer100.0100.099.60.099.8100.0100.00.098.096.480.80.099.398.893.50.097.272.9
GDN-299.899.286.240.8100.099.683.651.677.457.620.69.892.485.563.534.180.468.9
Muir (no γ{\bm{\gamma}} )100.0100.099.699.898.299.893.839.284.080.441.610.494.193.478.349.888.678.9
Muir (static γ{\bm{\gamma}} )100.0100.0100.099.6100.0100.099.041.096.280.643.813.698.793.580.951.491.181.2
Muir (ours)100.0100.0100.0100.0100.099.896.232.088.267.437.87.096.189.178.046.387.777.4
Appendix
Table 13: Complete synthetic associative recall accuracy on the six RULER benchmarks at 1.34B\bm{1.34\text{B}} parameters. (Top) Single-needle retrieval on S-NIAH-1/2/3, and their 3-task mean. (Bottom) Multi-hop VT, multi-value (MV-NIAH) and multi-query (MQ-NIAH) retrieval, and suite-wide averages across multi-needle retrieval (MV+MQ), all six tasks at 1K–4K,1\text{K}\text{--}4\text{K}, and all evaluated lengths ( 1K–8K1\text{K}\text{--}8\text{K} ).
  Extraction (%) ↑\uparrowPassage Q/A ( 2K2\text{K} Cut; %) ↑\uparrow6-Task Avg. (%) ↑\uparrow
 ModelSWDEFDASQuADTQANQDROP2K2\text{K} Cut1K1\text{K} Cut2K2\text{K} Trunc.
166MTransformer29.8039.5130.4744.9117.8316.2029.7931.4123.33
GDN-212.379.9025.5644.2512.1317.0120.2021.6019.41
Muir (no γ{\bm{\gamma}} )12.008.9027.3846.8013.9718.9721.3422.2920.41
Muir (static γ{\bm{\gamma}} )10.507.6325.2645.8512.4817.8719.9322.0118.96
Muir (ours)10.975.2727.6846.2713.4018.7820.3921.7819.51
Appendix
Table 14: Complete real-world extractive recall accuracy on the six Based/JRT benchmarks across all four model scales. Per-task accuracies are reported under the official 2K2\text{K} context-cut JRT protocol ( 12,07712{,}077 validation examples across structured attribute extraction and passage question answering; the 1.34B1.34\text{B} 2K2\text{K} context-cut average is reported in Table 2 ), alongside the 6-task average under 1K1\text{K} context-cut truncation and un-cut 2K2\text{K} left-truncation.
  PPL ↓\downarrowΔtail\Delta_{\mathrm{tail}} (nats) ↓\downarrow
 Model4K4\text{K}8K8\text{K}16K16\text{K}32K32\text{K}64K64\text{K}128K128\text{K} 
Contiguous FineWeb-Edu Long Documents
166MTransformer11.5550.97199.85473.31728.84895.74++ 4.5349  ± \,\pm\, 0.0798
GDN-211.6911.4411.3211.2911.2811.28++ 0.0051  ± \,\pm\, 0.0018
Muir (no γ{\bm{\gamma}} )11.4911.2511.1411.1111.1011.10++ 0.0051  ± \,\pm\, 0.0015
Muir (static γ{\bm{\gamma}} )11.3911.1411.0311.0010.9910.99++ 0.0076  ± \,\pm\, 0.0020
Appendix
Table 15: Zero-shot length extrapolation from 4K\bm{4\text{K}} to 128K\bm{128\text{K}} on contiguous FineWeb-Edu long documents and packed control streams. For each evaluation window T∈{4K,8K,16K,32K,64K,128K},T\in\{4\text{K},8\text{K},16\text{K},32\text{K},64\text{K},128\text{K}\}, we report aggregate token perplexity alongside the paired tail NLL differential \Delta_{\mathrm{tail}}\mathrel{\mathchoice{\vbox{\hbox{\displaystyle:}}}{\vbox{\hbox{\textstyle:}}}{\vbox{\hbox{\scriptstyle:}}}{\vbox{\hbox{\scriptscriptstyle:}}}{=}}\mathrm{NLL}_{[127\mathrm{K},128\mathrm{K})}(T=128\text{K})-\mathrm{NLL}_{[127\mathrm{K},128\mathrm{K})}(T=4\text{K}) (mean ±\pm paired SE in nats across NdocsN_{\mathrm{docs}} sequences).
  RULER S-NIAH-1 Accuracy (%) ↑\uparrowAvg. (%) ↑\uparrow
 Model ( 1.34B1.34\text{B} )1K1\text{K}2K2\text{K}4K4\text{K}8K8\text{K}16K16\text{K}32K32\text{K}64K64\text{K}128K128\text{K}256K256\text{K}512K512\text{K}1M1\text{M}1–32K1\text{--}32\text{K}1K–1M1\text{K--}1\text{M}
Zero-ShotTransformer (YaRN)100.0100.090.090.678.868.145.6OOMOOMOOMOOM87.9–
GDN-2100.098.880.040.013.811.934.423.815.010.010.057.439.8
Muir (no γ{\bm{\gamma}} )100.0100.0100.0100.098.173.846.246.926.215.611.995.365.3
Muir (static γ{\bm{\gamma}} )100.0100.0100.099.4100.098.898.198.895.691.990.099.797.5
Muir (ours)100.0100.0100.0100.0100.085.626.222.517.513.112.597.661.6
Appendix
Table 17: Zero-shot and post-cooldown RULER S-NIAH-1 passkey retrieval from 1K\bm{1\text{K}} to 1M\bm{1\text{M}} tokens at 1.34B\bm{1.34\text{B}} parameters. Complementing the depth-by-length heatmaps in Figure 8 , each context length T∈{1K,…,1M}T\in\{1\text{K},\dotsc,1\text{M}\} evaluates N=160N=160 examples stratified uniformly across 1010 needle-depth bins under greedy decoding, both zero-shot ( 4K4\text{K} ) and after the 64K64\text{K} continual cooldown ( 0.315B0.315\text{B} FineWeb-Edu tokens); unchunked Transformer (YaRN) prefill exceeds HBM (OOM) at T≥128K.T\geq 128\text{K}. Under Zero-Shot, the bottom two rows freeze dynamic readout in Muir at inference time to γstat{\bm{\gamma}}^{\mathrm{stat}} ( ztq←0{\bm{z}}^{q}_{t}\leftarrow{\bm{0}} ) or its held-out text mean γˉ=E[γt].\bar{\bm{\gamma}}=\mathbb{E}[{\bm{\gamma}}_{t}].
   Gate Ranks & FFN ( d=1,536d=1{,}536 )FineWeb-Edu ValWikiText
 Model ( 406M406\text{M} )Paramsrfr_{f}rbr_{b}rgr_{g}rwr_{w}rqr_{q}dffnd_{\mathrm{ffn}}TokenΔNLL\Delta\text{NLL}Token( 4K4\text{K} )
  (M)      NLL ↓\downarrow(mnat) ↓\downarrowPPL ↓\downarrowPPL ↓\downarrow
BaseTransformer409.7–––––6,1442.2424++ 7.09.4221.94
GDN-2409.7128dd128dd–4,3522.23540.09.3521.76
OursMuir (no γ{\bm{\gamma}} )409.7646412825605,6322.2230−- 12.49.2421.16
Muir (static γ{\bm{\gamma}} )409.7646412825605,6322.2191−- 16.39.2021.24
Appendix
Table 18: Fixed-FFN gate-compression, dense dynamic-readout, and unconstrained wide-FFN controls at the 406M\bm{406\text{M}} scale. Complementing the multi-scale rank ablations in Table 2 , we evaluate fixed-FFN gate compression ( dffn=4,352d_{\mathrm{ffn}}=4{,}352 ), dense dynamic-readout ( rq=d,r_{q}=d, dffn=3,584d_{\mathrm{ffn}}=3{,}584 ), and non-iso-parameter wide-FFN ( dffn=5,504,d_{\mathrm{ffn}}=5{,}504, 448.7M448.7\text{M} ) controls at the 406M406\text{M} scale ( L=11,L=11, d=1,536,d=1{,}536, 100B100\text{B} tokens); ΔNLL\Delta\text{NLL} is reported in millinats relative to GDN-2, with best and second-best iso-parameter models in bold and underlined .
 Kept SVD Rank rr (Validation ΔNLL\Delta\text{NLL} vs. r=128r=128 in nats ±\pm SE ↓\downarrow )
Gater=0r=0 (bias)r=16r=16r=32r=32r=48r=48r=64r=64r=96r=96
Decay ff ( αt{\bm{\alpha}}_{t} )++ 0.5801  ± \,\pm\, 0.2147++ 0.0259  ± \,\pm\, 0.0084++ 0.0043  ± \,\pm\, 0.0018++ 0.0020  ± \,\pm\, 0.0015−- 0.0007  ± \,\pm\, 0.0015−- 0.0001  ± \,\pm\, 0.0004
Erase bb ( bt{\bm{b}}_{t} )++ 0.1193  ± \,\pm\, 0.0154++ 0.0039  ± \,\pm\, 0.0004++ 0.0016  ± \,\pm\, 0.0009++ 0.0006  ± \,\pm\, 0.0008++ 0.0006  ± \,\pm\, 0.0004++ 0.0000  ± \,\pm\, 0.0001
Write ww ( wt{\bm{w}}_{t} )++ 0.4454  ± \,\pm\, 0.2248++ 0.0486  ± \,\pm\, 0.0048++ 0.0185  ± \,\pm\, 0.0022++ 0.0108  ± \,\pm\, 0.0015++ 0.0064  ± \,\pm\, 0.0010++ 0.0013  ± \,\pm\, 0.0007
Query qq ( γt{\bm{\gamma}}_{t} )++ 0.5639  ± \,\pm\, 0.2352++ 0.0296  ± \,\pm\, 0.0057++ 0.0118  ± \,\pm\, 0.0023++ 0.0077  ± \,\pm\, 0.0013++ 0.0045  ± \,\pm\, 0.0013++ 0.0015  ± \,\pm\, 0.0008
Output gg ( gt{\bm{g}}_{t} )++ 2.7725  ± \,\pm\, 0.2189++ 0.7489  ± \,\pm\, 0.2328++ 0.4824  ± \,\pm\, 0.2483++ 0.3492  ± \,\pm\, 0.2225++ 0.2472  ± \,\pm\, 0.1783++ 0.0376  ± \,\pm\, 0.0209
Appendix
Table 19: Post-hoc SVD gate-rank truncation sensitivity on the trained 406M\bm{406\text{M}} uniform-rank checkpoint. For each gate m∈{f,b,g,w,q},m\in\{f,b,g,w,q\}, we replace its rank- 128128 factor Wm↓Wm↑{\bm{W}}_{m\downarrow}{\bm{W}}_{m\uparrow} across all 1111 layers with its best rank- rr SVD approximation (retaining the additive bias cm{\bm{c}}_{m} ) and measure \Delta\text{NLL}(r)\mathrel{\mathchoice{\vbox{\hbox{\displaystyle:}}}{\vbox{\hbox{\textstyle:}}}{\vbox{\hbox{\scriptstyle:}}}{\vbox{\hbox{\scriptscriptstyle:}}}{=}}\mathrm{NLL}(r)-\mathrm{NLL}(128) (in nats ±\pm paired SE) on held-out FineWeb-Edu sequences.
DiagnosticWithin-Head StatisticScaleGDN-2MuirEffectTheory
Full row-law log-Pearson rrr(log⁡10mj,log⁡10mj∗)r(\log_{10}{m}_{j},\log_{10}{m}_{j}^{\ast})166M0.9510.9510.932/0.9160.932/0.916 pooledProp. 3.4
Decay-only log-Pearson rrr(log⁡10mj,−log⁡10(1−λj))r(\log_{10}{m}_{j},-\log_{10}(1-{\lambda}_{j}))166M0.8400.803−0.13-0.13 ( 0.869/0.8270.869/0.827 pooled)Rem. 3.5
Static gain alignmentrs(∣γj∣,ρj−1)r_{s}(|{\gamma}_{j}|,{\rho}_{j}^{-1})166M–++ 0.710100%100\% heads  ⁣> ⁣0\!>\!0Cor. 3.8
Readout power-law slope−dlog⁡γˉj/dlog⁡ρj-\mathrm{d}\log\bar{\gamma}_{j}/\mathrm{d}\log{\rho}_{j}1.34B0.000.44p∈(0,1)p\in(0,1)Cor. 3.8
Readout row-RMS spreadp90/p10\mathrm{p90}/\mathrm{p10} of γˉjρj\bar{\gamma}_{j}{\rho}_{j}1.34B3.83 ×\times3.10 ×\times−- 19.1%Cor. 3.8
Appendix
Table 20: Supplementary stationary row-energy and readout-scaling diagnostics across trained GDN-2 and Muir checkpoints. Complementing Table 3 , row-law ablations and static-gain alignment are evaluated across live channels ( ρj≥10−10{\rho}_{j}\geq 10^{-10} ) on the 166M166\text{M} GDN-2 and Muir (static γ{\bm{\gamma}} ) checkpoints, and dynamic-gain scaling and readout row-RMS spread are evaluated across all 2424 layers ( 384384 heads) at 1.34B1.34\text{B} (with γˉj=1\bar{\gamma}_{j}=1 in GDN-2); values report the median across heads.
  Validation ΔNLL\Delta\text{NLL} ( nats±SE\text{nats}\pm\text{SE} ) Across Token Positions
 InterventionModelMean EraseOverallEarlyMidLate
  ( 834M834\text{M} )μˉ\bar{\mu}[0,4,096)[0,4{,}096)[0,16)[0,16)[16,256)[16,256)[256,4,096)[256,4{,}096)
FreezeFreeze write gateGDN-20.733++ 0.8325  ± \,\pm\, 0.0581++ 0.5799  ± \,\pm\, 0.1467++ 0.9517  ± \,\pm\, 0.1275++ 0.8261  ± \,\pm\, 0.0566
Freeze erase gateGDN-20.672++ 0.0574  ± \,\pm\, 0.0066++ 0.0153  ± \,\pm\, 0.0143++ 0.0523  ± \,\pm\, 0.0152++ 0.0578  ± \,\pm\, 0.0062
Isotropic erase bt←μt1{\bm{b}}_{t}\leftarrow{\mu}_{t}{\bm{1}}GDN-20.732++ 0.0844  ± \,\pm\, 0.0044++ 0.0189  ± \,\pm\, 0.0116++ 0.0658  ± \,\pm\, 0.0157++ 0.0859  ± \,\pm\, 0.0044
Freeze write gateMuir0.610++ 0.3845  ± \,\pm\, 0.0368++ 0.2534  ± \,\pm\, 0.0767++ 0.4877  ± \,\pm\, 0.1309++ 0.3786  ± \,\pm\, 0.0328
Appendix
Table 21: Causal gate-freezing and erase-logit shift interventions on trained 834M\bm{834\text{M}} checkpoints. (Top) Validation ΔNLL\Delta\text{NLL} and mean erase rate μˉ\bar{\mu} when replacing the dynamic write gate wt{\bm{w}}_{t} or erase gate bt{\bm{b}}_{t} with its stream mean ( wˉ\bar{\bm{w}} or bˉ\bar{\bm{b}} ) or isotropic key-space projection ( μt1{\mu}_{t}{\bm{1}} ) across all layers. (Bottom) Additive erase-gate logit shifts ztb←ztb+δb{\bm{z}}^{b}_{t}\leftarrow{\bm{z}}^{b}_{t}+\delta_{b} for δb∈{−1.0,−0.5,+0.5},\delta_{b}\in\{-1.0,-0.5,+0.5\}, reported overall ( [0,4,096)[0,4{,}096) ) and stratified by token position bucket ( [0,16),[0,16), [16,256),[16,256), and [256,4,096)[256,4{,}096) ).
  MV-NIAH Recall (%) ↑\uparrow5-Hop VT Recall (%) ↑\uparrow
 Model ( 1.34B1.34\text{B} )V1V_{1}V2V_{2}V3V_{3}V4V_{4}Avg.H1H_{1}H2H_{2}H3H_{3}H4H_{4}H5H_{5}Avg.
1KTransformer84.887.285.666.080.925.85.86.06.65.29.9
GDN-267.057.062.074.465.128.434.449.461.664.647.7
Muir (no γ{\bm{\gamma}} )85.646.049.057.459.529.236.049.057.250.844.4
Muir (static γ{\bm{\gamma}} )57.455.257.259.657.413.217.220.823.820.219.0
Muir (ours)82.484.486.691.686.248.046.061.269.275.460.0
Appendix
Table 23: Write-order MV-NIAH and hop-resolved variable tracking recall at 1.34B\bm{1.34\text{B}} parameters. All results are evaluated over N=500N=500 examples per task at 1K1\text{K} (top) and 4K4\text{K} context lengths (bottom). MV-NIAH columns V1–V4V_{1}\text{--}V_{4} order the four target values bound to the shared key from oldest write ( V1V_{1} ) to most recent write ( V4V_{4} ), and VT columns H1–H5H_{1}\text{--}H_{5} order the five target variables from earliest root binding ( H1,H_{1}, VAR_1 = V ) to latest alias ( H5,H_{5}, VAR_5 = VAR_4 ).
More in this topic

Delta-Rule Linear Attention

Delta-rule linear attention updates a compact recurrent memory to retrieve and revise sequence information with linear-time processing.

See all 39 papers

Explore similar work

CardsList

  1. Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

    May 21, 2026Ali Hatamizadeh, Yejin Choi, Jan KautzLong-Context Language ModelingRecurrent Transformers

    NVIDIA

  2. Dual DeltaNet: Value-Side Memory Reorganization for Linear Attention

    Oct 28, 2026Yixiao Qian, Yaxin Zhang, Song Chen +3Linear AttentionDelta-Rule Linear Attention

    College of Control Science and Engineering, Zhejiang University · Polytechnic Institute, Zhejiang University · Department of Mathematics, National University of Singapore +1

  3. DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention

    Oct 30, 2026Zhangchi Feng, Harry Jiang, Yangyifei Luo +2Long-Context RetrievalLinear Attention

    SpeedAI · CCSE, School of Computer Science and Engineering, Beihang University