PaperlayerView exact revision on arXiv

Source-linked AI summary

Reasoning with Latent Thoughts: On the Power of Looped Transformers

Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, Sashank J. Reddi

arXiv:2502.17416v1cs.CLcs.AIcs.LG

TL;DR

The paper studies whether reasoning requires large depth rather than large parameter counts, and investigates looped Transformers as a way to obtain depth through weight sharing. It combines synthetic experiments, language-model evaluations, theoretical analysis, and regularization, finding that looped models can match or exceed deeper non-looped models on reasoning despite fewer parameters, while exhibiting latent-thought and effective-depth behavior. The paper also notes that its tests cover only a subset of reasoning problems.

  • Problem

    The paper asks how to solve reasoning problems that require large depth but not necessarily many parameters.

  • Method

    The paper studies k-layer Transformers looped L times, compares them with iso-parameter and iso-FLOP non-looped baselines, and analyzes their expressivity, latent thoughts, and regularization.

  • Results

    Looped models nearly match or outperform deeper non-looped models on many reasoning tasks and can show reasoning performance close to or above iso-FLOP models despite fewer parameters and worse perplexity.

  • Takeaways & Limitations

    Looping provides a parameter-efficient route to depth, with effective-depth scaling and latent-thought behavior connected to chain-of-thought reasoning.

  • Takeaways & Limitations

    The experiments test only a subset of reasoning problems, leaving whether the results extend to other forms such as multimodal and common-sense reasoning as an open question.

Abstract

from arXiv · show

Large language models have shown remarkable reasoning abilities and scaling laws suggest that large parameter count, especially along the depth axis, is the primary driver. In this work, we make a stronger claim -- many reasoning problems require a large depth but not necessarily many parameters. This unlocks a novel application of looped models for reasoning. Firstly, we show that for many synthetic reasoning problems like addition, $p$-hop induction, and math problems, a $k$-layer transformer looped $L$ times nearly matches the performance of a $kL$-layer non-looped model, and is significantly better than a $k$-layer model. This is further corroborated by theoretical results showing that many such reasoning problems can be solved via iterative algorithms, and thus, can be solved effectively using looped models with nearly optimal depth. Perhaps surprisingly, these benefits also translate to practical settings of language modeling -- on many downstream reasoning tasks, a language model with $k$-layers looped $L$ times can be competitive to, if not better than, a $kL$-layer language model. In fact, our empirical analysis reveals an intriguing phenomenon: looped and non-looped models exhibit scaling behavior that depends on their effective depth, akin to the inference-time scaling of chain-of-thought (CoT) reasoning. We further elucidate the connection to CoT reasoning by proving that looped models implicitly generate latent thoughts and can simulate $T$ steps of CoT with $T$ loops. Inspired by these findings, we also present an interesting dichotomy between reasoning and memorization, and design a looping-based regularization that is effective on both fronts.

1 INTRODUCTION

The paper argues that many reasoning problems need substantial computational depth but not proportionally many parameters, motivating looped Transformers with shared weights. It presents theoretical and empirical evidence that looping supports reasoning, including latent-thought behavior, downstream reasoning gains, and looping-inspired regularization.

  • Motivation and core claim: Many reasoning problems require depth but not necessarily parameters, and can be solved via looped models.Looped models iteratively apply the same few-parameter function to the input.
  • Language modeling: For language modeling, looped models show an inductive bias toward reasoning performance despite worse perplexity and memorization than an iso-FLOP non-looped model.Their reasoning performance can approach or exceed the iso-FLOP baseline despite L times fewer parameters and worse perplexity.
  • Latent thoughts and CoT: Looped models generate latent thoughts and can theoretically simulate chain-of-thought reasoning, with multiple latent thoughts generated per iteration.The paper connects looping to inference-time scaling through effective depth and CoT-style iterative reasoning.
  • Regularization: The paper proposes looping-inspired regularization to encourage conventional models to inherit looped models’ reasoning bias while retaining free parameters for perplexity.The regularizer makes corresponding blocks’ weights close using cosine similarity.
  • Experimental framing: The paper compares a k-layer model looped L times with iso-parameter k-layer and iso-FLOP kL-layer non-looped baselines.The looped model uses weight sharing, while the iso-FLOP baseline has the same effective depth and L times more parameters.
  • Synthetic reasoning tasks: On synthetic addition, p-hop induction, and GSM-style math tasks, iso-FLOP looped models can nearly match or outperform non-looped models despite using fewer parameters.Theoretical results support these findings by showing that many such problems admit iterative algorithms solvable with nearly optimal depth.

