The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Abstract: Transformers process tokens in parallel but are temporally shallow: at position $t$, each layer attends to key-value pairs computed based on the previous layer, yielding a depth capped by the number of layers. Recurrent models offer unbounded temporal depth but suffer from optimization instability and historically underutilize modern accelerators. 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. We show that the architecture can emulate both (i) a conventional Transformer and (ii) token-to-token recurrent updates under mild assumptions, while avoiding optimization instability. Naively, prefill/training appears bandwidth-bound with effective arithmetic intensity near $1$ because keys and values are revealed sequentially; we give an exact tiling-based algorithm that preserves the mathematical computation while reducing HBM traffic from $Θ(N^2)$ to $Θ(N\log N)$, increasing effective arithmetic intensity to $Θ(N/\log N)$ for sequence length $N$. On 150M and 300M parameter C4 pretraining, Recurrent Transformers improve cross-entropy over a parameter-matched Transformer baseline and achieve the improvement with fewer layers (fixed parameters), suggesting that recurrence can trade depth for width, thus reducing KV cache memory footprint and inference latency.

• The paper introduces a layerwise recurrent mechanism that computes persistent key-value pairs within each layer to enhance temporal depth and model expressivity.
• It employs an IO-aware tiling algorithm to reduce memory bandwidth and achieve near-linear latency scaling with longer sequences.
• Empirical results on synthetic benchmarks and language modeling tasks show improved performance and efficiency compared to standard Transformers.

The Recurrent Transformer: Enhancing Temporal Depth and Efficient Decoding

Introduction and Motivation

The Recurrent Transformer (RT) architecture proposes a fundamental modification to the standard Transformer, aiming to address its bounded temporal depth and improve inference efficiency. In standard Transformers, each layer at position $t$ exclusively attends to key–value (KV) pairs derived from the previous layer; the effective depth along the sequence is thereby capped by the number of layers. This restriction motivates the search for models that provide unbounded per-layer temporal depth akin to recurrent neural networks (RNNs), but without their characteristic optimization instability and hardware inefficiency.

RT achieves this by introducing layerwise recurrence: each layer computes persistent KV pairs from its own output activations, not just the previous layer's inputs. This enables later sequence positions to attend to representations already updated by same-layer attention and MLP computation, thereby enhancing the temporal expressivity within each layer. Unlike previous feedback or memory-augmented Transformers, RT maintains per-layer, per-token KV memories, simplifying efficient training and inference.

Figure 1: One layer of the Recurrent Transformer; persistent key-value pairs correspond to layer outputs and serve subsequent attention computations, in contrast to vanilla Transformers.

Architectural Innovations

Persistent and Temporary Key–Value Mechanism

RT departs from the vanilla Transformer by distinguishing between two KV types at each position: a temporary pair (from the current layer's input, used only for the current timestep to avoid circular dependencies) and a persistent pair (from the layer's output, used by all subsequent positions within the same layer). This circularity resolution is integral for stability and mathematical correctness.

Notably, the persistent KVs are computed after both attention and MLP updates, thus encoding deeper sequence-wide information. This allows RT layers to perform more complex, iterative computations on sequence data within a single layer, a paradigm inaccessible to standard attention blocks.

Representational Power

RT is representationally strict: the authors demonstrate that any width-$d'$ Transformer of arbitrary depth can be simulated by a width-3$d'$ RT under appropriate parameterization, and that RTs can also mimic token-to-token RNN recurrence under specified local attention concentrations. This positions RT as a strict generalization of both standard Transformers (full-permutation memory) and recurrent state-space models (iterative, positionwise computation), while avoiding fixed-size memory bottlenecks.

Theoretical and Empirical Training Stability

Training RNNs is historically fraught with vanishing/exploding gradients due to long sequential computation chains. In RT, the computation graph introduces multi-hop and direct attention paths between positions, promoting stable and non-degenerate gradient propagation. The coexistence of both one-hop (direct) and multi-hop (iterative) paths ensures robust information flow even across extended sequences, and the application of layer-normalization (RMSNorm) prior to KV projection empirically enhances training stability.

The authors formally analyze gradient dynamics in a simplified RT layer, showing that as long as the largest eigenvalue of the value-projection matrix (scaled by residual factor) is less than one, exploding gradients are avoided.

