Bridging the Gap Between Latent and Explicit Reasoning with Looped Transformers — part 3/6

Cleaned-text code points [31384, 46850). Bibliography omitted here; source.md includes it.

At 3B, LOTUS-aux performs comparably to LOTUS and well above the CODI + SIM-CoT baseline, which uses an auxiliary decoder of the same size as the base LM. The gains therefore come from the looped padded backbone with parallel CoT supervision, regardless of the routing. This parity holds only at 3B: at smaller scales LOTUS-aux degrades, whereas direct LM-head supervision (LOTUS) stays robust across scale.

##### LOTUS reasons 2.5 × 2.5\times faster than explicit CoT.

The thought phase dominates the difference between methods ( Table 2 ): CoT decodes its chain autoregressively, while LOTUS compresses the same “thinking” into R {R} parallel latent iterations, making its thought 2.5 × \mathbf{2.5\times} faster than CoT and 1.2 × \mathbf{1.2\times} faster than SIM-CoT 4 4 4 We measure latency on a single NVIDIA H100 NVL at batch size 1 1 with greedy decoding, split into query-prefill, thought, and answer phases. SIM-CoT and CODI both replace each explicit CoT step with a single latent token decoded sequentially. SIM-CoT additionally adds LoRA adapters (rank 128 128 ) on top of CODI, which inflates every phase. ( 2.1 × \mathbf{2.1\times} faster than CoT in total 5 5 5 LOTUS’s prefill and answer phases are just the standard LM operations ( Figure 2 ), so they match explicit CoT in Table 2 ( 16.7 16.7 vs. 15.9 15.9 ms and 31.5 31.5 vs. 29.5 29.5 ms); the looped latent thought is the only phase that differs. ). CODI is faster than LOTUS in the thought phase because it decodes only a single latent per step, rather than LOTUS’s K ​ c = 150 {K}{c}\!=\!150 parallel positions, but comes at the cost of accuracy. LOTUS is nonetheless only modestly slower ( 133 133 vs. 88 88 ms) because it consumes all positions in parallel. The inference-time width sweep in Table 8 further separates sequential depth from parallel width: with R {R} fixed, increasing the latent prefix from 6 6 to 300 300 positions changes thought latency by only 30 30 ms.

##### Speedup grows to 6.9 × 6.9\times on natural-language CoT traces.

The main GSM8K-Aug setting uses compact math-expression steps, while natural-language CoT requires explicit models to decode much longer reasoning in natural language. In the 3B natural-language CoT stress test ( Table 3 ), LOTUS is on par with explicit CoT in accuracy ( 68.13 68.13 vs. 68.41 68.41 ) while cutting thought-phase latency from 963.6 963.6 ms to 140.8 140.8 ms with a 6.9 × 6.9\times speedup. On this natural-language setting, LOTUS ( 68.13 % 68.13\% GSM8K) far exceeds the latent baselines PCCoT ( 47.6 % 47.6\% ), CODI ( 55.9 % 55.9\% ), and KaVa ( 60.0 % 60.0\% ) reported by Kuzina et al. (2026) ( Table 14 ), staying within variance of explicit CoT ( Table 3 ).

### 4.3 Ablation Studies

We ablate each component of LOTUS to isolate its contribution. Unless noted, all ablations use Llama-3.2-3B-Instruct , evaluate on GSM8K test set ( Cobbe et al., 2021 ) .

#### 4.3.1 Latent Supervision Design

LOTUS supervises the latent blocks via a direct readout ( ℒ step {\mathcal{L}_{\mathrm{step}}} ) of the post-loop latents 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through the main LM head f head {f_{\mathrm{head}}} . We compare it against five alternatives, holding the rest fixed ( K = R = 6 {K}{=}{R}{=}6 , c = 25 {c}{=}25 , also keeping ℒ ans {\mathcal{L}_{\mathrm{ans}}} ): (i) no latent supervision (only ℒ ans {\mathcal{L}_{\mathrm{ans}}} ); (ii) CODI ( Shen et al., 2025 ) , distilling a single pre-answer latent onto the teacher CoT model’s hidden state; (iii) LOTUS (per-iter), reading 𝒉 ( t ) {{{\bm{h}}}^{(t)}} through f head {f_{\mathrm{head}}} at each iteration t t rather than the post-loop 𝒉 ( R ) {{{\bm{h}}}^{({R})}} ; (iv) LOTUS-aux, reading 𝒉 ( t ) {{{\bm{h}}}^{(t)}} through a separate auxiliary decoder g ϕ {g_{\phi}} per iteration (default, Section 3.4 ); and (v) LOTUS-aux (post-loop), reading 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through g ϕ {g_{\phi}} instead.

