Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Recurrent Transformer computes each layer's key-value not from the previous layer but from the same layer's output, making later tokens see representations that have already passed through attention and MLP. The result is a structure with layer-wise independent recurrent memory, while keeping the autoregressive decoding cost of a standard Transformer.

Under weak assumptions, this structure can mimic the behavior of an existing Transformer, and conversely, by restricting attention to only previous positions, it can also implement token-wise recurrent updates. It creates more multi-hop paths, alleviating the long-path problem of RNN-style models, and training is stable when combined with normalization before key-value computation and depth-wise residual scaling.

Since key/value values are revealed sequentially during training and prefill, a naive implementation runs into a bandwidth bottleneck, but an exact tiling algorithm reduces HBM traffic from Θ(N^2) to Θ(N log N) and raises effective arithmetic intensity up to Θ(N/log N). In C4 pretraining with 150M and 300M parameters, cross-entropy improved over Transformer baselines with the same parameter count, and the model achieved performance with fewer layers, showing potential to reduce KV cache memory and inference latency.