Triadic linear attention improves long context modeling with 3d memory
Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
Machine LearningComputation and Language
Summary
Long-context sequence models need to remember a lot of information efficiently. The authors improve on existing methods by expanding memory from flat lists to three-dimensional structures, allowing these models to keep track of much more detail without much extra cost. They tested this new approach and found it helps models better understand and recall long sequences of text. This can lead to smarter language models that handle lengthy conversations or documents more effectively.
What this means in practice
- •For machine learning engineers: Build language models that better recall long documents by using triadic linear attention to expand memory efficiently.
- •For speech recognition developers: Improve transcription accuracy for long audio by integrating triadic linear attention into sequence models handling extended contexts.
Authors
Oliver Sieberling, Bharat Runwal, David Jin, Ryan Chin, Rameswar Panda, Yoon Kim
Abstract
Recurrent neural networks (RNNs) compress the historical context into a memory state of fixed size, thus allowing for constant-time inference. The memory state size is a crucial factor in their performance, as exemplified by the strong performance and resurgence of linear attention, which extends the vector-valued hidden states of ordinary RNNs to matrix-valued hidden states. Crucially, linear attention does so in a parameter-efficient way, in particular by using an outer product of the key and value vectors to write to the matrix-valued hidden state. We generalize this construction and propose triadic linear attention, which writes the triadic outer product of a key, a second key, and a value, into a third-order (i.e., 3D) tensor state, and reads from it by contracting both key axes with two queries. An $E$-dimensional second key thus yields an $E$-fold increase in state size while adding only two projections. Triadic linear attention is compatible with data-dependent forgetting, the delta rule, and chunkwise-parallel training. Applied to Gated DeltaNet and scalar-gated linear attention, triadic linear attention substantially improves long-context language modeling and recall, outperforming alternatives that enlarge the state.