RunningTensor improves memory for fast sequence processing models
RunningTensor: Generalizing Linear Attention to Higher-Order Recurrent States
Machine LearningArtificial Intelligence
Summary
Many tools that handle sequences of information quickly use a kind of memory stored as a matrix, which limits how much detail they can remember. The authors created RunningTensor, which lets this memory be a higher-dimensional shape, allowing more complex information to be stored without slowing down. They tested it with synthetic tasks and found it better than existing fast methods, and it also helped with real language understanding and retrieval tasks. This suggests using bigger memories in such models can improve performance.
What this means in practice
- •For machine learning engineers: Build fast sequence models that remember more complex patterns using higher-order memory tensors with RunningTensor.
- •For information retrieval developers: Improve retrieval accuracy in language-based search systems by integrating recurrent models with higher-order tensor states.
Authors
Luca Herranz-Celotti, Vincent Guigue
Abstract
Linear attention and state-space models provide linear-time sequence modeling, but their recurrent memory remains a second-order tensor (a matrix), limiting the order of interactions that can be represented in the state. We introduce the RunningTensor, which generalizes this memory to an order-$o$ tensor, updated by a rank-1 outer product and read by contracting against $o-1$ vector queries. Order $2$ recovers linear attention; we study order $3$ as a proof of concept, retaining both recurrent and parallel forms while remaining linear in sequence length $T$ and improving working memory capacity from $\mathcal{O}(W^2)$ to $\mathcal{O}(W^o)$. On synthetic multi-query associative recall, RunningTensor outperforms linear-attention and SSM baselines. After pretraining, it also improves performance on language-understanding and non-synthetic retrieval tasks, suggesting that higher-order recurrent state can provide useful additional memory capacity beyond matrix-valued state.