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

Authors: Tushaar Gangavarapu¹, Gil Tabak², Clemens JS Schaefer²
¹University of Texas at Austin, ²Google

Summary (Overview)

• Muir is a novel gated delta-rule linear attention architecture that fuses all recurrence gates (decay, erase, write, output, and a new query readout gain) into a single sensitivity-allocated low-rank projection, replacing the dense gate matrices in Gated DeltaNet-2 (GDN-2).
• The paper provides theoretical analysis showing that state updates are 13× more sensitive to write perturbations than erase perturbations (due to key sparsity from SiLU activation and ℓ₂-normalization), and derives a closed-form stationary row-energy law that explains why slow-decaying state rows accumulate much larger norms than fast rows.
• A matched-filter readout gate (combining static channel profile with dynamic per-token modulation) re-weights state rows, reducing write-time erasure and extending memory horizon—improving multi-hop retrieval, extractive recall, and length extrapolation.
• A boundary-pivoted chunkwise kernel enables numerically stable training on wide-matrix-unit accelerators (TPUs) by avoiding cumulative decay division overflow and 3D decay tensors, using divide-and-conquer block decomposition with leaf-level Neumann series inversion.
• Across 166M–1.34B parameters, Muir lowers perplexity by 1.8–2.0% over GDN-2 and 2.5–3.6% over Transformer, improves commonsense accuracy by up to 1.9 points, extrapolates 2,560× from 4K to 10M-token streams without perplexity drift, and maintains near-100% passkey retrieval to 1M tokens.

Introduction and Theoretical Foundation

Background and Motivation

Linear attention compresses context into a fixed-size key-value state for linear-time training and constant-memory inference. However, a finite-dimensional state admits only finitely many orthogonal key directions, forcing cross-key interference. Delta-rule recurrences mitigate this by erasing the current prediction for an incoming key before writing a new value.

The paper identifies three bottlenecks in channel-wise gated delta models:

1. Parameter allocation: GDN-2 uses full-rank $d \times d$ projections for erase and write gates, consuming $2d^2$ parameters (one-sixth of the $12d^2$ layer budget), shrinking the FFN by over 25%.

2. State readout: Slow-decaying state rows accumulate much larger norms than fast rows, dominating unweighted readout and forcing stronger write-time erasure that prematurely overwrites stored associations.

3. Parallel training: Channel-wise decay breaks commutativity with rank-one erasure, making standard chunkwise algorithms numerically unstable at the large chunk sizes needed for wide matrix units.

Theoretical Foundation

GDN-2 State Update (Equation 1):

$$\boldsymbol {S} _ {t} = \left(\mathbf {I} - \boldsymbol {k} _ {t} ^ {\top} (\boldsymbol {b} _ {t} \odot \boldsymbol {k} _ {t})\right) \operatorname{Diag} (\boldsymbol {\alpha} _ {t}) \boldsymbol {S} _ {t - 1} + \boldsymbol {k} _ {t} ^ {\top} (\boldsymbol {w} _ {t} \odot \boldsymbol {v} _ {t}), \quad \boldsymbol {o} _ {t} = \boldsymbol {q} _ {t} \boldsymbol {S} _ {t}$$

where $\alpha_t$ is the decay gate, $b_t$ the erase gate, $w_t$ the write gate, and $g_t$ the output gate.

Key theoretical results:

Proposition 3.1 (State sensitivity asymmetry): The expected squared Frobenius sensitivity to erase perturbations is:

$$\mathbb {E} _ {\delta \boldsymbol {b}} \left[ \| \mathrm{D} _ {\boldsymbol {b}} \boldsymbol {S} _ {t} [ \delta \boldsymbol {b} ] \| _ {F} ^ {2} \right] = \sum_ {j = 1} ^ {d _ {k}} k _ {t, j} ^ {2} \| \hat {\boldsymbol {S}} _ {t, j,:} \| _ {2} ^ {2}$$

while for write perturbations:

$$\mathbb {E} _ {\delta \boldsymbol {w}} \left[ \| \mathrm{D} _ {\boldsymbol {w}} \boldsymbol {S} _ {t} [ \delta \boldsymbol {w} ] \| _ {F} ^ {2} \right] = \| \boldsymbol {v} _ {t} \| _ {2} ^ {2}$$

Since keys concentrate energy on $\eta_{\text{eff}}(k_t) \ll d_k$ active channels (participation ratio ≈ 5–7 out of 128), erase perturbations have limited effect, while writes scale all channels.

Proposition 3.4 (Stationary row-energy law):

$$m _ {j} ^ {*} = \rho _ {j} ^ {2} = \frac {c _ {j} \sigma _ {v} ^ {2}}{1 - \lambda _ {j} + 2 \lambda _ {j} \bar {b} _ {j} c _ {j}}$$

