One-sentence Summary

Researchers from Microsoft Research, ETH Zürich, KRAFTON, and Ludo Robotics propose LOTUS, a looped padded Transformer that processes K latent blocks in parallel for R iterations with cross-entropy supervision on gold chain-of-thought step tokens, bridging the gap between latent and explicit reasoning at the $3B$3B scale while cutting thought-phase latency by $2.5 \times -6.9 \times$2.5×−6.9× and producing interpretable, CoT-aligned latent representations.

Key Contributions

• LOTUS is the first latent chain-of-thought method to close the performance gap to explicit chain-of-thought at the 3B scale, achieving this on Llama-3.2-3B-Instruct for GSM8K.
• It reduces thought-phase latency by 2.5× to 6.9× compared to explicit chain-of-thought and surpasses the out-of-domain average on math benchmarks.
• Post-loop latent states are interpretable: projecting them through the base language model head recovers gold reasoning steps and alternative valid chains, and ablations confirm the looped backbone and parallel gold-token supervision are essential.

Introduction

Scaling inference compute through chain-of-thought (CoT) reasoning has become a primary method for improving language model performance, but generating each reasoning token sequentially incurs high latency. Latent reasoning offers a more efficient alternative by performing intermediate computation in continuous hidden states instead of discrete tokens, condensing many steps into fewer model evaluations. Small-scale models benefit from this approach, yet past roughly 1B parameters no existing latent method matches explicit CoT accuracy on math tasks, and the gap widens with model size. The authors identify two common failures: prior methods generate latent thoughts autoregressively, which retains the expensive sequential bottleneck of CoT, and they lack direct, position-aligned CoT supervision, allowing latent representations to drift and destabilizing training. To overcome both, the authors present LOTUS, a looped padded Transformer that refines a fixed budget of latent token blocks over a small number of parallel passes and supervises these latents directly against gold CoT tokens via the base model’s own language modeling head. This parallel design eliminates the sequential dependency while ground-truth CoT alignment prevents latent drift. LOTUS is the first latent reasoning method to close the in-domain accuracy gap with explicit CoT on GSM8K at the 3B scale, exceeding CoT on out-of-domain math benchmarks and reducing thought-phase latency by up to 6.9 times on verbose natural-language reasoning.

Method

To transform a standard language model into an efficient latent reasoner, the authors propose LOTUS (Looped Transformers with parallel supervision on latents). The core design grounds latent computation directly in gold chain-of-thought (CoT) tokens using two main ingredients: a padded latent prefix processed by a looped transformer and parallel cross-entropy supervision on exact gold CoT tokens.

The input construction begins with a padded latent prefix. For a question $Q$Q, $S$S CoT steps $T_1, \ldots, T_S$T1​,…,TS​, and an answer $A$A, the model constructs an input sequence containing $Q$Q, a learnable beginning-of-thought token $\langle \mathrm{BoT} \rangle$⟨BoT⟩, $K$K blocks of $c$c learnable latent tokens $\langle \mathrm{lat} \rangle$⟨lat⟩, an end-of-thought token $\langle \mathrm{EoT} \rangle$⟨EoT⟩, and the answer $A$A. The block budget $K$K and per-block width $c$c are fixed hyperparameters, ensuring the latent region can align with the CoT steps across examples.

During the looped latent computation, the base causal language model $f_{\theta}$fθ​ first processes the prefix $[Q, \langle \mathrm{BoT} \rangle]$[Q,⟨BoT⟩] to populate a key-value cache $\mathcal{C}_{\mathrm{pre}}$Cpre​. The model then iterates $R$R times over the latent embeddings $\boldsymbol{E} \in \mathbb{R}^{Kc \times d}$E∈RKc×d. The hidden states at the latent positions after iteration $t$t are computed as:

$$\begin{array}{l} \boldsymbol{h}^{(0)} = f_{\boldsymbol{\theta}}(\boldsymbol{E} \mid \mathcal{C}_{\mathrm{pre}}), \\ \boldsymbol{h}^{(t)} = f_{\boldsymbol{\theta}}\left(\boldsymbol{E} + \boldsymbol{h}^{(t-1)} \mid \mathcal{C}_{\text{pre}}\right), \quad t = 1, \dots, R. \end{array}$$h(0)=fθ​(E∣Cpre​),h(t)=fθ​(E+h(t−1)∣Cpre​),t=1,…,R.​

