Lattice: Learning to Efficiently Compress the Memory

Mahdi Karami, Razvan Pascanu, Vahab Mirrokni

Keywords: Sequence models, Memory Compression, RNN, Transformer, State Space Models

Abstract:

Attention mechanisms have revolutionized sequence learning but suffer from quadratic computational complexity. This paper introduces Lattice, a novel recurrent neural network (RNN) mechanism that leverages the inherent low-rank structure of K-V matrices to efficiently compress the cache into a fixed number of memory slots, achieving sub-quadratic complexity.
We formulate this compression as an online optimization problem and derive a dynamic memory update rule based on a single gradient descent step. The resulting recurrence features a state- and input-dependent gating mechanism, offering an interpretable memory update process.
The core innovation is the orthogonal update: each memory slot is updated exclusively with information orthogonal to its current state, hence incorporating only novel, non-redundant data, which minimizes the interference with previously stored information.
The experimental results show that Lattice achieves the best perplexity compared to all baselines across diverse context lengths and model sizes.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 9658