where $\rho_j$ is the stationary row RMS norm, $c_j$ the key energy concentration, $\sigma_v^2$ the value variance, $\lambda_j$ the mean squared decay, and $\bar{b}_j$ the mean erase rate.

Proposition 3.6 (Matched-filter readout): The optimal readout gain for retrieval lag $\ell$ is:

$$\gamma_ {j} ^ {*} (\ell) = \kappa_ {\ell} \frac {\bar {a} _ {\ell , j}}{m _ {j} ^ {*}} \propto \frac {\lambda_ {j} ^ {\ell / 2}}{\rho_ {j} ^ {2}}$$

Proposition 3.10 (Erase-readout substitution): Matched-filter readout suppresses interference at query time, shifting the equilibrium erase rate downward and extending the effective memory horizon of every state channel.

Methodology

Muir Architecture

Fused Low-Rank Gate Projection (Equation 9):

$$\boldsymbol {h} _ {t} = \boldsymbol {x} _ {t} \boldsymbol {W} _ {\downarrow} \in \mathbb {R} ^ {1 \times R _ {\mathrm{tot}}}, \quad \boldsymbol {z} _ {t} ^ {m} = \boldsymbol {h} _ {t} ^ {m} \boldsymbol {W} _ {m \uparrow} + \boldsymbol {c} _ {m} \text { for } m\in\{f,b,g,w,q\}$$

Gate ranks $(r_f, r_b, r_g, r_w, r_q) = (64, 64, 128, 128, 128)$ with $R_{\text{tot}} = 512 = 2 \times 256$ for MXU tile alignment.

Gate parameterizations:

$$\log \boldsymbol {\alpha} _ {t} = - \exp (a _ {\mathrm{log}}) \operatorname{softplus} (\boldsymbol {z} _ {t} ^ {f}), \quad \boldsymbol {b} _ {t} = 2 \sigma (\boldsymbol {z} _ {t} ^ {b}), \quad \boldsymbol {w} _ {t} = \sigma (\boldsymbol {z} _ {t} ^ {w}), \quad \boldsymbol {\gamma} _ {t} = \boldsymbol {\gamma} ^ {\mathrm{stat}} \odot 2 \sigma (\boldsymbol {z} _ {t} ^ {q})$$

Iso-parameter FFN reallocation: Setting $d_{\text{ffn}}(R_{\text{tot}}) = 4d - R_{\text{tot}}$ maintains the $12d^2$ per-layer budget while widening the FFN from 5,888 to 7,680 channels at $d = 2048$ (+30.4%).

Boundary-Pivoted Chunkwise Kernel

The chunkwise solve requires forming the triangular system matrix $T = I + L_{\text{strict}}$ with entries $(L_{\text{strict}})_{t,s} = \langle \hat{k}_t, \pi_{t,s} \odot k_s \rangle$. The kernel uses:

1. Divide-and-conquer block decomposition (Equation 11):

$$\boldsymbol {T} _ {\mathrm{blk}} ^ {- 1} = \left[ \begin{array}{c c} \boldsymbol {T} _ {\mathcal {L}} ^ {- 1} & \boldsymbol {0} \\ \hline - \boldsymbol {T} _ {\mathcal {R}} ^ {- 1} \boldsymbol {T} _ {\mathrm{cross}} \boldsymbol {T} _ {\mathcal {L}} ^ {- 1} & \boldsymbol {T} _ {\mathcal {R}} ^ {- 1} \end{array} \right]$$

2. Boundary-pivoted decay splitting without division:

$$\boldsymbol {\pi} _ {t, s} = \exp (\boldsymbol {\Lambda} _ {t} - \boldsymbol {\Lambda} _ {s}) = \underbrace {\exp (\boldsymbol {\Lambda} _ {t} - \boldsymbol {\Lambda} _ {p})} _ {\boldsymbol {\pi} _ {t, p}} \odot \underbrace {\exp (\boldsymbol {\Lambda} _ {p} - \boldsymbol {\Lambda} _ {s})} _ {\boldsymbol {\pi} _ {p, s}}$$

3. Leaf-level Neumann series (Equation 12):

$$\pmb {T} _ {\mathrm{leaf}} ^ {- 1} = \prod_ {j = 0} ^ {\log_ {2} n _ {0} - 1} \big (\mathbf {I} + (- \pmb {L} _ {\mathrm{leaf}}) ^ {2 ^ {j}} \big)$$

4. Custom backward pass using saved $T^{-1}$ and $Z$:

$$\mathrm{d} \left[ \begin{array}{c c} \hat {\boldsymbol {V}} & \hat {\boldsymbol {K}} _ {\odot \pi} \end{array} \right] = \boldsymbol {T} ^ {- \top} \mathrm{d} \boldsymbol {Z}, \quad \mathrm{d} \boldsymbol {L} _ {\text { strict }} = - \operatorname{stril} \left(\mathrm{d} \left[ \begin{array}{c c} \hat {\boldsymbol {V}} & \hat {\boldsymbol {K}} _ {\odot \pi} \end{array} \right] \boldsymbol {Z} ^ {\top}\right)$$

