From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step

Thank you for the response. From the technical report of GPT-4 (https://arxiv.org/pdf/2303.08774v5), it seems like GPT-3.5 achieved 57.1 on GSM-8K with 5-shot No-CoT, while GPT-4 got 92.0 with 5-shot CoT (Table 2, last row). For Mistral, it looks like GSM-8K scores are consistent with 5-shot maj@1 greedy sampling. Look forward to the rest of your responses.

We thank the reviewer for the feedback.

Re: Internalized CoT models have worse performance compared to explicit CoT prompting

We believe that in addition to accuracy, inference efficiency is an equally important dimension of performance. For instance, o1-preview requires about 30 seconds of inference time to generate a CoT chain, while our approach can significantly reduce latency, as shown in Figure 2(b).

Additionally, our work explores the scientific question of what smaller models can achieve through alternative optimization strategies. For example, our method enables a small GPT-2 model to directly solve 20×20 digit multiplication, a feat not achievable by prior methods, including o1.

Re: Comparison with No CoT baselines

To clarify, we finetuned all models, including No CoT baselines, for comparable durations (24 hours or until convergence). If the reviewer suggests first finetuning No CoT models on CoT examples and then transitioning to No CoT finetuning, we have already conducted this experiment on GSM8K with Mistral 7B. Specifically, we finetuned the model for one epoch using CoT supervision before transitioning to No CoT finetuning. This approach resulted in worse performance (0.36 vs. 0.38 accuracy), likely due to the big formatting differences between CoT and No CoT settings.

Re: Synthetic training data

We chose intermediate CoT steps as training data because they align naturally with the reasoning process for arithmetic tasks, enabling systematic internalization. To our knowledge, our approach is the first to enable a small GPT-2 model to directly solve 20×20 digit multiplication while achieving a 25x inference speedup compared to explicit CoT (Figure 2(b)). If the reviewer has specific suggestions for simpler methods (e.g., deterministic transformations of problem statements), we would be eager to explore them.

Re: Higher number of intermediate steps

While multiplication may appear synthetic, a 20-digit multiplication problem requires $O(20\times 20) = O(400)$ intermediate steps.

5. Re: Testing on larger models and more complex datasets

Unfortunately, resource constraints limited our ability to finetune larger models such as LLaMA 3.1. However, we note that scaling model size or training data does not directly address the challenges of the No CoT setting, and larger models typically incur higher inference costs. Implicit CoT offers a unique advantage in this dimension by reducing inference latency. Furthermore, we compared our approach against state-of-the-art models, such as GPT-3.5 and GPT-4, which also struggled with the No CoT setting.

Re: Q1: What are the motivations for internalizing CoT steps beyond latency?

Latency is indeed a critical dimension of model performance. Additionally, our approach demonstrates that alternative optimization techniques can discover solutions that are not achievable through direct training. For instance, our method enabled a small GPT-2 model to solve $20\times20$ multiplication directly—something prior methods could not accomplish.

Re: Q2: Implications of losing token-based 'working memory' for complex compositional tasks

One of our motivations is that for state of the art LLMs such as GPT-4o, there are likely hundreds of layers (GPT-3 has 96 layers), while most reasoning problems do not require hundreds of intermediate steps, and it seems to be a significant waste of resources to forward across hundreds of layers just to compute each single CoT token.

Re: Q3: Comparison with other synthetic data approaches (e.g., deterministic transformations)

We would be excited to explore such methods further. However, it is unclear how deterministic transformations, such as changing numerical values in problem statements, would achieve the same performance and latency benefits as our approach. If the reviewer could provide specific examples, we would gladly experiment with them.

Re: Q4: Why was the method not evaluated on larger models and more complex datasets?

Resource constraints limited our ability to conduct such experiments. We believe it is unreasonable to require evaluations on larger models and more datasets without clear hypotheses to test. Finetuning a 7B model, such as Mistral, is already computationally intensive.

We thank the reviewer again for their detailed feedback and hope our responses address their concerns.

We appreciate your thoughtful feedback and suggestions, which have greatly helped us refine and clarify our work. Below, we address your concerns and outline the additional experiments and analyses conducted in response to your review.

Re: Figure 3 Concerns About Shortcut Learning

As noted in our quick clarification, the accuracy in Figure 3 reflects exact match accuracy, where even a single digit deviation results in a zero score. To address your concern about potential shortcut learning, we have included a token-level accuracy plot in our revised paper (Figure 8 in Appendix F). This plot provides additional insight into the training dynamics.

Our analysis shows that for most of the training period (fewer than 112 CoT tokens removed), token-level accuracy across all variants of our approach remains above 90%. Toward the end of training, when all CoT tokens are removed, the token-level accuracy for the ablation variants drops to approximately 75% and does not recover, and our qualitative analyses found that they struggle with predicting the middle digits of the result. In contrast, the No CoT approach (represented by the brown dashed horizontal line) achieves only 61% token-level accuracy even after training for over 24 hours. This shows the effectiveness of our internalization approach, which maintains significantly higher token-level accuracy even in ablated settings.

Re: Analysis in Section 6.1 and the Need for Pretrained Model Probing

Thank you for this valuable suggestion. In response, we have added details on the probing experiment to Appendix D, including the architecture, training procedure, and evaluation metrics. We also conducted a mechanistic interpretability analysis on an explicit CoT model as a control experiment, presented in Appendix E (Figure 7).

Our findings reveal that even explicit CoT models can encode some CoT token information in their hidden states. For example, the first and last digits of partial products and partial sums are predictable at the earliest positions. However, these predictions rely on simple patterns (e.g., $a_1\times b_1$ for the lowest digit) and do not reflect full internalization of the partial products. In contrast, the hidden states of implicit CoT models encode information about all partial products, including the more challenging middle digits of the results, which are not predictable from explicit CoT hidden states. This confirms that the internalization training process does internalize intermediate computations within the model.

Re: Q1 - "No CoT" Setting and Prompt Sensitivity

For Mistral 7B, the No CoT setting was achieved through finetuning on the same dataset as our internalization approach, but without CoT supervision. This trained model got much better results than prompting alone. For GPT-3.5 and GPT-4, we used a five-shot prompting setup, which performed better than zero-shot prompting with instructions to omit CoT steps in our preliminary experiments.

The low performance in these settings reflects the inherent difficulty of direct prediction without CoT. To illustrate this challenge, we reference two simple examples where GPT-4o fails even when explicitly instructed to avoid CoT reasoning (see https://bit.ly/gpt4_system1_ogden and https://bit.ly/gpt4_system1).

We have clarified the baseline descriptions in Section 4.2 of the revised paper to make it more clear.

Re: Q2 - Comparison with Mistral 7B Using CoT Training

To further address your suggestion, we conducted an additional experiment where Mistral 7B was first finetuned for one epoch with explicit CoT supervision before transitioning to No CoT finetuning. However, this approach resulted in slightly lower final accuracy (0.36). We hypothesize that this is due to the significant formatting differences between CoT and No CoT settings and the lack of reasoning internalization during explicit CoT training. Explicit CoT supervision allows the model to delay reasoning until the CoT tokens, whereas internalization forces the reasoning to be embedded directly in the hidden states.

We thank you again for your detailed feedback and hope these additional results and clarifications address your concerns. Please let us know if further clarifications or experiments are needed.