Loop, Think, & Generalize: Implicit Reasoning in Recurrent-Depth Transformers

In plain words

Language models know many facts, but they often fail to chain two of them together when the question needs both, especially when the particular pair was never seen as a question during training. The authors ask whether a model whose layers are reused in a loop can do this better than a normal stack of layers, where each layer has its own weights and can only use the knowledge stored at its own depth.

They test this on a small, controlled puzzle: a made-up graph of entities and relations, where the model must follow a chain of links from a starting entity to an answer in one forward pass. A looped model does better on two kinds of new cases. The first is a new combination of facts it has never seen chained together. The second is a chain longer than any it was trained on, which works if the model is given more loops at inference time. Looping also has a downside: past a point, extra loops make the answer worse, which the authors call overthinking.

TL;DR

• Small GPT-2-style transformers are trained from scratch on synthetic multi-hop queries over a knowledge graph §3.1 §5.1 §6.1.
• A recurrent-depth model applies the same block R times; a vanilla model is the R = 1 case §4 §5.1.
• Recurrent models generalize to unseen fact combinations; vanilla models do not, and the ability appears in three grokking stages §5.2.
• Deeper chains are reached by raising inference-time recurrence, but the gain stops and reverses because of overthinking §6.2.

Problem

Transformers store facts in specific layers. A second-hop fact that is stored only in shallow layers cannot be read by deeper layers, because parameters are not shared across layers §1. Prior work reports that vanilla transformers can learn two-hop composition only through grokking, and that even then they do not generalize to combinations of facts that never appeared in a composition during training §2 §5. Vanilla models also fail to extrapolate to deeper chains than they were trained on §1 §3.2.

The task assumes a single forward pass with no written reasoning steps (implicit reasoning) §1. Recurrence is only evaluated on this synthetic setting; the authors state that results do not transfer directly to LLMs trained at scale app. A.

Method

Task. A directed knowledge graph has entities E and relations R. Atomic facts are triples (h, r, t). A k-hop inferred fact is a chain of k atomic facts through k − 1 unnamed intermediate entities §3.1. The induced set of k-hop facts is

\mathcal{I}_k(\mathcal{C}) = \{(h, r_1, \dots, r_k, t) \mid \exists\, i_1, \dots, i_{k-1},\ (h, r_1, i_1), \dots, (i_{k-1}, r_k, t) \in \mathcal{C}\} §3.1

The model receives the prefix <e_h><r_1>…<r_k> and must output the single token <e_t>. Only the next-token prediction at the final prefix position is supervised §3.1 §4.

Splits. The atomic set is partitioned into C_ID and C_OOD. Training uses all atomic facts plus the inferred facts of C_ID up to a maximum depth k_train §3.2. Three evaluations follow §3.2:

1. In-distribution: held-out inferred facts built from C_ID.
2. Systematic: inferred facts built from C_OOD, whose atomic facts never appear in any training composition.
3. Depth extrapolation: inferred facts with k > k_train, built from C_ID.

Model. The model is a decoder-only transformer. An embedding layer feeds a stack of L layers that is applied R times, so the effective depth is L × R. The output of each pass is the input of the next. Layer norm and a tied output projection give the logits. The embedding and output head share weights §4 fig. 1. Each entity and each relation has its own token §4.

h^{(r+1)} = f_\theta\!\left(h^{(r)};\, m\right), \qquad r = 0, \dots, R-1 §4

Here f_θ is the stack of L layers with shared parameters θ, and m holds the causal and padding masks. The final h^(R) goes through layer norm and the tied projection §4.

Zero initialisation. The output projections (c_proj) of attention and MLP are set to zero, so each recurrent pass starts as an identity map. The authors say this keeps the input-output Jacobian stable as R grows §4.

Training recurrence. Two options §4:

• Fixed: the same R for every batch.
• Dynamic: a new R for each batch, drawn from a clipped Poisson distribution

R \sim \mathrm{clip}\big(\mathrm{Poisson}(\lambda),\ R_{\min},\ R_{\max}\big),\qquad R_{\min}=2,\ R_{\max}=8 §4

The mean is λ = 4 for R_max = 8 app. J. The sampled R is not conditioned on the complexity of the example, because the authors want to model cases where complexity is unknown in advance §4.

Systematic-generalization experiment §5.1. Knowledge graph with |E| = 2000 entities and |R| = 200 relations, average out-degree 20. There are 40k atomic facts, split 95% C_ID and 5% C_OOD, plus 273.6k inferred facts for training. The in-distribution test set has 3k two-hop facts from C_ID; the OOD test set has nearly 2k two-hop facts from C_OOD. The model has L = 4 layers, uses absolute position embeddings, and is trained with fixed R ∈ {1, 2, 4, 8}. Dynamic recurrence is not used here §5.1.