Empirical Validation / Results

Language Modeling (Table 1, 1.34B scale)

Associative Recall

• Single-needle (S-NIAH-1/2/3): Muir holds 100% on S-NIAH-1 through 8K; all Muir variants mitigate the 4K degradation of GDN-2.
• Multi-hop (VT): Dynamic γₜ leads all models at every context length; static γ fails on varying hop lags.
• Multi-needle (MV/MQ-NIAH): Dynamic γₜ matches or outperforms both baselines across 1K–4K.
• Real-world extractive recall: Muir improves over GDN-2 on all six Based/JRT tasks, with dynamic γₜ driving 69% of the gain.

Length Extrapolation

• Zero-shot: Muir (dynamic γ) extrapolates from 4K to 10M-token streams (2,560×) without perplexity drift; Dynamic YaRN Transformer diverges past 8K.
• Passkey retrieval: Muir (static γ) reaches 90.0% zero-shot and 99.4% post-cooldown recall at 1M tokens (256× pretraining length); GDN-2 collapses past 4K–8K.

Training Throughput

• Muir matches GDN-2 at 4K (179.4 vs. 179.3 K tokens/s) and runs +1.0–1.6% faster at 16K–32K.
• Muir trains 1.64× and 2.74× faster than Transformer at 16K and 32K respectively.

Numerical Stability

• Boundary-pivoted kernel eliminates decay overflow (NaN) on fast-decaying channels.
• Condition number $\kappa_2(T) < 40$ for trained models (max across 1,536 chunk systems).
• Muir attains lower p90 $\kappa_2(T)$ than GDN-2 (7.7 vs. 11.1).

Mechanistic Validation (Table 3)

Theoretical and Practical Implications

Theoretical Implications

1. Sensitivity-allocated gating is a general principle: State updates in delta-rule recurrences have fundamentally different sensitivity to erase vs. write gates due to input normalization and activation sparsity. This asymmetry justifies asymmetric rank allocation—a principle that extends beyond gated delta rules to any linear recurrent model with channel-wise gating.

2. Matched-filter readout is a principled remedy for state imbalance: The stationary row-energy law provides a closed-form prediction for how multi-timescale decay creates energy imbalance, and the matched-filter framework (balancing signal survival against interference) provides the optimal readout. The fractional power-law exponent $p \in (0,1)$ (measured at $p = 0.35$) is theoretically derived, not empirically tuned.

3. Erase-readout substitution: Decoupling interference suppression from write-time erasure is a fundamental mechanism for extending memory horizon—readout filtering and write-time erasure are substitutable control mechanisms.

Practical Implications

1. Iso-parameter efficiency: Muir's fused low-rank gating reduces gate-projection parameters by 4.5×, enabling a 30.4% wider FFN at zero parameter cost—improving both recurrent memory and parametric knowledge capacity.

2. Hardware-aligned algorithm design: The boundary-pivoted chunkwise kernel demonstrates that numerical stability and hardware efficiency can be jointly achieved through divide-and-conquer block decomposition with leaf-level Neumann series, routing 98% of chunkwise arithmetic onto matrix units.

3. Long-context capability at constant memory: At 1.34B, Muir maintains a 25 MB float32 state for 10M-token streams, versus 2.1 TB bfloat16 KV cache for Transformer—a fundamental advantage for deployment in memory-constrained settings.

Conclusion

Muir addresses three bottlenecks in channel-wise gated delta-rule models through theory-guided architecture design:

1. Parameter efficiency: Fusing all gates into a sensitivity-allocated low-rank projection (with smaller decay/erase bottlenecks) widens the FFN by up to 30.4% at fixed parameters and FLOPs.

2. State readout: Matched-filter readout (static profile × dynamic modulation) balances state rows, reduces write-time erasure by 17%, improves multi-hop retention from 37% to 45%, and enables 2,560× length extrapolation without perplexity drift.

3. Parallel training: The boundary-pivoted chunkwise kernel eliminates numerical instability at large chunk sizes, matching GDN-2 throughput and training 2.74× faster than Transformer at 32K.

Limitations and future directions:

• Static and dynamic readout trade single-horizon retention against multi-lag adaptability under a fixed gate-rank budget.
• Pure-JAX kernel is slower than hand-tuned Splash Attention at short contexts (< 11K).
• Future work: predicting one lag per head to unify passkey recall and multi-lag routing; pairing sliding-window attention with static-readout Muir; fusing the chunkwise kernel into a single Pallas kernel; extending sensitivity-allocated gating to matrix-state test-time optimization models.