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

heading: 1 Introduction
text: The capabilities of larger language models (LLMs) tend to scale with their model size, leading to a substantial demand for memory and compute resources (Chowdhery et al., 2023; Wei et al., 2022a). Distilling knowledge from larger LLMs to smaller LLMs has been crucial for the efficient deployment of AI (Hinton et al., 2015; Phuong and Lampert, 2019). Chain-of-Thought (CoT) (Wei et al., 2022b) 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 (Ma et al., 2023; Magister et al., 2023; Li et al., 2023). .
Distilling Step-by-Step (DSS) (Hsieh et al., 2023) 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 (Wang et al., 2023b). 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 (Mueller et al., 2022).
To address this issue, we model the DSS using information bottleneck and investigate it from an information-theoretic viewpoint (Tishby and Zaslavsky, 2015). Subsequently, we formulate the DSS as an optimization problem to maximize mutual information (MI) of label prediction and ra-
6857
tionale generation tasks. However, estimating MI from finite data has historically been a difficult problem in both deep learning and information theory (McAllester and Stratos, 2020; Belghazi et al., 2018; Paninski, 2003).
In this study, we introduce a variational method to estimate the MI. We propose a practical yet effective auxiliary loss to quantify the shared information between the prediction and the generation tasks, thereby enhancing the alignment between the two tasks and facilitating the knowledge transfer from CoT. We conduct comprehensive experiments with two smaller types of T5 models (Raffel et al., 2020), T5-base (220M) and T5-small (60M), on four popular datasets. Furthermore, we provide detailed analysis in Section 5. Our main contributions are summarized below: • We reframe the MTL framework of DSS as a
MI estimation challenge, aiming to maximize the MI between label prediction and rationale generation tasks. To achieve this, we introduce a variational approach grounded in the IB principle for effective MI estimation. To the best of our knowledge, we present the first work of improving CoT distillation from an IB perspective. • Beyond establishing a theoretical foundation, we present a practical approach for MI estimation, incorporating a simple yet effective auxiliary loss to learning to maximize MI and enhance DSS. • Our methodology demonstrably outperforms existing benchmarks across multiple datasets, evidencing the efficacy of our approach in enhancing the reasoning capabilities of distilled models. • We conduct a systematic review of the relationship between label prediction and rationale gen-
eration tasks under MTL training, presenting both qualitative and quantitative analysis results. Armed with theoretical proofs and experimental results, we aim to lay the groundwork for future research on enhancing CoT distillation within an effective MTL framework, guided by principles from information theory.

heading: 2 Related Work
text: 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 (Hinton et al., 2015), KD has since been extended to a variety of applications, owwwing to its effective transfer of knowledge across models and tasks (Chen et al., 2021; Wang and Yoon, 2021; Sanh et al., 2019; Jiao et al., 2020; Luo et al., 2024; Go et al., 2023; Zhang et al., 2022b). A crucial yet open challenge is how to effectively transfer the knowledge. To address the issue, previous studies (Zhang et al., 2022c; Allen-Zhu and Li, 2023; Zhang et al., 2021) 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 (Caruana, 1997; Zhang and Yang, 2021). In recent years, MLT has
been broadly applied to NLP tasks (Worsham and Kalita, 2020; Zhang et al., 2023b; Liu et al., 2019). 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 (Kendall et al., 2018; Lin et al., 2019). Recently, KD has also been applied within MTL frameworks, achieving state-of-the-art results in various applications (Li and Bilen, 2020; Xu et al., 2023; Yang et al., 2022; Garner and Dux, 2023; Zhang et al., 2023a).
Information Bottleneck (IB) IB (Tishby and Zaslavsky, 2015; Slonim, 2002) 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 (Cover, 1999; Covert et al., 2023; Liu et al., 2009). Estimating MI is known to be challenging, and recent progress has been made towards learning-based variational approaches (Tian et al., 2020; Covert et al., 2023; Bachman et al., 2019; Tschannen et al., 2019; Belghazi et al., 2018; Diao et al., 2023). 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 (Alemi et al., 2016; Wang et al., 2019). 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 (Tian et al., 2021; Du et al., 2020; Wan et al., 2021) and NLP (Chen and Ji, 2020; Zhang et al., 2022a; Paranjape et al., 2020). 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.

heading: 3 Methodology
text: 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.

