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

AI summary

Overview

Research area: Efficient sequence modeling for large language models — specifically gated linear attention (delta-rule recurrences) and hardware-aware kernel design for wide-matrix-unit accelerators.

Technical level: Advanced. The paper assumes familiarity with linear attention, the delta rule, chunkwise-parallel training forms, and accelerator memory hierarchies (HBM, VMEM, MXU/VPU tiling).

Scope: The paper analyzes why channel-wise gated delta-rule state updates are more sensitive to writes than to erasure, and uses that analysis to design a new architecture (Muir) plus a numerically stable chunkwise kernel, evaluated from 166M to 1.34B parameters.

What This Paper Is About

Linear attention compresses a whole context into a fixed-size key-value state, but packing many keys into a finite state causes interference, so delta-rule models erase a key's old prediction before writing a new value. Gated DeltaNet-2 (GDN-2) made decay, erasure, and writes channel-wise and gave erasure and writes equally large full-rank projections, which the authors show is wasteful and which shrinks the feed-forward network by over 25% under a fixed parameter budget. Muir reallocates that parameter budget according to how sensitive the state actually is to each gate, re-weights state rows at readout, and adds a training kernel that stays stable at the large chunk sizes modern accelerators need.

Key Contributions

1. A sensitivity analysis of channel-wise gated delta recurrences. Proposition 3.1 derives directional derivatives of the state with respect to the erase gate and the write gate, plus their expected squared Frobenius sensitivities under isotropic perturbations: the write term is simply the squared value norm, while the erase term is weighted jointly by squared key entries and decayed state-row energy. Because SiLU activation and unit normalization concentrate key energy onto a sparse set of channels (quantified by a key participation ratio in [1, d_k]), the erase gate can be parameterized at much lower rank than the write gate.

2. A closed-form scaling law for state row norms and a matched-filter readout. The paper derives how slow- versus fast-decaying state rows differ in norm, shows that unweighted query readout lets high-norm slow rows dominate and forces stronger write-time erasure, and proposes a readout gate that re-weights state rows using a static channel profile combined with dynamic per-token modulation.

3. The Muir architecture. All recurrence and output gates are fused into one sensitivity-allocated low-rank projection, and the saved gate parameters are reinvested to widen the FFN within the same 12d² per-layer budget used for iso-parameter comparison.

4. A boundary-pivoted chunkwise kernel. Standard chunkwise algorithms divide keys by cumulative decay prefixes, which underflows or overflows at large chunks; the proposed kernel splits relative decays without division or full-chunk 3D decay tensors and solves the within-chunk triangular system through block matrix multiplications, keeping work on the wide matrix units.

Main Findings

• Erasure tolerates far lower rank than writing. The expected squared state sensitivity to a write perturbation is ‖v_t‖²₂, independent of the state, whereas the erase sensitivity is Σ_j k²_{t,j}‖Ŝ_{t,j,:}‖²₂ — it vanishes on inactive key channels and on empty or decayed rows. The authors report that the gradient of the contraction scalar μ_t is supported on η_eff(k_t) ≪ d_k active channels.

• Uneven row norms distort readout. State rows decay at different rates, so slow rows accumulate writes over hundreds to thousands of tokens and row norms span more than an order of magnitude; under unweighted readout these rows dominate and force stronger write-time erasure, which prematurely overwrites stored associations.

• Parameter waste in GDN-2. The dense erase and write projections W_b, W_w ∈ ℝ^{d×d} consume 2d² parameters, one-sixth of the 12d² layer budget, shrinking FFN width from 8,192 to 5,888 at d = 2,048 (d_ffn = 3d − 256), a reduction of over 25%.

• Lower perplexity and NLL across scales. Over 166M–1.34B parameters, Muir lowers perplexity by 1.8–2.0% over GDN-2 and 2.5–3.6% over Transformer; Figure 1 reports validation NLL reductions of 18.0–20.4 millinats over GDN-2 and 25.0–36.2 millinats over Transformer.

• Better commonsense and recall accuracy. Commonsense accuracy rises by up to 1.9 points, and extractive recall improves over GDN-2. At 1.34B, Muir outperforms both baselines on multi-hop and multi-needle recall.

• Long-context extrapolation without drift. Muir extrapolates 2,560× from its 4K training length to 10M-token streams with no perplexity drift, holding a 25 MB float32 state, while Dynamic YaRN Transformer diverges past 8K and would need a 2.1 TB bfloat16 KV cache at 10M.

• Passkey retrieval holds to 1M tokens. Static matched-filter readout gives near-100% passkey accuracy to 1M tokens, whereas GDN-2 collapses past 4K–8K.

• Throughput is preserved or improved. On 8 TPU7x chips at 0.5M tokens/step and 1.34B parameters, Muir matches GDN-2 throughput despite an extra readout gate and a 30.4% wider FFN, reaches GDN-2's 100B-token validation loss 1.2× faster, and trains 1.64× and 2.74× faster than Transformer at 16K and 32K (the abstract rounds the 32K figure to 2.7×).

