From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step — part 2/2

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

Overall, our results demonstrate that Stepwise Internalization is an effective method for enabling implicit CoT reasoning, offering a compelling trade-off between accuracy and speed. This makes it a valuable approach for tasks requiring both high performance and low latency.

## 6 Analysis

### 6.1 Accuracy-Speed Trade-off

One notable advantage of ICoT-SI is that it allows trading off accuracy with speed by internalizing different amounts of CoT tokens. At one extreme, when no tokens are internalized, ICoT-SI can recover explicit CoT performance. At the other extreme, when all tokens are internalized, we achieve implicit CoT, typically with much better performance than directly training a No CoT model.

Even when ICoT-SI is not completely successful due to model capacity limitations, such as on more challenging tasks where it cannot internalize all CoT steps, we can still leverage intermediate checkpoints to achieve a trade-off between accuracy and speed. For example, as shown in Figure 2(a) , on the 11 × 11 11\times 11 multiplication task with GPT-2 Small, even though the model cannot internalize all CoT steps, ICoT-SI is still able to achieve an accuracy of over 0.7 at a speed four times that of explicit CoT when part of the CoT tokens are internalized.

This trade-off curve illustrates the flexibility of ICoT-SI in balancing computational efficiency and model performance. By adjusting the number of internalized CoT tokens, users can optimize for either higher accuracy or faster inference depending on the requirements of their specific application.

### 6.2 Ablation Studies

Figure 3 plots the validation accuracy versus the schedule for the number of CoT tokens removed during training for the 7 × 7 7\times 7 multiplication task. This figure compares the full approach to several ablated variants. Even for the full approach, there are fluctuations in the curve, and the validation accuracy briefly drops to zero at one point during training but eventually recovers. However, the ablated variants do not fully recover when accuracy drops.

#### Removal smoothing.

As mentioned in Section 3 , adding a small random offset o o to the number of removed tokens is crucial when the loss function changes due to the removal of more CoT tokens. The distribution of o o is parameterized by a hyperparameter λ \lambda , as introduced in Section 3 . We use λ = 4 \lambda=4 throughout this work, resulting in the distribution shown in Figure 2(b) . In this distribution, 98% of the time, o = 0 o=0 , but about 2% of the time, one or more additional tokens are removed. As shown in Figure 3 , the “Without Removal Smoothing” curve fails to recover after the accuracy drops to zero at around s ⁡ ( t ) = 50 s(t)=50 , whereas the full approach does much better.

#### Resetting the optimizer.

Another important technique for stabilizing training is resetting the optimizer when more tokens are removed. This avoids large estimates of second-order derivatives and stabilizes training. In Figure 3 , the “Without Optimizer Reset” curve drops to zero around 100 steps and does not recover, showing the importance of resetting the optimizer during training.

#### Removal side.

In our main experiments, CoT tokens are removed from the beginning (left side). Removing CoT tokens from the right side performs significantly worse, as shown by the “Right-Side Removal” curve in Figure 3 . We suspect this is because internalizing tokens at the beginning is easier than internalizing tokens at the end. CoT tokens at the end depend on the earlier tokens, so internalizing them between the end of CoT and the beginning of the final answer, which only has a few positions, is more challenging. In contrast, internalizing tokens at the beginning allows distributing them across the entire input.

#### Number of tokens removed per epoch.

The number of tokens removed per epoch ( Δ \Delta ) significantly affects the training stability and speed. In the main experiments, we used Δ = 8 \Delta=8 , which removes 8 tokens per epoch. A higher Δ \Delta value leads to faster training but risks not converging, as the model may not be able to keep up with the rapid changes in the loss function. For instance, when using Δ = 16 \Delta=16 , the training fails to converge, as shown by the “Aggressive Removal” curve in Figure 3 . Conversely, a lower Δ \Delta value is more likely to result in successful training but at a slower pace. Future work could explore adaptive Δ \Delta schedules based on loss values to balance speed and stability more effectively.

## 7 Related Work

#### No CoT approaches.

