Source-linked AI summary
From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step
Yuntian Deng, Yejin Choi, Stuart Shieber
TL;DR
The paper investigates whether models can internalize explicit chain-of-thought reasoning instead of generating intermediate tokens. It progressively removes CoT steps during finetuning and reports strong results on multiplication and GSM8K, while noting training-cost and generalization limitations.
Problem
Explicit CoT can improve reasoning accuracy but may be computationally expensive, motivating study of whether models can internalize its reasoning steps.
Method
Stepwise Internalization begins with explicit-CoT training and gradually removes intermediate CoT tokens while finetuning the model toward implicit reasoning.
Results
The method achieves 0.99 accuracy on 9 × 9 multiplication with GPT-2 Small and over 0.50 accuracy on GSM8K with Mistral-7B without intermediate steps.
Takeaways & Limitations
Stepwise Internalization enables accurate implicit reasoning across multiplication and grade-school math while avoiding explicit intermediate outputs.
Takeaways & Limitations
Training becomes expensive for longer CoT chains because each token-removal stage requires finetuning, and explicit CoT still achieves higher accuracy.
Abstract
from arXiv · showhide
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
The paper asks whether language models can internalize explicit chain-of-thought reasoning in hidden states. It proposes gradually removing intermediate steps during finetuning, reporting strong results on multiplication and GSM8K while acknowledging limited task coverage.
- Motivation: Explicit CoT improves reasoning accuracy by decomposing complex problems into manageable intermediate steps, but long reasoning chains can be computationally expensive.The paper also notes that multi-digit multiplication remains challenging for language models despite being easy for calculators.
- Approach: Stepwise Internalization starts from explicit-CoT training, then progressively removes intermediate steps and finetunes the model until it can predict outputs implicitly.The method can also shorten reasoning chains while maintaining accuracy when full implicit reasoning is not yet possible.
- Novelty: The method offers a simpler alternative to hidden-state knowledge distillation for teaching models implicit reasoning.The comparison is specifically with a teacher-student approach that transfers explicit-CoT teacher states into student states.
- Results: Over 50% accuracy on GSM8K was achieved by Mistral 7B without explicit intermediate steps, exceeding GPT-4’s 44% score under direct-answer prompting.The result extends the evaluation beyond multiplication to grade-school math word problems.
- Limitations: The empirical evaluation is limited to multi-digit multiplication and grade-school math, leaving generalization to broader tasks and more diverse CoT traces for future research.The authors state that experiments on other tasks were out of scope because of computational-resource limitations.
2 Background: Implicit Chain-of-Thought Reasoning
Implicit CoT produces no explicit intermediate reasoning tokens during generation, while still using reasoning-step supervision during training. Stepwise Internalization implements this idea by progressively removing CoT tokens until only direct output prediction remains.
- Implicit CoT: Implicit CoT omits explicit intermediate reasoning steps during generation, distinguishing it from No CoT because reasoning steps remain available as training supervision.The model can instead use internal states as a scratchpad during inference.
- Explicit CoT: Explicit CoT uses intermediate steps both to supervise task learning and to provide an inference-time scratchpad.Implicit CoT replaces the latter function with the model’s internal states.
- Illustrative example: For 12 × 34, explicit CoT predicts intermediate calculation steps before producing the final answer 408.The example reverses digit order and intermediate steps for prediction consistency.
- Comparison: No CoT and implicit CoT directly predict the answer without intermediate steps, potentially accelerating inference for long reasoning chains at an accuracy cost.Their key distinction is whether intermediate steps provide supervision during training.
- Stepwise Internalization: Stepwise Internalization is presented as a simpler curriculum-learning approach that progressively removes CoT tokens during training.The process moves from explicit CoT supervision toward fully implicit reasoning.
3 Stepwise Internalization
Stepwise Internalization progressively removes intermediate CoT tokens while finetuning the model, aiming to internalize reasoning in hidden states. The method uses a removal schedule and smoothing techniques to manage training transitions.
- Removal schedule: The number of removed tokens can follow a predefined schedule; this work uses a linear schedule controlled by the per-epoch removal hyperparameter ∆.The schedule stops removing tokens once it exceeds the number of actual CoT tokens.
- Training stability: Training instability can arise when removing an additional CoT token abruptly changes the loss function and optimizer’s gradient estimates.The method resets the optimizer state whenever an additional CoT token is removed.
- Training stability: Removal Smoothing adds a small random non-negative offset to the scheduled token removal count, easing transitions between successive removal stages.With finite λ, the model sometimes removes more than s(t) tokens, smoothing the jump to s(t)+1.
- Stepwise Internalization: Stepwise Internalization starts from explicit CoT training and progressively removes intermediate reasoning tokens across training stages.The process eventually removes all intermediate steps for implicit CoT reasoning.
4 Experimental Setup
The evaluation tests Stepwise Internalization on increasingly difficult multiplication and GSM8K reasoning tasks, comparing it with direct training, explicit CoT, and implicit-CoT distillation baselines across model scales. Accuracy and inference speed are the principal evaluation metrics.
- Tasks and datasets: The experiments evaluate Stepwise Internalization on multi-digit multiplication and grade-school math reasoning.Multiplication includes 4-by-4, 5-by-5, 7-by-7, and 9-by-9 tasks; grade-school reasoning uses GSM8K.
- Baselines and models: The baseline suite includes No CoT, Explicit CoT, and ICoT-KD implicit reasoning via knowledge distillation.Explicit CoT uses five-shot prompting for GPT-3.5 and GPT-4 and finetuning for other models.
- Baselines and models: The study uses GPT-2, Phi-3 3.8B, and Mistral-7B to examine performance across model scales.The proposed implicit method is termed ICoT-SI.
- Evaluation metrics: Accuracy measures exact-match final-answer performance, while speed measures normalized examples per second during batch-size-1 inference.Speed is measured on an Nvidia H100 GPU relative to the corresponding No CoT model.
5 Results
Stepwise Internalization substantially improves implicit reasoning performance across multiplication and GSM8K, while offering a speed–accuracy trade-off relative to explicit CoT.
- Main results: 0.99 accuracy on 9 × 9 multiplication is achieved by GPT-2 Small with ICoT-SI, whereas No CoT fails on 4 × 4 multiplication.ICoT-SI also outperforms ICoT-KD on GPT-2 Small multiplication and matches No CoT speed.
- Main results: 0.56 accuracy on 5 × 5 multiplication is achieved by MathGLM-100M, compared with 0.90 by MathGLM-2B, while ICoT-SI solves up to 9 × 9 multiplication.A related arithmetic-specific method reaches 14 × 14 multiplication, but ICoT-SI is presented as more general.
- Main results: Over 0.50 accuracy is achieved by Mistral-7B on GSM8K without intermediate steps, compared with 0.44 for GPT-4 without intermediate steps.The result supports applying ICoT-SI beyond arithmetic tasks.
- Accuracy–speed trade-off: 0.68 accuracy with explicit CoT versus 0.51 with ICoT-SI is reported for finetuned Mistral-7B on GSM8K.Implicit CoT therefore remains less accurate than explicit CoT in this comparison.
- Accuracy–speed trade-off: 11 times faster inference is achieved by ICoT-SI than Explicit CoT on 9 × 9 multiplication, with comparable accuracy.Overall, the results emphasize a trade-off between accuracy and speed.
6 Analysis
ICoT-SI supports adjustable accuracy–speed trade-offs through partial internalization, while training stability depends on smoothing, optimizer resets, removal direction, and removal pace.
- Trade-offs: Over 0.7 accuracy at four times Explicit CoT speed is achieved on 11 × 11 multiplication when GPT-2 Small partially internalizes CoT tokens.Intermediate checkpoints remain useful when the model cannot internalize every reasoning step.
- Trade-offs: Users can adjust the number of internalized CoT tokens to optimize for higher accuracy or faster inference.This flexibility allows partial internalization to serve application-specific requirements.
- Ablation studies: 98% probability at offset o = 0 and 2% at o ≥1 define the removal-smoothing distribution used with λ = 4.Without smoothing, accuracy fails to recover after dropping to zero.
- Ablation studies: Resetting the optimizer stabilizes training; without it, accuracy drops to zero around 100 steps and does not recover.The reset avoids large estimates of second-order derivatives after more tokens are removed.
- Ablation studies: Removing CoT tokens from the beginning performs better than removing them from the right side.Beginning removal can distribute internalized tokens across the entire input, whereas end tokens depend on earlier tokens.
- Ablation studies: ∆ = 16 tokens removed per epoch fails to converge, while lower removal rates are more likely to train successfully but more slowly.The main experiments use ∆ = 8 to balance training speed and stability.
7 Related Work
Stepwise Internalization belongs to a broader line of work on reasoning without explicit intermediate outputs, but differs through progressive token removal and does not require teacher-model distillation.
- No CoT approaches: MathGLM solves multi-digit arithmetic without intermediate steps using sufficient demonstrations containing lower- and higher-digit arithmetic.Stepwise Internalization is compared with this no-CoT training approach for multi-digit multiplication.
- No CoT approaches: Filler tokens provide an alternative to CoT tokens for transformer reasoning and can improve model expressivity.The paper suggests combining this approach with Stepwise Internalization for more challenging tasks.
- Internalizing CoT: ICoT-KD uses explicit CoT during training but not generation, distilling reasoning into the student’s hidden states through a teacher model.Stepwise Internalization pursues the same implicit-reasoning goal without requiring a teacher model.
- Related methods: Stepwise Internalization progressively removes CoT tokens, whereas Context Distillation matches outputs with and without a scratchpad.Each internalization stage can be viewed as context distillation within a curriculum-learning setting.
- Related methods: Searchformer shortens search traces through sampled shorter trajectories, while Stepwise Internalization forces internalization by removing CoT tokens.The methods therefore differ in how shorter reasoning is obtained.
8 Limitations
The paper identifies training cost, instability, interpretability loss, and an accuracy trade-off as limitations of Stepwise Internalization.
- Training costs: Longer CoT chains increase training cost because each token-removal stage requires additional finetuning.For N-digit multiplication, the reasoning-chain length grows exponentially with N, making training increasingly expensive.
- Instability: Aggressive token-removal schedules can destabilize training and prevent recovery from accuracy drops.Lower ∆ values generally improve stability but require longer training.
- Interpretability: Implicit CoT models lose interpretable intermediate steps, although probing hidden states or combining implicit and explicit CoT may help.A mixed mode could trade off interpretability and latency according to task requirements.
- Accuracy: Explicit CoT achieves higher accuracy, while Stepwise Internalization trades some accuracy for lower latency.On 11 × 11 multiplication, the method retains reasonable accuracy and is several times faster despite not fully solving the task.
9 Conclusions and Future Work
The paper concludes that Stepwise Internalization can produce accurate implicit CoT reasoning by progressively removing intermediate tokens. Future work targets interpretability, mixed-mode reasoning, and scaling to larger models and training setups.
- 9 Conclusions and Future Work: Stepwise Internalization progressively removes intermediate CoT tokens and finetunes the model to internalize reasoning steps.The approach aims to achieve implicit CoT reasoning without outputting intermediate steps.
- 9 Conclusions and Future Work: The method achieves high accuracy on up to 9 × 9 multiplication and outperforms GPT-4 on GSM8K without outputting intermediate reasoning steps.Compared with explicit CoT, it can be up to 11 times faster while maintaining similar accuracies.
- 9 Conclusions and Future Work: Future work includes probing internal processes, combining implicit and explicit CoT, and scaling the method to larger models and broader training setups.The proposed mixed mode is intended to balance accuracy, latency, and interpretability.
- 9 Conclusions and Future Work: Experiments used task- and model-specific optimization settings, including AdamW, batch-size choices, learning rates, ∆ values, and bfloat16 for larger models.The reported configurations differ across multiplication and GSM8K experiments.
B Stability Issues for Aggressive Removal
Aggressive CoT-token removal schedules can make training unstable, with outcomes varying across runs even under identical configurations apart from the random seed.
- B Stability Issues for Aggressive Removal: Identical configurations can diverge: one run solved the task after all CoT tokens were removed, while another failed.The runs differed only in random seed, illustrating sensitivity to training dynamics.
C Additional Experiments
Additional experiments examined unchanged position IDs and a binary-tree CoT format as alternatives to the main setup.
- C Additional Experiments: Keeping position IDs unchanged produced more stable training but performance similar to the main approach, so it was not used in the experiments.After CoT removal, the next token directly reused the following position ID.
- C Additional Experiments: The binary-tree CoT format decomposes multiplication into N-digit-by-1-digit products and repeatedly merges results with sum operators.Its shorter description length could make the reasoning format easier for transformers to learn.