Teaching Small Language Models to Reason (2023)

Abstract: Chain of thought prompting successfully improves the reasoning capabilities of large language models, achieving state of the art results on a range of datasets. However, these reasoning capabilities only appear to emerge in models with at least tens of billions of parameters. In this paper, we explore the transfer of such reasoning capabilities to smaller models via knowledge distillation, also investigating model and dataset size trade-off. Specifically, we finetune a student model on the chain of thought outputs generated by a larger teacher model. Our experiments show that the proposed method improves task performance across arithmetic, commonsense and symbolic reasoning datasets. For example, the accuracy of T5 XXL on GSM8K improves from 8.11% to 21.99% and 18.42% when finetuned on PaLM 540B and GPT-3 175B generated chains of thought, respectively.

## introduction chain of thought (cot) prompting encourages language models (lms) to break down a reasoning task into a series of intermediate steps @xcite . they demonstrate that this prompting significantly increases the task accuracy of large language models (llms) across commonsense, symbolic and mathematical reasoning datasets. here, llms are models with at least tens of billions of parameters, such as palm 540b @xcite , gpt-3 175b @xcite , or ul2 20b @xcite . however, the reasoning capabilities of smaller lms do not improve with cot prompting, mostly producing illogical cot. notably, cot prompting even reduces the accuracy of models with less than 10 billion parameters. @xcite attribute this to abilities, such as semantic understanding and symbolic mapping, only emerging at larger scales. this leads us to our re-search question: can the reasoning capabilities of llms be transferred to smaller lms via finetuning? this work explores cot knowledge distillation @xcite from palm 540b @xcite and gpt-3 175b @xcite to different sizes of the smaller language model t5 @xcite , such as t5 xxl, xl and base, which have 11 billion, 3 billion and 220 million parameters, respectively. as a result of our work, we make two recommendations: (1) perform knowledge distillation by finetuning the student model on the cot generated by a large teacher model; and (2) generate the cot from an llm, as proposed by @xcite , but crucially provide the solution to the task in the few-shot prompt. we demonstrate that the proposed method improves task performance across arithmetic, commonsense and symbolic reasoning datasets irrespective of the teacher model used. for example, we show an accuracy increase from 8.11% to 21.99% and 18.42% on the gsm8k @xcite dataset when finetuning t5 xxl on palm 540b and gpt-3 175b generated cot data, respectively. ## related work this work is inspired by the seminal work of @xcite on cot prompting. they demonstrate that prefixing an input with 2-8 exemplars of cot reasoning encourages lms to do the same, reaching state-of-the-art performance on datasets such as gsm8k @xcite . @xcite show that task accuracy can be further improved by using self-consistency in cot prompting. selfconsistency samples cot reasoning paths from a model's decoder and returns the most consistent path by taking the majority vote. subsequently, chung et al. (2022) explore finetuning a flanbased @xcite version of palm on manually generated cot data. concurrent to our work, a small number of other works propose methods focused on cot student-teacher knowledge distillation. @xcite and @xcite also explore knowledge distillation with the difference of proposing diverse sampling and rationalization prompting, respectively. in contrast to their work, our work explores more teacher models and demonstrates both the effects of dataset and model size on accuracy. we also achieve a higher accuracy on common datasets, such as gsm8k, than @xcite . in contrast to our work, @xcite focus on training two models, one for problem decomposition and one for solving. yet differently, the focus of @xcite relies on producing markupand-mask explanations for open-book question answering. lastly, @xcite present one related experiment, however, we present a more indepth exploration on more datasets. to the best of our knowledge, our work is the first to extensively explore the improvement of the reasoning ability of small lms via knowledge distillation across multiple model architectures, and observing the effects of student model size and dataset size on accuracy. ## method we propose a two-step pipeline for cot knowledge distillation. the first step comprises annotating an existing supervised dataset with cot reasoning generated by a teacher model. to generate high quality data, we propose using llms, such as palm 540b or gpt-3 175b, as teachers, based on the finding that cot reasoning improves with model scale @xcite . specifically, we perform few-shot prompting with 8 exemplars on these models to generate cots. however, we make a key modification to the prompts proposed by @xcite . we adapt the few-shot prompts to provide the model with the target after posing the question and before providing example cot. this is based on the observation that providing this guidance allows llms to correct small mistakes in the cot. lastly, we remove all incorrect cot based on the target answer to prevent the student to learn from bad examples. the second step comprises finetuning a student model via teacher forcing @xcite ). the student is provided with the question as input, and the cot and answer as the target. as the model is trained on producing a cot during finetuning, prompting is not required. ## experimental setup we follow a similar experimental setup to @xcite , focusing on tasks covering arithmetic, commonsense and symbolic reasoning. ## discussion we demonstrate that finetuning larger lms on the cot data generated by llms of over 100 billion parameters can significantly improve task accuracy. even a small number of cot examples appear to suffice for this. however, such improvements appear to be task dependent. for example, the effects are limited for the strategyqa dataset, which can be attributed to the task requiring specific factual knowledge, which smaller lms may not have memorised due to their limited capacity. nevertheless, there is some performance improvement, which may be attributed to the model learning how to approach such tasks. moreover, the cot knowledge distillation pipeline presented allows to trade-off model and dataset size with accuracy. future work could explore improving the reasoning of small models in multi-task settings, as well as the generation of new training data using llms, rather than annotating existing datasets. ## conclusion this work explores cot knowledge distillation from llms of over 100 billion parameters to smaller lms. we propose a knowledge distillation pipeline consisting of two keys steps: (1) generate cot for existing datasets using llms and