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

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

After R {R} iterations of looping ( Section 3.2 ), we supervise each latent position ( i , j ) ({i},j) in the grid with the corresponding CoT step token. Let T i = ( T i , 1 , … , T i , c ) {T_{{i}}}=({T_{{i},1}},\dots,{T_{{i},{c}}}) be the tokenized CoT step i {i} padded or truncated to c {c} tokens. We apply a single batched cross-entropy that directly aligns each block to its target step: ℒ step = 1 N step ​ ∑ i = 1 K ∑ j = 1 c CE ⁡ ( f head ​ ( 𝒉 i , j ( R ) ) , T i , j ) , {\mathcal{L}_{\mathrm{step}}}\;=\;\frac{1}{{N_{\mathrm{step}}}}\sum_{{i}=1}^{{K}}\sum_{j=1}^{{c}}{\mathrm{CE}}\!\left({f_{\mathrm{head}}}({{{\bm{h}}}^{({R})}_{{i},j}}),\,{T_{{i},j}}\right), (3) where f head {f_{\mathrm{head}}} is the LM head and N step = ∑ i | T i | {N_{\mathrm{step}}}=\sum_{{i}}|{T_{{i}}}| is the total number of supervised (non-padding) CoT tokens. 1 1 1 Padding tokens used to pad CoT steps to length c {c} are ignored by the cross-entropy (and excluded from N step {N_{\mathrm{step}}} ).

##### Answer supervision loss.

ℒ ans {\mathcal{L}_{\mathrm{ans}}} is computed in a separate final forward of f 𝜽 {f_{{{{\bm{\theta}}}}}} that reuses the prefix cache 𝒞 pre {\mathcal{C}}_{\text{pre}} and inserts the post-loop latent hidden states 𝒉 ( R ) {{{\bm{h}}}^{({R})}} at the latent rows: 𝒛 = f 𝜽 ​ ( [ ⟨ EoT ⟩ , A ] | [ 𝒞 pre , 𝒉 ( R ) ] ) , {{{\bm{z}}}}\;=\;{f_{{{{\bm{\theta}}}}}}\!\left([{\langle\texttt{EoT}\rangle},\,{{A}}]\;\middle|\;[{\mathcal{C}}_{\text{pre}},\,{{{\bm{h}}}^{({R})}}]\right), (4) where 𝒛 m ∈ ℝ d {{{\bm{z}}}}_{m}\in{{\mathbb{R}}}^{{d}} collects the resulting hidden state at each answer-suffix position m ∈ { 0 , 1 , … , | A | − 1 } m\in\{0,1,\dots,|{{A}}|-1\} ( m = 0 m=0 being the trailing ⟨ EoT ⟩ {\langle\texttt{EoT}\rangle} position, which predicts A 1 A_{1} ). We denote hidden states here by 𝒛 {{{\bm{z}}}} , distinct from the looped forward latents 𝒉 ( t ) {{{\bm{h}}}^{({t})}} , to mark that they come from different sources. Applying the LM head f head {f_{\mathrm{head}}} , standard next-token cross-entropy on the answer suffix gives: ℒ ans = 1 | A | ​ ∑ m = 0 | A | − 1 CE ⁡ ( f head ​ ( 𝒛 m ) , A m + 1 ) . {\mathcal{L}_{\mathrm{ans}}}\;=\;\frac{1}{|{{A}}|}\sum_{m=0}^{|{{A}}|-1}{\mathrm{CE}}\!\left({f_{\mathrm{head}}}({{{\bm{z}}}}_{m}),\,A_{m+1}\right). (5) Given a step supervision weight λ step {\lambda_{\mathrm{step}}} , the full objective is ℒ = ℒ ans + λ step ​ ℒ step . {\mathcal{L}}\;=\;{\mathcal{L}_{\mathrm{ans}}}+{\lambda_{\mathrm{step}}}\,{\mathcal{L}_{\mathrm{step}}}. (6)