2 LOOPED MODELS ON SIMPLE REASONING TASKS

The paper evaluates looped Transformers on synthetic reasoning tasks, testing whether repeated computation can provide effective depth with fewer parameters. Across addition, p-hop induction, and i-GSM, looped models outperform equally parameterized shallow models and approach deeper non-looped baselines.

  • Task design: The experiments cover n-ary addition, p-hop induction, and synthetically generated grade-school math problems.These tasks are procedurally constructed and are intended to provide insights into reasoning that has algorithmic structure.
  • Looped models: A looped model applies the same sequence-to-sequence function repeatedly, with looping independent of the architecture used for that function.The study uses the vanilla looping construction, denoting repeated application as f^(L).
  • Addition: For addition, looped models nearly match the 12-layer non-looped baseline and substantially outperform iso-parameter shallow baselines.The comparison evaluates different operand counts and model widths under matched parameter or compute settings.
  • Addition: A 1-layer network looped 12 times solves the addition task while using 1/12th of the baseline’s parameters.This result supports the claim that the task primarily requires depth rather than additional parameters.
  • p-hop induction: On p-hop induction, reasonably deep looped models perform as well as the baseline while using substantially fewer parameters.The task varies hop count between 16 and 32 and tests recursive backtracking over sequences.
  • i-GSM: On i-GSM, looped models significantly outperform iso-parameter models and perform as well as non-looped iso-flop models.The problems are symbolic arithmetic computation graphs with depth restricted to 4.

3 LANGUAGE MODELING WITH LOOPED MODELS

The experiments compare looped language models with iso-parameter and iso-compute non-looped baselines across perplexity, memorization, and reasoning evaluations. Looping generally sacrifices perplexity and memorization performance but provides a stronger inductive bias toward reasoning, with benefits that scale with effective depth and connect to latent-thought computation.

  • Experiments: The evaluation covers perplexity and four downstream groups—closed-book QA, open-book QA, math word problems, and reasoning primitives—spanning 19 tasks.Closed-book QA primarily measures memorization, while open-book QA, math problems, and reasoning primitives assess contextual inference or reasoning.
  • Experiments: Looped models improve over iso-parameter baselines but have worse perplexity than the 24-layer iso-flop baseline because they use fewer parameters.The looped models cover roughly 34–50% of the perplexity gap between the iso-parameter and iso-flop baselines.
  • Downstream results: The 12-layer model looped twice reaches 34.3 on math word problems versus 29.3 for the 24-layer baseline despite using 50% of its parameters and having worse perplexity.For k ≥6, looped models can almost match the 24-layer baseline on math word problems.
  • Downstream results: Looped models outperform the 24-layer baseline on reasoning primitives for every tested layer budget despite having 24/k× fewer parameters.Reasoning primitives are synthetic contextual tasks, so memorization cannot account for this result.
  • Inductive bias: At equivalent log perplexity, looped models perform similarly on closed-book QA but better on open-book QA and math word problems.The isoplot comparison therefore separates memorization-oriented tasks from tasks requiring more reasoning.
  • Scaling behavior: Accuracy increases with effective depth for both looped and non-looped models with diminishing returns, while looping has higher relative benefits on reasoning tasks and reaches 1.19x for reasoning primitives.The study fits accuracy as a function of effective depth and compares the impact of looping with adding parameters.
  • Connection to chain-of-thought: Looped models can generate multiple latent thoughts per iteration and theoretically simulate T steps of chain-of-thought reasoning with T loops.This connection motivates using looping as an inference-time scaling mechanism.

