Language Model Memory and Memory Models for Language

TL;DR

This work investigates how language-model memories store input information and introduces encoder–decoder memory architectures that enable near-arbitrary information retrieval with parallelizable processing. It demonstrates that causal language-model embeddings retain little information, while autoencoder-based memories form rich, reconstructible representations; combining causal and copy objectives or curriculum training yields accurate memory formation. The study provides metrics for memory fidelity (entropy-based and token-accuracy measures) and shows that frozen encoders can achieve high throughput, while pretrained LLM decoders can be repurposed to process memory embeddings with manageable trade-offs. These findings have implications for improving retrieval, long-context reasoning, and efficient inference by replacing sequences with compact, information-rich memories. Overall, the results suggest that carefully designed memory models with combined objectives can surpass the information limitations of standard causal embeddings, unlocking scalable memory-enabled inference for future language systems.

Abstract

The ability of machine learning models to store input information in hidden layer vector embeddings, analogous to the concept of `memory', is widely employed but not well characterized. We find that language model embeddings typically contain relatively little input information regardless of data and compute scale during training. In contrast, embeddings from autoencoders trained for input regeneration are capable of nearly perfect memory formation. The substitution of memory embeddings for token sequences leads to substantial computational efficiencies, motivating the introduction of a parallelizable encoder-decoder memory model architecture. Upon causal training these models contain information-poor embeddings incapable of arbitrary information access, but by combining causal and information retention objective functions they learn to form and decode information-rich memories. Training can be further streamlined by freezing a high fidelity encoder followed by a curriculum training approach where decoders first learn to process memories and then learn to additionally predict next tokens. We introduce the perspective that next token prediction training alone is poorly suited for accurate memory formation as the objective itself is non-invertible, motivating the use of combined objective functions for models where the entire input is not exposed.

Figures (13)

• Figure 1: Information retention experimental approach (left) and example training runs (right).
• Figure 2: Memory Model Architecture and $n_{ctx}=256$ per chunk, $s=4$ chunk causal training characteristics on FineWeb. Mixers are $d_m=512$ for encoders, $d_m=1024$ for decoders and Transformers $d_m=256$ and $d_m=512$ for compute equivalence.
• Figure 3: (a) Copy and (b) Blank Copy experimental schematics
• Figure 4: Memory Model Information Content and Copy Training. (a) Memory model training on for the copy task. (b) $d_m=256, n_{ctx}=512$ Transformer and $d_m=512, n_{ctx}=256$ Mixer memory model encoder information, (c) Memory Transformer training on the blank copy task, (d) Frozen encoder copy training, all curriculum trained except CLM encoder and Memory encoder. All models $d_m=512, n_l=16, n_{ctx}=256, s=4$ unless otherwise noted and copy 512 tokens.
• Figure S1: Mixer and Transformer information retention validation training curve examples on FineWeb. Transformers are $d_m=512, h=4$ and mixers $d_m=1024$ non-headed, both with $n_{ctx}=512, n_l=16$.