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

📌 Highlights

Muir combines a mechanistically motivated low-rank gate allocation, matched-filter readout, FFN reallocation, and accelerator-aware chunkwise execution. The authors report improvements over GDN-2 and Transformer across language modeling, commonsense reasoning, retrieval, and training efficiency.

• Across 166M–1.34B parameters, the paper reports 18.0–20.4 millinats lower validation NLL than GDN-2 and 25.0–36.2 millinats lower than Transformer.
• At 1.34B, Muir improves multi-hop and multi-needle retrieval over both baselines and extrapolates from a 4K training length to 10M-token streams without perplexity drift.
• Static matched-filter readout maintains near-perfect passkey retrieval to 1M tokens in the reported S-NIAH-1 evaluation.
• The kernel reports 24× lower forward-and-backward VPU compute, 12% lower total kernel FLOPs, and 98% of chunkwise arithmetic routed to the MXU under the evaluated configuration.
• The principal caveats are limited model scale, theoretical assumptions that may not hold on natural language, and hardware-specific systems evidence.

🎯 Introduction

The paper starts from the standard trade-off between Transformer attention and linear attention. Softmax self-attention has quadratic training cost in sequence length and a cache that grows linearly at inference, whereas linear-attention models compress history into a fixed-size key–value state. That compression creates cross-key interference because a finite-dimensional state cannot preserve arbitrarily many orthogonal directions.

Gated delta rules reduce interference by erasing the current prediction for an incoming key before writing its value. GDN-2 extends this design with channel-wise decay, erase, and write gates, but the paper identifies three bottlenecks: dense erase and write projections reduce FFN capacity, unequal decay rates create highly unequal state-row norms, and large-chunk parallel training becomes numerically difficult on accelerators with wide matrix units and slower vector units.

Muir's objective is therefore not simply to replace attention with another recurrence. It is to allocate model capacity according to how gates influence the state, correct readout imbalance using a matched filter, reinvest the parameter savings in the FFN, and provide a training kernel that makes the resulting recurrence efficient at large chunk sizes.

🔬 Methodology

The base recurrence decays the prior state row-wise, erases the current key prediction, and writes a gated value. Muir adds a fifth gate γ_t that scales the normalized query only at readout. This keeps the key unit-normalized for stable erasure while allowing the model to attenuate or amplify state rows according to token context and expected retrieval horizon.

The paper's sensitivity analysis shows that erase-gate changes are weighted by k_t,j² and the existing row energy, while write-gate changes act directly on v_t. This motivates a lower-rank erase projection and a higher-rank write projection. Muir fuses all five gates through one down-projection and five up-projections, with ranks (64, 64, 128, 128, 128), and uses channel biases so low-rank gates retain independent operating points.

Under an iso-parameter budget, the fused projection reduces gate cost and widens the FFN. At d = 2,048, the reported primary configuration uses d_ffn = 7,680, compared with 5,888 for GDN-2. The training implementation expresses each chunk as a triangular system, splits cumulative log-decay at interval boundaries so every exponential factor remains in (0, 1], recursively inverts block-triangular systems, and applies a custom backward pass.

📊 Experiments

The evaluation covers held-out FineWeb-Edu validation NLL, WikiText and LAMBADA perplexity, zero-shot commonsense benchmarks including PIQA, HellaSwag, WinoGrande, ARC-e, ARC-c, OpenBookQA, and SIQA, plus 5-shot WinoGrande and TriviaQA. The main table spans 166M, 406M, 834M, and 1.34B parameters after 100B training tokens. Retrieval evaluations include synthetic single-needle, 5-hop variable-binding, multi-needle, and multi-value tasks; RULER tasks; six real-world document-extraction tasks under the Based metric; S-NIAH-1 passkey retrieval; and long-context tests up to 10M tokens.

The authors report that Muir lowers validation NLL by 1.8–2.0% relative to GDN-2 and 2.5–3.6% relative to Transformer, raises commonsense accuracy by up to 1.9 points, and improves extractive recall over GDN-2. In the gate and readout ablation, the reported Muir row reaches 1.9973 FineWeb-Edu NLL at 1.34B, 16.41 WikiText perplexity, 10.00 LAMBADA perplexity, 53.84 LM-Eval accuracy, 34.60 Based recall, and 19.32 TriviaQA accuracy.

