From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step — part 1/2
Cleaned-text code points [0, 15900). Bibliography omitted here; source.md includes it.
# From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step
###### Abstract
When leveraging language models for reasoning tasks, generating explicit chain-of-thought (CoT) steps often proves essential for achieving high accuracy in final outputs. In this paper, we investigate if models can be taught to internalize these CoT steps. To this end, we propose a simple yet effective method for internalizing CoT steps: starting with a model trained for explicit CoT reasoning, we gradually remove the intermediate steps and finetune the model. This process allows the model to internalize the intermediate reasoning steps, thus simplifying the reasoning process while maintaining high performance. Our approach enables a GPT-2 Small model to solve 9-by-9 multiplication with up to 99% accuracy, whereas standard training cannot solve beyond 4-by-4 multiplication. Furthermore, our method proves effective on larger language models, such as Mistral 7B, achieving over 50% accuracy on GSM8K without producing any intermediate steps.
## 1 Introduction
A prevalent approach to improving the performance of language models (LMs) to perform complex reasoning tasks is chain-of-thought (CoT) reasoning, in which the LM generates explicit intermediate reasoning steps before arriving at a final answer [ 14 , 19 ] . This method allows models to break down complex problems into simpler, manageable parts, thereby improving the accuracy of their final predictions. However, this explicit reasoning process can be computationally expensive, especially when the reasoning chain is long [ 6 ] . Additionally, using explicit intermediate steps might not align with the intrinsic computational strengths of LMs [ 12 ] : for instance, multi-digit multiplication is very easy for calculators but remains challenging for GPT-4 [ 20 ] .
In this work, we examine the possibility of internalizing the reasoning process in the model’s hidden states. We propose an approach, Stepwise Internalization, which begins with a model trained for explicit CoT reasoning. We then gradually remove the intermediate steps and finetune the model, forcing it to internalize the reasoning process. Once all intermediate steps are internalized, we achieve a model capable of full implicit CoT reasoning. Moreover, even in cases where the model does not have the capacity for full implicit CoT reasoning, this method still allows for shortening the reasoning chain while maintaining accuracy.
Our approach is an alternative to the approach proposed by Deng et al. [6] , which shares the goal of implicitly reasoning using the hidden states of transformers instead of relying on explicit CoT tokens. To teach the model to use hidden states for reasoning, that method employs a teacher model that performs explicit CoT reasoning, and then distills the teacher’s hidden states into the student model’s hidden states. In comparison, our approach is much simpler yet more effective.
Our approach demonstrates significant improvements over standard training methods. For instance, a GPT-2 Small model trained with Stepwise Internalization on multiplication can solve even 9-by-9 multiplication problems nearly perfectly, while standard training without CoT struggles even with 4-by-4 multiplication. Furthermore, our method scales effectively to larger models, such as the Mistral 7B model [ 10 ] , achieving over 50% accuracy on the GSM8K dataset of grade-school math word problems [ 5 ] , without producing any explicit intermediate steps, outperforming the much larger GPT-4 model without chain-of-thought reasoning, which only scores 44% when prompted to directly generate the answer.
It is important to note that our empirical evaluation focuses on specific reasoning tasks like multi-digit multiplication and grade-school math problems. While our results show the potential of Stepwise Internalization in these contexts, and the simplicity of the method makes it applicable to chain-of-thought approaches in a wide range of tasks, further research is needed to explore its efficacy across a broader range of tasks and more diverse CoT traces. Due to limitations in available computational resources, experiments on other tasks are out of scope for this work. This paper aims to lay the groundwork for this new approach and highlight its promise, while acknowledging that its full generalization is still under investigation.
The contributions of our work are as follows: First, we introduce Stepwise Internalization, a simple method for implicit CoT reasoning. Second, we demonstrate the effectiveness of internalizing intermediate hidden states via Stepwise Internalization. Third, we provide empirical results showing the superior performance of models trained with Stepwise Internalization on different reasoning tasks and model scales. Our code, data, and pretrained models are available at https://github.com/da03/Internalize_CoT_Step_by_Step .
## 2 Background: Implicit Chain-of-Thought Reasoning
Implicit chain-of-thought reasoning (implicit CoT, or ICoT) is a concept introduced by Deng et al. [6] , where during generation, the language model does not produce explicit intermediate reasoning steps in words. It is distinct from not using chain-of-thought reasoning (No CoT), in that explicit reasoning steps are allowed during training, enabling the ICoT model to learn the underlying reasoning approach from the supervision provided on the reasoning process. The key insight of Deng et al. [6] is that intermediate reasoning steps serve two purposes in explicit CoT: they provide supervision during training to facilitate learning the task [ 14 ] , and they act as a scratchpad during inference to assist in solving the task [ 19 ] . However, the latter purpose can be fulfilled by utilizing the internal states of the model instead of explicit tokens.
As an illustrative example, consider using a language model to solve a multi-digit multiplication problem, such as 12 × 34 12\times 34 . (The actual input reverses the digit order as 2 1 * 4 3 for consistency with Deng et al. [6] .) In the long multiplication algorithm, 12 × 34 12\times 34 is broken into: 12 × 4 + 12 × 30 = 48 ⏟ reversed: 84 + 360 ⏟ reversed: 063 . 12\times 4+12\times 30=\underbrace{48}_{\text{reversed: }84}+\underbrace{360}_{\text{reversed: }063}. In explicit CoT, the model is trained to predict these intermediate steps 8 4 + 0 6 3 before predicting the final answer 8 0 4 ( 408 408 reversed). Predicting these intermediate steps facilitates the model’s ability to solve the task. (The intermediate steps are also reversed to make it easier for the model to predict [ 17 ] .)
In both No CoT and implicit CoT settings, the model needs to directly predict the answer 408 408 from the input, bypassing the intermediate steps. This approach can make inference much faster for long reasoning chains, albeit at the cost of accuracy.
The primary difference between implicit CoT and No CoT lies in the use of intermediate reasoning steps as supervision during training. In the work of Deng et al. [6] , a knowledge distillation approach was employed to distill explicit reasoning into implicit reasoning within the hidden states. This method involves training a teacher model to perform explicit CoT reasoning and then transferring this knowledge to a student model, which internalizes the reasoning process within its hidden states.
In the present work, we propose a far simpler yet more effective approach based on a kind of curriculum learning that we call Stepwise Internalization, which we detail in the next section.
## 3 Stepwise Internalization
Stepwise Internalization is a method designed to achieve implicit chain-of-thought reasoning by gradually removing intermediate reasoning steps during training. We define the input as x x , the intermediate steps as z = z 1 , z 2 , ⋯ , z m z=z_{1},z_{2},\cdots,z_{m} , and the final output as y y . A language model with parameters θ \theta is first trained using the following loss function: min θ − log P θ ( y , z 1 : m ∣ x ) , \min_{\theta}-\log P_{\theta}(y,z_{1:m}\mid x), where z 1 : m z_{1:m} denotes the sequence of intermediate steps z 1 , z 2 , ⋯ , z m z_{1},z_{2},\cdots,z_{m} .
At each step t t of the training process, we remove (up to) s ( t ) s(t) tokens from the intermediate steps z z . The updated loss function then becomes: min θ − log P θ ( y , z 1 + min ( s ( t ) , m ) : m ∣ x ) . \min_{\theta}-\log P_{\theta}(y,z_{1+\min(s(t),m):m}\mid x).
There are multiple ways to parameterize s ( t ) s(t) . For instance, it might be based on a threshold of the loss value or a predefined schedule similar to learning rate schedulers used in optimizers. In this work, for simplicity, we use a linear schedule for removing tokens: s ( t ) = ⌊ Δ t T ⌋ , s(t)=\left\lfloor\Delta\frac{t}{T}\right\rfloor, where T T is the total number of steps per epoch, and Δ \Delta is a hyperparameter controlling how many CoT tokens are removed per epoch. (Once s ( t ) s(t) exceeds the number of actual chain-of-thought tokens, all tokens are removed.)
During initial experiments, we observed instability in the training process due to changes in the loss function over time. This instability arises for two primary reasons:
First, the optimizer commonly used in training language models, such as AdamW [ 11 , 13 ] , maintains estimates of second-order gradients. A sudden change in the loss function, caused by the removal of one more CoT token, results in abrupt changes in the second-order gradients. To address this issue, we reset the optimizer’s state whenever an additional CoT token is removed.
Second, even if a model fits perfectly to the current loss when s s tokens are removed, transitioning to the next stage, where s + 1 s+1 tokens are removed, leads to a significant increase in the loss, as the model is not yet trained for this new setting. To mitigate this issue, we introduce a technique which we term “Removal Smoothing”, where we add a small random offset to the original number of tokens to remove s ( t ) s(t) , such that: s ( t ) ∗ = s ( t ) + o , s(t)^{*}=s(t)+o, where o o is a random variable with support of non-negative integers ℤ ≥ 0 \mathbb{Z}_{\geq 0} , and its distribution is parameterized by another hyperparameter λ \lambda : P ( o ) ∝ exp ( − λ o ) . P(o)\propto\exp(-\lambda o). When λ = ∞ \lambda=\infty , o = 0 o=0 and we recover the version without Removal Smoothing. However, when λ < ∞ \lambda<\infty , the model is trained to remove more than s ( t ) s(t) tokens at step t t with a small probability, which helps smooth the transition into the next stage of removing s ( t ) + 1 s(t)+1 tokens, reducing the abrupt jumps in the loss function.
Figure 1 illustrates the high-level idea of the Stepwise Internalization approach. The training process consists of multiple stages, where the model progressively learns to internalize reasoning steps by removing tokens from the CoT at each stage, eventually achieving implicit CoT reasoning.
## 4 Experimental Setup
### 4.1 Data
We evaluate our proposed Stepwise Internalization method on two reasoning tasks following Deng et al. [6] : multi-digit multiplication and grade-school math reasoning.
#### Multi-digit multiplication.
We use two of the most challenging arithmetic tasks from BIG-bench [ 3 ] : 4-by-4 multiplication and 5-by-5 multiplication, as described by Deng et al. [6] . Given the effectiveness of Stepwise Internalization on these tasks, we extend our evaluation to 7-by-7 and 9-by-9 multiplication. The complexity of multiplication tasks grows significantly with the number of digits, as the program length grows quadratically with the number of digits [ 7 ] . We use the scripts and setup from Deng et al. [6] to generate synthetic training data for our main experiments 1 1 1 Following Deng et al. [6] , K K -by K K multiplication only considers K K -digit numbers but not lower digits. .
#### Grade school math.
We use the GSM8K dataset [ 5 ] , with the augmented training data provided by Deng et al. [6] . Detailed dataset statistics are provided in Table 1 .
### 4.2 Baselines and Models
We compare our method to the following baselines: • No CoT: Models directly trained without chain-of-thought supervision.
• Explicit CoT: Models finetuned or prompted with explicit chain-of-thought reasoning [ 14 ] . We use 5-shot prompting for GPT 3.5 and GPT-4 but full finetuning for other models.
• ICoT-KD: The implicit chain-of-thought via knowledge distillation method proposed by Deng et al. [6] .
Our proposed method, implicit chain-of-thought via Stepwise Internalization, is termed ICoT-SI. To verify the effectiveness of our approach across different model scales, we use pretrained models GPT-2 [ 16 ] , Phi-3 3.8B [ 1 ] , and Mistral-7B [ 10 ] .
### 4.3 Evaluation
Because the premise for implicit chain-of-thought methods is to approach the speed of no chain-of-thought and the accuracy of explicit chain-of-thought, we use two main evaluation metrics: First, we evaluate the accuracy of each method on the respective tasks of generating the final output. Second, we compare the inference speed of each method to the No CoT baseline. We measure speed, in examples per second, on an Nvidia H100 GPU with a batch size of 1. For ICoT-KD, we directly take numbers from Deng et al. [6] . However, due to hardware differences, we recompute speed relative to No CoT when speed numbers from ICoT-KD are not available.
## 5 Results
Table 2 presents the main results, where we compare Stepwise Internalization to various baselines.
#### Stepwise Internalization is effective.
Compared to other methods that do not output intermediate steps, Stepwise Internalization (ICoT-SI) proves to be highly effective. For example, ICoT-SI enables a GPT-2 Small model to solve the 9 × 9 9\times 9 multiplication problem with an accuracy of 0.99, whereas the No CoT method fails on even 4 × 4 4\times 4 multiplication. Additionally, ICoT-SI outperforms Implicit CoT via Knowledge Distillation (ICoT-KD); while ICoT-KD fails to solve 5 × 5 5\times 5 multiplication using a GPT-2 Small model, ICoT-SI can solve up to 9 × 9 9\times 9 multiplication. Also, while ICoT-KD is slightly slower than No CoT due to the additional emulator model, ICoT-SI has the same speed as No CoT 2 2 2 The speed of ICoT-SI in Table 2 is not always 1.00 due to randomness in hardware speed. .
When compared to existing literature, ICoT-SI is also competitive. For instance, at a similar model size, MathGLM-100M [ 20 ] can only solve 5 × 5 5\times 5 multiplication with an accuracy of 0.56. Even with 2 billion parameters, MathGLM-2B can solve 5 × 5 5\times 5 multiplication with an accuracy of 0.90. Although another related work [ 17 ] is able to train a GPT-2 Small model to solve up to 14 × 14 14\times 14 multiplication, the method proposed in that work is specific to arithmetic tasks, whereas ICoT-SI is more general.
ICoT-SI enables the internalization of CoT reasoning in a general way, making it applicable to tasks beyond arithmetic, such as grade-school math problems. For example, on the GSM8K dataset, ICoT-SI achieves a new state-of-the-art accuracy for models not using any intermediate steps. It finetunes the Mistral-7B model to achieve over 0.50 accuracy, whereas even GPT-4 can only achieve 0.44 without using intermediate steps.
#### Stepwise Internalization lags behind explicit CoT in accuracy but is faster.
In terms of accuracy, implicit CoT methods still lag behind explicit CoT. For instance, a finetuned Mistral-7B model can achieve an accuracy of 0.68 on GSM8K with explicit CoT but ICoT-SI only got 0.51. However, implicit CoT methods offer significant speed advantages. For example, on the 9 × 9 9\times 9 multiplication task, ICoT-SI is comparable in accuracy to Explicit CoT but is 11 times faster during inference.