Investigating Mysteries of C o T -Augmented Distillation (2024)

Abstract: Eliciting chain of thought (CoT) rationales - sequences of token that convey a “reasoning” process has been shown to consistently improve LLM performance on tasks like question answering. More recent efforts have shown that such rationales can also be used for model distillation: Including CoT sequences (elicited from a large “teacher” model) in addition to target labels when fine-tuning a small student model yields (often substantial) improvements. In this work we ask: Why and how does this additional training signal help in model distillation? We perform ablations to interrogate this, and report some potentially surprising results. Specifically: (1) Placing CoT sequences after labels (rather than before) realizes consistently better downstream performance – this means that no student “reasoning” is necessary at test time to realize gains. (2) When rationales are appended in this way, they need not be coherent reasoning sequences to yield improvements; performance increases are robust to permutations of CoT tokens, for example. In fact, (3) a small number of key tokens are sufficient to achieve improvements equivalent to those observed when full rationales are used in model distillation.

## introduction chain of thought (cot) reasoning-i.e., generating tokens which communicate step-by-step "thinking"-can (sometimes dramatically) improve model performance on reasoning tasks @xcite . in the context of model distillation @xcite , recent work has elicited such rationale chains from massive llms (e.g., ) to augment data with which to fine-tune much smaller (<2b parameters) task-specific models. figure 1 illustrates this distillation approach: the student model is trained to generate the rationales in addition to the target token(s). this simple cot-augmented distillation strategy consistently and sometimes dramatically improves the performance of student models @xcite . for example, @xcite used rationales from gpt-3 (175b) to teach a comparatively tiny student lm (opt-1.5b) to produce similar "reasoning" token sequences at inference time. they show an average increase in task accuracy of 12.4% across three commonsense reasoning datasets. @xcite adopted a similar approach to finetune gpt-2 (large; 774m) on grade-school math datasets with improvements of 8.23% on gsm8k and 16.20% on svamp. beyond commonsense reasoning, @xcite achieved sota results (+6.23 absolute gain in micro-f1, on average) with a distilled model for relation extraction by exploiting cot rationales. in this work we ask: why does distillation with cot augmented targets consistently improve the performance of distilled lms? one might naively suspect that the student model benefits from learning to mimic the relevant "reasoning" process. but we find that it is not the case that student models benefit from "reasoning" at inference time. rather, consistent with contemporaneous work @xcite , we observe that placing cot sequences after target tokens for distillation actually improves student performance (compared to when cot is pre-prended to labels). this means the student model need not bother generating its "reasoning" at test time, as the label will be generated ahead of this anyway. further, we find that rationale grammatically is not necessary; one can shuffle rationale tokens and/or include only "important" tokens from chains of thought during distillation and still realize performance benefits equivalent to those observed when using the full rationales. through ablations with three small student lms (gpt-2, phi-1.5, and gemma-2b), we report the following, sometimes counter-intuitive, findings regarding cot-augmented distillation and how rationales benefit student models. we summarize our key findings as ice is likely easier to glide over, thereby reducing the friction, hence the answer is b. ## experimental design cot-augmented distillation entails eliciting rationales from a large teacher model and using these as additional training signal for a small student model. rationales here comprise the logical steps taken to reach a response from a given input. 1 these are inserted into distillation training targets, and the student model is in this way taught to generate reasoning in addition to labels. this has been shown empirically to provide (sometimes dramatic) performance benefits @xcite . but why? what accounts for the success of cot-augmented distillation? in this work we investigate the following questions about the role of cot-rationales in distillation. (rq1) does the placement of the reasoning chain relative to the target label (pre-or post-) matter? relatedly, might observed performance gains owe to simply allowing the student model additional compute during inference? (rq2) must rationales feature logical and coherent "chain-of-thought" reasoning, or could we, e.g., scramble the ordering of tokens and still observe improvements? finally, (rq3) could we realize the same benefits in distillation using only a handful of key tokens from rationales, rather complete reasoning sequences? to answer these questions empirically, we establish baseline student lm performance, and then compare this to ablated variants of cot-augmented models. we use a fixed icl prompt (appendix b) with the input and target label to elicit a possible rationale for each instance in a dataset. we use mistral-7b-instruct @xcite as the teacher model and gpt-2 @xcite , gemma-2b @xcite and phi-1.5 @xcite as student models. note that one could instead replace mistral-7b-instruct with gpt-4 (or any other llm capable of generating cot-style rationales in icl settings) as the teacher model. see appendix b for the prompt used to elicit rationales for training instances of all datasets used in our work. following prior related work @xcite , we select three commonsense reasoning datasets: commonsenseqa @xcite , openbookqa @xcite , and quarel @xcite . each dataset provides an input consisting of a question, and a predefined set of answer choices. the target labels are the correct answer choices (appendix a). with a learning rate of 3e-5, batch size of 4 for com-monsenseqa and openbookqa, and 8 for quarel, with a maximum input length of 512, maximum output length of 256. we evaluated checkpoints every 500 steps with early stopping (@xmath0, @xmath1.02). because we are only interested in measuring relative performance of fine-tuned models across ablations (as opposed to necessarily realizing sota performance), we left the remaining hyperparamters to their default values. ## conclusions we have investigated why and under what circumstances does cot-augmented distillation improve student model performance. specifically, we evaluated the degree to which the following aspects contribute to the observed gains realized in cotaugmented distillation. some of these findings corroborate and deepen observations made in contemporaneous work, e.g., models can benefit from additional compute at inference time @xcite , and cotaugmentation fares best when rationales are placed after the target labels @xcite . we have not fully characterized the mechanism by which cot augmentation aids distillation, but we have ruled out some explanations and provided empirical insights into when and how cot augmentation provides useful signal to student models.