Distilling Reasoning Capabilities into Smaller Language Models (2023)

Abstract: 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: https://github.com/kumar-shridhar/Distiiling-LM .

## introduction large language models (llms) have demonstrated strong performance on a variety of reasoning tasks @xcite @xcite . one particularly interesting strategy for prompting these models is chainof-thought (cot), which has been shown to elicit reasoning abilities in llms by asking the model to incorporate intermediate reasoning steps while solving a problem @xcite ## methodology the setting we consider consists of a data set d, where each problem p i is accompanied by a final answer a i that can be reached by several steps of reasoning. the task of solving the problem using a model ψ is to predict an answer â = ψ(p ) such that â = a. we consider different data scenarios where intermediate annotations of the solution may be available in different forms (e.g., step-by-step, as a semantic decomposition by subquestions) or may not be present. depending on the availability of annotations, we propose different approaches to augment the training of a small model on d by a robe takes 2 bolts of blue fiber and half that much white fiber. how many bolts in total does it take? ## results and discussion can our framework improve the reasoning capabilities of smaller models? ## ablation studies in this section, we describe additional analyses regarding specific components of the framework we propose, as well as negative results that we obtained with alternative strategies. how good are the sub-questioning capabilities of a smaller model? we investigate in more detail the ability of a small model to decompose a problem by generating meaningful subquestions. we fine-tuned gpt-2 large on the gpt-3 generated subquestions provided in the gsm8k dataset. we then evaluated the quality of the generated questions in terms of bleu score @xcite , bert f 1 score @xcite , and by measuring for how many problems the number of questions generated by gpt-2 (#q) matches the number of gpt-3 annotated questions for a given problem. we found that the fine-tuned gpt-2 predicted an incorrect number of subquestions for the majority of problems (see eliminating the need for a subquestion module. we have experimented with an alternative training solution that does not involve a question-generation model. this strategy aims to improve the supervision for fine-tuning a small model through subquestioning, but without relying on the presence of subquestions at test time. the procedure consists of training the student model to generate the entire chain of steps leading to an intermediate answer. that is, when the sub-question q (1) is asked, the model is trained to generate the answer s (1) , but when q (j) is asked, the model is trained to generate the chain of thought reasoning {s (1) , s (2) , . . . , s (j) } (instead of just s (j) ). this eliminates the need for the intermediate subquestions at inference time, as the model is trained to implicitly decompose the main problem into smaller reasoning steps. however, this method gpt-2 no subq subq with qg small 2.70 5.98 medium 7.20 11.57 large 8.18 17.89 : faye was placing her pencils into rows with number0 pencils in each row. she had number1 packs of pencils each one having number2 pencils. how many rows could she make? gt solution: <</ + number0 number1 @xmath0 ## conclusion the chain-of-thought style of step-by-step reasoning has proven to be very effective for reasoning in llms. in this work, we propose ways to distill these reasoning capabilities into smaller models and suggest ways to further improve them by explicitly asking stepwise questions. we demonstrate the effectiveness of our proposed methodology on three popular multi-step reasoning datasets, and discuss cases where one method should be preferred over the other for different datasets.