[论文] Multi-State Linear Attention with Generalized Second-Order Routing

论文概要

研究领域: NLP / 注意力机制
作者: Hee Seung Hwang, Sanghyuk Chun, Xindi Wu, Olga Russakovsky, Zhiwei Deng
发布时间: 2026-10-27
arXiv: 2610.37645

中文摘要

自注意力在长序列建模中起关键作用，但其二次时间复杂度从根本上限制了LLM的扩展。多状态线性注意力（MSLA）通过将上下文分割为段并在各自的记忆状态中压缩来提高效率，但MSLA仅使用query与每段keys的平均对齐度来聚合信息，忽略了对齐度在段内的分布。为更全面建模段与query的相关性，本文提出MERLIN（MomEnt Routing for multi-state LINear attention），跟踪每段keys的一阶和二阶矩。利用keys的矩，query可以恢复其与每段对齐的均值和方差，实现更有内容感知的聚合。在Gated DeltaNet-1.3B和Mamba-2-780M的多状态变体上，MERLIN在长上下文任务上一致优于现有聚合方法：RULER needle-in-a-haystack和长上下文QA上提升+8.3%（Gated DeltaNet）和+7.0%（Mamba-2），LongBench多样子任务上提升+13.5%和+12.0%，甚至超过了消耗4倍内存的基线。

原文摘要

Self-attention plays a key role in long-sequence modeling, but its quadratic time complexity fundamentally limits the scaling of Large Language Models. Recently, Multi-State Linear Attention (MSLA) mechanisms have been proposed to make sequence modeling more efficient by splitting the context into segments and compressing each in its own memory state. MSLA aggregates the information extracted from each memory state using only the mean alignment between a query and each segment's keys, ignoring how the alignment is distributed within the segment. In order to model the relevance of a segment to a query more comprehensively and incorporate the variance of the alignment in the aggregation, we propose MERLIN (MomEnt Routing for multi-state LINear attention), a method that tracks the first and second moments of each segment's keys. Using the moments of keys, a query can recover both the mean and the variance of its alignment with each segment to perform a more content-aware aggregation of per-state readouts. Experiments on multi-state variants of Gated DeltaNet-1.3B and Mamba-2-780M show that MERLIN consistently outperforms existing aggregation methods on various long-context tasks. In particular, MERLIN outperforms baselines on RULER needle-in-a-haystack and long-context QA (+8.3% for Gated DeltaNet-1.3B and +7.0% for Mamba-2-780M) and diverse subtasks in LongBench (+13.5% and +12.0%). It even outperforms baselines that consume 4x more memory. MERLIN provides a simple and general routing mechanism for improving long-context capabilities of MSLA models.

自动采集于 2026-10-29

#论文 #arXiv #AI #小凯