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

Harsh Kohli; Srinivasan Parthasarathy; Huan Sun; Yuekun Yao

Source-linked AI summary

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

Harsh Kohli, Srinivasan Parthasarathy, Huan Sun, Yuekun Yao

arXiv:2604.07822v2Apr 9, 2026cs.CLcs.AIcs.LG

TL;DR

The paper asks why transformers with stored factual knowledge struggle with implicit multi-hop composition. It studies recurrent-depth transformers on controlled synthetic reasoning tasks and finds systematic generalization and depth extrapolation, while excessive recurrence causes overthinking.

Problem
Transformers store substantial parametric knowledge but often struggle to compose it for implicit multi-hop reasoning and unfamiliar or deeper combinations.
Method
The paper trains recurrent-depth and vanilla transformers from scratch on synthetic multi-hop knowledge-graph tasks, varying recurrence and generalization setting.
Results
Recurrent-depth transformers show systematic generalization through three-stage grokking and extrapolate to greater reasoning depths as inference-time recurrence increases, unlike vanilla transformers.
Takeaways & Limitations
Shared-weight recurrence supports compositional generalization in implicit reasoning, with additional inference iterations enabling deeper reasoning.
Takeaways & Limitations
Excessive inference-time recurrence causes overthinking, degrading predictions and limiting generalization to extremely deep compositions.

Abstract

from arXiv · show

We study implicit reasoning, i.e. the ability to combine knowledge or rules within a single forward pass. While transformer-based large language models store substantial factual knowledge and rules, they often fail to compose this knowledge for implicit multi-hop reasoning, suggesting a lack of compositional generalization over their parametric knowledge. To address this limitation, we study recurrent-depth transformers, which enables iterative computation over the same transformer layers. We investigate two compositional generalization challenges under the implicit reasoning scenario: systematic generalization, i.e. combining knowledge that is never used for compositions during training, and depth extrapolation, i.e. generalizing from limited reasoning depth (e.g. training on up to 5-hop) to deeper compositions (e.g. 10-hop). Through controlled studies with models trained from scratch, we show that while vanilla transformers struggle with both generalization challenges, recurrent-depth transformers can effectively make such generalization. For systematic generalization, we find that this ability emerges through a three-stage grokking process, transitioning from memorization to in-distribution generalization and finally to systematic generalization, supported by mechanistic analysis. For depth extrapolation, we show that generalization beyond training depth can be unlocked by scaling inference-time recurrence, with more iterations enabling deeper reasoning. We further study how training strategies affect extrapolation, providing guidance on training recurrent-depth transformers, and identify a key limitation, overthinking, where excessive recurrence degrades predictions and limits generalization to very deep compositions.

1 Introduction

Transformers store substantial parametric knowledge but struggle to compose it for implicit multi-hop reasoning and unfamiliar or deeper combinations. The paper studies recurrent-depth transformers, finding systematic generalization and depth extrapolation alongside a grokking dynamic and an overthinking limitation.

Transformers often fail to flexibly compose stored knowledge for implicit reasoning within a single forward pass.
Recurrent-depth transformers iteratively apply shared layers, enabling more flexible access to and composition of parametric knowledge.
The study uses controlled synthetic datasets to examine systematic generalization and depth extrapolation in models trained from scratch.
Recurrent-depth transformers show systematic generalization through a three-stage transition from memorization to in-distribution generalization and then systematic generalization.
Recurrent-depth transformers enable depth extrapolation as inference-time recurrence increases, but excessive recurrence causes overthinking and degrades performance.

2 Related Work

Prior work finds that looped transformers can improve algorithmic length generalization and emulate iterative learning procedures. This paper focuses on controlled multi-hop composition, emphasizing training dynamics, systematic composition, and inference-time extrapolation.

Looped transformers have been studied as architectures that emulate learning algorithms such as gradient descent, neural networks, and decision trees.
Looped transformers offer superior length generalization on algorithmic tasks including parity and binary addition.
The paper’s controlled setting attributes generalization differences to architectural choices and studies systematic composition and multi-hop extrapolation with increased recurrence at inference time.