4 LOOPING-INSPIRED REGULARIZATION

The paper introduces weight regularization that encourages consecutive blocks in a deep transformer to resemble repeated loop iterations. This preserves the reasoning-oriented inductive bias of looping while retaining free parameters for perplexity optimization.

  • Regularization method: The regularizer encourages consecutive k-layer blocks to have similar weights using cosine similarity across parameter groups.It is added to the standard cross-entropy objective to make an L-layer model approximately looped.
  • Regularization method: λreg = 0 gives standard training, whereas λreg →∞ converges to a fully looped model; intermediate values produce approximately looped models.For a (4 ⊗6) setting, the authors use k = 4, L = 24, and λreg = 10.
  • Regularization effect: After training, cosine similarity between successive blocks is around 0.98 or higher for all parameter groups.This verifies that the regularization produces the intended similarity between blocks.
  • Results: Regularized models retain a strong reasoning-oriented inductive bias and improve on reasoning tasks without hurting perplexity.The isoplots show similar behavior to baseline on closed-book QA but stronger behavior on open-book QA and reasoning problems.

5 THEORETICAL ANALYSIS FOR LOOPED MODELS

The theoretical analysis explains why looped Transformers can achieve substantial reasoning depth with few parameters, establishing constructions for group composition, p-hop induction, and latent-thought CoT simulation.

  • Definitions: A looped Transformer is defined by repeatedly applying the same Transformer block before the output layer.The model is pθ,T = OUTPUTθOUTPUT ◦ (TBθTB)^T ◦ EMBEDθTE,θPE.
  • Group composition: A one-layer Transformer looped ⌈log2 n⌉ times computes composition of n elements from any finite group.The construction uses constant precision and matches the best known depth upper bound for non-looped models.
  • Layer simulation: Looped Transformers can simulate non-looped Transformers with repeated layers using fewer parameters while preserving depth.Theorem 5.2 provides the general simulation result under bounded activations and a bounded number of distinct layers.
  • p-hop induction: A constant-layer Transformer with log(p) loops suffices to solve p-hop induction, matching the non-looped layer lower bound.This gives a theoretical basis for the strong performance of small looped models on p-hop tasks.
  • CoT simulation: With L loops, a looped Transformer can simulate L steps of another Transformer’s chain-of-thought computation.The construction uses dummy tokens and masking so loop iterations reproduce the outputs of sequential CoT steps.

6 RELATED WORK

The related work situates looped Transformers among studies of reasoning, depth, computational representation, chain-of-thought, and parameter-efficient or adaptive computation.

  • Reasoning and representation: Prior work studies Transformers’ reasoning abilities, computational representation, and chain-of-thought behavior.The cited literature includes synthetic reasoning tasks, computational problems, and formal analyses of CoT reasoning.
  • Chain-of-thought: Figure 4 presents CoT as a looped process producing one thought token per iteration and contrasts it with parallel latent thoughts.The right panel describes theoretical CoT simulation through appropriately masked updates.
  • Looped models: Earlier research connected looping with supervised learning, BERT pretraining, adaptive computation, and extrapolation on algorithmic tasks.These studies motivate examining looping specifically for reasoning.
  • Inductive bias: Recent work on inductive biases suggests that architectural choices can disproportionately affect reasoning performance during pretraining.The paper relates this line of inquiry to stacking and its hypothesized connection with looped models.

7 CONCLUSIONS, LIMITATIONS AND FUTURE WORK

The paper concludes that looping supports parameter-efficient reasoning and may bias language models toward reasoning over memorization, while identifying broader validation as future work.

  • Conclusions: Looped models solve many studied reasoning problems with fewer parameters and disproportionately improve reasoning relative to memorization.The conclusion frames these findings as a new direction for looped models in reasoning.
  • Limitations and future work: The evidence is limited to a subset of reasoning problems, leaving multimodal and common-sense reasoning as open questions.The authors also identify formalizing reasoning problems as an important future direction.
  • Future work: The authors suggest looped models as a possible route toward more efficient inference-time scaling for deeper reasoning.This proposal is motivated by the paper’s connections among effective depth, latent thoughts, and CoT.

