Teaching Small Language Models Reasoning through Counterfactual Distillation (2024)

Abstract: With the rise of large language models (LLMs), many studies are interested in transferring the reasoning capabilities of LLMs to small language models (SLMs). Previous distillation methods usually utilize the capabilities of LLMs to generate chain-of-thought (CoT) samples and teach SLMs via fine-tuning. However, such a standard distillation approach performs poorly when applied to out-of-distribution (OOD) examples, and the diversity of the generated CoT samples is insufficient. In this work, we propose a novel counterfactual distillation framework. Firstly, we leverage LLMs to automatically generate high-quality counterfactual data. Given an input text example, our method generates a counterfactual example that is very similar to the original input, but its task label has been changed to the desired one. Then, we utilize multi-view CoT to enhance the diversity of reasoning samples. Experiments on four NLP benchmarks show that our approach enhances the reasoning capabilities of SLMs and is more robust to OOD data. We also conduct extensive ablations and sample studies to understand the reasoning capabilities of SLMs.

## introduction large language models (llms) have demonstrated remarkable performance in a wide range of downstream tasks @xcite @xcite . recent works have shown that chain-of-thought (cot) can elicit reasoning capabilities in llms by asking the model to incorporate intermediate reasoning steps while solving a problem @xcite . however, the efficacy of prompt-based cot methods is restricted to very large models (beyond 10b parameters) @xcite . due to the substantial computational resources or expensive api calls required for accessing llms that support cot, various studies have delved into distilling the * corresponding author: yin zhang. the weasel was becoming a problem, it kept getting into the chicken eggs kept in the what? question a:forest b:barn c:public office d:freezer weasels are known for their small size and agility, making them adept at entering buildings or enclosures where they can find food. in this context, the mention of a weasel becoming a problem by getting into where the eggs were kept strongly suggests a scenario where the eggs are stored in a barn. if the weasel kept getting into the chicken eggs, it would likely be in the freezer. chicken eggs are often stored in the freezer to preserve their freshness and prevent them from spoiling. reasoning ability of llms into slms @xcite . existing works focus on cot distillation, a method that utilizes the cot rationales of llms as supervision for training slms, rather than just labels @xcite @xcite . the slms improve their reasoning capabilities by imitating the reasoning process of the llms. although the cot distillation method has been proven to be effective, it has the following weaknesses: in figure 1 , the rationales generated by the slm indicate that it has learned the correlation between "egg kept" and "freezer" through cot distillation. however, during inference, it completely ignores the context and directly links these two phrases together, lacking a deep understanding of the text semantics. to address these issues, we propose to enhance the standard cot distillation process from two aspects respectively. in order to mitigate the reliance of slms on spurious correlations @xcite @xcite , we propose counterfactual data augmentation to encourage slms to better learn the causal structure of the text. specifically, we utilize llms to edit task instances by adjusting causal relationships in the instances or modifying text segments that are crucial for label assignments. first, we use offthe-shelf language processing tools and llms to obtain the parts of the text that need to be edited. then, prompt engineering and in-context learning are applied with llms to generate and filter counterfactual data, which closely resembles the original input but with altered answers. additionally, we employ the multi-view cot to introduce more diversified knowledge for slms, including both positive view cot (pvc) and negative view cot (nvc). pvc denotes the standard cot generated by llms when answering questions, while nvc requires llms to generate refutational rationales for each incorrect option, explaining why the current option is considered erroneous. we conduct experiments on four question answering tasks that require knowledge-intensive reasoning. experiments demonstrate that: (1) both counterfactual data augmentation and multi-view cot are beneficial to improving model performance. (2) on the same distribution dataset, the proposed approach significantly outperforms stan-dard cot distillation, with an average improvement of 11.43%. (3) across various parameter scales (ranging from 120m to 770m) and model structures (from decoder-only to encoder-decoder) for small models, our method consistently shows improved performance. (4) compared to the baseline model, our approach demonstrates robust generalization on ood data. furthermore, extensive experiments indicate that our method enhances the reasoning capabilities of slms. ## method the core idea of our method is to disrupt the causal relationship in the original text and reconstruct it using llm. since the answers to the newly generated questions change, slm is able to learn the causal differences between similar texts through multi-view cot distillation, thereby enhancing its reasoning capabilities. the overview of our approach is illustrated in figure 2 . in this section, we elaborate on our method and discuss the motivation behind it. ## experiments we study how small models can learn to reason better on four multi-step reasoning datasets: common-senseqa (csqa) @xcite , quarel @xcite , arc @xcite and qasc @xcite . the arc dataset is divided into two parts: the challenge set and the easy set. in our experiments, we combine these two subsets for model training. since the test labels of csqa and qasc datasets are not public, we use the official development set as our test set. in this paper, we utilize the gpt-3.5-turbo api foot_0 to generate reasoning paths. in the experiment, gpt-2 @xcite , opt @xcite , and gpt-neo @xcite are used as small models, and the batch size is set to 64 during training, with a total of 20 epochs. we use huggingface transformers and pytorch for the implementation. ## conclusion in this paper, we propose a distillation approach based on counterfactual data augmentation and multi-view cot, aiming to enable llms to instruct small language models in reasoning. the experimental results demonstrate that our method outperforms baseline models across various scales and architectures of small models, and exhibits strong generalization and robustness on ood data. moreover, our method shows broad application potential across a variety of tasks.