The Surprising Effectiveness of Shared Memory in Looped Transformers

2026-10-01

AI Summary

This study addresses the linear growth of KV cache with computational steps in recurrent Transformers by proposing a pre-trained shared memory mechanism. The method writes to the cache only during the initial recurrence, while subsequent steps directly read from it and retain merely a short context window, thereby substantially compressing memory footprint. Through large-scale pretraining based on recurrent predictive Transformers and their hybrid architectures, we find that shared memory not only avoids performance degradation but actually enhances model capabilities, establishing a new quality-memory Pareto frontier. At parameter scales ranging from 150M to 1B, this approach reduces perplexity by 1.12–1.82 while decreasing context memory consumption by 76%–79%.

Abstract

Looped Transformers apply the same layers several times per token, adding compute to improve quality without more parameters. Each recursion, however, writes its own key-value cache, so memory still grows with compute. Inference-time techniques can shrink this cache at a cost in quality. We pretrain looped language models to share memory: only the first recursion writes a cache, and later recursions read it while keeping a short window of their own. Surprisingly, we find that sharing memory does not cost quality and instead improves it. At 150M-1B parameters, our Looped Prediction Transformer (LPT) and its hybrid variant set a new quality-memory frontier for looped models: with five recursions, the hybrid lowers validation perplexity on FineWeb-Edu by 1.12-1.82 relative to a same-size standard Transformer while using 76-79% less context memory. Through an extensive analysis, we investigate why memory sharing helps. Shared and local memory develop different representations, and later recursions attend mostly to the shared memory, which also acts as a gradient highway to the first recursion.

Problem

• Looped Transformers
• Shared Memory
• Key-Value Cache
• Language Models
• Memory Efficiency

Innovation

• Looped Transformers
• Shared Memory
• Key-Value Cache
• Language Model Pretraining
• Gradient Highway