Mixed Distillation Helps Smaller Language Models Reason Better (2024)

Abstract: As large language models (LLMs) have demonstrated impressive multiple step-by-step reasoning capabilities in recent natural language processing (NLP) reasoning tasks, many studies are interested in distilling reasoning abilities into smaller language models (SLMs) via fine-tuning. Previous distillation methods usually utilize the capabilities of LLMs to generate chain-of-thought (CoT) samples to teach SLMs. However, this distillation approach performs poorly in certain scenarios due to the limitations of CoT. In this work, we introduce a novel Mixed Distillation (MD) framework, distilling multiple step-by-step reasoning abilities into SLMs. First, we leverage LLMs to generate multiple step-by-step reasoning rationales by sampling automatically. Then, we create high-quality, well-balanced mixed thought data and design a novel multi-task loss to help SLMs better learn and adaptively activate multiple step-by-step reasoning. Our extensive experiments demonstrate that MD enhances both single-path (using either CoT or PoT) and multi-path (using both CoT and PoT) reasoning abilities of SLMs during inference across reasoning tasks. Notably, a single model generated by MD exceeds the comprehensive performance of an ensemble of two individual CoT and PoT distilled models. Mistral-7B using MD can achieve remarkable improvements of 87.5%, 74.0% and 77.1% on SVAMP, GSM8K and ASDIV, respectively, outperforming the teacher model, GPT-3.5-Turbo. We hope our work provides insight into SLMs’ multiple step-by-step reasoning abilities.