Depth-extrapolation experiment §6.1. Knowledge graph with |E| = 200 entities and |R| = 10 relations, average out-degree 10, with a permutation constraint (below) app. B. The dataset has 2k atomic facts and 15k k-hop facts for each k ∈ [2, 40], and each model is tested on 750 held-out k-hop facts §6.1. The model uses no positional embeddings (NoPE). Fixed R ∈ {1, …, 8} and a dynamic variant are compared §6.1.

Permutation constraint app. B. Each relation is a random permutation of the entity set, so every entity has exactly one outgoing edge per relation. Without this, a model could read the answer from the last few relations alone. The pseudocode below follows app. B.

# Permutation knowledge graph  [app. B]
# 200 entities, 10 relations; each relation is a bijection on entities

def build_atomic_facts(entities, relations):
    facts = []
    for relation in relations:
        # Random permutation pi_r over the entity indices  [app. B]
        permutation = random_permutation(entities)
        for entity in entities:
            facts.append((entity, relation, permutation[entity]))
    return facts  # 2000 atomic facts for 200 x 10

# Sample one k-hop example  [app. B]
# facts_by_head[e] lists (relation, tail) pairs leaving entity e

def sample_k_hop(all_facts, facts_by_head, k):
    # Start from one atomic fact chosen at random  [app. B]
    head, relation, tail = random_choice(all_facts)
    relations = [relation]
    for hop in range(k - 1):
        # Follow one outgoing edge, chosen uniformly  [app. B]
        relation, tail = random_choice(facts_by_head[tail])
        relations.append(relation)
    # Intermediate entities are not part of the example  [app. B]
    return head, relations, tail  # prefix = head + relations; target = tail

Curriculum §6.1. Training is easy-to-hard over hop depth. It starts with atomic and 2-hop facts and adds 3-hop facts once the model reaches 95% accuracy on the held-out 2-hop test split. The process repeats for each k, and each stage keeps all earlier facts to avoid forgetting. The largest k the model can fit this way is its learnable recursion depth §6.1.

Inference. Two ways to set the number of iterations §6.2:

• Fixed: use a chosen R at test time; scaling it past the training value allows depth extrapolation.
• Adaptive halting (r*): stop per example once the prediction has settled. The stop rule uses the KL divergence between successive iterations and the entropy of the current prediction:

\mathrm{KL}\big(p_t \,\|\, p_{t-1}\big) < \varepsilon_{\mathrm{KL}} \quad\text{and}\quad H(p_t) < H_{\mathrm{thresh}} §6.2

The paper gives ε_KL = 0.01 and H_thresh = 3.00 app. C. For the code, see Paper vs code.

# Adaptive halting (r*)  [§6.2]
# Input:  prefix tokens; upper bound max_recurrence
# Output: logits used for the prediction; number of iterations used

def adaptive_predict(prefix, max_recurrence, eps_kl, entropy_thresh):
    hidden = embed(prefix)
    previous_logp = None
    for t in range(max_recurrence):
        # One more pass through the shared layers  [§4]
        hidden = apply_shared_layers(hidden)
        logits = readout(hidden)  # (B, V), last prefix position
        logp = log_softmax(logits)
        if t > 0:
            # KL(p_t || p_{t-1}) between successive predictions  [§6.2]
            kl = sum(exp(logp) * (logp - previous_logp))
            # Entropy of the current prediction  [§6.2]
            entropy = -sum(exp(logp) * logp)
            if kl < eps_kl and entropy < entropy_thresh:
                return logits, t + 1
        previous_logp = logp
    # Limit reached without a halt: use the last pass
    # [code: gpt_utils_extrapolation.py#L591-L596]
    return logits, max_recurrence

Logit lens and causal analysis §5.2 app. E. The authors project each layer and each iteration through the output head to see when the bridge entity and the target entity become decodable. They also run activation patching.

Hyperparameters

Values stated in the paper, with the code default where it differs.

Results

Author-reported. The text gives few numbers, and the main accuracy curves are in figures, so the values below are the ones stated in the text. Setups vary per row.

Paper vs code

Code: github.com/OSU-NLP-Group/Loop-Think-Generalize, commit bb22b192977b1b136fde985c8e7a2392d172d97a, Apache-2.0. Discrepancies below were checked by reading the cited lines.