heading: 3.1 Preliminaries
text: Under the DSS framework, the prefixes [PREDICT] and [EXPLAIN] will be prepended to the input text, TEXT, for tasks corresponding to label prediction and rationale generation, respectively. In the label prediction task, given the input [PREDICT] + TEXT along with predictive labels Y, a representation feature V, where V ∈ Rd, is trained using Y. Similarly, in the rationale generation task, the input [EXPLAIN] + TEXT and rationale label R guide the training of a representation feature Z, where Z ∈ Rd, using R.
Our goal is to distill CoT knowledge from larger LLMs to smaller LLMs models. To achieve this, based on the basis of IB (Tishby and Zaslavsky, 2015; Zhang et al., 2022a; Wang et al., 2019), we model the DSS as following:
I(Z;Y ) = ∫ p(z, y) log p(z, y)
p(z)p(y) dzdy. (1)
where sampling observations z ∼ Z and v ∼ V. Here p(·) represents the probability distribution.
To encourage CoT distillation to focus on the information represented in label Y, we propose using IB to enforce an upper, bound Ic, on the information flow from the representation features V to the representation features Z. This is achieved by maximizing the following objective:
max I(Z;Y ) s.t. I(Z;V ) ≤ Ic. (2) By employing Lagrangian objective, IB allows Z to be maximally expressive about Y while being maximally compressive regarding the input data, as follows:
CIB = I(Z;V )− βI(Z;Y ) (3)
where β is the Lagrange multiplier. Clearly, Eq. 3 demonstrates the trade-off optimization between high mutual information and high compression (Zhang et al., 2022a; Alemi et al., 2016). In our scenario, given a predefined small student model, the compression ratio is fixed. Therefore, we formulate the CoT distillation as an optimization problem:
max I(Z;V ) (4)
Due to symmetric property of MI, I(Z;V ) = I(V ;Z). CoT distillation can also enhance rationale generation task with the label knowledge. This is validated in Section 5.

heading: 3.2 Variational Bounds of MI
text: We rewrite I(Z;V ) of Equation 4 as follows:
I(Z;V ) = Ep(z,v) [ log
p(v|z) p(v)
] (5)
According to (Poole et al., 2019; Covert et al., 2023), a tractable variational upper bound can be established by introducing a variational approximation q(v) to replace the intractable marginal p(v), demonstrated by:
I(Z;V ) = Ep(z,v) [ log
p(v|z)q(v) p(v)q(v)
]
= Ep(z,v) [ log
p(v|z) q(v)
]
−KL(p(v)||q(v))
(6)
here KL[·||·] denotes Kullback-Leibler divergence. The bound is tight when q(v) = p(v). Consequently, KL(p(v)||q(v)) is equal to KL(p(v)||p(v)), which becomes zero. Therefore, we can derive at the following inequality:
I(Z;V ) ≤ Ep(z,v) [ log
p(z|v) p(v)
] (7)
We can then express the MI from Eq. 5 as the follows: Ep(z,v) [ log
p(z|v) p(v)
] = ∑ p(z, v) log
p(v|z) p(v)
= ∑ p(z|v)p(v) log p(v|z) − ∑
p(v)p(z|v) log p(v) (8)
Assuming that p(v) is uniform distribution for maximal entropy (Schröder and Biemann, 2020), the term ∑ p(v)p(z|v) log p(v) is considered as a constant. This also allows for the omission of p(v) in ∑ p(z|v)p(v) log p(v|z). By combining Eq. 5 and Eq. 8, then maximization of I(Z;V ) can be expressed as:
max I(Z;V ) = maxEp(z,v) [ log
p(z|v) p(v)
]
∝ max ∑ p(z|v) log p(v|z) = max(− ∑
p(z|v) log 1 p(v|z))
= min( ∑
p(z|v) log 1 p(v|z))
= min( ∑
CE(z|v, v|z)) (9)
here CE represents cross entropy. Accordingly, CoT distillation in Eq. 4 is transformed into the problem outlined in the above equation. This problem can be addressed with a learning-based method. To tackle this issue, we have developed a new MI loss that minimizes cross-entropy of representation features of the rationale generation (p(z|v)) and representation features of the label prediction (p(v|z)), effectively maximizing MI during CoT distillation process.

