Improve Student’s Reasoning Generalizability through Cascading Decomposed C o T s Distillation (2024)

Abstract: Large language models (LLMs) exhibit enhanced reasoning at larger scales, driving efforts to distill these capabilities into smaller models via teacher-student learning.Previous works simply fine-tune student models on teachers’ generated Chain-of-Thoughts (CoTs) data. Although these methods enhance in-domain (IND) reasoning performance, they struggle to generalize to out-of-domain (OOD) tasks.We believe that the widespread spurious correlations between questions and answers may lead the model to preset a specific answer which restricts the diversity and generalizability of its reasoning process.In this paper, we propose Cas cading Decomposed Co Ts D istillation (CasCoD) to address these issues by decomposing the traditional single-step learning process into two cascaded learning steps. Specifically, by restructuring the training objectives—removing the answer from outputs and concatenating the question with the rationale as input—CasCoD’s two-step learning process ensures that students focus on learning rationales without interference from the preset answers, thus improving reasoning generalizability. Extensive experiments demonstrate the effectiveness of CasCoD on both IND and OOD benchmark reasoning datasets

## introduction recent developments in llms have brought remarkable improvements in reasoning tasks via cot prompting @xcite . however, these great reasoning capabilities are often associated with more parameters @xcite , which is not practical to emergent in smaller language models (slms). existing works @xcite @xcite try to make the reasoning capabilities isolated and distilled to student slms by simply fine-tuning figure 1 : (a) empirical results of standard cot distillation (std-cot) and directly fine-tuning on answer labels without cots (answer sft) on one in-domain @xcite and the other four out-of-domain benchmark reasoning datasets. (b) in the given example, the semantic similarity between "swimsuit" in the question and "swim" in the answer demonstrates a high level of match, which could allow the model to predict the answer using simple keyword matching or certain rules. on teacher llms generated cots data, known as standard cots distillation @xcite . although the method effectively leverages the llms' cots to boost the reasoning performance of student models on seen tasks, it does not ensure effective reasoning in ood settings, leading to weak generalization on unseen tasks. our pioneer study demonstrates that, as shown in figure 1 (a), when using the same ind training dataset, student models developed via the method std-cot perform better on ind tasks but significantly worse on ood tasks compared to models fine-tuned directly with question-answer pairs. the surprising findings indicate that students' cots do not effectively transfer to new domains and these slms seem to be more adept at learning to predict answers directly from questions. we attribute these issues to the spurious correlations between questions and answers that are commonly found in implicit reasoning tasks @xcite @xcite , as illustrated in figure 1 (b ). the std-cot approach requires models to learn both the rationale and the answer in a single step, where the learned spurious correlations in training stage can adversely affect the quality of rationale generation during inference. that is to say, upon reading a question, student models may fastly, unconsciously, and automatically formulate a "preset answer" @xcite , which in turn may lead them to implicitly reduce the token generation space when producing cot. this results in diminished diversity and generalizability of their rationales. in this paper, we propose cascading decomposed cots distillation (cascod), a straightforward yet effective method to address these issues. specifically, we decompose the traditional singlestep learning process of std-cot into two cascaded learning steps: a rationale learning step and an answer learning step. in the rationale learning step, the training objective, with the answer removed, is defined as 2 : @xmath0. in the answer learning step, we concatenate the question with the target output from the rationale learning step and use this combined input for the answer learning step, setting the training objective as q, @xmath1. this cascading two-step learning configuration mitigates the capture of spurious correlations between questions and answers during the training phase, ensuring that students focus on learning rationales without interference from the preset answers. furthermore, the inference phase execution pipeline is aligned with the training phase; the model first generates a rationale when given a question, and then, based on the question-rationale pair, predicts the final answer, further alleviating potential reasoning biases caused by spurious correlations. extensive experiments demonstrate that cascod outperforms the baselines on both ind and ood benchmark reasoning datasets ( §4.3). besides, we validate the generalizability of cascod across different models, model sizes, and training data sizes 2 q: the question, r: the rationale, a: the answers. ## related works cot capability of language models. llms have demonstrated a wide array of capabilities in numerous natural language processing tasks, underscored by various studies @xcite . one notable manifestation of this is the cot prompting technique @xcite , which facilitates models in articulating a series of deductive reasoning steps. this method has substantially enhanced llms' problem-solving abilities, as evidenced in several works @xcite @xcite . despite these advancements, the effectiveness of cot prompting notably diminishes in smaller models @xcite . research by chung et al. (2022) indicates that with targeted training on cot data via instruction tuning, slms can unlock cot capabilities. in our study, we show that slms' cot performance can be further enhanced by decomposing the standard cot distillation process into two cascaded learning steps. distilling knowledge from llms. numerous studies @xcite @xcite have explored the knowledge distillation from advanced llms like chatgpt @xcite . these efforts typically focus on student small language model ...how many ...? we start ... let's add ... 1 + 2= 3 . theref ore, the is 3. [bos] answer ...how many ...? we start ... let's add ... the is 3. [eos] answer ignored output single-step learning loss student small language model ...how many ...? ra on ale: we ... let's add ... 1 + 2 = 3 . [bos] ...how many ...? we ... let's add ... 1 + 2 = 3 . [eos] ignored output ra onale learning step loss ques on chain-of-thoughts ra onale student small language model ...how many ...? ... [bos] ...how many ...? ... [eos] ignored output answer learning step loss theref ore, the is 3. answer 3. answer delimiter standard cots dis lla on cascading decomposed cots dis lla on decompose cascade ra on ale: ra on ale: ra on ale: theref ore, 1 + 2= 3 . theref ore, the is answer distilling a broad range of abilities via instruction tuning on extensive and varied datasets @xcite @xcite . however, our work is aimed at distilling the cot reasoning capabilities from llms, in line with magister et al. (2023); ho et al. (2023), who propose standard cots distillation by directly finetuning slms on teacher llms' cots. fu et al. (2023) expands on this by using various reasoning data formats for specializing domain-specific slms. wang et al. (2023c) distill slms via learning from self-reflection and feedback in an interactive, multi-round paradigm with teacher llms. hsieh et al. (2023) propose to learn the rationale and answers as separate goals for optimizing. li et al. (2022) propose learning the entire cots and the single answers to enhance the reasoning of students. based on these, liu et al. ( 2023 ) introduce an additional distillation objective, self-evaluation, aiming for slms to assess the accuracy of their cots akin to llms' evaluative processes. and chen et al. (2024) maximize the mutual information between objectives. @xcite use in-family and out-family teachers to generate more cots for fine-tuning students. however, these methods are affected by the spurious correlations due to their isolated optimization objectives. in contrast, we reorganize the training objectives to effectively mitigate this issue. ## methodology in this section, we introduce our new distillation method that decomposes the single-step leanring process of standard cots distillation into two cascaded learning steps, as illustrated in figure 2 . formally, the standard cots distillation objective @xmath2, rationale step learning with the objective @xmath3, @xmath4. below we first describe how to extract cots from teacher llms in §3.1. then we describe the standard cots distillation method and discuss its limitations in §3.2. finally, we provide a detailed presentation of our method in §3.3. ## experiments in this section, we conduct extensive experiments and analysis to evaluate the effectiveness of our method across both ind and ood datasets. ## conclusion we propose a simple yet effective cots distillation method cascod to address the issue of questionanswer spurious correlations that previous cots distillation methods suffer from. specifically, we decompose the traditional single-step learning process into two cascaded learning steps and restructure their training objectives. extensive experiments show that cascod significantly outperforms the baselines across both ind and ood datasets. further analysis reveals that cascod is robust to model size, training data size, different models, and weights and can lead to faithful student models.