• Halting KL threshold. Paper: ε_KL = 0.01 app. C. The adaptive eval calls evaluate_model_test_adaptive_metrics without passing eps_kl, so the wrapper's default of 0.001 is used inference_extrapolation_adaptive.py gpt_utils_extrapolation.py. The inner forward_adapative has its own default of 0.01 gpt_utils_extrapolation.py, but the wrapper overrides it.
• Halting entropy. A second variant forward_adapative_ent uses a threshold H0 + H1/seq_len, with H0 = 3.0 and H1 = 12.0 gpt_utils_extrapolation.py gpt_utils_extrapolation.py, and eps 0.5. The adaptive eval script does not call it. Paper: fixed H_thresh = 3.00 app. C.
• First-iteration halting. The paper does not say whether one iteration may halt. The code checks KL and entropy only from t = 1, so at least two iterations are always used gpt_utils_extrapolation.py.
• Curriculum failure rule. Paper: if the threshold is not reached, training terminates and (k+1)-hop data is never introduced §6.1. Code: training at a stage stops only on reaching the threshold or the cumulative cap of 100001 epochs train_extrapolation.py train_extrapolation.py. The outer loop then moves to k+1 regardless train_extrapolation.py. After the cap, each later stage trains for one epoch.
• Systematic weight decay. Paper: 0.01 for all experiments app. C. The systematic trainer defaults to 0.1 train_systematicity.py. The extrapolation trainer uses 0.01 train_extrapolation.py.
• Systematic batch size. Paper: 512 app. C. The trainer's parser default is 128 train_systematicity.py. The README says the defaults reproduce the paper, which conflicts with the stated 512; the code does not settle which is right [README].
• Zero initialisation. Paper: c_proj zeroed for all experiments §4. The extrapolation trainer scales c_proj by c_scale, default 0.0, which zeroes it train_extrapolation.py gpt_utils_extrapolation.py. The systematic trainer imports gpt_utils_systematicity, where no c_proj scaling appears (checked by grep), so zero init does not apply there.
• LR schedule. Paper: linear warmup of 2000 steps app. C. Code: linear warmup then linear decay to 0 over len(dataloader) × num_epochs steps train_extrapolation.py. An optional per-stage decay exists but is off by default train_extrapolation.py.
• Max recurrence in adaptive eval. The paper's active text does not give a cap. The script default is 16 inference_extrapolation_adaptive.py.
• Systematic data. README: the 2-hop data come from the external GrokkedTransformers pipeline, not from this repository [README]. The ratio 7.2 is inferred facts to atomic facts in C_ID: 273.6k / 38k §5.1 [README].

No other discrepancies found in the paper text and the code files I read.

Pitfalls

• Readout slot. Batches are right-padded to max_len (50 by default). With the default pred_pos = last_token, the prediction is read from the last column, index max_len − 1, which is a padding slot for any query shorter than that gpt_utils_extrapolation.py gpt_utils_extrapolation.py train_extrapolation.py. Training uses the same layout, so the setup is consistent. The paper says the supervised position is the final prefix position §3.1 §4. Reproduce the padding exactly, or the results will not match.
• Stop rule uses the test split. The curriculum advances when accuracy on the current held-out k-hop test split exceeds 95% train_extrapolation.py. The paper calls it held-out §6.1. Stage selection and evaluation share one split, so reported "learnable depth" is optimistic. Use a separate validation split if you need an unbiased number.
• Curriculum cap. The cap is cumulative across stages, and the code moves on to the next hop level after the cap. The paper's stated rule (training stops, no higher hops) does not hold in the code. See Paper vs code.
• Default init is unstable. With the default Gaussian init for c_proj, results are mixed across seeds, and some runs degrade with more recurrence app. F. The zero-init is what the main results depend on.
• Shortcuts without the permutation constraint. On a non-permutation graph, a model reached high OOD accuracy at 80 hops by reading a short suffix of the relation sequence, not by composing all hops app. K. Use the permutation construction, or your deep-chain results will mean something else.
• Overthinking. For the dynamic model, more than about 15 inference iterations stops improving OOD accuracy §6.2. The margin peaks and then declines with more iterations; the peak is lower for harder chains §6.2 fig. 8. Adaptive halting does not fix this. The authors say the input-injection fix from prior work does not solve it here §6.2.
• Seed variance. Fixed-recurrence results are shown for two seeds and dynamic for three app. I. One seed was not affected by overthinking but also showed no scaling app. F. Expect variance.
• Compute. Grokking to 4 hops took over 1.3M steps §6.2. Budget for long runs at the first stage, not just the last.
• Scope. Synthetic tokens, a tiny vocabulary, and no natural-language surface variation. The authors do not claim the results carry over to LLMs trained at scale app. A.

Lineage

• Builds on: Universal Transformers and ALBERT for weight sharing across depth app. D; looped transformers for length generalization (Fan et al., 2024) §2; the looped regularizer of Saunshi et al. §2 §4; the easy-to-hard curriculum for k-hop tasks, and 2-hop grokking results, both from prior work §6.1 §5.
• Adaptive halting by KL divergence on the output distribution comes from Geiping et al. (2025), which the paper extends with an entropy term §6.2.
• The overthinking name comes from Bansal et al. (2022) §6.2.
• Compared against: vanilla transformers of matched and larger depth §5.2 app. G.
• Superseded by: not stated.