🤖 AI Summary
This study addresses the limitation of linear attention in modeling long contexts due to its reliance on fixed-size matrix states. To overcome this, we propose Ternary Linear Attention, which generalizes the hidden state from a second-order matrix to a third-order tensor. The state is updated via the third-order outer product of the key, a secondary key, and the value, and retrieved through dual-query contraction. Efficient optimization is achieved by integrating a data-dependent forgetting mechanism, the Delta rule, and chunkwise parallel training. This approach significantly expands memory capacity with minimal parameter overhead. When incorporated into architectures such as Gated DeltaNet, it substantially enhances long-context language modeling and recall capabilities, outperforming existing state-expansion methods.
📝 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.