Towards Compressive and Scalable Recurrent Memory

TL;DR

Elastic Memory introduces a novel, scalable recurrent memory architecture based on the HiPPO framework, enabling efficient long-context processing and outperforming existing methods in memory usage and speed.

Contribution

The paper presents Elastic Memory, a new memory architecture that applies optimal online compression and flexible sampling, improving long-context modeling beyond current recurrent memory approaches.

Findings

Outperforms baselines on 32k+ long-context datasets

Uses 16x less memory than Memorizing Transformer at equal parameters

Faster and more effective than Melodi at larger scales

Abstract

Transformers face a quadratic bottleneck in attention when scaling to long contexts. Recent approaches introduce recurrent memory to extend context beyond the current window, yet these often face a fundamental trade-off between theoretical principles and practical scalability. To address this, we introduce Elastic Memory, a novel memory architecture grounded in the HiPPO framework for online function approximation. Elastic Memory treats historical sequence as samples from continuous signals, applying optimal online compression to encode them into a fixed-size memory state. For retrieval, we propose a flexible \textit{polynomial sampling} mechanism that reconstructs a history summary from this compressed state. Elastic Memory consistently outperformed baselines on long-context (32k+) datasets across three domains. With equal parameters, it beat Memorizing Transformer by 16x memory and outperformed Melodi at all memory sizes, even when Melodi had 30% more parameters. When scaling model size, Elastic Memory stayed ahead of all baselines and was significantly faster than Melodi at 4x size. Furthermore, its decoupled design allows for injecting inductive biases at test-time to boost performance.