##### Key properties of the objective.

The supervision in Equation 6 has three properties: (i) it is direct : the per-block target T i {T_{{i}}} is scored through the same LM head used to produce the answer; (ii) it is parallel : all K {K} blocks are supervised simultaneously; and (iii) it is post-loop : the readout reads 𝒉 ( R ) {{\bm{h}}}^{({R})} after the final iteration, rather than at every iteration. We ablate the design choice in Section 4.3 , and compare against prior latent reasoning methods in Appendix C .

##### Inference.

At inference, LOTUS runs the looped forward and then decodes the answer. We compute the KV cache of the question Q {{Q}} once, iterate the loop R {R} times to obtain the post-loop latents 𝒉 ( R ) {{{\bm{h}}}^{({R})}} , and then decode the answer A {{A}} autoregressively through the base LM head while conditioning on those latents (same as Figure 2 only without computing the losses). The reasoning is carried entirely by the parallel latent blocks, so the only sequential decoding is over the short answer suffix, which is the source of the latency gains reported in Section 4.2 .

##### No latent decoding for answer generation.

Note that the latent reasoning steps are never read out during answer generation. They are projected to the token space through the LM head only to align them with CoT reasoning during training with ℒ step {\mathcal{L}_{\mathrm{step}}} . The answer is always generated from the latent hidden states, both in the final-forward ℒ ans {\mathcal{L}_{\mathrm{ans}}} pass ( Equation 4 ) and at inference.

#### 3.3.2 Roles of ℒ step {\mathcal{L}_{\mathrm{step}}} and ℒ ans {\mathcal{L}_{\mathrm{ans}}}

A natural worry about the LOTUS objective ( Equation 6 ) is that it supervises each latent position independently : Equation 3 scores every gold CoT token in parallel, with no autoregressive factorization of the chain. Why should independent per-position targets train a model that produces a globally coherent answer? We give a lens that resolves this, distinguishing the role of the two losses and explaining why both are needed.

##### Parallel chain likelihood.

The step loss ℒ step {\mathcal{L}_{\mathrm{step}}} supervises a parallel chain likelihood (PCL) over the readout positions. Maximizing the per-position probability p 𝜽 ​ ( T i , j ∣ Q ) p_{{{{\bm{\theta}}}}}({T_{{i},j}}\mid{Q}) of each gold token at its latent position induces the chain likelihood p 𝜽 PCL ​ ( T ∣ Q ) = ∏ i = 1 K ∏ j = 1 c p 𝜽 ​ ( T i , j ∣ Q ) , p_{{{{\bm{\theta}}}}}^{\mathrm{PCL}}({T}\mid{Q})=\prod_{{i}=1}^{{K}}\prod_{j=1}^{{c}}p_{{{{\bm{\theta}}}}}({T_{{i},j}}\mid{Q}), (7) rather than the autoregressive conditionals p 𝜽 ​ ( T i , j ∣ Q , T < i , T i , < j ) p_{{{{\bm{\theta}}}}}({T_{{i},j}}\mid{Q},{T}_{<{i}},{T_{{i},<j}}) . The factorization treats the readout as conditionally independent across positions, but the latent states themselves are not independent: the looped Transformer computes them jointly over the padded workspace, so the dependence the factorization drops is carried by the shared latent computation rather than by the loss.

##### Complementary roles.