• Ablation attribution. Per Figure 1(a), the sensitivity-allocated low-rank gates plus FFN expansion drive the bulk of the validation NLL gain over GDN-2, while the dynamic matched-filter readout drives the gain on extractive recall.

• Mechanistic checks support the theory. The authors report that state updates are confirmed to be more write- than erase-sensitive, that row norms closely follow the derived scaling law, and that the learned readout balances rows, lowers required write-time erasure, and retains associations over longer contexts.

Methodology in Plain English

The approach is analysis-first. The authors begin from the GDN-2 update rule and ask a simple question with calculus: if you nudge the erase gate or the write gate, how much does the memory state actually move? The derivative for writes scales the incoming value directly, so every channel is affected every step. The derivative for erasure is filtered twice — once by how much of the key's energy sits on that channel, and once by how much stored energy that row already holds — so most erase channels rarely matter. A second piece of analysis tracks how row norms grow over long sequences given per-channel decay rates, explaining why some rows dominate the readout.

Those two results drive the design. Instead of giving erasure and writing equally wide matrices, Muir gives erasure a low-rank, biased projection (so different state rows still get distinct baseline erase rates) and fuses all gates into one low-rank projection. The parameters freed up go into a wider FFN. On the readout side, queries are multiplied by a learned gate that mixes a fixed per-channel profile with a per-token component, evening out the contribution of slow and fast rows.

The training kernel is an engineering response to a numerical problem: computing within-chunk interactions in the standard way requires dividing keys by cumulative decay products, which underflow (the paper notes π can reach 2⁻²⁵⁶ ≈ 10⁻⁷⁷ at chunk size 256) on fast-decaying channels, while the stable alternative — materializing a full chunk-by-chunk-by-channel decay tensor, about 32 MB per head in float32 at C = 256 — runs on the slow vector unit. The boundary-pivoted kernel avoids both, keeping the triangular solve in matrix multiplications that map onto 256×256 matrix units.

Evaluation compares Muir against GDN-2 and Transformer at matched parameter budgets from 166M to 1.34B, using perplexity/NLL, commonsense accuracy, extractive recall, synthetic multi-hop and multi-needle retrieval, and passkey retrieval, with throughput measured on TPU7x Ironwood chips (192 GB HBM at 7,380 GB/s, 128 MB VMEM, 2,307 TFLOP/s bfloat16 per chip). The provided text is truncated partway through Section 3, so the synthetic memory experiments of §3.2, the full training data and hyperparameter details, and the appendix proofs are not reported in the supplied content.

Why This Matters

Impact on research. The paper argues against a symmetry assumption baked into recent gated delta-rule designs: if erasure and writing affect the state differently, allocating identical full-rank matrices to each is a misallocation. It pairs this with a hardware argument — that kernels tuned for 16×16 TensorCores do not transfer to accelerators with 256×256 matrix units and a vector unit roughly 64× slower on bfloat16 — and shows architecture and kernel must be co-designed.

Real-world applications.

• Long-document question answering and summarization over contexts far beyond the training length.
• Retrieval-augmented and in-context retrieval workloads, where extractive recall over many needles matters and the KV cache of attention is prohibitive.
• Multi-hop reasoning over long inputs, where the 1.34B model is reported to beat both baselines.
• Streaming or session-long assistants needing constant-memory inference (25 MB state versus a projected 2.1 TB bfloat16 cache at 10M tokens).

Industry relevance. Constant-memory inference and linear-time training directly reduce serving memory and compute cost for long contexts, and the reported throughput parity with GDN-2 at a 30.4% wider FFN means the quality gains are not paid for in step time on TPU-class hardware.

Future Directions

• Scaling beyond 1.34B. The paper's largest reported model is 1.34B; whether the sensitivity-based parameter allocation holds at frontier scale is not established in the provided text.
• Portability of the kernel. The boundary-pivoted chunkwise design targets wide-matrix-unit accelerators such as TPU7x and Trainium; behaviour on GPU-style hardware with 16×16 TensorCores, and the best chunk size there, is an open question.
• How the rank allocation should be chosen. The analysis motivates low-rank erasure, but whether the rank split between gates should be fixed, tuned per layer, or learned is not resolved in the content provided.
• Composition with other efficiency methods. Whether the freed FFN capacity and the matched-filter readout combine cleanly with Mixture-of-Experts, quantization, or hybrid attention layers is untested here.

Target Audience

Researchers and engineers working on linear attention, state-space models, and gated delta-rule architectures; kernel and compiler engineers optimizing sequence mixing for TPU-class and other wide-matrix-unit accelerators; and practitioners evaluating whether constant-memory models can replace attention for long-context retrieval workloads. Readers without background in delta-rule recurrences or accelerator memory hierarchies will find Sections 2 and 3 demanding, though the introduction and Figure 1 convey the main results at a higher level.

Authors’ 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.