Transformer language model improves generalization with recurrent memory

Persistent Recurrent Memory Between Transformer Layers - Improves Language Model Generalization

Computation and Language

Summary

Language models sometimes struggle to learn well from small amounts of data without overfitting. The authors added a small extra memory component between layers of a transformer model that keeps track of information over time. This change made the model better at predicting text it hadn’t seen before, reducing errors significantly. They showed this improvement came from the new memory setup itself, not from any extra tasks or tricks. The memory also seemed to help the model understand story position, which normal transformers handle less effectively.

What this means in practice

Authors

Eduardo Novaes Hering

Abstract

We introduce a simple architectural modification to decoder-only transformers: a persistent recurrent state that observes hidden representations via cross-attention, updates itself through a GRU, and modulates subsequent processing via gated addition. Inserted between the lower and upper halves of a 6-layer transformer, this module adds only 3.7\% additional parameters while reducing evaluation loss from $2.438 \pm 0.004$ to $1.743 \pm 0.018$, corresponding to a 28.5\% reduction on held-out language modeling data. The improvement is statistically significant across 5 random seeds ($p < 0.01$) and corresponds to reduced overfitting (generalization gap 0.12 vs 0.26). Through controlled ablations, we demonstrate that the improvement stems entirely from the persistent memory topology, not from auxiliary self-prediction objectives. A model with identical topology but no auxiliary loss performs equivalently, while a random auxiliary loss provides no benefit. Representation probing reveals that the persistent state encodes narrative position (52\% vs 33\% chance level)---information that standard attention maintains less efficiently. Our results suggest that bridging transformer layers with a lightweight recurrent memory is a simple, effective approach to improving generalization in small-scale language models.