3 Task Formulation

The paper formulates implicit reasoning as multi-hop traversal over a synthetic knowledge graph and separates in-distribution generalization, systematic generalization, and depth extrapolation. These challenges differ in whether compositions are unseen, atomic facts are compositionally held out, or reasoning depth exceeds training.

3 Task Formulation: The task formulation categorizes in-distribution generalization, systematic generalization, and depth extrapolation as three challenges in synthetic implicit reasoning.
3.1 Task Definition: The knowledge graph contains entities as nodes and relations as edges, with atomic facts represented as triplets (h, r, t).
3.1 Task Definition: A k-hop inferred fact chains k atomic facts from a head entity to a final tail through k−1 intermediate entities.
3.1 Task Definition: Given a head and relation sequence, an autoregressive decoder predicts the final tail while implicitly retrieving intermediate entities.
3.2 Generalization Challenges: The inferred-fact set I_k(C) contains relation chains whose atomic facts all belong to the complete fact set C.
3.2 Generalization Challenges: Training includes all atomic facts and inferred facts up to maximum depth k_train, using a partition C = C_ID ∪ C_OOD to define generalization settings.
3.2 Generalization Challenges: In-distribution generalization tests held-out compositions from C_ID, whereas systematic generalization tests compositions from atomic facts in C_OOD never used in training compositions.
3.2 Generalization Challenges: Depth extrapolation evaluates C_ID compositions deeper than k_train, requiring rules to be applied iteratively beyond observed training depths.

4 Recurrent-Depth Transformer

The model reuses a shared transformer stack for recurrent computation, allowing effective depth and inference-time computation to increase without retraining. Training uses fixed or dynamically sampled recurrence, with dynamic recurrence designed for settings where task complexity is unknown.

Model architecture: A decoder-only transformer reuses L shared layers for R recurrent iterations, producing effective depth D = L × R.
Model architecture: Increasing R at inference time allocates more computation without changing the architecture or retraining parameters.
Model architecture: The final recurrent representation is normalized and projected through tied output weights to produce vocabulary logits, with supervision only at the tail position.
Zero-initializing attention and feed-forward output projections makes each recurrent block an identity mapping initially and stabilizes the input-output Jacobian.
Stopping strategies: Fixed recurrence uses one iteration count for all training instances, whereas dynamic recurrence samples iteration counts independently for each batch.
Stopping strategies: The setup uses R_min = 2 and R_max = 8 and avoids oracle complexity matching, reflecting cases where task complexity and required computation are unknown.

5 Systematic Generalization

The study tests systematic generalization on two-hop knowledge-graph reasoning and finds that recurrent-depth transformers can combine unfamiliar facts, unlike vanilla transformers. This ability develops through staged training dynamics and is reflected in intermediate representations.

Experimental setup: The study evaluates systematic generalization on two-hop inferred facts composed from atomic knowledge, using held-out in-distribution and out-of-distribution splits.The dataset contains 40k atomic facts and 273.6k inferred training facts, with nearly 2k OOD two-hop test facts.
Results: Recurrent-depth transformers generalize systematically, whereas the vanilla transformer completely fails on unfamiliar atomic-fact combinations.Even R = 2 achieves non-trivial generalization, while R = 1 is equivalent to the four-layer vanilla transformer.
Results: R = 4 converges in 2k epochs versus 7k for R = 2, with the acceleration also observed in wall-clock time.The comparison concerns convergence of recurrent-depth models under fixed training recurrence.
Training dynamics: Systematic generalization follows three stages: training-set overfitting, in-distribution generalization, and later OOD generalization after near-perfect in-distribution accuracy.The final transition occurs much later than training overfitting, around 10^4 versus 10^2 epochs.
Mechanistic analysis: Logit-lens analysis shows a progression from memorized target prediction to bridge decoding and finally successful OOD second-hop composition.Vanilla transformers can recover OOD bridge entities but fail at the second hop; recurrent models succeed only in Stage 3.

6 Depth Extrapolation