heading: 3.3 Training Loss
text: The training loss is given by
Ltotal = α1Lprediction + α2Lgeneration + α3LCE (10)
where α1, α2 and α3 are regularization parameters, all of which are non-negative. Lprediction represents the loss of the label prediction task, and Lgeneration represents the loss of the rationale generation task. Both are general cross-entropy loss as defined in (Hsieh et al., 2023).
According to the last line of Equation 9, we define the our MI loss as
LCE = l(f(Z), f(V)) (11)
f represents our proposed mutual information (MI) loss module, and l denotes the cross-entropy loss. As shown in Figure 1, the MI loss module consists of softmax and max reduction layers. The softmax function separately calculates the distributions for the outputs of the vocabulary spaces in the label prediction and rationale generation tasks. Subsequently, a max reduction operation is employed to reduce the dimensionality of the predicted outputs from both tasks to a uniform dimension for the loss calculation. Specifically, in the label prediction task, dimensions are reduced from Rm×d to R1×d, and in the rationale generation task, from Rn×d to R1×d.

heading: 4 Experiments

heading: 4.1 Experimental Setting
text: Datasets. We conducted the experiments on four widely-used benchmark datasets across three different NLP tasks: e-SNLI (Camburu et al., 2018) and ANLI (Nie et al., 2020) for natural language inference; CQA (Talmor et al., 2018) for commonsense question answering; and SVAMP (Patel et al., 2021) for arithmetic math word problems. We
used rationale generated by PaLM 540B (Chowdhery et al., 2023), which were collected and opensourced by (Hsieh et al., 2023)1.
Setup. Based on CoT properties and the comparative experimental study in (Hsieh et al., 2023), our work adopted T5-base (220 million) and T5small (60 million) to the student models. α1 and α2 were set as 0.5 and α3 is set as 0.1. We trained our models on one A100 GPU with 80G memory. For T5 base model, the training time for the full-size four dataset was approximately 14.4 hours. For T5 small model, the training times was approximately 8.6 hours.
Baselines. We compare our work with the stateof-the-art DSS (Hsieh et al., 2023) by running their open-sourced code and include two other baseline reported in their work: (1) Standard Finetune, which involves using the prevailing pretrainthen-finetune paradigm to finetune a model with ground-truth labels through standard label supervision. (Howard and Ruder, 2018). (2) Single-task, which finetunes the model using both of the label and non-CoT rationale as supervision .
Evaluation Settings. Following the DSS work (Hsieh et al., 2023), we adopt the accuracy as the performance metrics across all four datasets. Higher accuracy indicates better results. Besides accuracy, we also adopt Expected Calibration Errors (ECE) and Average Confidence Scores to evaluate calibration of the T5-base model. A lower ECE and higher Average Confidence Scores indicate better performance. We adopt GPT-4 to evaluate Quality of CoT examples and subjective analysis. Please refer to our codes for more details.

heading: 4.2 Results
text: Experiments of T5-base Model. We present our experimental results for the T5-base model in Table 1. In single-task training, the rationale and label are concatenated into a single sequence, which is treated as the target in training models (Hsieh et al., 2023). Our proposed method consistently achieves better performance than standard fine-tuning and single-task methods across all datasets. Compared to DSS, our method outperforms DSS on ANLI, CQA, and SVAMP, and achieves nearly the same accuracy on e-SNLI.
1Data and DSS code are from https://github.com/ google-research/distilling-step-by-step.
Experiments of T5-small Model. The experimental results for the T5-small model are shown in Table 2. The patterns of the results are similar to those of T5-base. Our proposed method consistently achieves better performance than standard finetuning across all dataset. Compared to DSS, our method outperforms DSS on ANLI, CQA and SVAMP, and is just 0.2% less accuracy on e-SNLI.
Distillation with LLM Labels. We conducted an experiment on e-SNLI and ANLI datasets with T5-base model to evaluate the effect of label quality. We distilled the student models using labels generated by 540B PaLM instead of the ground truth. The results are shown in Table 3. Comparing Table 1 and Table 3, we observe the label quality affects the distillation results in both methods. Even With the noisy LLM labels, our model still outperforms DSS on both datasets.
Distillation with smaller datasets. To evaluate the performance of our models on smaller datasets, we distilled T5-base and T5-small models on various sizes of four datasets and compared to DSS method. The results are shown in Figure. 2 and 3 respectively.

