Effectiveness of Chain-of-Thought in Distilling Reasoning Capability from Large Language Models (2025)

Abstract: Chain-of-Thought (CoT) prompting is a widely used method to improve the reasoning capability of Large Language Models (LLMs). More recently, CoT has been leveraged in Knowledge Distillation (KD) to transfer reasoning capability from a larger LLM to a smaller one. This paper examines the role of CoT in distilling the reasoning capability from larger LLMs to smaller LLMs using white-box KD, analyzing its effectiveness in improving the performance of the distilled models for various natural language reasoning and understanding tasks. We conduct white-box KD experiments using LLMs from the Qwen and Llama2 families, employing CoT data from the CoT-Collection dataset. The distilled models are then evaluated on natural language reasoning and understanding tasks from the BIG-Bench-Hard (BBH) benchmark, which presents complex challenges for smaller LLMs. Experimental results demonstrate the role of CoT in improving white-box KD effectiveness, enabling the distilled models to achieve better average performance in natural language reasoning and understanding tasks from BBH.

## introduction reasoning is the process of using logic, evidence, and knowledge to make sense of information, draw conclusions, solve problems, and make decisions. while reasoning is considered as one of the essential capabilities, it tends to emerge only in large language models (llms) at a certain scale @xcite . smaller llms, typically with fewer than tens of billions of parameters, have limited reasoning capability @xcite . improving the reasoning capability of llms, especially of smaller ones, is an important research topic with various practical applications @xcite . indeed, improving the reasoning capability for smaller llms enables real-time responsiveness, de-* correspondence: cong-thanh.do@toshiba.eu livering faster processing than larger models-an essential factor for interactive applications utilising natural language. these more efficient llms facilitate edge deployment on devices with limited resources, addressing the computational demands of larger models. furthermore, better reasoning helps smaller models understand nuanced language and maintain context throughout extended natural language interactions, even in the presence of linguistic imperfections @xcite . finally, they offer reduced operational cost for inference, making them more accessible, and empower natural language processing systems to be more proactive in problem-solving and task completion. several approaches have been developed to improve the reasoning capabilities of llms. for example, techniques like few-shot prompting @xcite and supervised fine-tuning (sft) @xcite can improve the capabilities of llms, including reasoning, in specific domains when training data are available. chain-of-thought (cot) prompting is a popular method used to elicit the reasoning capabilities of llms @xcite . cot, which involves a series of intermediate reasoning steps, significantly improves the capability of llms to perform complex reasoning tasks @xcite . knowledge distillation (kd) @xcite is a promising technique to transfer knowledge, including reasoning capability, from larger llms to smaller llms. this technique is frequently employed to condense the knowledge stored in larger neural networks into smaller counterparts, thereby reducing computational resource requirements and improving inference speed without substantial performance sacrifices. more recently, cot has been leveraged in kd to transfer multi-step reasoning capability from a larger llm to a smaller one @xcite . in @xcite , a mentor model, which is an intermediate-sized, task-specific fine-tuned model, was used to augment additional cot annotations and provide soft labels for the student model during reasoning distillation. while cot enhances complex reasoning in llms, integrating it with white-box kd-where the teacher model's internal workings are fully visible during distillation-presents opportunities and challenges that require deeper investigation. in this paper, we examine the role of cot in the distillation of reasoning capability from larger llms to smaller llms using white-box kd. we demonstrate the effectiveness of cot in improving kd performance, enabling the distilled models to achieve better average performance in various natural language reasoning and understanding tasks from the big-bench-hard benchmark @xcite . ## related work cot distillation is a popular approach which is used to distill reasoning capability from larger llms to smaller llms @xcite @xcite @xcite @xcite . the main idea in these works is to prompt a very large teacher llm, such as gpt-4 (openai, 2024), to solve complex questions via zero-shot cot prompting @xcite , and then use the reasoning samples to fine-tune a smaller student llm. this approach is applied with black-box kd @xcite , where only the texts generated by the teacher llms are accessible. white-box kd approach @xcite ) leverages full access to the teacher llm's out-put probabilities or intermediate hidden states during distillation. this approach enables the student llm to learn not just the final outputs but the underlying reasoning process-leading to improved performance on complex language understanding and reasoning tasks @xcite . with the availability of more open-source llms, whitebox kd is becoming increasingly valuable for both academia and industry, as it provides good potential to improve performance. however, most studies on white-box kd have focused on traditional small language models @xcite , and the application of cot in white-box kd has yet to be widely explored. ## kd+cot: white-box knowledge distillation with chain-of-thought in this paper, we examine the role of cot in the distillation of reasoning capability from larger llms to smaller llms using white-box kd, where the output probabilities of the teacher llm can be accessed during the distillation process @xcite . we assume that the teacher and student llms have tokenizers of the same size n and the same vocabulary, to simplify the computation of the distillation loss between probability distributions. we use kullback-leibler (kl) divergence as the only loss in the distillation loss to keep the kd step pure with only information flowing from the teacher model @xcite . the kd+cot approach for distilling reasoning capability from the teacher llm to the student llm is illustrated in fig. 1 . during the kd+cot process, training instances with rationales are forwarded through both the teacher and student llms. the output probabilities produced at the output layers of these models are used to compute the distillation loss. the error derivative of the distillation loss is then back-propagated through the student llm to update its weights and minimise the distillation loss. the key difference between kd+cot and vanilla white-box kd is that the former incorporates rationales in the training data, while the latter does not. figure 2 illustrates examples of training instances from the cot-collection dataset @xcite , showcasing integrated step-by-step reasoning rationales. incorporating these rationales into the training data for kd is expected to elicit the teacher llm's reasoning capability during training, effectively distilling it to the student llm by minimising the distillation loss between the models' output probabilities. we use two families of llms for the experiments, qwen @xcite and llama2 @xcite . these models are chosen for their proficiency in few-shot learning and their unique vocabulary coverage. the qwen-1.8b and qwen-7b models which have 1.8 and 7 billion parameters, respectively, were proposed by alibaba cloud. these models are used as the student and teacher llms, respectively, in the qwen-based kd experiments. qwen-1.8b and qwen-7b are transformer-based llms, which were pre-trained on a large volume of data, including web texts, books, codes, etc. the qwen tokenizer consists of @xmath0, 936 tokens. in llama2-based kd experiments, the llama2-13b-chat is used as the teacher model, and the llama2-7b and tinyllama are used as the student models. llama2 is an auto-regressive language model that employs an optimised transformer architecture and was trained on 2 trillions tokens of data. the llama2-13b-chat model, which consists of 13 billion parameters, is a fine-tuned version of the llama2-13b model. it utilises sft and reinforcement learning with human feedback (rlhf) to align with human preferences for helpfulness and safety. the llama2-7b model has 7 billion parameters, while the tinyllama model @xcite ) has 1.1 billion parameters and was pre-trained on 1 trillion tokens of data. tinyllama is built on the architecture and tokenizer of llama2 which consists of @xmath1, 000 tokens. figure 2 : examples of training data instances with integrated rationales that lead to the answers. these data are from the cot-collection dataset @xcite . ## conclusion this paper examined the effectiveness of cot in distilling reasoning capability from a teacher llm to a student llm using white-box kd, where the distillation loss between their output probabilities is minimised. experimental results demonstrated that cot improves the effectiveness of white-box kd, even when the vanilla white-box kd does not outperform the baseline. the results also suggest that cot elicits the teacher llm's reasoning capability, which are then effectively distilled into the student model during the distillation process. in summary, the kd+cot approach enhances the distilled models, improving their overall performance in bbh natural language reasoning and understanding tasks over the student models while maintaining their efficiency advantage over the teacher models in terms of inference speed.