PaperlayerView exact revision on arXiv

Source-linked AI summary

Bridging the Gap Between Latent and Explicit Reasoning with Looped Transformers

Ying Fan, Anej Svete, Kangwook Lee

arXiv:2606.31779v2cs.LGcs.CL

TL;DR

Existing latent reasoning methods fall behind explicit CoT beyond 1B parameters, while explicit CoT requires sequential token generation. LOTUS pairs a looped padded Transformer with parallel supervision on gold CoT tokens, bridging the gap at 3B and reducing thought-phase latency. Its latent readouts recover gold reasoning steps, and the approach remains scoped to math benchmarks with fixed block and loop hyperparameters.

  • Problem

    Existing latent reasoning methods underperform explicit CoT on math reasoning beyond 1B parameters, even as explicit CoT’s sequential token generation makes reasoning costly.

    Source: §1 · p. 1

  • Method

    LOTUS processes padded latent blocks with a looped Transformer and applies parallel cross-entropy supervision aligning each latent position to its gold CoT-step token.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    On LLAMA-3.2-3B-INSTRUCT, LOTUS bridges the in-domain GSM8K gap to explicit CoT, surpasses it on the out-of-domain math average, and cuts thought-phase latency by 2.5ˆ.

    Source: §6 · p. 13

  • Takeaways & Limitations

    Post-loop latents can be read through the base LM head to recover gold reasoning steps and surface alternative valid intermediate steps.

    Source: §1 · pp. 2–3

    Evaluation is on math benchmarks, and fixed K, c, and R mean chains longer than K steps require autoregressive completion of the tail.

    Source: §6 · p. 13

Abstract

from arXiv · show

Language models typically reason via explicit chain-of-thought (CoT), generating intermediate steps token-by-token. Latent CoT offers an alternative: it performs multi-step reasoning in the model's hidden states, replacing decoded tokens with continuous representations for greater efficiency. However, existing latent CoT methods underperform explicit CoT beyond 1B parameters, and the gap widens with scale. Looped, or recurrent-depth, Transformers, which reuse their weights to increase computation depth without adding parameters, are a natural fit for latent reasoning. We therefore ask whether looped Transformers can bridge this gap. We answer affirmatively with a simple recipe: a looped padded Transformer that processes K latent blocks in parallel for R iterations, with a cross-entropy loss on each latent position's gold CoT-step token, similar to explicit CoT supervision. We instantiate it as LOTUS (Looped Transformers with parallel supervision on latents). LOTUS is, to our knowledge, the first latent-CoT method to bridge the gap to explicit CoT at the 3B scale, while cutting thought-phase latency by 2.5x-6.9x from compact math expressions to natural language. Projecting LOTUS's post-loop latents through the base LM head recovers the gold reasoning steps and even surfaces alternative valid intermediate steps, evidence that its latent space is interpretable and CoT-aligned. Ablations confirm that both the looped backbone and the parallel supervision on gold CoT tokens are essential. Code is available at https://github.com/yingfan-bot/lotus.

1 Introduction

Existing latent reasoning methods lose ground to explicit chain-of-thought (CoT) beyond 1B parameters, while sequential token generation makes CoT costly. LOTUS combines looped parallel computation with direct supervision on gold CoT tokens, bridging the gap at 3B while reducing thought-phase latency.

1 Introduction

  • Explicit CoT requires N sequential model evaluations for N reasoning tokens, while common latent methods also generate latent tokens autoregressively and retain a sequential bottleneck.Other latent approaches reduce autoregression with parallel iterations but use indirect supervision rather than position-aligned gold CoT targets.Source: §1 · pp. 1–2
  • Beyond 1B parameters, existing latent methods do not match explicit CoT on math reasoning, and their gap widens as model size grows.On small backbones, latent methods can match CoT accuracy.Source: §1 · p. 1
  • At 3B, LOTUS brings GSM8K accuracy within 1 point of explicit CoT and surpasses it on the out-of-domain math average, with thought-phase latency reduced by 2.5ˆ.On natural-language CoT, it matches explicit CoT while reducing thought-phase latency by 6.9ˆ.Source: Figure 1 · p. 2, §1 · pp. 2–3
  • LOTUS places K learnable latent blocks between question and answer, loops the base model over them, and aligns each latent position in parallel to its corresponding gold CoT-step token.The post-loop latents are supervised through the base model’s LM head; LOTUS-aux routes supervision through an auxiliary decoder.Source: §1 · pp. 2–3