##### Both the looped backbone and latent supervision are beneficial.

Without any latent supervision, the looped backbone alone ( 63.3 % 63.3\% ) already exceeds CODI + SIM-CoT ( 62.3 % 62.3\% , Table 1 ), which uses a single latent per step; the looped padded backbone is thus already beneficial on its own. Every supervised variant exceeds both no latent supervision and CODI-only, confirming that the gold CoT supervision contributes to the gain. LOTUS (post-loop direct, 70.0 % 70.0\% ) matches the best LOTUS-aux schedule (per-iter, 69.9 % 69.9\% ). The direct (LOTUS) and auxiliary decoder (LOTUS-aux) routings thus perform comparably at their best schedule.

##### The best supervision schedule differs between LOTUS and LOTUS-aux.

LOTUS performs best with post-loop readout ( 70.0 % 70.0\% vs. 68.2 % 68.2\% per-iter), which lets the early latent blocks refine fully until the end of the loop. LOTUS-aux instead works best with per-iteration readout ( 69.9 % 69.9\% vs. 68.4 % 68.4\% post-loop), which shortens the gradient path—possibly helpful because its extra decoder already lengthens it.

#### 4.3.2 Looped Architecture Design

##### Latent token budget per block c c .

We train separate models at c ∈ { 1 , 5 , 10 , 25 , 30 } c\in\{1,5,10,25,30\} to measure how the per-block budget affects accuracy, holding K = 6 K{=}6 fixed ( Table 7 ). Accuracy rises sharply from c = 1 c{=}1 ( 49.7 % 49.7\% ) to c = 5 c{=}5 ( 67.5 % 67.5\% ), then climbs only marginally before plateauing with c = 25 c{=}25 and c = 30 c{=}30 tied at 70.0 % 70.0\% . A single token per CoT step is too narrow for the direct-readout supervision, but moderate widths suffice.

##### Inference-time c c .

We sweep the inference-time budget c ∈ { 1 , 5 , 10 , 25 , 30 , 50 } c\in\{1,5,10,25,30,50\} on a single checkpoint trained at c = 25 c{=}25 ( Table 8 ). Reducing c c below the trained value hurts accuracy ( − 19 -19 points at c = 1 c{=}1 , − 11 -11 points at c = 5 c{=}5 , − 6 -6 points at c = 10 c{=}10 ), while exceeding it slightly helps before plateauing ( 70.5 % 70.5\% at both c = 30 c{=}30 and c = 50 c{=}50 ). Thought-phase latency increases only ∼ 30 {\sim}30 ms ( 111 111 to 141 141 ms) as c c varies by 50 × 50\times , since the latents are processed in parallel within the fixed R = 6 {R}{=}6 sequential steps. This contrasts with autoregressive CoT, whose latency grows linearly in the number of thought tokens.

##### Looped iterations R {R} .

We train separate models for R ∈ { 2 , … , 6 } {R}\in\{2,\dots,6\} , fixing K = 6 K{=}6 and c = 25 c{=}25 so that the latent prefix has K ​ c = 150 Kc=150 positions and only the depth of sequential refinement changes ( Table 6 ). Larger R {R} gives the model more refinement iterations before the readout. As a result, accuracy rises steeply with R {R} (from 14.6 % 14.6\% at R = 2 {R}{=}2 to 70.0 % 70.0\% at R = 6 {R}{=}6 ), with gains continuing through R = 5 {R}{=}5 ( 68.1 % 68.1\% ) before nearly saturating at R = 6 {R}{=}6 .

##### Inference-time R {R} .

Reusing the model trained at R = 6 {R}{=}6 , we vary R ∈ { 1 , … , 7 } {R}\in\{1,\dots,7\} at inference with the latent prefix fixed at 150 150 positions ( Table 6 ). Accuracy climbs with R {R} ( 22.7 % 22.7\% at R = 1 {R}{=}1 to 70.0 % 70.0\% at R = 6 {R}{=}6 ), peaks at the trained R = 6 {R}{=}6 , and dips slightly at R = 7 {R}{=}7 ( 69.3 % 69.3\% ). Running fewer iterations leaves the model less opportunity to refine, while running more than at training time brings no further gain.

## 5 Latent Representation Analysis