heading: 4.3 Ablation Study
text: Effectiveness of Difference Dimension Reduction Method In our proposed MI loss module, we utilize maximum reduction to align the dimensions of different features. Additionally, mean reduction serves as an alternative method for dimension reduction, based on the hypothesis that important features can represent better than average features. In Table 4, we present the results of two different layer of MI module. The results indicate the superiority of the MI module with maximum reduction.
Comparison with KL Divergence KL divergence loss has been extensively utilized in KD tasks,serving as a metric for assessing the similarity between two data distributions (Hinton et al., 2015; Zhang et al., 2022c; Gou et al., 2021; Husain et al., 2024). While KL divergence is widely applied in various KD scenarios, modeling DSS using IB framework has proven to be more accurate than using similarity measures, as discussed in Section 3. To validate our hypothesis, we conducted experiments on T5-base model across all four datasets. As shown in Table 5, our proposed method consistently outperforms the KL divergence approach, demonstrating superior performance.

heading: 5 Discussion

heading: 5.1 Analysis on T5 Calibration
text: Calibration measures the alignment between a model’s predicted accuracy and its confidence levels. Lee et al. (2022) introduced an innovative perspective on model distillation, positioning the teacher model not only as a source of knowledge but also as a tool for identifying mis-calibration during the training of the student model. This ability
to maintain calibration and make reliable predictions is crucial for downstream applications and has been the focus of prior studies (Chen et al., 2023; Lee et al., 2022; Jiang et al., 2021). Here, we apply the Expected Calibration Errors (ECE) and Average Confidence Scores to reflect the alignment between the model’s predicted probabilities and the actual outcomes, thereby gauging the reliability and certainty of its predictions. Despite the potential limitations inherent in these metrics, we still employ ECE in our experiments due to its simplicity and popularity, as in previous work on investigating the calibration quality of T5 (Chen et al., 2023; Lee et al., 2022).
We employ a 10-bin-based ECE metric and a softmax-based approach to compute average confidence scores from the test outputs across all four datasets. Given that e-SNLI and ANLI essentially represent the same task, we conduct an out-ofdomain experiment by testing the model checkpoint trained on one dataset with the test set of the other. This analysis gives us insights into how well our model generalizes across similar tasks and the robustness of its predictions in out-of-domain scenarios and to assess the calibration quality of the model more comprehensively.
Table 6 presents the results of the distilled model calibration evaluation. Overall, both models report lower ECE and confidence scores on SVAMP and e-SNLI, indicating that these two tasks are more challenging and models are less certain about their prediction. Lower ECE values from our MI-based distillation approach are presented for e-SNLI and ANLI, and their respective out-of-domain tests. Notably, our method achieves an ECE of 4.35 in eSNLI, significantly lower than DSS’s 8.54. However, in SVAMP and CQA, our method records higher ECE, indicating potential areas for improvement in these domains. The trade-off in calibration accuracy in specific tasks like SVAMP and CQA compared to DSS suggests future directions for refining our approach.
Regarding average confidence scores (Conf.), our method generally maintains competitive confidence levels, with notable improvements in eSNLI and ANLI. In e-SNLI, the confidence is lower (30.06) compared to DSS (34.33), which, combined with a lower ECE, suggests a more realistic confidence estimation. Conversely, in the out-of-domain scenarios for e-SNLI and ANLI, our method shows marginally higher confidence scores
than DSS, which, coupled with the lower ECE, indicates robustness in out-of-domain generalization.

heading: 5.2 Analysis on CoT Output

heading: 5.2.1 Quality of CoT Examples by GPT-4 Evaluation
text: We evaluate the quality of CoT examples using GPT-4, as it achieves the state-of-the-art human alignment performance and is used for text generation evaluation in previous work (Liu et al., 2023; Hsu et al., 2023; Wang et al., 2023a). Inspired by (Wang et al., 2023a), we ask GPT-4 to evaluate the quality of the provided CoT examples based on their coherency and relevancy to the input questions and answers. We randomly sample 50 CoT examples from the four datasets and ask GPT-4 to score based on a scale from 1 to 5, where 1 indicates completely incoherent and irrelevant responses, and 5 represents highly coherent, relevant, and helpful responses. For each sample, we run the same sample for four times to obtain self-consistency to measure the reliability of the responses. Table 7 presents the prompt we use for GPT-4 evaluation, average scores and standard deviation on the scores obtained over the four datasets. We report the scores on both provided CoT (“gold”) rationales and distilled model predicted rationales.
The effectiveness of our MI-based distillation
method is closely linked to the quality of CoT reasoning in the training data. When the CoT quality is high, as in SVAMP, a strong correlation is observed between the model’s label prediction accuracy and the quality of its generated CoT. However, this correlation weakens significantly when the CoT quality is low (e-SNLI), suggesting that the model struggles to align label prediction with coherent rationale generation under poor training conditions. Interestingly, with average-quality CoT data (ANLI), the performance gap between our MIbased distillation and DSS is minimal, suggesting that the effectiveness of our approach is particularly reliant on the presence of high-quality reasoning in the training data.