## introduction recent llms @xcite @xcite such as gpt-4 have demonstrated multiple step-by-step reasoning capabilities @xcite @xcite with chain-of-thought (cot) and program-of-thought (pot) previously unseen in slms. cot boosts reasoning by guiding llms to produce intermediate natural language steps @xcite , while pot stimulates reasoning by generating intermediate program code that can be executed by the python executor @xcite . however, deploying these advanced llms in real-world applications presents significant costs and computational demands @xcite . to address these challenges, distilling step-by-step reasoning capabilities from llms emerges as a resource-friendly and effective strategy. on the other hand, llms with multiple stepby-step reasoning can address distinct challenges. as shown in figure 1 , llms can solve the majority of problems (p ∧ c, ∼ p ∧ c, p ∧ ∼ c) through step-by-step reasoning. among them, a large part of the problems can be solved by both @xcite pot: def solution(): cot: paige and her mom planted 10 seeds in each flower bed. there are 45 flower beds in total. so they planted 10x45 = 450 seeds. the answer is 450. question: paige was helping her mom plant flowers and together they planted some seeds. they put 10 seeds in each flower bed. if there are 45 flowerbeds. how many seeds did they plant? cot (p ∧ c). moreover, each reasoning ability has its own advantages. pot excels in large-scale and intensive numerical calculations @xcite , whereas cot demonstrates proficiency in language comprehension, geometry, algebra, fraction calculations, and equation comparisons @xcite . specifically, only about 6% of the problems can be solved exclusively via cot (∼ p ∧c) across tasks. in contrast, pot exclusively addresses 31.98% of problems on gsm8k, and this figure exceeds 10% for other tasks (p ∧ ∼ c). we show more details in appendix d.1. meanwhile, compared with llms, slms struggle with generating effective intermediate steps due to limited knowledge, which makes their reasoning challenging @xcite @xcite . existing works focus on cot distillation, a method that utilizes the cot rationales of llms as supervision for training slms, which excel in generating natural language content @xcite @xcite @xcite @xcite . however, the sample data generated by llms is only based on a single reasoning path, cot, and slms can't learn pot and activate multiple stepby-step reasoning making it challenging for slms to develop robust reasoning abilities @xcite . in addition, previous works @xcite @xcite have proved that multi-task learning enhances model performance by involving various knowledge domains. therefore, we propose a novel distillation framework, mixed distillation (md), which combines multiple prompting techniques to create highquality, well-balanced mixed thought data within a novel multi-task learning approach, as shown in figure 2 . specifically, we utilize multiple prompting techniques with sampling to prompt llms to generate multiple step-by-step reasoning paths for each input, resulting in a collection of thought data. to help slms learn and adaptively activate multiple step-by-step reasoning capabilities, we further conduct quality control and categorization to create mixed thought data. then, we present a novel multi-task loss to fine-tuning slms, using consensus via majority vote to obtain answers with multiple prompting techniques. this approach allows slms to activate different step-by-step reasoning capabilities for different problems. we conduct extensive experiments on svamp, gsm8k, asdiv, and strategyqa reasoning tasks to validate the efficacy of md. as shown in figure 3 , the overall performance of slms is better than that of models using other distillation methods. notably, on svamp, llama2-7b achieves a significant 15% improvement over cot distillation and records an impressive 5% increase compared to pot distillation. the contributions of our work are as follows: ## related work multiple thoughts prompting techniques in llm recent work @xcite @xcite @xcite , focusing on eliciting the step-by-step reasoning process of llms, has validated its effectiveness in reasoning domains, such as svamp @xcite , gsm8k @xcite , asdiv @xcite and strat-egyqa @xcite . cot @xcite enhances reasoning by prompting llms to generate intermediate natural language thought steps. pot @xcite stimulates llm's reasoning ability by prompting them to generate intermediate code that can be executed by the python executor. zhang et al. (2023); li et al. (2023a) have shown the adaptability of combining natural language reasoning and program synthesis within prompt-based learning to effectively solve reasoning tasks. yue et al. (2023b) improve the mathematical reasoning ability of llms combining pot and cot by designing prompts. yue et al. (2023a) develop a cost-effective approach by combining pot and cot. however, previous work (yue et al., 2023b,a; imani et al., 2023) provided limited insights into pot and cot in improving slms by distilling llms. our md presents a novel multitask loss to help slms learn and adaptively activate both reasoning capabilities on specific tasks. knowledge distillation from llms knowledge distillation @xcite @xcite @xcite has proved its effectiveness in improving slms. some works @xcite @xcite @xcite leverage generative cot as a supervisory signal to fine-tune smaller task-specific models. however, previous works @xcite al., 2023; zhu et al., 2024) focus solely on llm-generated single reasoning path as the supervisory signal, without considering the capability of cot and pot in slms. our proposed md emphasizes the importance of pot with cot as supervisory signals and the slms can effectively learn and activate multiple step-bystep reasoning from high-quality and well-balanced mixed thought data. ## experiments in this section, we first prove that pot, as a supervisory signal, enhances the slms' reasoning capabilities (sec. 4.1). moreover, our findings emphasize the benefits of md, which enhances slms' capabilities in single-path reasoning and multiple-path reasoning, enabling slms to learn and adaptively activate step-by-step reasoning (sec. 4.2). we further conduct extensive experiments based on llama2-7b, codellama-7b, mistral-7b, t5-large, and gpt2-large, compared to other distillation methods (sec. 4.3). finally, we validate the generalizability of md (sec. 4.4). datasets our experiments primarily center on the following datasets: svamp @xcite , gsm8k @xcite , and asdiv @xcite . we extend our assessment to strate-gyqa @xcite , where we evaluate the capability of commonsense reasoning. more dataset details are provided in appendix a. baselines we evaluate md by comparing it with experiments using closed-source models (openai et al., 2023), open-source models (touvron et al., method svamp asdiv closed-source models gpt-3.5-turbo 82.0 75.8 open-source models llama2 (touvron et al., 2023) † 38.0 50.7 codellama (roziere et al., 2023) † 59.0 61.4 wizardmath (luo et al., 2023) † 57.3 59.1 single-path distillation llama2-7b + cot 64.5 (↑26.5) 63.2 (↑12.5) llama2-7b + pot 56.5 (↑18.5) 64.2 (↑13.5) ensemble single-path distillation llama-7b + cot w/ pot * 61.5 (↑23.5) 65.9 (↑15.2) mixed distillation (ours) llama2-7b-md + cot 65.0 (↑27.0) 64.2 (↑13.5) + pot 70.0 (↑32.0) 64.2 (↑13.5) + cot w/ pot * 74.5 (↑36.5) 68.9 (↑18.2) ## case study slms with md can better master two types of step-by-step reasoning abilities. in single-path reasoning, it mitigates the shortcomings of cot in solving large-scale and intensive numerical calculations, while also easing language comprehension challenges caused by insufficient knowledge for pot. we present the actual output of slms for reasoning tasks, as shown in figure 6 . for the question labeled 35624, it is clear that correct answers are obtained using md involving cot or pot. however, when using cot or pot distillation alone, errors occur in the reasoning process. specifically, the error in cot distillation is due to the inability to effectively compute 244*146, a common issue indicative of poor computational capability in cot. we show more details in appendix d.2. 6 conclusion in this paper, we introduce a novel framework md that distills the reasoning paths of cot and pot from llms into slms. our experimental results demonstrate that md enhances the slms' singlepath reasoning and multi-path reasoning, enabling slms to learn and adaptively activate step-by-step reasoning. comparative analysis and experimental results show that our md can effectively extract two different forms of capabilities, cot and pot from llm, to improve the reasoning ability of slms.