Source-linked AI summary
Distilling Reasoning Capabilities into Smaller Language Models
Kumar Shridhar, Alessandro Stolfo, Mrinmaya Sachan
TL;DR
CoT reasoning is powerful but typically requires very large language models, creating a need to transfer such capabilities to smaller models. The paper distills LLM-generated reasoning through CoT and Socratic CoT, whose subquestion decomposition guides smaller student models. Across GSM8K, StrategyQA, and SVAMP, the strategies improve smaller-model performance, including cases where GPT-2 large performs as well as a nearly 10× larger GPT-3 6B model.
Problem
CoT reasoning primarily works with very large models, while accessing CoT-capable LLMs requires significant computational resources or expensive API calls.
Method
The framework prompts an LLM for reasoning annotations and distills them into smaller models using CoT or Socratic CoT, which decomposes problems into subquestion-solution pairs.
Results
The distillation strategies improve smaller-model reasoning across GSM8K, StrategyQA, and SVAMP, with Socratic CoT yielding gains of up to 38% and GPT-2 large performing as well as nearly 10× larger GPT-3 6B.
Takeaways & Limitations
Socratic CoT can be an effective alternative to standard CoT for transferring reasoning capabilities into smaller models across multiple reasoning datasets.
Takeaways & Limitations
The study uses one LLM solution and a single prompt, leaving majority-voting over sampled pairs and broader prompt comparisons for future work; the models remain unsuitable for sensitive settings because of hallucination risk.
Abstract
from arXiv · showhide
Step-by-step reasoning approaches like chain of thought (CoT) have proved to be very effective in inducing reasoning capabilities in large language models. However, the success of the CoT approach is fundamentally tied to the model size, and billion parameter-scale models are often needed to get CoT to work. In this paper, we propose a knowledge distillation approach that leverages the step-by-step CoT reasoning capabilities of larger models and distills these abilities into smaller models. In this work, we propose an alternative reasoning scheme, Socratic CoT, that learns a decomposition of the original problem into a sequence of subproblems and uses it to guide the intermediate reasoning steps. We use Socratic CoT to train a combination of two small distilled models: a problem decomposer and a subproblem solver. In practice, given a new problem, the two distilled models work in sync to decompose and solve complex problems. On multiple reasoning datasets (GSM8K, StrategyQA, and SVAMP), our proposed distillation strategies boosts the performance of smaller models over 70% compared to the baselines. Finally, we investigate when Socratic CoT is an effective alternative to CoT, demonstrating cases where a much smaller model (GPT-2 large) can outperform a 10X larger model (GPT-3 6B). Our code is available here: https://github.com/kumar-shridhar/Distiiling-LM
1 Introduction
CoT reasoning is effective in large language models but generally depends on very large models, motivating distillation into smaller models. The proposed framework uses LLM-generated step-by-step or subquestion annotations to supervise smaller models, with Socratic CoT providing additional improvement.
- CoT elicits intermediate reasoning steps in LLMs but primarily works on models with hundreds of billions of parameters or broad task tuning.
- The method uses an LLM to generate step-by-step annotations, which provide additional supervision when fine-tuning smaller student models.
- The framework addresses the computational resources and API costs of accessing CoT-capable LLMs by transferring their reasoning abilities to smaller models.
- Socratic CoT decomposes each problem into subproblem-solution pairs, using a question generator and a question-answering model to guide intermediate reasoning.
- The framework evaluates gold, teacher-generated CoT, and Socratic CoT annotations across GSM8K, StrategyQA, and SVAMP.
- Supervision by CoT-decomposed examples improves smaller-model performance, while Socratic CoT provides further improvement, with gains of up to 40%.
2 Related Work
Related work spans prompting and decomposition methods for multi-step reasoning, knowledge distillation using teacher rationales, and question-based supervision for reasoning tasks.
- Decomposing Multi-Step Reasoning Tasks: Chain-of-thought prompting generates intermediate steps to improve reasoning in language models, while related methods sample multiple paths or use zero-shot prompts.
- Decomposing Multi-Step Reasoning Tasks: Modern multi-step reasoning approaches increasingly rely on large language models through prompting or in-context learning.
- Knowledge Distillation: Knowledge-distillation work trains student networks to mimic large teacher models and emphasizes instructions or rationales alongside final answers.
- Subquestioning as supervision: Prior subquestioning research uses information-seeking or clarification questions to supervise discovery, question answering, and commonsense reasoning.
3 Methodology
The framework distills LLM-generated reasoning into smaller models by decomposing problems into ordered subquestion-solution steps. It supports unified generation or an iterative setup with separate question-generation and question-answering models.
- Problem setup: The task is to predict a final answer from problems whose solutions may include intermediate step-by-step or subquestion annotations.The framework accommodates datasets with different annotation availability.
- Teacher annotation: A large language model generates intermediate reasoning steps for problems, retaining chains whose final solution matches the ground-truth answer.Incorrect chains are discarded and resampled up to three times before constructing the augmented training set.
- Socratic CoT: Socratic CoT represents each solution as an ordered sequence of subquestion-solution pairs generated from few-shot decompositions.Each pair corresponds to one intermediate step in the overall solution.
- Student architectures: The iterative strategy separates subquestion generation from intermediate answering into QG and QA models trained and used sequentially.QG produces subquestions from the problem, while QA solves them one by one using intermediate solutions during training.
- Student architectures: The unified strategy trains one student to generate the complete sequence of subquestion-solution pairs autoregressively.Its loss is computed for each pair leading to the final answer.
- Inference: At inference, the unified model directly predicts subquestions and answers, whereas the iterative pipeline generates subquestions before decoding each answer and parsing the final prediction.The final numerical answer is extracted from the last generated solution using simple heuristics.
4 Empirical Analysis
The empirical analysis evaluates the distillation strategies across three multi-step reasoning datasets and several annotation settings. The reported results indicate that CoT-based supervision improves smaller models, with Socratic CoT providing further gains.
- Datasets: The study evaluates smaller models on GSM8K, StrategyQA, and SVAMP under varied annotation scenarios.These datasets cover mathematical word problems, factual binary questions, and settings with different intermediate-supervision availability.
- Annotation settings: GSM8K provides manually annotated step-by-step solutions and a Socratic version containing intermediate subquestion-solution pairs.The authors report that the subquestions were generated by prompting GPT-3 and reproduce them with a BERT F1 score of 95%.
- Annotation settings: StrategyQA supplies supporting facts and decomposed questions, but these annotations are not always aligned, preventing a GSM8K-like setup.The dataset requires a distinct treatment because a fact need not answer a corresponding subquestion.
- Annotation settings: For AsDiv/SVAMP, GPT-3 generates intermediate subquestions and solutions because the training data contains final answers only.The generated solutions serve as CoT, while generated subquestion-solution pairs serve Socratic CoT.
- Experimental setup: The experiments use GPT-2 variants as students and GPT-3 175B as the teacher for decomposing complex problems.Models are trained with the Huggingface library, equal iteration budgets, periodic validation, and teacher forcing.
- Evaluation: Final-answer accuracy is the evaluation metric for GSM8K, SVAMP, and StrategyQA.Performance is computed from the final answer produced by the student model.
5 Results and Discussion
Across GSM8K, StrategyQA, and SVAMP, the supervision strategy affects how effectively smaller models perform multi-step reasoning. Socratic CoT often improves over standard step supervision, while dataset structure and model size influence which strategy works best.
- Overall results: Smaller models improve when fine-tuned with reasoning supervision across all three considered datasets.The study compares Answer Only, GT Steps, GT Facts, CoT, SocCoT, SocGT, Iterative, and Unified strategies for GPT-2 models.
- Step-by-step solutions: GPT-2 Large with Socratic CoT performs as well as GPT-3 6B fine-tuned with human-annotated data.The comparison uses GPT-2 Large with 774M parameters and a nearly 10x larger model.
- Supporting facts: On StrategyQA, GPT-3-generated supporting facts hurt GPT-2 Large accuracy relative to ground-truth facts: 58.07 vs. 60.51.Using provided subquestions with LLM-generated answers instead produced 60.31 vs. 52.02 for GPT-2 Medium.
- Answers only: On SVAMP, both CoT and Socratic CoT improve performance, with Socratic CoT more effective for smaller models but CoT preferred for the larger model.This dataset contains only final answers, so LLMs generate both intermediate steps and subquestions.
- Prompting: GPT-3 prompting with subquestioning boosts GSM8K accuracy by over 40% compared with standard CoT prompting.The experiment uses a single-shot prompt in which GPT-3 first decomposes the problem and then solves the subproblems.
6 Ablation Studies
The ablation studies examine subquestion generation quality, guidance mechanisms, and whether explicit subquestions are needed at inference. Guidance improves generated-question quality, while removing the question-generation module substantially degrades performance.
- Subquestion generation: GPT-2 Large generated an incorrect number of subquestions for most GSM8K problems.The evaluation compared BLEU, BERT F1, and question-count matching against GPT-3 annotations.
- Guided generation: Conditioning question generation on predicted intermediate-solution equations improved BLEU, BERT F1, and question-count matching.The guidance model predicts equations from the original problem, and its outputs condition the question-generation model.
- Guided generation: The guidance mechanism helped generate more meaningful subquestions even when its predicted equations often did not yield the correct solution.Figure 6 reports overall GPT-2 student accuracy for answer prediction and question generation with and without guidance.
- Removing question generation: An alternative that trains one student to generate complete reasoning chains without test-time subquestions caused significant performance degradation.The approach implicitly decomposes problems by generating cumulative reasoning steps, but Table 5 compares it unfavorably with the question-generation setup on GSM8K.
- Qualitative analysis: The ablations also include qualitative GPT-2 predictions on GSM8K and SVAMP problems.These examples are reported in Figures 5 and 7.
7 Conclusion
The paper distills LLM reasoning capabilities into smaller models and augments this process by explicitly asking stepwise questions. Across three multi-step reasoning datasets, it demonstrates the methodology and identifies cases where CoT or Socratic CoT should be preferred.
- Conclusion: The proposed methods distill chain-of-thought reasoning capabilities into smaller models.The paper also explores explicit stepwise questions as a further improvement.
- Conclusion: The methodology is evaluated on GSM8K, StrategyQA, and SVAMP.These are described as popular multi-step reasoning datasets.
- Conclusion: The paper discusses when CoT or Socratic CoT should be preferred for different datasets.The conclusion frames the choice between methods as dataset-dependent.
Limitations
The limitation discussed is that distillation uses only one solution from the LLM, leaving multi-sample distillation and broader prompt comparisons for future work.
- Limitations: The distillation experiments use only one LLM-generated solution per problem.The paper notes that sampling multiple subquestion-solution pairs and majority voting could provide an alternative.
- Limitations: The CoT and Socratic CoT comparison uses a single prompt because of computational constraints.The authors suggest that using up to eight prompts might produce a fairer comparison and better results.
Ethical Considerations
The authors caution that the improved smaller models are not yet powerful enough for sensitive settings such as education and remain prone to hallucination.
- Ethical Considerations: The models are not yet powerful enough for sensitive settings such as education.The authors advise users to apply caution when using them.
- Ethical Considerations: The generative models, including those developed in this work, are prone to hallucination.This risk motivates careful use of the released code and model checkpoints.