The synthetic experiments use 20 seeds. In the stationary single-head task, the row-energy law matches empirical energies across a 52.4× dynamic range with Spearman and log-Pearson correlations of 1.0. At a geometric horizon of mean H = 27, the horizon filter shifts optimal erase strength from 0.60 to 0.50, while full inverse-RMS whitening worsens retrieval. In the bimodal lag experiment, the lag-conditioned oracle changes correlation with the decay factor from −0.98 at lag 2 to +0.83 at lag 80. In learned associative recall, dynamic readout improves long-lag recall by 4.18 points over unit readout and 2.00 points over dense gates, while reducing final-timestep state energy by 45.7% relative to unit readout.

For long-context and systems evaluation, Muir extrapolates from 4K training to 10M-token streams with a constant 25MB float32 state and no reported perplexity drift, maintains near-perfect S-NIAH-1 retrieval to 1M tokens with static readout, matches GDN-2 throughput on 8 TPU7x chips at 0.5M tokens per step, reaches GDN-2's 100B-token validation loss 1.2× faster, and trains 2.7× faster than Transformer at 32K. These are strong results, but they measure different quantities: NLL, accuracy, recall, state memory, throughput, and time-to-target should not be treated as interchangeable efficiency claims.

fig:toys

Caption: 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 μ* from 0.60 to 0.50, with the Corollary 3.8 filter matching the static oracle. (b) On a 50/50 mixture of lags ℓ ∈ {2, 80} at μ = 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 and beats the static oracle on long and mixed horizons. (c) Low-rank erase gates r_b outperform dense gates at every bottleneck r...

Why It Matters: This figure directly tests the proposed matched-filter mechanism and the asymmetric rank allocation before the full language-model experiments.

fig:arch

Caption: 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, x_t feeds linear-Conv-SiLU pathways for (q_t, k_t, v_t) alongside a single fused low-rank gate projection that produces all five gates (log α_t, b_t, g_t, w_t, γ_t). (c) Asymmetric gate ranks (r_f, r_b, r_g, r_w, r_q) = (64, 64, 128, 128, 128) sum to R_tot = 512 = 2 × 256, which fills two 256-column MXU tiles in the fused down-projection and aligns (f, b) and (g, w, q) on 128-lane VMEM tile boundaries for batched up-projections. (d) At d = 2,048, replac...

Why It Matters: It shows how the theoretical sensitivity argument becomes a concrete architecture and how hardware tile alignment is part of the design.

fig:kernel

Caption: 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 × 256 MXU, an unfused 256 × 128 right-hand side V̂ achieves only 50% utilization due to zero-padded lanes, whereas the fused solve [U; Y] = T⁻¹[V̂; K̂ ⊙ π] achieves 100% utilization. (c) At chunk size C = 256, leaf size 8, and d_k = d_v = 128, replacing naive C × C × 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× and total kernel FLOPs by 12%, routing 98% of chunkwise arith...

Why It Matters: The kernel is necessary for turning the architectural savings into practical training throughput on wide matrix-unit hardware.

fig:nll_curves

Caption: Muir consistently lowers validation NLL relative to GDN-2 and Transformer throughout 100B-token pretraining. Held-out FineWeb-Edu validation token NLL difference relative to GDN-2 across 166M–1.34B parameters as a function of training tokens. All three Muir variants improve over GDN-2 within the first 10B tokens and reach a 12–20 mnat reduction by 100B tokens, while static and dynamic matched-filter readout provide additional reductions over Muir without readout gain.

Why It Matters: It indicates that the language-modeling improvement appears early in training and persists through the reported 100B-token horizon.

fig:recall

Caption: 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 variants mitigate the 4K degradation of GDN-2 and extrapolate to 8K, where unextended RoPE degrades beyond its 4K training window. (b) On 5-hop variable-binding chains, static readout fails to match the varying hop lags, whereas dynamic γ_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 matches or outperforms both GDN-2 and Transformer across 1K–4K. (e) Averaged across the RULER tasks by needle depth and context length, Muir achieves th...

Why It Matters: This is the main retrieval-focused comparison and separates the contribution of dynamic readout from the language-modeling gains.

fig:longctx_niah