2 Preliminaries

The preliminaries distinguish explicit CoT’s decoded intermediate tokens from latent reasoning in hidden states, and introduce looped padded Transformers and related latent methods.

2 Preliminaries

  • Explicit CoT generates intermediate reasoning steps before the answer; training scores all CoT tokens in parallel, but inference generates them sequentially.N denotes the total number of CoT tokens across all steps.Source: §2 · p. 3
  • Latent reasoning performs intermediate computation in hidden states instead of decoded CoT tokens; prior work describes greater information capacity for latent steps on graph reachability.The cited example contrasts linearly many Transformer latent steps with quadratically many explicit CoT steps.Source: §2 · p. 3
  • Looped padded Transformers combine repeated application of shared backbone weights with learnable padding tokens that add parallel computation at each iteration.The passage characterizes the combination as more capable than either component alone.Source: §2 · p. 3
  • Related latent methods differ in how they generate or supervise latents: Coconut, CODI, and SIM-CoT decode autoregressively, while PCCoT and KaVa refine fixed latent budgets in parallel.Their supervision ranges from answer-only training and teacher distillation to auxiliary decoding or compressed teacher key–value caches.Source: §2 · p. 3

3 Method

LOTUS uses a padded latent workspace processed by a recurrent-depth Transformer and directly supervises post-loop latents against gold CoT tokens. Answer supervision complements this step loss, while LOTUS-aux provides an alternative auxiliary-decoder route.

3 Method

  • LOTUS grounds latent computation in exact gold CoT tokens and uses a looped padded Transformer, addressing the need for structured supervision in parallel reasoning.The method reads this supervision through the base model’s LM head; LOTUS-aux routes it through an auxiliary decoder.Source: §3 · pp. 3–4

3.1 Padded Latent Prefix

  • The padded latent prefix contains K blocks of c learnable latent tokens, delimited by learnable start and end tokens, with each block intended to represent one CoT step.K is fixed across examples and chosen to cover the step counts in the data; thus K is at least S on almost all examples.Source: §3.1 · p. 4

3.2 Looped Latent Computation

  • LOTUS runs the base model for R iterations to refine latent embeddings while attending to a cached question prefix; the recurrence is unrolled to fixed depth with gradients through all iterations.The answer suffix and trailing end token do not enter the loop.Source: §3.2 · pp. 4–5

3.3.1 Objective

  • After looping, LOTUS scores each block against its corresponding gold CoT step in parallel, then uses a separate answer-prediction pass conditioned on the post-loop latents.At inference, it computes the question cache once, loops to obtain latents, and decodes only the answer suffix autoregressively.Source: §3.3.1 · p. 5

3.3.2 Roles of Lstep and Lans

  • The step loss supplies per-position gold-token coverage, while answer loss favors jointly computed latent states that support the correct answer.The looped Transformer computes latent states jointly even though the step-loss readout factorizes across positions; dropping either loss degrades gold-chain likelihood in the reported analysis.Source: §3.3.2 · p. 6

3.4 LOTUS-aux: An Auxiliary Decoder Variant

  • LOTUS-aux replaces direct base-LM-head supervision with an auxiliary decoder that reads a latent block and scores its gold CoT step under teacher forcing.All CoT positions are scored in one forward pass; the autoregressive factorization affects the loss, not the parallel computation.Source: §3.4 · pp. 6–7

4 Experiments

Experiments compare LOTUS with explicit CoT and latent baselines across three backbones and evaluate accuracy, latency, and design choices. LOTUS stays close to explicit CoT at 3B while reducing thought-phase latency; ablations examine supervision and looped-architecture choices.

4.1 Setup

  • Experiments use GSM8k-Aug training, GSM8K and three out-of-domain benchmarks, plus a natural-language GSM8K-Aug stress test and three backbones.The backbones are GPT-2 (124M), LLAMA-3.2-1B-INSTRUCT, and LLAMA-3.2-3B-INSTRUCT; the main comparison uses six latent blocks and six loop iterations.Source: §4.1 · pp. 7–8, §4 · p. 7

