The Recurrent Transformer: Greater Effective Depth and Reduced KV-Cache Footprint
Costin-Andrei Oncescu ⋅ Depen Morwani ⋅ Samy Jelassi ⋅ Alexandru Meterez ⋅ Mujin Kwun ⋅ Sham Kakade
Abstract
Autoregressive foundation models operating under memory and bandwidth constraints benefit directly from shallower networks, but reducing layer count ordinarily sacrifices quality. We introduce the *Recurrent Transformer*, a simple architectural change where *each layer* attends to key-value pairs computed off its own activations, yielding layerwise recurrent memory while preserving standard autoregressive decoding cost. This adds temporal depth within a layer and enables a different depth-width tradeoff at fixed parameter count. On 300M-parameter C4 pretraining, a 6-layer Recurrent Transformer improves validation cross-entropy over 12- and 24-layer Transformer baselines while reducing cached key-value state by approximately $27\%$ and $50\%$, respectively. The gain persists in a compute-controlled comparison where a reduced-width RT uses fewer parameters and no more dense-block FLOPs than the Transformer baselines. Naive training and prefill have effective arithmetic intensity $\Theta(1)$ because persistent keys and values are revealed sequentially; an exact tiled schedule preserves the mathematical computation while reducing HBM traffic from $\Theta(N^2)$ to $\Theta(N\log N)$ and increasing effective arithmetic intensity to $\Theta(N/\log N)$. The architecture can emulate both conventional self-attention and token-to-token recurrent updates while retaining direct attention paths. Together, these results position layerwise recurrence as a foundation-model design with a reduced decoding footprint and a practical training/prefill schedule.
Chat is not available.
Successful Page Load