heading: 5.2.2 Case Studies on the Output Rationale
text: We performed case studies on SVAMP and e-SNLI as illustrated in Figure 4 and 5. In the SVAMP example, the question asks the difference in the number of kids Julia played with from Monday to Tuesday, with specific numbers provided for Monday, Tuesday, and Wednesday. DSS generates an incorrect explanation, which contradicts the given question, resulting in to a wrong answer. Conversely, our method correctly identifies the comparison needed between the number of kids Julia played with on Monday and Tuesday, leading to the correct answer. Notably, our generated CoT reason-
Prompt for GPT 4 Evaluation
Given an input pair of a question and an answer of a taskname task, how good is the given Chain-of-thought example? From 1-5, where 1 is completely incoherent and irrelevant, 2 is somewhat incoherent and irrelevant, 3 is coherent, relevant but not helpful, 4 is somewhat helpful, and 5 is helpful and it explains the answer well.
Average Scores and Standard Deviation
ing is identical to the golden one, demonstrating that by precisely grasping the rationale, our approach effectively resolves the math problem. We also show the evaluation results (score and reasoning) from GPT-4, where our method gains a top score of 5 and DSS gains only a mere score of 1. This example showcases that the high-quality CoT generated by our method enhances problemsolving capabilities in math tasks like SVAMP.
Another example (Figure 5) is from e-SNLI, where the task is to identify whether the hypothesis is entailment, contradiction, or neutral, based on the given premise and hypothesis. Although both our method and DSS generate the correct label output, it is worth noting that, the CoT of our method
points out the relationship between the premise and the hypothesis, while DSS only restates the hypothesis without providing any extra explanation or connecting the hypothesis to the premise. Our generated rationale also gains a higher score than DSS. A higher-quality rationale tends to facilitate more accurate label prediction, thereby enhancing overall task performance.

heading: 6 Conclusion
text: 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
our method and shed light on aspects of language model distillation and CoT applications.

heading: 7 Limitation
text: Our comparative analysis primarily focuses on the Distilling Step-by-Step (DSS) framework, which serves as our main benchmark. This concentrated comparison, while valuable for a deep understanding of DSS’s nuances and our advancements over it, constitutes a limitation of our work. Specifically, our analysis does not extend to a broader range of knowledge distillation methods currently employed in the field, focusing exclusively on T5 and not including other LLMs like Mistral and Llama2, Llama3 model family. This focus may overlook the potential insights and contrasts that could emerge from evaluating our approach against a wider array of distillation techniques. Future research could benefit from a more expansive comparative study, incorporating diverse methodologies to fully contextualize our findings within the broader landscape of knowledge distillation practices. This broader comparison would not only validate the efficacy of our method in various settings but also illuminate areas for further refinement and innovation.
However, it is important to note that our contribution lies in providing an in-depth analysis from both theoretical and practical viewpoints to enhance the CoT distillation process. Our work delves into the intricacies of utilizing mutual information to
improve distillation outcomes, offering significant advancements in understanding and applying CoT distillation techniques.

heading: 8 Ethical Issues
text: In this paper, we carefully considered the ethical implications in line with the ACL code of ethics. We evaluated the potential dual-use concerns, ensuring our research serves to benefit society and does not cause inadvertent harm. Our methodology and applications were thoroughly assessed for fairness, non-discrimination, and privacy, particularly in the context of data handling and model outputs. We also ensured our study did not expose any negative impact on individuals and groups. Moreover, we did not engage in academic dishonesty and adhered to high-quality processes and product standards in our professional work. We include this detailed discussion of these ethical considerations, affirming our commitment to responsible and beneficial computational linguistics research.

heading: Acknowledgements
text: We thank reviewers for providing helpful feedback.

heading: Future Work
text: Future research could benefit from a more expansive comparative study, incorporating diverse methodologies to fully contextualize our findings within the broader landscape of knowledge distillation practices. This broader comparison would not only validate the efficacy of our method in various settings but also illuminate areas for further refinement and innovation. However, it is important to note that our contribution lies in providing an in-depth analysis from both theoretical and practical viewpoints to enhance the CoT distillation process. Our work delves into the intricacies of utilizing mutual information to
improve distillation outcomes, offering significant advancements in understanding and applying CoT distillation techniques