Recurrent-depth transformers extrapolate to deeper multi-hop reasoning by scaling recurrence, but training strategy and excessive inference iterations constrain this ability.

6.1 Experiment Setup: Curriculum training progressively introduces higher-hop data after each held-out accuracy threshold is reached, defining the model’s learnable recursion depth.Training stops before exposing facts beyond the model’s capability.
6.2 In-Distribution Generalization: Higher training recurrence increases learnable recursion depth without adding parameters, while dynamic recurrence extends it further than fixed recurrence.More iterations are not always beneficial: R=7 and R=8 both learn up to 16-hop tasks.
6.2 In-Distribution Generalization: Models first require prolonged training on low-hop tasks, then rapidly generalize to more complex in-distribution compositions once the compositional rule is discovered.Figure 6 tracks cumulative gradient updates required to first generalize at each hop complexity.
6.3 Depth Extrapolation: Increasing inference-time recurrence enables progressively harder depth extrapolation, but this benefit appears only when training recurrence exceeds four iterations.At the training recurrence, models struggle beyond observed complexity; additional inference iterations alleviate this limitation.
6.3 Depth Extrapolation: With identical training data through 12-hop, R=6 extrapolates to 14-hop whereas R=8 reaches 19-hop, and dynamic recurrence also reaches 19-hop.These results isolate training iteration strategy from exposure to more complex training data.
6.3 Depth Extrapolation: Dynamic recurrence should be preferred over fixed strategies with the same maximum iteration budget when training data is sufficiently complex.It effectively exploits the extrapolation range by increasing learnable recursion depth.
6.3 Depth Extrapolation: Excessive inference recurrence causes overthinking: for dynamic models, iterations beyond 15 do not improve OOD performance and can degrade predictions.Across models, logit margin rises to a peak and then declines; dynamic recurrence has slower margin decay than fixed recurrence.
6.3 Depth Extrapolation: Adaptive halting allocates inference-time iterations according to hop complexity, although output distributions may change little while entropy remains high.Figure 5 marks adaptive halting with r∗; Figure 9 compares iteration allocation against hop count.

7 Conclusion

The paper finds that recurrent-depth transformers support systematic generalization and depth extrapolation for implicit multi-hop reasoning, while acknowledging important limits in task scope and real-world transfer.

Conclusion: Recurrent-depth transformers address systematic generalization and depth extrapolation in controlled implicit multi-hop reasoning experiments.The models are trained from scratch on synthetic tasks, enabling architectural effects to be isolated.
Conclusion: Systematic generalization emerges through three stages, progressing from rote memorization to generalizable solutions.The broader literature reports that vanilla transformers struggle with compositional generalization as complexity increases.
Conclusion: Depth extrapolation is enabled by scaling inference-time recurrence, although latent overthinking limits performance on highly complex tasks.The dataset uses permutation-based knowledge graphs with bijective relations and sampled multi-hop traversals.
Limitations: The experiments use small, structured tasks and therefore do not establish immediate transfer to frontier language models or natural-language reasoning.The setup abstracts away surface variation, underspecification, distractors, distribution shift, and large-scale training factors.
Experimental setting: The study uses a recurrent block of four transformer layers for most experiments, with larger recurrent blocks and vanilla models examined separately.Training uses embedding dimension 768, 12 attention heads, AdamW, learning rate 10^-4, weight decay 0.01, and 2000 warmup steps.
Related work: Related work links recurrent-depth models to latent reasoning studies that probe hidden-state trajectories and intermediate representations at LLM scale.Reported analyses disagree about evidence for latent reasoning, while some metrics distinguish correct from incorrect trajectories.

E Causal analysis of systematic generalization

Causal analysis shows that systematic generalization depends on recovering the first-hop bridge early enough to support the second-hop composition, a condition recurrence helps satisfy.

Figure 11: Figure 11 compares bridge-mediated composition across effective depth in the systematicity setup.The relevant comparison concerns when bridge information becomes available and whether it supports the final composition.
Panel A: On in-distribution data, the bridge becomes decodable at r1 in shallow layer 5, whose activations causally affect final-answer prediction.This establishes an early-access condition for the first-hop rule in the composition.
Panel B: On out-of-distribution data, vanilla transformers recover the bridge only at layer 7+ and never make the target decodable.The late bridge recovery leaves insufficient depth for using the first-hop rule in the second-hop composition.