This view makes the division of labor between the two losses precise. First, the support inclusion supp ⁡ ( q ⁡ ( T ∣ Q ) ) ⊆ ∏ i = 1 K ∏ j = 1 c supp ⁡ ( q ⁡ ( T i , j ∣ Q ) ) \mathrm{supp}\!\left({q}({T}\mid{Q})\right)\;\subseteq\;\prod_{{i}=1}^{{K}}\prod_{j=1}^{{c}}\mathrm{supp}\!\left({q}({T_{{i},j}}\mid{Q})\right) (8) shows that every gold chain lies inside the Cartesian product of per-position gold token supports. 2 2 2 Minimizing ℒ step {\mathcal{L}_{\mathrm{step}}} , the cross-entropy to these marginals, pulls the model’s PCL toward that product, i.e. minimizes KL ( ∏ i , j q ( T i , j ∣ Q ) ∥ p 𝜽 PCL ( T ∣ Q ) ) \mathrm{KL}\big(\prod_{{i},j}{q}({T_{{i},j}}\mid{Q})\,\big\|\,p_{{{{\bm{\theta}}}}}^{\mathrm{PCL}}({T}\mid{Q})\big) . So ℒ step {\mathcal{L}_{\mathrm{step}}} plays a support-coverage role: it makes each position place mass on the right gold tokens, without needing to reproduce the joint distribution—which LOTUS never samples, since the answer is decoded directly from the jointly computed latent states. Second, ℒ ans {\mathcal{L}_{\mathrm{ans}}} supplies the global selection pressure that PCL alone lacks: because the answer is trained while conditioning on the entire latent configuration, gradients favor jointly computed hidden states that can actually support the correct answer. The two losses are thus complementary by construction—coverage from ℒ step {\mathcal{L}_{\mathrm{step}}} , joint selection from ℒ ans {\mathcal{L}_{\mathrm{ans}}} —which is exactly the behavior we verify empirically in Section 5.3 , where dropping either loss degrades the gold chain likelihood.

To further probe how LOTUS compares to modeling the autoregressive chain likelihood, we introduce an autoregressive decoder variant in Section 3.4 and compare the results in Section 4.2 .

### 3.4 LOTUS-aux: An Auxiliary Decoder Variant

An alternative route for CoT supervision explored in existing work is autoregressive chain likelihood: an auxiliary autoregressive decoder conditioned on the latent blocks scores the CoT tokens under teacher forcing. This creates a trade-off: in exchange for modeling the chain autoregressively, CoT supervision no longer flows directly through the answer-generating base LM head, and teacher forcing introduces an extra train/inference mismatch.

We instantiate this routing on our looped padded backbone as LOTUS-aux. LOTUS-aux reuses the same looped forward and final forward as LOTUS ( Figure 2 ). The only difference is the supervision in the looped forward, which is routed through an auxiliary decoder rather than read directly through the base LM head. The auxiliary decoder mirrors SIM-CoT ( Wei et al., 2025 ) , but instead of generating latents autoregressively and supervising a single latent per CoT step, LOTUS-aux produces all K {K} latent blocks in parallel through the looped backbone and supervises a full c {c} -token block per step. We describe this routing below.

##### Latent supervision via the auxiliary decoder.

