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

Muir. Linear Attention with Sensitivity-Allocated Gating and Matched-Filter Readout. By Tushaar Gangavarapu and two others, from the University of Texas at Austin and Google. Published on the fifth of November.

Linear attention squeezes context into a key-value state of fixed size, which buys linear-time training and inference in constant memory. Gated delta-rule models manage that state in three moves: they decay it, erase what it currently predicts for an incoming key, and write a new value. Gated DeltaNet-2, or GDN-2, generalizes this with channel-wise decay and with erase and write gates that are decoupled and full rank.

The authors find two imbalances in that arrangement. First, the state is far less sensitive to erasure than to writing. Key activation and unit normalization concentrate a key's energy on a sparse subset of channels, while a value write scales every channel, so the erase gate does not need full rank. Second, state rows that decay slowly build up much larger norms than rows that decay quickly. Those rows dominate a readout that treats every row equally, which forces the model into stronger erasure at write time just to keep interference down.

Muir answers each imbalance directly. All the gates are fused into a single low-rank projection whose capacity is allocated by sensitivity. A matched-filter readout gate re-weights state rows when the state is read, so the loud slow rows no longer drown out the rest. The parameters saved by the low-rank gating are reinvested in a wider feed-forward network.

There is also a training obstacle. Saturating wide matrix units calls for large chunk sizes, and at those sizes the standard chunkwise algorithms become numerically unstable. The authors design a parallel update that keeps training both stable and efficient at that scale.

Models span one hundred sixty-six million to one point three four billion parameters. Across that range, Muir lowers perplexity relative to GDN-2 by between one point eight and two percent, and relative to a Transformer by between two point five and three point six percent. Commonsense accuracy rises by as much as one point nine points, and extractive recall improves over GDN-2.

At the largest size, Muir beats both baselines on multi-hop and multi-needle recall. Trained at a length of four thousand tokens, it extrapolates to streams of ten million tokens, a factor of two thousand five hundred and sixty, with no drift in perplexity. With static readout, passkey retrieval stays near one hundred percent out to one million tokens.

None of this costs speed. Throughput matches GDN-2. Muir reaches its validation loss at one hundred billion tokens one point two times faster, and at a context of thirty-two thousand tokens it trains two point seven times faster than a Transformer.

Muir's gains rest on measuring where a gated state is actually sensitive and spending parameters only there. Erasure gets less capacity, readout learns which rows to trust, and the budget that frees up goes into a wider feed-forward network.