SIM-CoT: Supervised Implicit Chain-of-Thought — part 2/4

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

After constructing the implicit latents z 1 : K z_{1:K} , the model switches to explicit decoding to generate the final answer. Let W o ∈ ℝ | 𝒱 | × d W_{o}\in\mathbb{R}^{|\mathcal{V}|\times d} be the output projection (LM head). With teacher forcing on the partial answer a < t a_{<t} , the generation is h T + K + t \displaystyle h_{T+K+t} = H θ ​ ( U ( K ) ⊕ e ⁡ ( a < t ) ) , \displaystyle=H_{\theta}\!\big(U^{(K)}\oplus e(a_{<t})\big), (2) p θ ( a t ∣ x , z 1 : K , a < t ) \displaystyle p_{\theta}(a_{t}\mid x,z_{1:K},a_{<t}) = softmax ​ ( W o ​ h T + K + t ) a t , \displaystyle=\mathrm{softmax}\!\big(W_{o}\,h_{T+K+t}\big)_{a_{t}}, (3) p θ ( a ∣ x , z 1 : K ) \displaystyle p_{\theta}(a\mid x,z_{1:K}) = ∏ t = 1 L a p θ ( a t ∣ x , z 1 : K , a < t ) . \displaystyle=\prod_{t=1}^{L_{a}}p_{\theta}(a_{t}\mid x,z_{1:K},a_{<t}). (4)

### 3.4 Training-time Decoder and Step-level Supervision

During training, a decoder p ϕ p_{\phi} (architecturally identical to the LLM) takes only the k k -th implicit latent z k z_{k} as conditioning signal and autoregressively generates the k k -th textual step s k = ( y k , 1 , … , y k , L k ) s_{k}=(y_{k,1},\ldots,y_{k,L_{k}}) . This provides step-level supervision that directly grounds z k z_{k} to its corresponding reasoning content: p ϕ ( s 1 : K ∣ z 1 : K ) = ∏ k = 1 K p ϕ ( s k ∣ z k ) = ∏ k = 1 K ∏ t = 1 L k p ϕ ( y k , t ∣ z k , y k , < t ) . p_{\phi}(s_{1:K}\mid z_{1:K})=\prod_{k=1}^{K}p_{\phi}(s_{k}\mid z_{k})=\prod_{k=1}^{K}\prod_{t=1}^{L_{k}}p_{\phi}\!\big(y_{k,t}\mid z_{k},\,y_{k,<t}\big). (5)

Parameterization. For step k k , the decoder is conditioned on the implicit latent z k z_{k} obtained from the LLM. Since z k z_{k} does not correspond to any token in the vocabulary, it is not included in the loss calculation. Instead, z k z_{k} is injected as an additional prefix vector that initializes the decoder’s hidden state for step generation. Concretely, the decoder input sequence is U k dec = [ z k ; e ⁡ ( y k , 1 ) , … , e ⁡ ( y k , L k ) ] , U^{\text{dec}}_{k}=\big[\,z_{k}\,;\,e(y_{k,1}),\ldots,e(y_{k,L_{k}})\,\big], where e ⁡ ( ⋅ ) e(\cdot) denotes the embedding function of the LLM shared between both models. During training with teacher forcing, the decoder predicts each token y k , t y_{k,t} autoregressively: p ϕ ​ ( y k , t ∣ z k , y k , < t ) = softmax ​ ( W dec ​ h k , t dec ) y k , t , p_{\phi}(y_{k,t}\mid z_{k},y_{k,<t})=\mathrm{softmax}\!\big(W^{\text{dec}}\,h^{\text{dec}}_{k,t}\big)_{y_{k,t}}, where h k , t dec h^{\text{dec}}_{k,t} is the decoder hidden state at position t t and W dec W^{\text{dec}} is the LM head of the decoder.

The training loss for step k k is then ℒ step , k = − ∑ t = 1 L k log p ϕ ( y k , t ∣ z k , y k , < t ) , \mathcal{L}_{\text{step},k}=-\sum_{t=1}^{L_{k}}\log p_{\phi}(y_{k,t}\mid z_{k},y_{k,<t}), which supervises only the textual step tokens. The decoder is used exclusively for this supervision during training and is discarded at inference.

### 3.5 Objectives

