Recurrent Looped Transformer

AI summary

Overview

Research area: Natural Language Processing / sequence modeling and algorithmic reasoning with Transformers.

Technical level: Advanced. The paper assumes familiarity with Transformer architecture, attention, key–value caches, state tracking formalisms, and backpropagation through time.

Scope: The paper introduces the Recurrent Looped Transformer (RLT), an architecture that couples a parallel causal encoder with a recurrent decoder, and evaluates how encoder–decoder depth allocation and feedback frequency affect length generalization on six algorithmic tasks.

What This Paper Is About

A standard Transformer applies a fixed number of layers to every token, so its computation depth does not grow with sequence length. That is a problem for tasks like parity or permutation tracking, where the model must update an internal state at every input. The paper proposes splitting an eight-layer stack into a causal encoder that processes known tokens in parallel and a recurrent decoder that merges each token's encoder output with the previous token's final decoder state, so the computation path grows with sequence length at a fixed per-token cost.

Key Contributions

1. Architecture. RLT combines encoder-memory reuse with feedback through the full decoder at every token. The appendices specify training by backpropagation through the full recurrent history, cache reuse across turns, and exact policy replay.
2. Length generalization. Trained on at most 40 bits, RLT splits 5+3 and 7+1 reach 100 ± 0% parity accuracy at 256 bits in every seed, compared with 50.07 ± 1.63% for the Transformer. On swap-based S5 at 256 operations (eight times the training length), 4+4 reaches 97.30 ± 2.76% final-state accuracy versus 0.85 ± 0.30%. After 5,000 training steps, 6+2 reaches 93.36 ± 5.69% on flat mod-5 expressions of length 63, compared with 33.20 ± 2.33%.
3. Depth allocation. The best split depends on the task: parity generalizes best with one or three decoder layers, whereas swap-based S5 accuracy at 256 operations rises with decoder depth, from chance with one decoder layer to 97% with four.
4. Feedback interval. Without feedback, parity and swap-based S5 stay near chance from length 64 at every split. With four-token chunks, RLT-2 keeps 64-bit parity at 98.99 ± 1.66% but lowers swap-based S5 at 64 operations from 100% to 19.60 ± 6.54%.

Main Findings

• Parity extrapolation: At 256 bits, RLT splits 5+3 and 7+1 retain 100 ± 0% accuracy, compared with 50.07 ± 1.63% for the eight-layer Transformer. The advantage appears during learning: at step 500, 6+2 reaches 99.44 ± 0.98% validation accuracy versus 48.48 ± 0.53%.
• Larger-depth series: In a separate seed-42 sixteen-layer series, RLT-1 8+8, 9+7, 11+5, and 16+0 reach 100% at 256 bits, compared with 49.41% for Transformer 16.
• Permutation tracking: At 256 operations, 4+4 reaches 97.30 ± 2.76% final-state accuracy versus 0.85 ± 0.30% for the Transformer. At 512 operations, 4+4 still reaches 55.70 ± 25.78%, while the Transformer stays at 0.85 ± 0.30%. Splits 7+1 and 8+0 are already near the uniform reference at 256 operations despite high training-length accuracy.
• Modular arithmetic: On flat expressions of length 63, 6+2 reaches 93.36 ± 5.69%, compared with 33.20 ± 2.33% for the Transformer. On bracketed expressions of length 64, 5+3 reaches 67.97 ± 2.91% versus 46.71 ± 1.21%. Accuracy falls on longer expressions, and several flat mod-5 splits vary widely across seeds.
• Feedback removal: RLT-0 remains near chance at 64-bit parity (50.23 ± 3.80%), and removing feedback drops parity and swap-based S5 to chance at every split.
• Chunked feedback: At 64 bits, RLT-1 reaches 100 ± 0% parity and chunk4 reaches 98.99 ± 1.66%. Swaps-S5 is more sensitive: at 64 operations, RLT-1 reaches 100 ± 0% and chunk4 19.60 ± 6.54%. Chunk4 7+1 and 8+0 are already near chance at 64 bits. Chunk4 5+3 reaches 79.92 ± 10.89% at 48 operations and 50.16 ± 18.77% at 64.
• Addition and standard S5: All models reach 100% teacher-forced accuracy on the training-width addition validation set, but accuracy drops beyond eight digits; eight-layer means at 32 digits range from 14.89% to 16.84%. Standard S5 remains difficult, with final-state accuracy near the 1/120 reference across feedback variants.
• Speed: In seed-42 CPU timing at 4+4, chunk4 training steps run 2.27 times as fast as RLT-1 steps, and RLT-0 steps 4.17–4.30 times as fast.

