论文

Muir:低秩门控加匹配滤波读出的线性注意力新架构

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

精选理由

一个改进 Gated DeltaNet-2 的线性注意力新架构,长文本外推到千万 token 还不掉点,训练还比 Transformer 快 2.7 倍,做长上下文的值得看看实现思路。

论文提出 Muir,一种线性注意力架构,核心发现是状态更新对擦除的敏感度远低于写入,因此擦除门可以用低秩参数化。Muir 将所有门融合为按敏感度分配的低秩投影,并用匹配滤波读出门重新加权状态行,把省下的参数投入更宽的 FFN。在 166M-1.34B 参数规模上,Muir 相比 GDN-2 降低 1.8-2.0% 的困惑度,相比 Transformer 降低 2.5-3.6%。1.34B 版本从 4K 训练长度外推到 1000 万 token 流而没有困惑度漂移,静态读出下在 100 万 token 内保持接近 100% 的 passkey 检索准确率,同时训练速度比 32K 上下文的 Transformer 快 2.7 倍。

原文 · arXiv cs.LG

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

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.