Training involves two complementary cross-entropy losses: one for supervising the textual steps through the decoder, and one for supervising the final answer through the base LLM.

Step-level supervision. For each implicit latent z k z_{k} , the decoder p ϕ p_{\phi} generates the corresponding reasoning step s k = ( y k , 1 , … , y k , L k ) s_{k}=(y_{k,1},\ldots,y_{k,L_{k}}) . Since z k z_{k} is not a vocabulary token, the loss is computed only over the textual step tokens: ℒ step = − ∑ k = 1 K ∑ t = 1 L k log p ϕ ( y k , t ∣ z k , y k , < t ) . \mathcal{L}_{\text{step}}=-\sum_{k=1}^{K}\sum_{t=1}^{L_{k}}\log p_{\phi}\!\big(y_{k,t}\mid z_{k},y_{k,<t}\big). (6) This loss grounds each latent z k z_{k} to a specific reasoning step, ensuring that the latent sequence carries fine-grained semantics.

Answer supervision. After K K implicit steps, the LLM F θ F_{\theta} switches back to explicit decoding to generate the final answer a = ( a 1 , … , a L a ) a=(a_{1},\ldots,a_{L_{a}}) . We optimize the standard language modeling loss: ℒ ans-lm = − ∑ t = 1 L a log p θ ( a t ∣ x , z 1 : K , a < t ) . \mathcal{L}_{\text{ans-lm}}=-\sum_{t=1}^{L_{a}}\log p_{\theta}\!\big(a_{t}\mid x,z_{1:K},a_{<t}\big). (7)

Total objective. The overall loss is a weighted sum: ℒ = λ step ​ ℒ step + λ lm ​ ℒ ans-lm . \mathcal{L}=\lambda_{\text{step}}\,\mathcal{L}_{\text{step}}+\lambda_{\text{lm}}\,\mathcal{L}_{\text{ans-lm}}. (8) Gradients from ℒ step \mathcal{L}_{\text{step}} propagate through the decoder into the latent representations z 1 : K z_{1:K} and further into the LLM (via Eq. equation 1 ), shaping the hidden states to encode step-level reasoning. Meanwhile, ℒ ans-lm \mathcal{L}_{\text{ans-lm}} trains the base model to produce the final answer directly, so the decoder can be discarded at inference time without affecting efficiency. Implementation details, inference procedures, and diagnostic analyses are provided in Appendix F .

## 4 experiment

### 4.1 Experimental Setup

Training Data. We follow previous works ( Deng et al., 2024 ; Hao et al., 2025 ) to use the GSM8k-Aug dataset Deng et al. (2024) for training implicit CoT models. The GSM8k-Aug expands the original GSM8k training set ( Cobbe et al., 2021 ) to 385k examples by using GPT-4 for data generation. To facilitate implicit CoT training, the GSM8k-Aug removes the reasoning chain of natural language, preserving only a sequence of structured mathematical expressions. Each expression is logically linked to the previous step, as illustrated by the example: <<12*3=36>><<9*2=18>><<17*2=34>><<36+18+34=88>> .

Evaluation Benchmarks. We report results on the GSM8k-Aug test set ( Cobbe et al., 2021 ) , which serves as our in-domain (ID) evaluation benchmark. To further evaluate mathematical reasoning under a distribution shift, we also evaluate models on three out-of-domain (OOD) benchmarks: (1) SVAMP ( Patel et al., 2021 ) , a dataset of grade-school arithmetic word problems that introduces simple variations to assess robustness; (2) GSM-Hard ( Gao et al., 2022 ) , a modified version of the GSM8k test split where numbers are replaced with larger magnitudes to increase problem difficulty; and (3) MultiArith ( Roy & Roth, 2015 ) , a subset of MAWPS ( Koncel-Kedziorski et al., 2016 ) consisting of multi-step arithmetic word problems. Please refer to the Appendix F for more details.

Implementation Details. We follow the training setup of previous works ( Hao et al., 2025 ; Shen et al., 2025b ) , and adopt consistent hyperparameter choices for GPT-2, LLaMA 1B/3B/8B. Detailed configurations, such as learning rates, curriculum strategies, are provided in Appendix E .

### 4.2 Main Results