We analyze the learned latent representations along three axes. Together, these analyses show that LOTUS’s accuracy gains are accompanied by structured latent computation: the post-loop latents carry readable CoT token signal, assign nontrivial mass to unseen valid alternatives, and rely on both gold CoT supervision and answer supervision to form coherent reasoning states. Section 5.1 measures the likelihood assigned to the gold CoT under different readouts, Section 5.2 tests whether the latents place mass on unseen but valid reasoning chains, and Section 5.3 ablates the roles of gold CoT and answer supervision in shaping the representation. Extended block-to-step decode examples are deferred to Section F.1 .

### 5.1 Likelihood of the Gold Chain under Different Readout Options

Recall that 𝒉 ( R ) {{{\bm{h}}}^{({R})}} denotes the post-loop latents and 𝒛 ~ ( t ) {\widetilde{{{\bm{z}}}}^{(t)}} the aux decoder’s output latents for t ∈ { 1 , … , R } t\in\{1,\dots,{R}\} . We evaluate four readouts in Table 9 : (i) LOTUS, reading 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through the base LM head; (ii) LOTUS-aux, reading 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through g head {g_{\mathrm{head}}} ; (iii) LOTUS-aux, reading 𝒛 ~ ( t ) {\widetilde{{{\bm{z}}}}^{(t)}} through g head {g_{\mathrm{head}}} under teacher forcing (TF) that matches training (each block’s c {c} latents paired with that block’s gold step tokens); and (iv) LOTUS-aux, reading 𝒛 ~ ( t ) {\widetilde{{{\bm{z}}}}^{(t)}} through g head {g_{\mathrm{head}}} but under free-running (FR), 6 6 6 LOTUS-aux (FR) gives a fairer reference to LOTUS than LOTUS-aux (TF): at inference the gold CoT is unavailable, so LOTUS-aux must feed its decoder its own generated tokens, exactly the FR setting. The TF row instead feeds the gold prefix the model would not have at test time. feeding the aux decoder’s own generated token in place of the gold tokens. 7 7 7 Here we report greedy decoding results. Replacing them with temperature- T T sampling ( T ∈ { 0.5 , 0.7 , 1.0 , 1.5 } T\in\{0.5,0.7,1.0,1.5\} ) moves FR NLL by ≤ 0.05 \leq 0.05 and top- 1 1 by ≤ 1 \leq 1 pp.

##### Metrics.

We report the negative log-likelihood (NLL) of the gold CoT token sequence T {T} at the latent positions. 8 8 8 Note that the 𝒉 ( R ) {{{\bm{h}}}^{({R})}} readout scores the gold CoT under the PCL ( Section 3.3.2 ) factorization and 𝒛 ~ ( t ) {\widetilde{{{\bm{z}}}}^{(t)}} under an autoregressive one, so NLL is only directly comparable under the same factorization and only for reference across factorizations. Mathematically, PCL NLL has a higher lower bound than the autoregressive NLL (details in Appendix E ). We also report average P ⁡ ( gold ∈ top-1 ) {P(\text{gold}\in\text{top-1})} and P ⁡ ( gold ∈ top-5 ) {P(\text{gold}\in\text{top-5})} against the full 128 ​ K 128\text{K} -vocabulary softmax. All metrics are averaged over the 1,319 1{,}319 GSM8K test questions.

##### LOTUS’s post-loop latent 𝒉 ( R ) {{{\bm{h}}}^{({R})}} contains the gold CoT via a direct readout.

Reading LOTUS’s post-loop latent 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through f head {f_{\mathrm{head}}} assigns the gold CoT a low 3.07 3.07 NLL and 70.9 % 70.9\% top- 1 1 accuracy ( Table 9 ), edging out LOTUS-aux (FR) at 64.5 % 64.5\% . The teacher-forced LOTUS-aux reader does better still ( 1.56 1.56 NLL, 84.3 % 84.3\% top- 1 1 ), but with the caveat that it uses the gold prefix as input, so we read it as an upper bound rather than a direct comparison. Still, the results confirm that the latents carry enough information to reconstruct the gold CoT chain.

##### LOTUS-aux has CoT signal in both 𝒉 ( R ) {{{\bm{h}}}^{({R})}} and 𝒛 ~ ( t ) {\widetilde{{{\bm{z}}}}^{(t)}} .

