Multi-state linear attention improves long document understanding
Multi-State Linear Attention with Generalized Second-Order Routing
Computation and Language
Summary
Understanding very long sequences of text is hard and slow for many AI models because they compare every part of the text with every other part. The authors worked on a method called multi-state linear attention that breaks the text into chunks and keeps track of each chunk separately. They improved the way those chunks are combined by not just looking at average importance but also noticing how varied that importance is inside each chunk. This new method, called MERLIN, helps AI models better find relevant information in long texts and perform better on tasks like answering questions from long documents.
What this means in practice
- •For ai model developers: Enhance large-scale language models to handle longer text inputs more effectively and efficiently with improved memory use.
- •For natural language processing engineers: Build question-answering systems that better identify relevant information within very long documents using moment-based attention routing.
Authors
Hee Seung Hwang, Sanghyuk Chun, Xindi Wu, Olga Russakovsky, Zhiwei Deng
Abstract
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.