KronLSTM improves long context memory with stable linear attention
Kronecker Delta Rule for Long Context Linear Attention
Machine LearningArtificial Intelligence
Summary
Long language sequences are hard for computers to remember without mixing up old information. The authors developed KronLSTM, a method that updates a computer's memory carefully so it doesn't forget or mess up previous memories. This approach keeps the memory steady, making it better for understanding long texts. Experiments show that KronLSTM often works better than other similar techniques, especially when dealing with very long information.
What this means in practice
- •For natural language processing teams: Improve language models to handle longer documents with less memory interference and more stable performance using KronLSTM.
- •For software engineers building ai assistants: Create AI assistants capable of remembering and revising long user conversations more accurately over time employing KronLSTM's memory update method.
Authors
Anamaria-Roberta Hartl, Levente Zólyomi, Michael List, Günter Klambauer, Sebastian Böck, Sepp Hochreiter
Abstract
We propose KronLSTM, a linear attention architecture based on the Kronecker Delta Rule, a novel update rule for associative matrix memories. KronLSTM extends the xLSTM family with a memory-revision mechanism that corrects targeted associations while modifying the memory only along a single direction in matrix space. This minimally invasive memory revision reduces interference with previously stored information compared with delta-rule updates, making it particularly well suited for long-context language modeling. The resulting memory dynamics can be mathematically described as a Markov chain that converges to a steady-state distribution, while the expected norm of the memory matrix remains uniformly bounded throughout the entire input sequence. Hence, KronLSTM ensures stable memory dynamics without requiring additional normalization and stabilization mechanisms for the memory state. Experiments show that KronLSTM is competitive with, and often outperforms, existing linear attention architectures, including models of the (Gated) DeltaNet family, with particularly strong performance on long context benchmarks. These results establish the Kronecker Delta Rule as an effective and analytically tractable foundation for stable associative memory in recurrent linear attention architectures. Our code (https://github.com/NX-AI/KronLSTM) and models (https://huggingface.co/collections/NX-AI/kronlstm) are open-source.