As shown in Figure 3 , the auxiliary decoder g ϕ {g_{\phi}} is an extra decoder model that, at each loop iteration t t , reads latent block 𝒉 ( t ) t {{{\bm{h}}}^{(t)}}_{t} (we select block i = t i=t at iteration t t , the c {c} post-loop latents of the t t -th block) as a c {c} -token prefix and predicts the gold CoT step T t = ( T t , 1 , … , T t , | T t | ) {T_{t}}=(T_{t,1},\dots,T_{t,\lvert{T_{t}}\rvert}) under teacher forcing, where | T t | \lvert{T_{t}}\rvert is the step’s token length. Let 𝒛 ~ ( t ) = ( 𝒛 ~ 0 ( t ) , … , 𝒛 ~ | T t | − 1 ( t ) ) {\widetilde{{{\bm{z}}}}^{(t)}}=({\widetilde{{{\bm{z}}}}^{(t)}_{0}},\dots,{\widetilde{{{\bm{z}}}}^{(t)}_{\lvert{T_{t}}\rvert-1}}) denote the output hidden states at each input CoT position m ∈ { 0 , 1 , … , | T t | − 1 } m\in\{0,1,\dots,\lvert{T_{t}}\rvert-1\} ( m = 0 m{=}0 being the last latent position 𝒉 t , c ( t ) {{{\bm{h}}}^{(t)}_{t,{c}}} , which predicts T t , 1 T_{t,1} , and m > 0 m{>}0 the gold token T t , m T_{t,m} ). Mirroring the final-forward pass of Equation 4 , we generate 𝒛 ~ ( t ) {\widetilde{{{\bm{z}}}}^{(t)}} with 𝒛 ~ ( t ) = g ϕ ( [ 𝒉 t , c ( t ) , T t , 1 , … , T t , | T t | − 1 ] | 𝒉 t , 1 ( t ) , … , 𝒉 t , c − 1 ( t ) ) , {\widetilde{{{\bm{z}}}}^{(t)}}\;=\;{g_{\phi}}\!\left([\,{{{\bm{h}}}^{(t)}_{t,{c}}},\;T_{t,1},\dots,T_{t,\lvert{T_{t}}\rvert-1}\,]\;\middle|\;{{{\bm{h}}}^{(t)}_{t,1}},\dots,{{{\bm{h}}}^{(t)}_{t,{c}-1}}\right), (9) where the first c − 1 {c}-1 latent positions 𝒉 t , 1 ( t ) , … , 𝒉 t , c − 1 ( t ) {{{\bm{h}}}^{(t)}_{t,1}},\dots,{{{\bm{h}}}^{(t)}_{t,{c}-1}} are the conditioning prefix. The teacher-forced next-token cross-entropy on the gold CoT gives: ℒ step aux = 1 N step ​ ∑ t = 1 K ∑ m = 0 | T t | − 1 CE ⁡ ( g head ​ ( 𝒛 ~ m ( t ) ) , T t , m + 1 ) , {\mathcal{L}_{\mathrm{step}}^{\mathrm{aux}}}\;=\;\frac{1}{N_{\text{step}}}\sum_{t=1}^{{K}}\sum_{m=0}^{\lvert{T_{t}}\rvert-1}{\mathrm{CE}}\!\left({g_{\mathrm{head}}}({\widetilde{{{\bm{z}}}}^{(t)}_{m}}),\;T_{t,m+1}\right), (10) where N step = ∑ t = 1 K | T t | N_{\text{step}}=\sum_{t=1}^{{K}}\lvert{T_{t}}\rvert is the total number of supervised CoT tokens. The supervision itself remains parallel: under teacher forcing the gold prefix is fed at every position, so all CoT positions are scored in a single forward pass with no sequential generation. In other words, the autoregressive factorization affects only the loss, not the computation.

##### Per-iteration supervision and the full objective.

Equation 10 reads block i = t i=t at iteration t t (per-iteration, rather than at the final iteration R R ). We validate this choice in Section 4.3.1 , where per-iteration readouts outperform reading all blocks at the final iteration R {R} for the auxiliary decoder routing. The answer loss ℒ ans {\mathcal{L}_{\mathrm{ans}}} from Section 3.3 is unchanged, so the full LOTUS-aux objective is ℒ aux = ℒ ans + λ step ​ ℒ step aux {\mathcal{L}^{\mathrm{aux}}}={\mathcal{L}_{\mathrm{ans}}}+{\lambda_{\mathrm{step}}}\,{\mathcal{L}_{\mathrm{step}}^{\mathrm{aux}}} . In our experiments g ϕ {g_{\phi}} is initialized with a copy of the base LM and trained together with the base LM. Training and architectural details are in Appendix D .

##### Inference.

The auxiliary decoder g ϕ {g_{\phi}} is used only at training time. At inference, LOTUS-aux generates the final answer autoregressively from the post-loop latents exactly as LOTUS does ( Figure 2 (b)), so the inference efficiency gain of LOTUS also applies to LOTUS-aux.