Baselines. We compare our SIM-CoT against five representative baselines: (1) CoT-SFT : Supervised fine-tuning (SFT) on CoT-annotated data, where the model is trained to generate explicit intermediate reasoning steps followed by the final answer. (2) No-CoT-SFT : Supervised fine-tuning on direct answers only, without producing intermediate steps. (3) iCoT ( Deng et al., 2024 ) : A curriculum learning method based on “Stepwise Internalization,” which injects CoT reasoning patterns into the model’s internal representations, enabling it to produce more accurate direct answers during inference. (4) Coconut ( Hao et al., 2025 ) : A curriculum learning approach that gradually replaces explicit reasoning steps with implicit tokens until the reasoning process becomes fully implicit. This method has shown strong empirical performance and serves as a primary baseline in our experiments. (5) CODI ( Shen et al., 2025b ) : A distillation-based method where explicit CoT acts as the teacher and implicit CoT as the student. By aligning the last hidden states of the full reasoning trajectory, CODI effectively internalizes knowledge and alleviates catastrophic forgetting.

In-Domain Math Benchmark Results. Table 1 (first column) reports GPT-2 results on GSM8k-Aug. SIM-CoT outperforms SFT-CoT and is the first training-based approach where implicit CoT surpasses explicit CoT. With GPT-2 using Coconut as the backbone, it achieves a + 2.1 +2.1 point improvement over SFT-CoT. It also exceeds other training-based implicit reasoning models; for example, on Coconut, it improves by + 8.2 +8.2 points, a relative gain of 22.4 % 22.4\% . Moreover, when applied on top of CODI—the current SOTA implicit reasoning method—SIM-CoT yields an additional + 0.6 +0.6 point improvement.

Table 2 (first column) shows the results when CODI is used as the backbone. In this setting, our method achieves a substantial + 3.4 +3.4 point improvement. Furthermore, we are the first to achieve performance comparable to SFT-CoT on LLaMA-1B, reaching 96 % 96\% of its accuracy. Given that prior studies ( Xu et al., 2025 ; Shen et al., 2025b ) reported that curriculum learning in larger models leads to catastrophic forgetting and that homogeneous knowledge harms model training ( Zhao et al., 2023 ) , we choose CODI as the backbone because its KL-regularized objective constrains the training not to deviate too far from the original model distribution, thereby alleviating catastrophic forgetting.

Out-of-Domain Math Benchmark Results. To evaluate the robustness of our method, we train on GSM8k and evaluate on out-of-domain datasets (GSM-Hard, MultiArith, and SVAMP). From the third column of Table 1 , we observe that SIM-CoT consistently outperforms SFT-CoT, with an average improvement of + 4.3 +4.3 points when using Coconut as the backbone. From the third column of Table 2 , our method further improves upon the current SOTA implicit reasoning method CODI by + 1.0 +1.0 point. Moreover, when scaling model size from GPT-2 to LLaMA-1B, SIM-CoT enlarges the performance gap against iCoT, Coconut, and other baselines.

We attribute the robustness of SIM-CoT to its step-level implicit supervision. Unlike SFT-CoT, which forces the model to mimic deterministic natural language annotations, and unlike CODI, which applies trajectory-level alignment to a coarse-grained reasoning path, our method introduces a moderate form of supervision. This design ensures the plausibility of each reasoning step while preserving the diversity of reasoning trajectories, thereby improving generalization to unseen inputs.

Inference Efficiency. In terms of inference speed, our method maintains the same efficiency as other implicit reasoning approaches on both GPT-2 and LLaMA-1B. On GPT-2, SIM-CoT not only surpasses SFT-CoT on both in-domain and out-of-domain benchmarks, but also achieves a 2.3 × 2.3\times and 2.2 × 2.2\times speedup on Coconut, respectively. On LLaMA-1B, SIM-CoT remains comparable to SFT-CoT in accuracy while delivering 1.9 × 1.9\times and 1.7 × 1.7\times speedups on in-domain and out-of-domain benchmarks, respectively. These results demonstrate the effectiveness of our approach in retaining or even enhancing the performance of explicit CoT while substantially reducing inference cost.

### 4.3 Ablation Studies

