Extending LLM Context via Associative Recurrent Memory

Closed-source LLMs support context windows of up to 1M tokens and beyond, while smaller open-source models still have a limited window of 32K-128K tokens. However, for some tasks, it is necessary to use a small local LLM to avoid any data leaks; some of these tasks require domain-specific knowledge and long-context understanding. To address this gap, we introduce two domain-specific long-context datasets, ManyTypes-long and GovReport-long, and present a practical recipe for extending short-context LLMs using the Associative Recurrent Memory Transformer (ARMT) architecture. Finally, we analyze the associative memory in trained ARMT models and show that associative memory primarily benefits from representations in the middle and upper layers of the transformer, allowing us to reduce the size of the models by removing redundant associative memory from other layers. Our results demonstrate an effective approach to enabling long-context capabilities in small, privacy-preserving LLMs for domain specific tasks.