Efficient Training via IO-Aware Tiling

A foundational bottleneck in recurrent-time architectures is low arithmetic intensity (FLOPs per byte), leading to bandwidth-bound training when naively accumulating over ever-growing KV prefixes. RT addresses this using a tiling algorithm inspired by FlashAttention, reorganizing the training schedule to process available queries and KVs in blocks. This reduces high-bandwidth memory (HBM) traffic from $\Theta(N^2)$ to $\Theta(N \log N)$ for sequence length $N$, and improves arithmetic intensity to $\Theta(N/\log N)$. This is critical for scaling on modern accelerators.

Figure 3: Tiling schedule used for efficient attention computation; tiles of key–value pairs are reused across multiple queries to increase arithmetic intensity.

The empirical results show latency scaling nearly linearly with sequence length under this schedule, in stark contrast to the quadratic growth seen with naive recurrent implementations.

Figure 2: One-layer forward-pass latency with respect to sequence length. Tiled RT is near-linear, outperforming naive recurrent and matching or exceeding vanilla Transformer performance at long contexts.

Experimental Validation

Synthetic Benchmarks

On the MAD suite and sequence-copying diagnostics, RT decisively outperforms single-layer Transformers on all but the compression task. These benchmarks specifically stress the ability to perform iterative, depth-intensive computation—domains where classical Transformers are theoretically known to be weak.

Figure 4: Sequence-level accuracy on synthetic diagnostics. RT surpasses Transformer performance on all but the compression task.

Figure 5: Token-level accuracy on the same synthetic tasks, demonstrating the recurrent mechanism's effectiveness even in granular settings.

Language Modeling and Depth–Width Tradeoff

RT was pretrained (C4 dataset) at 150M and 300M parameter scales and compared to parameter-matched Transformer baselines. Results indicate a consistent cross-entropy improvement for RT at both 6 and 12 layers, with RT-6L models performing similarly to or better than Transformer-12L (with appropriate width scaling). This suggests that increased effective temporal depth can compensate for shallower stacking and reduce KV cache footprint and inference latency.

Figure 6: C4 pretraining loss curves for a 300M parameter model—RT achieves lower cross-entropy with fewer layers.

Other ablations confirm the importance of RMSNorm for stability; removal leads to significant training degradation.

Figure 7: Ablating RMSNorm results in significantly worse validation loss, highlighting normalization's importance for deep recurrent architectures.

Practical Considerations and Implementation

Several engineering optimizations underpin the practicality of RT:

• Batch Size vs. MLP Efficiency: Sequential MLP application reduces compute intensity unless larger batches are feasible.
• CUDA Graphs: Required to minimize kernel launch overhead due to frequent, small kernel invocations.
• Cache Locality: Position-major KV memory layout is essential to take full advantage of tiled computation.
• Custom Backward Passes: In-place KV cache updates and careful scheduling permit parallel and memory-efficient backward computation.

Further advances (e.g., custom kernels or blockwise recurrence) could yield additional gains.

Implications and Future Directions

RT shifts the effective depth–width tradeoff in neural sequence modeling by exposing deeper in-layer computation, and leverages architectural and IO-aware algorithmic changes for practical deployment. This tradeoff reduces memory usage and potentially increases parallelism at inference and training time. The enhanced representational power could allow RT-based models to capture temporally extended dependencies more economically than canonical Transformers, with implications for both scaling laws and downstream performance as model and data scales increase.

Potential extensions include recurrence over blocks for even greater efficiency, further hardware co-design, and comprehensive analysis of scaling laws as recurrence is integrated at various granularity levels.

Conclusion

The Recurrent Transformer provides a principled architectural mechanism for augmenting temporal depth in Transformers while maintaining per-token, per-layer memory and ensuring computational feasibility via tiling. Numerical results substantiate improved learning and inference efficiency, as well as enhanced representational capacity for depth-intensive tasks. The work paves the way for new architectures that bridge the expressivity of recurrent models and the practical efficiency of attention mechanisms, warranting further research into large-scale pretraining, task-specific adaptation, and deeper theoretical analysis of effective depth in neural sequence models.