## 4 Experiments

Our experiments investigate four questions: how accurate LOTUS is, how efficient it is at inference, how much each design choice contributes, and whether its latent space is interpretable and CoT-aligned. Section 4.1 introduces the experimental setup, Section 4.2 presents accuracy and latency results against explicit CoT and prior latent baselines, Section 4.3 ablates the supervision design, latent block width c {c} , loop depth R {R} , and inference-time robustness to changing c {c} and R {R} without retraining, and Section 5 analyzes the learned latent representations. Full training and evaluation details are in Appendix D .

### 4.1 Setup

##### Datasets and backbones.

Following SIM-CoT ( Wei et al., 2025 ) , we train on GSM8k-Aug ( Deng et al., 2023 ) with 385k training samples and report in-domain accuracy on the GSM8K test set ( Cobbe et al., 2021 ) , together with three out-of-domain benchmarks: GSM-Hard, MultiArith, and SVAMP. We additionally train on a natural-language version of GSM8K-Aug ( Wu et al., 2025 ) as an efficiency stress test, where each reasoning step is a full-sentence rationale rather than a compact math expression. We use three backbones in increasing size: GPT-2 (124M) ( Radford et al., ) , Llama-3.2-1B-Instruct , and Llama-3.2-3B-Instruct ( Grattafiori et al., 2024 ) .

##### Methods.

We report LOTUS ( K = 6 {K}{=}6 blocks of c = 25 {c}{=}25 tokens on the Llama backbones, c = 13 {c}{=}13 on GPT-2 , and R = 6 {R}{=}6 looped iterations) 3 3 3 K {K} is chosen to cover the target maximum CoT step count: GSM8K solutions have at most six steps for 99 % 99\% of examples, so K = 6 {K}{=}6 assigns one block per step and falls back to autoregressive completion otherwise. and LOTUS + CODI , which adds CODI’s single-position distillation loss. The auxiliary decoder variants LOTUS-aux and LOTUS-aux + CODI ( Section 3.4 ) share the same configuration. We compare with Explicit CoT and No-CoT as reference, and with CODI ( Shen et al., 2025 ) and CODI + SIM-CoT ( Wei et al., 2025 ) for each backbone, which share the sequential latent reasoning budget ( R = 6 {R}{=}6 ) with LOTUS and differ only in the number of latent tokens per step. Latent methods that use a different sequential compute budget: Coconut ( Hao et al., 2025 ) , Coconut + SIM-CoT (with 10 sequentially generated latent thought in total), and the parallel-latent methods PCCoT ( Wu et al., 2025 ) and KaVa ( Kuzina et al., 2026 ) (with 3 sequential steps and 24 latent tokens in total) are compared separately in Table 13 and Table 14 ( Appendix C ).

### 4.2 Main Results

##### LOTUS matches explicit CoT at scale, where prior latent methods fall behind.

LOTUS scales well from GPT-2 to Llama-3.2-3B-Instruct ( Table 1 ). Prior latent methods match explicit CoT at small scale but fall behind as the backbone grows: CODI + SIM-CoT is level with explicit CoT at GPT-2 ( 42.6 42.6 vs. 42.7 42.7 ) yet trails it by 9.2 9.2 points at 3B, and KaVa ( Kuzina et al., 2026 ) widens similarly ( 1.9 1.9 to 5.8 5.8 points, Table 13 ). LOTUS instead stays within ∼ 1.5 {\sim}1.5 points of explicit CoT in-domain at each scale, and its gap does not widen. While the smaller-scale results are rather clustered, the decisive comparison is at Llama-3.2-3B-Instruct : LOTUS surpasses explicit CoT on the out-of-domain average and brings the in-domain GSM8K gap to within 1.5 1.5 points. LOTUS + CODI further tightens it to within 1 1 point.

##### The two routings match at 3B, but only direct LOTUS stays robust at smaller scales.