Methodology in Plain English

The researchers take an eight-layer stack and divide it into two parts: an encoder that reads known tokens in parallel and builds token representations plus key–value memory, and a decoder that runs token by token. At each position, a gated merge combines the encoder output with the previous token's final decoder output, then the decoder attends to the encoder prefix and to its own recent activations (a sliding window of eight) before predicting the next token. Prompt and response tokens follow the same transition.

They compare five allocations of eight layers — 4+4, 5+3, 6+2, 7+1, and 8+0 — with an eight-layer decoder-only Transformer. RLT-1 has 26.10–28.73M parameters versus 25.31M for the Transformer, so equal layer counts do not match parameters or compute. All eight-layer accuracy comparisons use seeds 42, 43, and 44, a shared training stream, global batch size 512, microbatch size 32, and AdamW with (β1, β2) = (0.9, 0.95), weight decay 0.1, and gradient clipping at norm 1. The learning rate warms up to 1e-4 over 200 steps and decays to 5e-6.

Parity, addition, and S5 use 2,000 steps; the mod-5 comparisons use 5,000 steps. Parity trains on 3–40 bits; S5 trains on 32 operations using all 120 permutations or the identity and ten single transpositions. Flat mod-5 expressions use +, −, and × with multiplication precedence and odd training lengths 3–39; bracketed expressions use trees of lengths 3–40. Evaluation uses 1,024 shared test examples per length, and addition uses 256 shared pairs per test width. Checkpoint selection uses the minimum in-distribution validation example loss, with exact ties selecting the earliest step; test scores do not enter selection.

They also vary the feedback interval: RLT-1 feeds back every token (B=1), RLT-2 shares one feedback state across each chunk of B tokens, and RLT-0 removes the feedback (B=∞). Training uses teacher forcing and cross-entropy, with full backpropagation through time; TBPTT128 covers every training sequence, so no gradients are truncated.

Why This Matters

The work argues that a fixed-depth Transformer cannot, unless TC⁰ = NC¹, track compositions of S5 permutations over arbitrarily long sequences, and shows a concrete architectural fix that keeps per-token cost fixed while letting the computation path grow with sequence length. It also treats the allocation of depth between encoding and recurrent updates as an explicit design variable rather than a fixed choice.

Real-world applications:

• Long-horizon state tracking: tasks where a system must maintain and update a state across long input streams, such as parsing or formal reasoning pipelines.
• Algorithmic and symbolic reasoning: arithmetic, modular arithmetic, and permutation composition beyond trained lengths, relevant to verifier-style or calculator-assisted systems.
• Inference-time efficiency: the chunked RLT-2 variant lets known tokens run in parallel, which the paper frames as useful for pretraining with a large chunk size and adapting toward per-token feedback later.
• Reinforcement learning pipelines: the appendices describe exact policy replay, in which a sampled response's history is replayed to rebuild the state under current parameters.

Industry relevance: the chunk-size trade-off is directly practical — larger B shortens the sequential decoder path (chunk4 training steps run 2.27 times as fast as RLT-1 steps at 4+4), while smaller B extrapolates better on parity and swap-based S5. The paper suggests pretraining can begin with a large B and reduce it mid-training and post-training toward B=1.

Future Directions

• Feedback schedules: the paper states that changing B changes the model's computation, so inference and RL replay must use the chunk size of the final training stage. It explicitly notes that its experiments train each chunk size from scratch and do not test a schedule that anneals B during training.
• RL performance: the authors state that their experiments measure supervised algorithmic performance and that RL performance is not evaluated, leaving policy-gradient evaluation as an open question.
• Depth allocation selection: because the best split depends on the task (parity favors one or three decoder layers; swap-based S5 improves with four), how to choose or adapt a split per task remains unresolved.
• Task scope: standard S5 remains difficult, with final-state accuracy near the 1/120 reference across feedback variants, and addition accuracy drops beyond eight digits — areas where the approach has not demonstrated gains.

Target Audience

Researchers and engineers working on Transformer architecture, length generalization, and algorithmic reasoning, particularly those interested in recurrent or looped designs, encoder–decoder memory reuse, and the trade-off between sequential feedback and token parallelism. Practitioners building models for long-context or state-tracking tasks will also benefit from the depth-allocation and feedback-interval ablations.

Authors’ 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%.