F Experiments with default initialization

Default Gaussian initialization produces unstable recurrence behavior: increasing training recurrence is usually inconsistent, while five random seeds show substantially different scaling and overthinking patterns.

Default initialization: Increasing recurrence from R ∈ {1, . . . , 7} does not consistently improve in-distribution or out-of-distribution generalization under default initialization.Only the R = 7 model shows strong inference-time scaling and robustness to latent overthinking.
Default initialization: The R = 5 model shows no out-of-distribution generalization when inference-time recurrence increases.Most other runs instead show performance degradation attributed to latent overthinking.
Figure 12: Figure 12 reports results for training recurrences R ∈ {1, . . . , 7} and dynamic recurrence under default initialization.The figure provides the run-level context for comparing recurrence strategies.
Random seeds: Five random-seed runs of the R = 5 model vary substantially in inference-time scaling and susceptibility to latent overthinking.Seed 1 is stable without scaling, seed 2 is stable with scaling, and seeds 3–5 combine scaling with degradation at higher recurrence.

G Experiments with larger vanilla transformers

Larger vanilla transformers improve in-distribution generalization with depth, but still lack out-of-distribution generalization and flexible inference-time computation compared with recurrent-depth models.

Larger vanilla transformers: 8-, 16-, and 24-layer vanilla transformers often exceed the 2-hop or 3-hop in-distribution generalization of the 4-layer baseline.This agrees with prior findings that greater vanilla depth can improve in-distribution generalization.
Larger vanilla transformers: Vanilla models remain weaker than recurrent-depth models with the same effective depth on in-distribution generalization.The comparison isolates the advantage of recurrence under matched effective depth.
Larger vanilla transformers: The larger vanilla models still show no out-of-distribution generalization to more complex samples.Increasing layer count therefore does not reproduce the reported systematic-generalization behavior.
Depth extrapolation: Vanilla transformers lack recurrent-depth models' flexibility to scale computation according to sample complexity.Figure 14 presents the depth-extrapolation results for the larger vanilla models.

H Additional experiments with matched-hop data

Matched-hop experiments show that higher train-time recurrence supports stronger out-of-distribution generalization when recurrence is scaled at inference time. Results are reported across training curricula up to 8-, 10-, and 12-hop compositions.

H Additional experiments with matched-hop data: Higher train-time recurrence produces stronger OOD generalization through inference-time scaling.This trend holds for models trained on data up to 8-hop, 10-hop, and 12-hop compositions.

I Results with different random seeds

Across random seeds, fixed-recurrence models generally improve ID and OOD generalization with higher train-time recurrence, while dynamic recurrence remains robust to overthinking and supports reliable adaptive halting. Additional analyses show that apparent deep generalization can arise from shallow suffix-based shortcuts rather than full multi-hop composition.

Fixed recurrence: Across two random seeds, higher train-time recurrence generally increases ID and OOD generalization, although the r = 8 model is an exception in one run.The experiments use inference-time recurrence up to 40 iterations.
Dynamic recurrence: Across three seeds, dynamic recurrence maintains generally high ID generalization and OOD extrapolation while remaining robust to latent overthinking.Adaptive halting r* reliably stops recurrence when further computation is unnecessary.
Extrapolation trends: Generalization ratios start near 1 after 2-hop training and increase as curriculum training exposes models to more complex compositions.The ratio compares maximum OOD generalization with maximum ID generalization under inference-time scaling.
Extrapolation trends: Neither larger models nor larger maximum train-time recurrence consistently improves the generalization ratio across the tested settings.The study varies model depth and maximum dynamic-recurrence training settings.
Shortcut analysis: Models showing strong performance up to 80 hops after training to 40 hops were using suffix shortcuts rather than resolving the full compositional chain.Activation patching concentrates causal effects near the relation-sequence tail, where trailing relations can identify the final entity.