Learning to Maximize Mutual Information for Chain-of-Thought Distillation (2024)

Abstract: Knowledge distillation, the technique of transferring knowledge from large, complex models to smaller ones, marks a pivotal step towards efficient AI deployment. Distilling Step-by-Step (DSS), a novel method utilizing chain-of-thought (CoT) distillation, has demonstrated promise by imbuing smaller models with the superior reasoning capabilities of their larger counterparts. In DSS, the distilled model acquires the ability to generate rationales and predict labels concurrently through a multi-task learning framework. However, DSS overlooks the intrinsic relationship between the two training tasks, leading to ineffective integration of CoT knowledge with the task of label prediction. To this end, we investigate the mutual relationship of the two tasks from Information Bottleneck perspective and formulate it as maximizing the mutual information of the representation features of the two tasks. We propose a variational approach to solve this optimization problem using a learning-based method. Our experimental results across four datasets demonstrate that our method outperforms the state-of-the-art DSS. Our findings offer insightful guidance for future research on language model distillation as well as applications involving CoT. Codes are available at https://github.com/xinchen9/cot_distillation_ACL2024.

## introduction the capabilities of larger language models (llms) tend to scale with their model size, leading to a substantial demand for memory and compute resources @xcite . distilling knowledge from larger llms to smaller llms has been crucial for the efficient deployment of ai @xcite . chain-of-thought (cot) @xcite distillation represents a pivotal advance in the quest to endow smaller language models with the sophis-ticated reasoning capabilities of their larger counterparts. by distilling complex thought processes into more compact models, this approach aims to democratize access to advanced natural language understanding and reasoning across a wider array of computational resources @xcite @xcite . . @xcite ) introduces a cot distillation method that guides smaller models using rationales from llms within a multi-task learning (mtl) framework, training them for both label prediction and rationale generation tasks. this framework simultaneously optimizes the model for two related objectives on the same input, enhancing its chain-of-thought learning by sharing representations between the two tasks, thereby improving overall performance efficiently. while dss brings out the benefits of reducing computational costs, it suffers from the same problem as the conventional mtl framework, that is the difficulty in effectively connecting the prediction and generation tasks. the intricacies inherent in training models within the mtl framework can undermine the effectiveness and reliability of the dss process @xcite . despite the successful setup of an mtl framework in dss, where the tasks of label prediction and rationale generation are intrinsically related, the current configuration may not optimally capture and maximize the mutual knowledge between these tasks. furthermore, llms are prone to producing hallucinations and inconsistent rationales, which potentially mislead the student model toward incorrect answers and cause conflicts in mtl that destruct student model learning @xcite . ## related work we present an overview of previous work across three areas related to our study: knowledge distillation, multi-task learning, and information bottleneck. knowledge distillation (kd) originally designed to train small models by leveraging the extensive knowledge of larger models @xcite , kd has since been extended to a variety of applications, owwwing to its effective transfer of knowledge across models and tasks @xcite @xcite @xcite @xcite . a crucial yet open challenge is how to effectively transfer the knowledge. to address the issue, previous studies @xcite @xcite have extracted various features and designed auxiliary loss functions to enhance kd. our work focuses on improving the model by acquiring mutual knowledge to address both label prediction and rationale generation tasks. multi-task learning (mtl) by exploiting commonalities and differences among relevant tasks, mtl can enhance learning efficiency and improve prediction accuracy by learning multiple objectives from a shared representation @xcite . in recent years, mlt has been broadly applied to nlp tasks @xcite @xcite . however, some studies have identified that training multiple tasks trained simultaneously could lead to conflicts among them, making it challenging to optimize the performance of all tasks simultaneously @xcite . recently, kd has also been applied within mtl frameworks, achieving state-of-the-art results in various applications @xcite @xcite @xcite . information bottleneck (ib) ib (tishby and zaslavsky, 2015; @xcite provides a powerful statistical tool to learn representation to preserve complex intrinsic correlation structures over high dimensional data. as a general measure of the dependence between two random variables, mi is also widely used in deep learning to effectively represent feature dependencies @xcite @xcite . estimating mi is known to be challenging, and recent progress has been made towards learning-based variational approaches @xcite @xcite @xcite . another challenge associated with the ib principle is the optimization process, which involves a trade-off between achieving concise representation and maintaining strong predictive capabilities @xcite . consequently, optimizing ib becomes a complex task that heavily depends on the problem formulation and the provision of an effective optimization solution. recent studies have applied ib to solve complex machine learning problems both in computer vision @xcite @xcite and nlp @xcite @xcite . in this paper, we formulate our cot distillation problem with mtl training pipeline using ib method, and provide a learning-based solution to optimize ib for our cot distillation, as detailed in section 3. ## methodology this section begins with an introduction to preliminaries of ib. following this, we formulate our cot distillation idea within the ib framework and propose a learning approach to optimize mi. ## conclusion in this paper, we re-investigate the dss framework from an information-theoretic perspective. we model it using information bottleneck and propose to strengthen it by maximizing the mutual information between rationale generation and label prediction tasks. the proposed learning-based method can automatically optimize the cot distillation and bolster the reasoning ability of the distilled smaller models. both our qualitative and quantitative analysis demonstrate the rationale behind