Several works in the literature focus on training language models to solve arithmetic tasks without outputting intermediate steps. MathGLM [ 20 ] demonstrated that with sufficient training data, including both lower-digit and higher-digit arithmetic task demonstrations, a 2 billion parameter LM can solve multi-digit arithmetic tasks without any intermediate steps. Compared to this work, Stepwise Internalization achieves higher accuracy in solving multi-digit multiplication with much smaller models, likely due to leveraging chain-of-thought supervision during training. Another notable work by Shen et al. [17] showed that by mixing lower-digit and higher-digit multiplication demonstrations, even a GPT-2 Small can learn up to 14-digit multiplication. However, Stepwise Internalization does not require specially prepared training data with mixed task difficulties. Additionally, Stepwise Internalization is theoretically applicable to any reasoning task with CoT reasoning steps, as demonstrated by its effectiveness on grade-school math problems.

Also relevant is the work of Pfau et al. [15] , which shows that transformer language models can reason using filler tokens as an alternative to CoT tokens. They showed reasoning using these filler tokens improves a language model’s expressivity. Our approach has the potential to be combined with their approach to solve even more challenging tasks.

#### Internalizing CoT.

Our work is closely related to that of Deng et al. [6] (ICoT-KD), which introduced the task of implicit CoT reasoning. ICoT-KD allows using CoT during training but not during generation, and it implements this via knowledge distillation to internalize the reasoning steps within hidden states. Compared to ICoT-KD, Stepwise Internalization has three advantages: First, it is simpler to implement as it does not require a teacher model. Second, while ICoT-KD internalizes reasoning into a single “column” of states (corresponding to the final input position), Stepwise Internalization allows the model to internalize reasoning across all input positions. Lastly, Stepwise Internalization achieves better accuracy compared to ICoT-KD.

Our work is also related to Context Distillation [ 18 ] , which trains a model to produce the same output when conditioned on a scratchpad versus without it. Each stage of Stepwise Internalization can be viewed as a form of context distillation, where one CoT token is distilled into the model’s internal states. Stepwise Internalization extends Context Distillation into a curriculum learning setting.

Another relevant work is Searchformer [ 12 ] , which first trains a transformer to imitate A* search and then finetunes it on sampled shorter search traces. This allows the model to perform searches using fewer steps than those provided during training. While Searchformer relies on sampling to find shorter traces, Stepwise Internalization forces the model to internalize steps by removing CoT tokens.

## 8 Limitations

#### Training costs.

One limitation of the proposed approach is its high training cost due to the finetuning required when removing each set of CoT tokens. As discussed in Section 6.2 , removing CoT tokens too fast leads to non-convergence. Therefore, the longer the CoT chain, the longer the training duration. For tasks like N N -digit multiplication, where the reasoning chain length grows exponentially with N N , training becomes expensive as N N increases.

#### Instability.

Another practical issue we observed is the instability of training with aggressive Δ \Delta values. For example, Figure 4 in Appendix B shows a case where the model could not recover from a drop in accuracy. Using lower Δ \Delta values generally leads to more stable training, but at the cost of longer training time. Identifying and addressing unstable dynamics early on, potentially by restarting training as suggested by Hu et al. [9] , could be a valuable improvement.

#### Interpretability.

Similar to existing work on No CoT and implicit CoT training, models trained using our approach lose interpretable intermediate steps. However, it might be possible to interpret the internal hidden states of these models using probing techniques [ 2 , 8 ] . Additionally, combining implicit and explicit CoT training could allow users to choose between interpretability and latency, providing flexibility based on the requirements of future tasks.

#### Accuracy.

