Looped transformer

Overview

A looped transformer reuses one stack of layers. Its output is fed back in as its input, and the same weights run again. The number of passes R sets the effective depth, which is the layer count times R 2604.07822 §4. A vanilla transformer is the case R = 1.

Two questions drive the work on this design. The first is whether reusing weights lets a model combine facts it could not combine at a fixed depth. In a vanilla stack, a fact stored in shallow layers cannot be read by deeper layers that have different weights 2604.07822 §1. The second is whether more passes at inference can reach longer chains than the model was trained on 2604.07822 §1 2604.07822 §6.1.

The name also covers the original work on looped transformers, which neither card covers, so it is only named here: Looped Transformers. The sources below use it as a baseline, and their own reproductions are the numbers cited.

Variants

Fixed recurrence

Every batch and every input uses the same R. The paper trains a four-layer model with R in {1, 2, 4, 8} on a synthetic multi-hop task 2604.07822 §5.1. Reported results on that task:

• R = 1 fails on the out-of-distribution (OOD) split, where the atomic facts never appear in a training composition 2604.07822 §5.2.
• R = 2 reaches non-trivial OOD accuracy, with no number given in the text. R = 4 reaches OOD generalization in about 2k epochs, against about 7k for R = 2 2604.07822 §5.2 fig. 3.

For depth extrapolation, the fixed model trains on chains of up to a set hop count and is tested on longer ones. With matched data (training up to 12 hops, testing 12 to 24), R = 6 generalizes to 14 hops and R = 8 to 19 hops 2604.07822 §6.2. With the paper's own curriculum, R = 6 reaches 17 hops and R = 8 reaches 24 hops 2604.07822 §6.2. Raising R at inference past the training value is what allows the extra depth 2604.07822 §6.2. Card: Loop, Think, & Generalize.

Dynamic recurrence

Each batch draws its own R from a Poisson distribution with mean 4, clipped to [2, 8] 2604.07822 §4 2604.07822 app. J. The sampled R does not depend on the example's complexity, because the authors model cases where complexity is unknown in advance 2604.07822 §4.

The paper reports that the dynamic model matches R = 8 on the matched-data setup, at 19 hops. In the main setup it exceeds R = 8, but the text gives no number 2604.07822 §6.2. The same source reports overthinking: beyond about 15 inference iterations, OOD accuracy stops improving, and the margin peaks and then declines. The peak is lower for harder chains 2604.07822 §6.2 fig. 8. Card: Loop, Think, & Generalize.

Stochastic guidance with deep supervision

This variant keeps the loop but makes each latent step stochastic. A learned Gaussian term is added to the high-level state at each transition. Its mean and variance come from the state, so the run can take different paths 2605.19376 §2.1. Training uses a posterior that sees the target and a prior that does not, with a KL term between them 2605.19376 §2.2. Inference drops the target and draws N trajectories, then returns the most frequent answer 2605.19376 §2.3 2605.19376 alg. 2.

Training follows deep supervision. A supervision step runs several transitions, gradients flow only through the last one, and the state is detached before the next step 2605.19376 §2.2. The paper calls this biased but memory efficient 2605.19376 app. A.2.

Reported on Sudoku-Extreme, with the paper's own reproduction of each baseline: N = 1 accuracy is 86.6 for this variant, against 61.7 for its looped baseline and 83.6 for its TRM baseline. At N = 20 accuracy is 96.1 2605.19376 tab. 8. The ablation removes each part in turn. Without the stochastic guidance, Sudoku drops to 83.6 and N-Queens drops to 66.8. Stochasticity without guidance gives 96.0 on Sudoku and 62.1 on N-Queens. Guidance without stochasticity gives 0.0 on both 2605.19376 tab. 3. Card: Generative Recursive Reasoning.

When to use which

Implementation sketch

The fixed and dynamic variants share one loop. The stochastic variant adds a noise head to each transition, which is not shown here.

# Looped forward pass, fixed or sampled recurrence  [arxiv:2604.07822 §4]
# Input: prefix tokens; R, or None to sample it per batch

def looped_forward(prefix, R=None, lam=4, r_min=2, r_max=8):
    # Dynamic: a new R per batch, clipped Poisson  [arxiv:2604.07822 §4]
    if R is None:
        R = clip(poisson(lam), r_min, r_max)
    hidden = embed(prefix)
    for _ in range(R):
        # Same shared layers each pass; output feeds the next  [arxiv:2604.07822 §4]
        hidden = shared_layers(hidden)
    # Answer read at the final prefix position  [arxiv:2604.07822 §3.1]
    return readout(layer_norm(hidden))

Three details from the cards matter for a faithful reimplementation. The output projections of attention and MLP are zeroed at start, so each pass begins as an identity map 2604.07822 §4. The synthetic task must use a permutation construction, or a model can read the answer from a short suffix of the relations and look like it composes all hops 2604.07822 app. K 2604.07822 app. B. The padded batch layout also decides which column holds the prediction, so reproduce it exactly (Loop, Think, & Generalize, Pitfalls).

Contradictions

None known between the two sources. They use different tasks, and neither reports the other's measure, so their numbers are not comparable. The looped baseline in the second source is its own reproduction, not a number from the first source.

Within the second source, the paper and its code disagree on several training details, including the halting label and the weight decay used for ARC-AGI. These are listed on that card (Paper vs code), not resolved here.

Papers

• Loop, Think, & Generalize: Implicit Reasoning in Recurrent-Depth Transformers: fixed and dynamic recurrence on a synthetic knowledge-graph task, with depth extrapolation and overthinking.
• Generative Recursive Reasoning: stochastic latent guidance on a looped reasoner, trained with variational inference and majority-voted at test time.