Caption: Static matched-filter readout extends passkey retrieval to 1M tokens both zero-shot and after a 64K cooldown. S-NIAH-1 accuracy across 10 needle-depth bins and context lengths 1K–1M at 1.34B scale under zero-shot evaluation from 4K pretraining and pure-text continual cooldown at 64K. Cooldown extends Dynamic YaRN from 8K to 32K before early-depth accuracy drops at 64K and unchunked prefill exceeds HBM at 128K, whereas recurrent horizons change little under pure-text cooldown: GDN-2 collapses past 4K–8K, Muir without readout gain fades past 32K, dynamic readout holds through 32K before tapering at 64K, and static readout maintains near-perfect recall to 1M tokens in both regimes.

Why It Matters: It provides the clearest evidence that static matched-filter readout can preserve a single-lag retrieval signal far beyond the training context.

🔮 Conclusion

The paper presents a coherent co-design: sensitivity analysis motivates asymmetric gate ranks, stationary row-energy analysis motivates matched-filter readout, parameter savings widen the FFN, and the boundary-pivoted kernel makes large-chunk training practical on wide matrix-unit hardware.

The reported evidence supports improvements in language modeling, downstream reasoning, synthetic and real-world retrieval, long-context extrapolation, and training throughput. The strongest practical result is that static readout can preserve single-lag passkey retrieval to 1M tokens while keeping the recurrent state fixed at the reported size.

The results should be interpreted as evidence for a promising linear-attention design, not as proof that recurrent state can universally replace softmax attention. The paper leaves open how the method behaves at larger scale, on other accelerators, and on broader long-document workloads.

🛠️ Future Research Improvements

The most important next experiment is scale: test whether the reported rank proportions and FFN reallocation remain optimal as hidden size, depth, and parameter count grow beyond 1.34B. A useful study would separate gains from architecture, optimizer, kernel, and training-token allocation.

The theoretical assumptions should be stress-tested on natural-language activation statistics. In particular, measure whether key participation ratios, row-energy laws, and lag-dependent optimal gains remain stable across domains, languages, tokenization schemes, and instruction-tuned workloads.

The long-context evaluation would benefit from more real-document tasks at millions of tokens, variable-lag retrieval, multi-document interference, and streaming updates without a special cooldown phase. Hardware studies should also compare GPU and non-TPU implementations using matched batch sizes, precision, and memory constraints.

A practical research direction is adaptive rank allocation: use observed gate sensitivity, row-energy spread, or retrieval horizon to allocate rank per layer or head instead of fixing the same rank pattern throughout the network.

🏭 Potential Industry Use Scenarios

Muir is a plausible candidate for long-document language models where constant-size recurrent state is valuable: streaming summarization, persistent document assistants, log and event-stream modeling, and retrieval over sequences longer than the training window. The reported 25MB float32 state at 1.34B and 10M-token evaluation suggest a memory profile that could be attractive for continuous processing.

Its matched-filter readout may be useful in systems that must retrieve facts at heterogeneous time horizons, such as multi-hop document analysis, multi-needle extraction, and passkey-like stateful workflows. The dynamic readout is particularly relevant when the appropriate channel weighting changes with the retrieval lag.

The kernel co-design is relevant to accelerator-focused training stacks. Fused gate projections, MXU-aligned ranks, and recursive chunkwise solves could reduce training bottlenecks on hardware where matrix units are much faster than vector units.

Deployment decisions should remain conditional on workload validation: the paper establishes promising benchmark behavior, but production use would require latency, fault tolerance, quantization, checkpoint portability, and broad real-data evaluations.

💬 Critical Analysis

The paper's strongest aspect is the alignment between mechanism and measurement. The sensitivity derivatives motivate the asymmetric gate ranks; the stationary energy law is checked against trained checkpoints; matched-filter readout reduces measured readout energy spread; and the synthetic experiments show that full whitening is not automatically beneficial. This is more persuasive than presenting the architecture only through end-task scores.

The main technical risk is that the proposed law is derived under simplified stochastic assumptions while the final model operates on correlated, nonstationary language representations. The checkpoint diagnostics support the direction of the theory, but they do not establish that the approximations remain accurate across every layer, domain, or retrieval horizon.

The systems contribution is substantial but hardware-specific. The reported 24× VPU-compute reduction and 98% MXU arithmetic share are meaningful for the evaluated TPU7x configuration, yet the wider FFN, extra query gate, recursive inversion, and custom backward pass may have different trade-offs on GPUs or smaller accelerators.

For builders, the fairest reading is that Muir offers a strong blueprint for capacity-efficient linear attention with explicit long-context retrieval mechanisms. Reproduction should preserve metric direction, distinguish perplexity from recall and throughput, and separately validate architectural gains from kernel and hardware gains.