Recurrent Looped Transformer

Yifan Zhang; Jichen Feng; Shihan Qin

cs.CLOct 6, 2026

Recurrent Looped Transformer

Authors: Yifan Zhang, Jichen Feng, Shihan Qin

Organizations: Princeton University · University of Pennsylvania

Abstract

State tracking requires an update at every input, but the depth a Transformer applies to each token is fixed regardless of sequence length. We introduce the Recurrent Looped Transformer (RLT), which splits its layers between a parallel causal encoder and a recurrent decoder. At each token, the decoder merges the encoder output with the previous token's final decoder state, so the computation path grows with sequence length at a fixed per-token cost. On six algorithmic tasks, we compare five splits of eight layers with an eight-layer Transformer over three seeds. Trained on at most 40 bits, two RLT splits generalize parity to 256 bits with 100% accuracy in every seed, while the Transformer stays at chance. On swap-based $S_5$ permutation tracking at eight times the training length, RLT reaches 97% final-state accuracy versus under 1% for the Transformer, and accuracy increases with decoder depth. On modular arithmetic beyond the training lengths, RLT reaches up to 93% versus 33% for the Transformer. Ablations show that these gains depend on the feedback: removing it drops parity and swap-based $S_5$ to chance at every split. Updating the feedback once per four-token chunk lets known tokens in a chunk run in parallel and keeps 64-bit parity at 99%, while permutation tracking depends on per-token feedback: chunking lowers length-64 swap-based $S_5$ from 100% to 20%.

Figures & tables

Panel of Figure 2 — Figure 2: RLT-1 architecture for one token update. Computation flows upward; each outlined stack repeats for the indicated depth. The encoder output supplies both the gated merge and the global KV projection. With $G=1$ , all decoder layers use their own queries to read the same cached, projected global KV. Each decoder layer constructs separate SWA KV from its own input, attends to its local window, then reads global memory and applies an FFN. Residual paths bypass each attention or FFN sublayer. The final decoder output $s_{t}$ feeds the next token’s merge, while each layer retains its own SWA cache. Blue arrows carry encoder features and global memory, green arrows carry local KV, and orange arrows carry recurrent hidden states.

Panel of Figure 3 — Figure 3: Recurrent path length and block count per token. For $L_{E}=L_{D}=48$ , the path traverses $48t$ decoder blocks after $t$ tokens. Each token evaluates 96 encoder and decoder blocks. Orange counts decoder blocks along the recurrent path; blue counts blocks in both stacks per token.

Panel of Figure 4 — Figure 4: RLT-2 architecture for one chunk. Computation flows upward; the stacks repeat for the indicated depths. Each position $t\in I_{k}$ uses its own gate to merge its encoder output with the preceding boundary state $h_{k-1}$ . State normalization and projection are shared across the chunk. Each decoder layer applies causal SWA, prefix-restricted encoder-memory attention, and an FFN, with residual connections around each sublayer. For $G=1$ , all decoder layers read the same projected encoder KV; each layer retains its own SWA KV across chunk boundaries. Every output $y_{t}$ predicts the next token; only $y_{kB}$ at a complete boundary becomes the next feedback state. A partial chunk preserves $h_{k-1}$ during continuation. Blue arrows carry encoder features and memory, green arrows carry decoder KV, and orange arrows carry chunk-level feedback.

Panel of Table 1 — Table 1: Execution costs at matched stack dimensions. $T$ counts known input tokens, $K=\lceil T/B\rceil$ , and $N$ counts generated tokens; all three variants share the $L_{E}$ -stage causal encoder prefill. Counts omit common readout and encoder-memory projection work. RLT-2 reuses the normalized and projected boundary state; gates depend on each token. Backward counts cover the decoder dependency graph, with the encoder’s backward pass common to all variants. Counts describe dependencies and arithmetic; wall-clock speed depends on the implementation and hardware.

Explore similar work

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
Apr 23, 2026Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3Transformer Architectures Kv-Cache Management
Harvard University
FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates
Sep 24, 2026Wanqi Yang, Shiwei LiuBlock Sparse Flash Attention Transformer Architectures
ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems Tübingen AI Center
Refit the Probe: Single-Direction Ablation Is Not a Necessity Test
May 30, 2026Yuhang JiangLayer-Wise Sequence Modeling
Independent Researcher

Apr 23, 2026, cs.LGJ/K move · Enter open · S save

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi, Alexandru Meterez+2

Harvard University

Transformer Architectures Kv-Cache Management

Open paper AskPDF arXiv

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)

Θ(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.

Abstract

Explore similar work

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates

Refit the Probe: Single-Direction Ablation Is Not a Necessity Test