This finite-unroll recurrence allows the model to refine the latent embeddings jointly over the padded workspace while attending to the cached question context.

As shown in the figure below:

The training objective combines step-aligned CoT supervision and answer supervision. After $R$R iterations, the step CoT supervision loss $\mathcal{L}_{\text{step}}$Lstep​ aligns each latent position $(i,j)$(i,j) in the grid with the corresponding CoT step token $T_{i,j}$Ti,j​ using a single batched cross-entropy through the base model's LM head $f_{\text{head}}$fhead​:

$$\mathcal{L}_{\text{step}} = \frac{1}{N_{\text{step}}} \sum_{i=1}^{K} \sum_{j=1}^{c} \mathrm{CE}\left(f_{\text{head}}\left(\boldsymbol{h}_{i,j}^{(R)}\right), T_{i,j}\right),$$Lstep​=Nstep​1​i=1∑K​j=1∑c​CE(fhead​(hi,j(R)​),Ti,j​),

where $N_{\text{step}}$Nstep​ is the total number of supervised CoT tokens. The answer supervision loss $\mathcal{L}_{\text{ans}}$Lans​ is computed in a separate final forward pass that reuses the prefix cache and inserts the post-loop latent hidden states $\boldsymbol{h}^{(R)}$h(R). The model predicts the answer suffix $A$A autoregressively:

$$\mathcal{L}_{\text{ans}} = \frac{1}{|A|} \sum_{m=0}^{|A|-1} \mathrm{CE}(f_{\text{head}}(\boldsymbol{z}_m), A_{m+1}),$$Lans​=∣A∣1​m=0∑∣A∣−1​CE(fhead​(zm​),Am+1​),

where $\boldsymbol{z}_m$zm​ represents the hidden states at the answer-suffix positions. The full objective is $\mathcal{L} = \mathcal{L}_{\text{ans}} + \lambda_{\text{step}} \mathcal{L}_{\text{step}}$L=Lans​+λstep​Lstep​. This parallel chain likelihood approach ensures coverage of gold tokens at each position, while the answer loss provides global selection pressure for jointly computed hidden states.

To explore autoregressive chain likelihood, the authors also introduce LOTUS-aux, which routes CoT supervision through an auxiliary decoder $g_{\phi}$gϕ​ instead of the base LM head.

The auxiliary decoder variant is illustrated in the following diagram:

In LOTUS-aux, at each loop iteration $t$t, the auxiliary decoder reads the latent block $\boldsymbol{h}_{t}^{(t)}$ht(t)​ as a prefix and predicts the gold CoT step $T_t$Tt​ under teacher forcing. The output hidden states $\widetilde{\boldsymbol{z}}^{(t)}$z(t) are generated by:

$$\widetilde{\boldsymbol{z}}^{(t)} = g_{\phi}\left([\boldsymbol{h}_{t,c}^{(t)}, T_{t,1}, \dots, T_{t,|T_t|-1}] \mid \boldsymbol{h}_{t,1}^{(t)}, \dots, \boldsymbol{h}_{t,c-1}^{(t)}\right).$$z(t)=gϕ​([ht,c(t)​,Tt,1​,…,Tt,∣Tt​∣−1​]∣ht,1(t)​,…,ht,c−1(t)​).

The step loss for this variant is computed as:

$$\mathcal{L}_{\text{step}}^{\text{aux}} = \frac{1}{N_{\text{step}}} \sum_{t=1}^{K} \sum_{m=0}^{|T_t|-1} \mathrm{CE}\left(g_{\text{head}}(\tilde{\boldsymbol{z}}_m^{(t)}), T_{t,m+1}\right).$$Lstepaux​=Nstep​1​t=1∑K​m=0∑∣Tt​∣−1​CE(ghead​(z~m(t)​),Tt,m+1​).

The auxiliary decoder is used exclusively during training to model the chain autoregressively, while inference for both LOTUS and LOTUS-aux relies solely on the post-loop latents to decode the final answer, preserving the efficiency gains of parallel latent reasoning.

Experiment