Surprisingly, even though 𝒉 ( R ) {{{\bm{h}}}^{({R})}} is never directly supervised to predict the gold CoT through g head {g_{\mathrm{head}}} in LOTUS-aux, we can still read 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through g head {g_{\mathrm{head}}} with a 25.8 % 25.8\% top- 5 5 probability. Greedy decoding also surfaces some meaningful gold tokens, indicating that the CoT content is carried by the latents themselves. In fact, despite being supervised through different heads, LOTUS and LOTUS-aux have a similar effect of making the post-loop latents 𝒉 ( R ) {{{\bm{h}}}^{({R})}} put mass on the CoT steps, with the caveat that LOTUS-aux readouts are less formatted (which is expected since it is not supervised directly). We provide side-by-side 𝒉 ( R ) {{{\bm{h}}}^{({R})}} readout examples for both models in Appendix F .

### 5.2 Mass on Unseen Valid Chains

Beyond the gold chain in the test set, we test whether LOTUS and LOTUS-aux place mass on unseen valid reasoning chains rather than only memorizing the seen trace in the training set. We use the same four readout configurations as in Section 5.1 .

##### Setup.

Given each question, we compare its ground-truth reasoning chain against an unseen-but-valid alternative. Let G G and U U denote the sets of intermediate numbers appearing uniquely in the ground-truth and the unseen reasoning chain, respectively. A random control set serves as a baseline. We report each set’s average per-token negative log-likelihood (NLL) at the position where it surfaces most strongly (lower means more confident). We also report P ⁡ ( U ∈ top- ​ k ) P(U\in\text{top-}k) , the fraction of U U numbers that surface within the top- k k readout at any latent position. See Section F.2 for the full construction and metric details.

##### Findings.

Across all four readouts the ordering G ≪ U ≪ G\!\ll\!U\!\ll\! Random holds ( Table 10 , lower NLL = = surfaces more): each reader assigns much lower NLL to a ground-truth-only number than to an unseen-but-valid one, which in turn sits far below the random control. Crucially, the U U signal is concentrated rather than diffuse: P ⁡ ( U ∈ top- ​ k ) P(U\in\text{top-}k) columns confirm the unseen-valid intermediates genuinely surface in the top- k k (e.g., for LOTUS through the base LM head, 15.3 % 15.3\% are top- 1 1 and 64.0 % 64.0\% fall within top- 5 5 , despite never being trained on). Section F.3 and Figure 4 give a path-level example and visual grid: two valid chains reach the same answer ( 108 108 ) through disjoint intermediates, and the unseen-path numbers ( 24 24 , 18 18 ) surface in the post-loop readout despite never appearing in the trained chain, question, or answer.

### 5.3 Disentangling the Two Losses in LOTUS

We now ask what each of LOTUS’s two supervision losses contributes to the latent representation observed above. The theory in Section 3.3.2 ascribes complementary roles to ℒ step {\mathcal{L}_{\mathrm{step}}} and ℒ ans {\mathcal{L}_{\mathrm{ans}}} . We compare three models that differ in supervision: LOTUS (both losses), ℒ step {\mathcal{L}_{\mathrm{step}}} only, and ℒ ans {\mathcal{L}_{\mathrm{ans}}} only.

##### Gold-chain NLL.

Reading the 𝒉 ( R ) {{{\bm{h}}}^{({R})}} through the base LM head f head {f_{\mathrm{head}}} , Table 11 reports NLL, P ⁡ ( gold ∈ top-1 ) P(\text{gold}\in\text{top-1}) , and P ⁡ ( gold ∈ top-5 ) P(\text{gold}\in\text{top-5}) on the GSM8K test set, using the same protocol as Table 9 .

##### ℒ step {\mathcal{L}_{\mathrm{step}}} and ℒ ans {\mathcal{L}_{\mathrm{ans}}} play complementary roles.

LOTUS beats both ablations ( Table 11 ). ℒ step {\mathcal{L}_{\mathrm{step}}} -only, lacking the cross-block coupling and the answer grounding that ℒ ans {\mathcal{L}_{\mathrm{ans}}} supplies, is worst on every column. ℒ ans {\mathcal{L}_{\mathrm{ans}}} -only, without directly tying 𝒉 ( R ) {{{\bm{h}}}^{({R})}} to the gold CoT, lands nearer the gold tokens (NLL 5.97 5.97 , top- 5 5 26.5 % 26.5\% ) yet has a top- 1 1 rate similar to ℒ step {\mathcal{L}_{\mathrm{step}}} -only and stays far short of LOTUS (NLL 3.07 3.07 , top- 5 5 85.8 % 85.8\% ). The answer loss alone settles 𝒉 ( R ) {{{\bm{h}}}^{({R})}} into an answer-consistent neighborhood without aligning to the gold chain, matching the complementary roles posited in Section 3.3.2 .

## 6 Conclusion