4.2 Main Results

  • At LLAMA-3.2-3B-INSTRUCT, LOTUS nearly closes the in-domain GSM8K gap to explicit CoT and leads its out-of-domain average.Across scales, LOTUS stays within 1.5 points of explicit CoT in-domain; LOTUS + CODI tightens the 3B gap to within 1 point.Source: §4.2 · p. 8, Table 1 · p. 9
  • At 3B, LOTUS-aux performs comparably to LOTUS, but degrades at smaller scales while direct LM-head supervision remains robust.The 3B auxiliary-decoder variant performs well above CODI + SIM-CoT, which uses an auxiliary decoder of the same size as the base LM.Source: §4.2 · p. 8
  • LOTUS cuts thought-phase latency 2.5x versus explicit CoT on compact math traces and 6.9x on natural-language traces, with accuracy on par in the latter setting.On natural-language GSM8K, accuracy is 68.13 versus 68.41 for explicit CoT, while thought latency falls from 963.6 ms to 140.8 ms.Source: §4.2 · p. 8, Table 3 · p. 10

4.3 Ablation Studies

  • LOTUS’s prefill and answer phases match explicit CoT latency; only the looped latent thought phase differs.The reported prefill times are 16.7 versus 15.9 ms and answer times are 31.5 versus 29.5 ms for LOTUS and explicit CoT, respectively.Source: §4.3 · p. 8, Table 2 · p. 9

4.3.1 Latent Supervision Design

  • Supervised variants outperform both no latent supervision and CODI-only; direct LOTUS works best with post-loop readout, while LOTUS-aux works best per iteration.The best reported accuracies are 70.0% for direct post-loop LOTUS and 69.9% for per-iteration LOTUS-aux; the authors suggest the latter may benefit from a shorter gradient path.Source: §4.3.1 · p. 10, Table 4 · p. 10

4.3.2 Looped Architecture Design

  • Training accuracy rises from 49.7% at c=1 to 67.5% at c=5, then plateaus at 70.0% for c=25 and c=30.At inference with a c=25-trained checkpoint, reducing c hurts accuracy, while increasing it to 30 or 50 yields 70.5%; latency changes only about 30 ms across the sweep.Source: §4.3.2 · p. 10, Table 7 · p. 11, Table 8 · p. 11
  • With the latent prefix fixed at 150 positions, training accuracy rises from 14.6% at R=2 to 70.0% at R=6; inference accuracy peaks at the trained R=6.Inference accuracy increases from 22.7% at R=1 to 70.0% at R=6, then dips to 69.3% at R=7.Source: §4.3.2 · pp. 10–11, Table 5 · p. 10, Table 6 · p. 10

5 Latent Representation Analysis

LOTUS’s latent states carry readable gold-CoT signal and assign probability to unseen valid reasoning chains. The analyses also indicate that gold-step and answer supervision play complementary roles in shaping those states.

5.1 Likelihood of the Gold Chain under Different Readout Options

  • 3.07 NLL and 70.9% top-1 accuracy for LOTUS’s post-loop latents read through the base LM head show that they carry enough information to reconstruct the gold CoT.LOTUS-aux under free-running readout reaches 64.5% top-1; its teacher-forced result is an upper bound because it uses the gold prefix.Source: §5.1 · p. 12, Table 9 · p. 12
  • LOTUS’s post-loop latents also carry CoT signal under LOTUS-aux, despite not being directly supervised through its auxiliary head.Reading LOTUS-aux post-loop latents through that head gives a 25.8% top-5 probability; the readouts are less formatted than LOTUS’s, as expected without direct supervision.Source: §5.1 · p. 12

5.2 Mass on Unseen Valid Chains

  • Across four readouts, ground-truth-only numbers receive lower NLL than unseen-valid numbers, which in turn score better than random controls.For LOTUS through the base LM head, 15.3% of unseen-valid numbers appear in top-1 and 64.0% in top-5.Source: §5.2 · p. 12, Table 10 · p. 13

5.3 Disentangling the Two Losses in LOTUS

  • LOTUS outperforms both single-loss ablations, indicating that step and answer supervision contribute complementary structure to its latents.Answer-only reaches 5.97 NLL and 26.5% top-5, while step-only is worst on every reported column; answer loss alone does not align latents to the gold chain.Source: §5.3 · pp. 12–13, Table 11 · p. 13

6 Conclusion

The paper concludes that LOTUS approaches explicit CoT performance with parallel gold-CoT supervision on a looped padded Transformer, while producing readable and alternative valid reasoning signals in its latents. Its evaluation is limited to math benchmarks, and fixed latent-budget hyperparameters leave longer chains to autoregressive completion.

6 Conclusion

  • On LLAMA-3.2-3B-INSTRUCT, LOTUS bridges the in-domain gap to explicit CoT on GSM8K, exceeds CoT on the out-of-domain average, and cuts thought-phase latency by 2.5×.The authors report that the looped backbone, parallel gold-CoT supervision, sufficient block width, and loop depth are each necessary.Source: §6 · p. 13
  • Evaluation follows prior latent-reasoning work [Hao et al., 2025, Shen et al., 2025, Wei et al., 2025] in focusing on math benchmarks, so transfer to other domains remains open.Fixed K, c, and R must cover expected step counts; chains longer than K steps use autoregressive completion for the tail.Source: §6 · p. 13

Appendix

  • The appendices collect related-work discussion, notation, experimental details, PCL and autoregressive NLL bounds, and latent-representation analyses.Appendix F includes greedy readout examples and multi-path setup and examples.Source: Appendix · p. 21

A Related Work

  • The related-work discussion spans looped transformers, diffusion and padding approaches, latent CoT methods, and speculative decoding.It distinguishes LOTUS’s parallel latent reasoning from approaches that parallelize proposals within an explicit token trace.Source: A Related Work · pp. 22–23

C Closely Related Latent Reasoning Methods

The comparison situates LOTUS against latent methods that use autoregressive decoding or Jacobi-style parallel refinement. LOTUS’s fixed, step-aligned latent blocks and direct gold-step supervision accompany strong GSM8K results across compact-expression and natural-language settings.

C Closely Related Latent Reasoning Methods

  • LOTUS removes the auxiliary path used by CODI and SIM-CoT, supervising its own post-loop hidden states through the answer LM head at block granularity.LOTUS-aux retains an auxiliary autoregressive decoder but produces all blocks in parallel through the looped backbone.Source: C Closely Related Latent Reasoning… · p. 23
  • LOTUS fixes each latent block to a gold CoT step and supervises it against that step’s tokens, unlike PCCoT and KaVa’s unaligned latent positions and indirect supervision.PCCoT and KaVa use 24 latent tokens refined over 3 Jacobi iterations; LOTUS improves with loop depth up to 6, whereas PCCoT accuracy peaks at about 3 iterations and degrades beyond that.Source: C Closely Related Latent Reasoning… · pp. 23–24

Comparison with latent methods at other compute budgets

  • At 1B and 3B, LOTUS reaches 57.3 vs. 56.5 and 70.0 vs. 65.7 GSM8K accuracy against KaVa, respectively, while PCCoT and KaVa fall further behind as backbone size grows.The comparisons use results reported in the cited prior work; Coconut is evaluated only up to LLAMA-3.2-1B in Wei et al. [2025].Source: C Closely Related Latent Reasoning… · p. 25, Table 13 · p. 25

Natural-language CoT stress-test

  • On natural-language GSM8K, 3B LOTUS reaches 68.13% accuracy with 140.8 ms thought latency, versus explicit CoT’s 68.41% and 963.6 ms.LOTUS also scores 73.40 on SVAMP versus explicit CoT’s 71.93, and 16.27 on GSM-Hard versus 18.27; it leads the latent baselines on both tasks.Source: C Closely Related Latent Reasoning… · p. 25, Table 14 · p. 25

D Experimental Details

The experiments train and evaluate LOTUS on math reasoning with three backbone sizes, including a natural-language CoT stress test. Implementation choices set a fixed latent workspace and loop schedule, with longer-than-budget chains completed autoregressively.

D Experimental Details

  • The latent prefix contains Kc positions, with K=6 covering the reference CoT step count for 99% of GSM8k-Aug examples; longer problems use autoregressive completion for the unsupervised tail.For LLAMA backbones c=25 gives 150 positions, while GPT-2 uses c=13 for 78 positions.Source: D Experimental Details · p. 25
  • Training uses a staged curriculum that converts one additional CoT step into a latent block per Estage epochs, with loop count Re set by the loops-per-block ratio.Unless stated otherwise, Estage=1 and r=1, so the saturated model uses R=K=6; the loop-depth ablation varies r while holding K=6.Source: D Experimental Details · pp. 25–26
  • Models start from explicit-CoT checkpoints and are fine-tuned for 30 epochs with AdamW, gradient clipping, and a step-loss weight of 0.05.Checkpoints are selected by GSM8k validation accuracy; the final-epoch checkpoints perform similarly.Source: D Experimental Details · p. 26
  • The natural-language stress test keeps the 3B looped backbone, switches to natural-language GSM8k-Aug, increases c to 50, and sets λstep to 0.033.Other settings match the original GSM8k-Aug training setup.Source: D Experimental Details · p. 26
  • LOTUS-aux uses a full-size auxiliary-decoder copy of the base model; lightweight auxiliary decoders were unstable in training.The full-size choice matches SIM-CoT’s decoder architecture.Source: D Experimental Details · p. 26

E PCL Versus Autoregressive NLL Lower Bounds

PCL scores gold CoT tokens in parallel, but its best achievable NLL can exceed the autoregressive factorization when gold tokens remain dependent given the latents.

E PCL Versus Autoregressive NLL Lower Bounds

F.1 Greedy Readout Examples

Greedy-readout examples show that post-loop latents can surface correct and alternative reasoning steps, but may also yield plausible wrong answers. The discrete readouts are diagnostic projections, not the inference pipeline or a measure of reasoning quality.

F.1 Greedy Readout Examples

  • The examples show LOTUS-aux’s auxiliary decoder can produce initial gold steps and then repeat or loop on later values.Source: §F.1 · pp. 27–29
  • Janet’s ducks and James’s sprints are answered correctly, with gold or valid alternative steps surfacing in LOTUS’s latent readouts.For James, the readout surfaces the valid alternative step «3*60=180» before 540 stabilizes.Source: §F.1 · pp. 27–28
  • LOTUS and LOTUS-aux both miss Carla’s gold total of 160, producing 180 and 120 respectively after six loops.Their latents recover several relevant intermediates but fail to combine them into the correct total.Source: §F.1 · p. 28
  • For Melanie’s vacuum problem, LOTUS outputs 30 and LOTUS-aux outputs 36, although the gold answer is 18.LOTUS’s valid chain surfaces in B5, but its final readout commits to 30.Source: §F.1 · p. 29
  • The displayed A/B/C streams are diagnostic projections of post-loop latents; inference generates the final answer from continuous latents, not a discretized chain.Noisy streams can accompany correct answers, and a clean auxiliary-decoder chain is not required for correctness.Source: §F.1 · p. 29

F.2 Details for Multi-Path Analysis (Section 5.2)

The multi-path analysis compares readouts against trained-path and unseen-valid-path intermediates using filtered GSM8K examples and a per-number scoring procedure.

F.2 Details for Multi-Path Analysis (Section 5.2)

  • The training split is augmented from about 7.5K GSM8K problems to 385,620 examples, and the path bank uses original GSM8K questions also present in training.Source: §F.2 · p. 29
  • The path bank pairs each question’s training gold chain with an unseen but valid alternative sampled from up to ten verified paths.The alternatives come from rejectionsampling LLAMA-3.1-8B on original GSM8K training questions and are verified absent from training.Source: §F.2 · p. 29
  • Filtering leaves 341 questions, with 871 trained-only, 974 unseen-only, and 974 matched random-control numbers.The token-clean filter excludes unseen-path numbers sharing a sub-token with trained-path, question, or answer numbers.Source: §F.2 · p. 29
  • The analysis scores each target number by its best-position per-token NLL, then averages within number sets and across questions.Top-k columns instead report the fraction of unseen-path numbers recovered among the top-k candidates.Source: §F.2 · p. 29

F.3 Multi-Path Readout Example

A bus-distance question has two valid reasoning chains to the same answer, and LOTUS’s post-loop readout surfaces intermediates found only in the unseen chain.

F.3 Multi-Path Readout Example

  • Figure 4 aligns the trained path, unseen valid path, and latent positions where the only-U intermediates surface.Source: Figure 4 · p. 30, §F.3 · p. 30
  • The unseen valid chain computes daily distances first, using 24 and 18 as intermediates before reaching the shared answer 108.The trained chain instead computes hours across days before combining the distances.Source: §F.3 · p. 30
  • LOTUS’s post-loop readout surfaces unseen-path values 24 and 18, with probabilities 0.77 at top-1 position 142 and 0.14 at top-2 position 116, respectively.Both numbers are single tokens and absent from the trained path, question, and answer, so their surfacing is not a digit-overlap artifact.Source: §F.3 · p. 30

Related papers

Loading 2606.31779v2…