The evaluation investigates LOTUS, a latent reasoning model, across accuracy, efficiency, design choices, and interpretability using GPT-2 and LLAMA backbones on GSM8K and out-of-domain math benchmarks. LOTUS matches explicit chain-of-thought accuracy at scale while prior latent methods fall behind, and it achieves a 2.5x to 6.9x speedup in the thought phase by compressing reasoning into parallel latent iterations. Ablations show that both the looped backbone and latent supervision are beneficial, with accuracy saturating at moderate latent widths and loop depths, and that inference-time adjustments to width or depth are robust without retraining. Analysis of the latent representations confirms they carry readable chain-of-thought signal, place nontrivial mass on unseen valid reasoning chains, and rely on complementary roles of step and answer supervision to form coherent reasoning states.

LOTUS nearly matches explicit chain-of-thought in-domain accuracy at all tested scales and leads the out-of-domain average at the largest scale, while prior latent methods like CODI+SIM-CoT and KaVa see their accuracy gaps widen with model size. It achieves this with a 2.5× faster thought phase than explicit CoT, compressing reasoning into parallel latent iterations, and the speedup grows to 6.9× on natural-language traces. LOTUS stays within roughly 1.5 points of explicit CoT on GSM8K at every scale and surpasses explicit CoT on the out-of-domain average at the 3B scale, whereas CODI+SIM-CoT falls 9.2 points behind at 3B. LOTUS is 2.5× faster than explicit CoT in the thought phase, with speedup increasing to 6.9× on natural-language CoT traces while maintaining accuracy on par with explicit CoT.

LOTUS achieves substantially lower thought-phase latency than explicit chain-of-thought by compressing reasoning into parallel latent iterations, while maintaining competitive accuracy. The thought phase dominates the overall latency difference, with LOTUS being 2.5 times faster than explicit CoT and 1.2 times faster than SIM-CoT in this phase. CODI is faster in the thought phase due to decoding only a single latent per step, but this comes at the cost of reduced accuracy. LOTUS thought-phase latency is 133.0 ms, compared to 338.8 ms for explicit CoT and 162.7 ms for SIM-CoT. CODI achieves the lowest thought latency at 88.2 ms but sacrifices accuracy relative to LOTUS and explicit CoT. Total inference latency per example is 181.2 ms for LOTUS versus 384.2 ms for explicit CoT, a speedup of over 2 times. Query prefill and answer phases contribute relatively small and similar latency across methods, making the thought phase the primary differentiator.

LOTUS matches explicit chain-of-thought accuracy on a natural-language GSM8K stress test with a 3B-parameter model, while reducing thought-phase latency by 6.9 times. It far exceeds other latent reasoning methods, outscoring the next best baseline by over 8 percentage points. LOTUS cuts thought-phase latency from 963.6 ms to 140.8 ms (a 6.9× speedup) while staying within 0.3 points of explicit chain-of-thought accuracy. It outperforms all other latent reasoning baselines, with the best prior method, KaVa, trailing by 8.1 percentage points (60.0% vs. 68.13%).

Post-loop supervision of the final latent state through the main language model head achieves the highest test accuracy at 70.0%, outperforming per-iteration supervision and auxiliary decoder variants. Any form of LOTUS latent supervision yields a large gain over the answer-loss-only baseline (63.3%) and the CODI-only model (64.4%). Post-loop supervision of h(R) via the main LM head achieves 70.0% accuracy, higher than per-iteration supervision's 68.2%. LOTUS-aux with per-iteration supervision through an auxiliary decoder reaches 69.9%, nearly matching the best configuration. Both answer-loss-only (63.3%) and CODI-only (64.4%) baselines fall far short of any LOTUS variant.

Increasing the number of looped refinement iterations R consistently improves accuracy on GSM8K, with steep gains from R=2 to R=5 and near saturation at R=6. The model benefits from deeper sequential refinement of the latent prefix before readout. Accuracy rises sharply from 14.6% at R=2 to 68.1% at R=5, then plateaus near 70.0% at R=6. Gains continue through R=5 before nearly saturating, indicating that additional refinement iterations beyond five provide diminishing returns.

LOTUS, a latent reasoning method that compresses chain-of-thought into parallel iterative refinement, is evaluated on GSM8K and out-of-domain tasks against explicit chain-of-thought and prior latent approaches. It matches explicit chain-of-thought accuracy while achieving up to 6.9× faster thought-phase latency, and it substantially outperforms earlier latent methods whose accuracy gaps widen with scale. Ablations show that post-loop supervision of the final latent state and increasing refinement iterations (with diminishing returns beyond five) are critical to its performance.