Undoubtedly, explicit CoT still achieves higher accuracies compared to our approach to implicit CoT. However, our method enables a trade-off between latency and accuracy. Even on tasks it cannot fully solve without intermediate steps, such as 11 × 11 11\times 11 multiplication, it maintains reasonable accuracy while being several times faster than explicit CoT. Moreover, our results demonstrate the potential of leveraging hidden states for reasoning: even a GPT-2 Small model can be trained to solve 9 × 9 9\times 9 multiplication, despite having only 12 layers, far fewer than the number of reasoning steps in the CoT for 9 × 9 9\times 9 multiplication. When scaled to larger models with hundreds of billions of parameters and up to a hundred layers, such as GPT-3 [ 4 ] , they could potentially solve even more challenging reasoning tasks without explicit CoT steps.

## 9 Conclusions and Future Work

In this work, we introduced Stepwise Internalization, a novel approach for achieving implicit chain-of-thought reasoning in language models. By gradually removing intermediate CoT tokens and finetuning the model, we enable the internalization of reasoning steps incrementally. Our approach demonstrates significant improvements over existing methods, achieving high accuracy on up to 9 × 9 9\times 9 multiplication using GPT-2 Small and outperforming GPT-4 on GSM8K while not outputting any intermediate reasoning steps. Compared to explicit CoT methods, our approach can be up to 11 times faster while maintaining similar accuracies.

For future work, probing the internal processes as the model internalizes each reasoning step could provide insights into the learning mechanisms. Additionally, developing a mixed-mode approach that combines implicit and explicit CoT reasoning could potentially offer the best of both worlds, balancing accuracy, latency, and interpretability based on user preferences. Another promising direction is scaling Stepwise Internalization to larger models and more extensive training/pretraining setups, which could further enhance its effectiveness on a broader range of reasoning tasks.

## Acknowledgments and Disclosure of Funding

This work was supported by NSF grant DMS-2134012 and ONR grant N00014-24-1-2207. We also thank Harvard University FAS Research Computing for providing computational resources.

## Appendix A Hyperparameters

For all experiments, we use the AdamW optimizer [ 13 ] , with λ = 4 \lambda=4 and an effective batch size of 32 by default. For Phi-3 3.8B and Mistral 7B, we use a batch size of 16 with a gradient accumulation of 2. For the multiplication tasks, we use a learning rate of 5 × 10 − 5 5\times 10^{-5} and Δ = 8 \Delta=8 . For GSM8K, we use a learning rate of 5 × 10 − 5 5\times 10^{-5} and Δ = 1 \Delta=1 for GPT-2 Small and GPT-2 Medium, and a learning rate of 1 × 10 − 5 1\times 10^{-5} and Δ = 8 \Delta=8 for Phi-3 3.8B and Mistral 7B, with bfloat16 precision. Additionally, for GSM8K, we only consider sequences with 150 or fewer tokens for training and remove all CoT tokens when 39 or more tokens are scheduled to be removed. All experiments are run on a single H100 with 80GB of GPU memory for up to 200 epochs or 24 hours, whichever is reached first.

## Appendix B Stability Issues for Aggressive Removal

We found that using aggressive removal schedules (that is, bigger Δ \Delta values) can sometimes lead to unstable training dynamics. As one example, Figure 4 shows two different runs under identical configurations except for the random seed. One run was eventually able to solve the task after all CoT tokens were removed, whereas the other failed to solve the task after all CoT tokens were removed.

## Appendix C Additional Experiments

#### Keeping position IDs.

As CoT tokens are removed, the position where the final output starts changes. We tried a variant where position IDs remain unchanged, meaning the position ID of the next token is used directly after removing a CoT token. Although this approach was more stable during training, its performance was similar to the current approach. For simplicity, we did not use this variant in our main experiments.

#### Alternative CoT formats.

Different valid reasoning paths can lead to the correct final answer for the same problem. We explored using a binary tree formatted CoT chain for the multiplication problems. This format decomposes an N N -digit multiplication into a sequence of N N -digit-by-1-digit multiplication problems, merges the results using sum operators, and continues merging until the final sum is computed. This program has a shorter description length, potentially making it easier for transformers to learn [ 7 ] . However, its performance was similar to the current approach: for 9 × 9 9\times 9 multiplication using GPT-2 Small, it achieved 0.95 accuracy and failed on 11 × 11 11\times 11 multiplication.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .