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

1. Problem Statement

The paper addresses a fundamental limitation in Large Language Models (LLMs): the inability to perform implicit reasoning (reasoning within a single forward pass without explicit Chain-of-Thought). While LLMs store vast parametric knowledge, they struggle to compose this knowledge to solve novel multi-hop reasoning tasks.

Specifically, the authors identify two critical generalization failures in standard (vanilla) transformers:

1. Systematic Generalization: The inability to combine knowledge components (facts/rules) that were never seen together during training. For example, if a model learns $A \to B$ and $C \to D$ separately, it often fails to infer $A \to D$ if the path requires chaining unseen combinations.
2. Depth Extrapolation: The inability to generalize to reasoning depths (number of hops) significantly deeper than those observed during training (e.g., training on 5-hop tasks but failing on 10-hop tasks).

The authors hypothesize that vanilla transformers fail because knowledge is distributed across specific layers; deeper layers cannot access facts stored in shallower layers without parameter sharing or iterative processing.

2. Methodology

To address these limitations, the authors propose and study Recurrent-Depth Transformers (also known as looped transformers).

Model Architecture

• Core Mechanism: Instead of a deep stack of unique layers, the model uses a single transformer block (with $L$ layers) that is applied iteratively ($R$ times) to the hidden state.
• Computation: The output of iteration $r$ becomes the input for iteration $r+1$. This allows the model to allocate more computation (increasing $R$) at inference time without changing the parameter count.
• Stabilization: To prevent instability from weight sharing in deep unrolling, the authors use zero-initialization for the output projection matrices of attention and feed-forward blocks. This ensures the model starts as an identity mapping.
• Stopping Strategies:
  • Fixed Iteration: A constant number of loops ($R$) for all inputs.
  • Dynamic Iteration: The number of loops is sampled per batch (e.g., from a Poisson distribution) to simulate variable complexity.
  • Adaptive Halting: A custom stopping criterion that combines KL-divergence (stability of output) and Entropy (confidence of prediction) to stop recurrence only when the model is both stable and confident.

Task Formulation

The authors use a controlled synthetic task based on a Knowledge Graph (KG):

• Atomic Facts: Triplets $(h, r, t)$ representing entities and relations.
• k-hop Inference: Predicting the tail entity $t$ given a head $h$ and a sequence of $k$ relations.
• Generalization Challenges:
  • Systematic: Testing on combinations of atomic facts never seen in training (OOD).
  • Depth Extrapolation: Testing on $k$-hop tasks where $k > k_{train}$.
• Dataset Design: To prevent "shortcut" learning (where models guess based on relation suffixes), they employ a permutation-based constraint where every relation acts as a bijection over entities, forcing the model to traverse the full chain.

3. Key Contributions & Findings

A. Systematic Generalization via "Grokking"

• Vanilla vs. Recurrent: Vanilla transformers fail completely on systematic generalization (OOD 2-hop tasks), regardless of training time. Recurrent-depth transformers succeed.
• Three-Stage Grokking Process: The authors identify a distinct training dynamic in recurrent models:
  1. Memorization: The model overfits training data; only training accuracy improves.
  2. In-Distribution Generalization: After prolonged training, the model learns the compositional rule for seen combinations (the "grokking" phenomenon).
  3. Systematic Generalization: Only after mastering in-distribution tasks does the model suddenly generalize to unseen (OOD) combinations.
• Mechanistic Analysis: Using Logit Lens, they show that in Stage 1, the model predicts the final answer without decoding intermediate "bridge" entities (rote memorization). In Stage 3, the model correctly decodes bridge entities at intermediate layers, indicating true compositional reasoning.

B. Depth Extrapolation via Inference-Time Scaling

• Unlocking Depth: Recurrent-depth models can generalize to depths far beyond training (e.g., training on 12-hop, testing on 24-hop) if inference-time recurrence ($R_{inference}$) is scaled up.
• Training Strategy Matters:
  • Dynamic Recurrence: Training with variable iteration counts (dynamic) yields better extrapolation than fixed recurrence, even with the same maximum iteration budget.
  • Prerequisite: Sufficient training-time iterations are required to "learn the rule" before inference-time scaling can be effective.
• The "Overthinking" Limitation: There is a critical trade-off. Increasing inference iterations improves performance up to a point, after which performance degrades. This is termed "Overthinking," where excessive recurrence causes the model to lose confidence or drift from the correct answer.
  • Observation: The logit margin (confidence) peaks and then declines as iterations increase.
  • Mitigation: The proposed Adaptive Halting (using both KL-divergence and Entropy) effectively stops recurrence before overthinking occurs, allocating compute proportional to task complexity.

4. Results Summary

• Systematic Generalization: Recurrent models (e.g., $R=4$) achieve near-perfect accuracy on OOD 2-hop tasks, while vanilla transformers ($R=1$) remain at chance level.
• Depth Extrapolation:
  • Models trained on up to 12-hop tasks successfully generalized to 24-hop tasks when inference iterations were increased.
  • Dynamic training strategies achieved the highest extrapolation limits (up to 22-hop) compared to fixed strategies.
• Efficiency: Adaptive halting reduced the average number of iterations required for complex tasks, proving that compute can be allocated dynamically based on difficulty.

5. Significance

• Architectural Insight: The paper demonstrates that weight sharing via recurrence is a powerful inductive bias for compositional reasoning, allowing models to overcome the "layer-depth" bottleneck of standard transformers.
• Training Dynamics: It provides a mechanistic explanation of "grokking" in the context of reasoning, showing that systematic generalization is a late-stage emergent property following memorization and rule discovery.
• Inference-Time Compute: It validates the hypothesis that scaling compute at inference time (via recurrence) can unlock capabilities not present at training time, provided the model is trained with appropriate strategies (dynamic recurrence) and stopping mechanisms (adaptive halting).
• Limitations Identified: The discovery of "overthinking" highlights that more computation is not always better; unbounded recurrence can degrade performance, necessitating intelligent stopping criteria.

In conclusion, the paper establishes recurrent-depth transformers as a promising architecture for implicit reasoning, offering a path to compositional generalization and depth extrapolation that vanilla transformers cannot achieve, provided specific training and inference strategies are employed.