A.1 SIMPLE REASONING SETUP DETAILS

The appendix describes the synthetic reasoning tasks, datasets, model configurations, and downstream evaluation plots used to assess looped models.

  • Synthetic tasks: The synthetic benchmark suite includes n-ary addition, p-hop induction, and procedurally generated i-GSM math problems.These tasks test arithmetic, multi-step token tracking, and symbolic grade-school mathematics.
  • p-hop induction: The p-hop dataset enforces valid p-hop chains, then shuffles them among filler tokens while preserving hop order.Training uses 4M examples, with 262k examples each for testing and validation.
  • Downstream evaluation: Figure 7 plots task-group evaluations against reversed validation log perplexity for an 8-layer baseline and an 8-layer model looped three times.Looping follows the baseline on closed-book QA but performs better at equivalent log perplexity on open-book QA and math word problems.

A.3 RESULTS FOR EACH TASK GROUP

This section provides individual-task results for closed-book QA, open-book QA, math word problems, and reasoning primitives, including looped and regularized models.

  • Tables 5–8 report individual-task results for closed-book QA, open-book QA, math word problems, and reasoning primitives.
  • The evaluations include looped models from Table 3 and models trained with regularization from Table 4.
  • Table 8 specifically presents downstream evaluations for reasoning primitives.

B THEORETICAL RESULTS

The theoretical preliminaries define transformer components, causal self-attention, feed-forward networks, finite-precision arithmetic, and loop-compatible binary representations.

  • The embedding layer maps token sequences to vectors by combining token and positional embeddings.
  • The transformer block uses causal self-attention with query, key, value, and output projections, alongside a fully connected feed-forward layer.
  • The output layer maps vector sequences to probability distributions over the vocabulary, while deterministic decoding selects the highest-probability token.
  • The looped-transformer framework is defined over transformer blocks applied repeatedly, with the stated attention definition using a causal mask.
  • The analysis assumes finite-precision fixed-point numbers, rounding, binary encodings, and finite-precision operations such as inner products, multiplication, and softmax.

B.2 PROOFS

The proofs construct transformers that simulate iterative computations by selectively activating layer-specific MLPs and attention heads, with group composition serving as a central example.

  • A depth counter can be implemented by adding one hidden dimension to the MLP.
  • Undesired MLP neurons are deactivated through control signals that identify the layer-specific parameters being used.
  • The construction activates only the desired attention head at each loop, while routing undesired heads to a zero-valued position.
  • Group composition: Group composition is computed iteratively by updating intermediate group elements across ⌈log2 n⌉ rounds and returning the final element.
  • Motivation: Constant-depth transformers cannot solve composition of S5 in one pass under NC1 ≠ TC0, whereas shallow transformers can simulate it with chain-of-thought.
  • Group composition: The transformer construction represents group elements and positions with signed binary vectors, then retrieves required values through attention and composes them with an MLP.

B.3 CONNECTION TO CHAIN-OF-THOUGHT REASONING

This section formalizes how looped transformers connect to chain-of-thought reasoning. It shows that looping can simulate multiple CoT steps while preserving the target transformer's final output.

  • Theorem B.5 states that an L-layer non-looped transformer with m CoT steps can be simulated by a looped transformer with L + O(1) layers and m loops.
  • The construction relies on MLP subroutines for argmax, decoding and embedding, control gating, position shifting, binary increment, and comparison.
  • The simulation uses CoT masking so that each loop updates the embedding at the next generated-token position while preserving earlier positions.
  • The constructed looped model stores position and loop-count information in additional embedding coordinates.
  • A masked looped transformer can itself be simulated without masking by increasing layers, embedding size, and hidden dimensions by constant factors.
  • After m loops, the looped transformer's output is identical to the non-looped transformer's output after m CoT steps.
Loading 2502.17416v1…