Ablation on the Number of Implicit Tokens. We study the effect of varying the number of implicit latents on GPT-2, comparing SIM-CoT with Coconut trained on GSM8k-Aug and evaluated on GSM8k-Aug, GSM-Hard, MultiArith, and SVAMP (Fig. 3 ). Following Coconut, each latent corresponds to two tokens. As shown in Fig. 5 , most problems involve two to six steps with a small proportion of harder cases, so we set the maximum number of implicit latents to 8. For each configuration, we report the best performance, and results show that SIM-CoT provides more stable training and achieves consistent gains over Coconut, indicating that step-level implicit supervision scales effectively with larger latent capacity.

Ablation on Scaling to Larger Backbones. To examine robustness and scalability, we extend experiments to larger LLaMA backbones, including LLaMA 3.2 3B and LLaMA 3.1 8B. Table 3 reports results on GSM8k-Aug (in-domain) and GSM-Hard, MultiArith, and SVAMP (out-of-domain). Overall, SIM-CoT scales effectively to larger backbones, consistently surpassing or matching explicit CoT on out-of-domain tasks while reducing reliance on trajectory-level supervision. For clarity, the detailed results and discussion are provided in Appendix A .

Ablation on Different Decoder Sizes. We replace the decoder of LLaMA 1B with larger ones from the same vocabulary family and evaluate on GSM8k-Aug, GSM-Hard, MultiArith, and SVAMP. A 1B-scale decoder yields consistent gains, whereas larger decoders (3B or 8B) slightly reduce accuracy; detailed discussion is in Appendix B .

Ablation on Soft Thinking. We also study the effect of integrating soft thinking ( Zhang et al., 2025b ; Wu et al., 2025 ) with both Coconut and SIM-CoT. For clarity, the detailed experimental setup, results, and analyses are provided in Appendix C .

Interpretability of Implicit Reasoning. Continuous thoughts in implicit reasoning models are not aligned with vocabulary tokens and cannot be directly decoded into human-readable text. We address this by reusing the training decoder to project and visualize the semantic meaning of each latent step as shown in Fig. 4 . We also analyze latent token distances under different configurations. Detailed descriptions, numerical results, and examples are provided in Appendix I .

## 5 Related Work

A large body of work has studied explicit chain-of-thought (CoT) prompting, including self-consistency ( Wei et al., 2022 ; Wang et al., 2023 ) , least-to-most prompting ( Zhou et al., 2023 ) , reflection-based reasoning ( Shinn et al., 2023 ; Madaan et al., 2023 ) , and the integration of external tools ( Yao et al., 2023b ) . Other work investigates step-level supervision to structure explicit reasoning ( Zheng et al., 2023 ; Wei et al., 2025a ) . While effective, explicit CoT increases inference cost with longer sequences and often produces redundant steps, limiting efficiency and reasoning diversity ( Li et al., 2025 ; Zhang et al., 2025b ; Xu et al., 2025 ) .

Implicit CoT aims to reduce output length while retaining multi-step reasoning. Prior work explores knowledge internalization ( Deng et al., 2024 ) , architectural modification ( Saunshi et al., 2025 ; Chen et al., 2025 ; Cheng & Van Durme, 2024 ; Su et al., 2025 ; Mohtashami et al., 2023 ; Geiping et al., 2025 ) , training-free latent construction ( Zhang et al., 2025b ; Wu et al., 2025 ) , and auto-regressive latent reasoning ( Xu et al., 2025 ; Tan et al., 2025 ) . Coconut applies answer-level supervision ( Hao et al., 2025 ) , and CODI uses trajectory-level distillation ( Shen et al., 2025b ) . Our work introduces step-level supervision, which distributes signals across latent steps and improves stability. See extended discussion in Appendix D .

## 6 Conclusion

We introduce SIM-CoT, a training-based implicit reasoning method with step-level supervision on latent tokens. On GPT-2, SIM-CoT outperforms the strong explicit baseline SFT-CoT, while also surpassing implicit baselines such as Coconut and CODI. When scaling to larger LLaMA backbones, the performance achieves consistent gains over existing implicit reasoning methods and maintains fast inference efficiency. Ablation studies further show that it improves training stability with more latent tokens and can benefit from integration with training-free techniques such as soft thinking. Distance analysis confirms that SIM-CoT produces latent representations that are diverse yet stable.

## APPENDIX

## Usage of Large Language Models

In this paper, we used LLMs only for minor language polishing and formatting, without generating ideas, analyses, or experimental results.

## Outline