SIKeD: Self-guided Iterative Knowledge Distillation for Mathematical Reasoning

We thank the reviewer for a careful examination of our work, and providing positive feedback for our research problem. We discuss the weaknesses and questions highlighted by the reviewer below.

Weaknesses

How SIKeD promotes strategy diversity

Although the method demonstrates effectiveness, the motivation requires further clarification. While iteratively forming new training datasets based on model outputs is an existing approach, the paper's contribution lies in showing this method can enable diverse reasoning strategies in smaller models. However, the underlying mechanism for how this approach promotes strategy diversity needs clearer explanation.

Knowledge distillation uses all data from the LLMs while self-distillation relies on smaller models data for training. SIKeD lies in the middle of the two extremes, where at each iteration, the smaller model generates multiple outputs (k=10) and a reward is applied to it (0 if incorrect, 1 if correct). The correct reward incentivizes the model to improve its performance over iterations and also to choose the right strategy for a given task which may not be the greedy strategy making the strategy choice more diverse (if needed). In a sense, optimizing for performance automatically leads to strategy diversity over iterations, since the smaller model cannot solve all the questions with the same strategy. The qualitative analysis presented in Fig. 7 shows that the smaller model changes its strategy to get the correct solution if needed. Essentially, in our setting, diversifying strategy is entangled with improved performance.

The experimental evaluation is currently limited to mathematical reasoning tasks. Exploring the effectiveness of the proposed method in other scenarios would provide valuable insights into its generalization capabilities.

We have discussed this in detail in the general comments section. In short, we wanted to demonstrate the self-iterative knowledge distillation approach where a smaller model can learn to pick the right strategy to solve a given task following LLM. We used Chain of Thought, subquestion decomposition and Program of Thought as our strategies to pick from. Essentially we were looking for a task that possesses the qualities of intermediate reasoning, can be decomposed into smaller tasks and can be represented in the form of a program. We found mathematical reasoning as the perfect fit for our use case and other reasoning tasks were hard to fit into the chosen strategies. This is why we limited ourselves to mathematical reasoning and also added that to the title to not oversell our idea. We are collecting a dataset in a different domain that can work across multiple strategies but that is for future work.

Below is a discussion on the questions raised by the reviewer.

Q1: Can the authors elaborate on the underlying reasons for smaller models' bias towards specific strategies? Does the preferred strategy vary?

Our experiments (Fig 5) suggests that the dominant strategy (measured in terms of the performance) is often preferred by the smaller model. This might be because of relatively more samples present in the training data. However, when we balance the samples per strategy (for ablation study), the model prefers CoT, possibly because of the pre-training bias.

Q2: Bias introduced due to discarding samples

The paper states that mixing LLM-generated data with self-generated outputs helps align smaller models with their learned knowledge (L83). However, discarding samples with mismatched outputs creates a different form of bias? Please clarify this.

Although we discard the mismatched outputs from the smaller model and utilize the correct generations from LLM, the KL divergence between the training data and SLM still decreases as all the data is not just from the LLM and mixing smaller model data reduces the KL divergence. This is also demonstrated in Fig. 2.

Q3: Please explain the mechanism by which mixing the data promotes diverse strategy selection in smaller models.

• Please refer to Weaknesses 1 above. We have addressed this concern in details.

Q4: How to judge if the generated r_i is correct? (L294)

Knowing if the rationale is correct or not is a strong limitation of all the reasoning tasks involving intermediate steps and is not a specific limitation for our work. This has been discussed in the past works like https://arxiv.org/pdf/2212.10001

Q5: Is it a typo in L211 & L215 that the same notation appears in different contexts?

L211 is the rolled out version of L215. We did not find any typo in those lines. Please let us know if there is anything specific you would like to point out.

Q6: Discussing how selection works for biased strategy

Could the authors clarify the strategy sampling process for smaller models? Specifically regarding L413, "if both CoT and PoT are sampled correctly, our biased strategy choice is PoT" - is this strategy determined by the smaller model's output? Additional details on the strategy sampling mechanism would be appreciated and helpful

The underlying assumption for using a biased strategy selection mechanism lies in the fact that for some datasets, a specific strategy might be more beneficial than the others (L408-410). Based on this, we create biased training datasets for all the three strategies. For example, if we are creating a biased dataset for PoT, we firstly select only the problems which can be solved via PoT. We ignore the CoT and L2M reasoning chains for these problems. For problems that cannot be solved via PoT, we select CoT or L2M chains at an instance level, but not both. This helps us to create a dataset biased on PoT.

Please let us know if all the concerns are addressed. We are happy to discuss more. Thank you.

We want to sincerely thank the reviewers for their time and effort in evaluating our paper. We would appreciate it if you could kindly confirm that the rebuttal was received and let us know if any additional steps or clarifications are required from our side. Your feedback is highly important to us, and we remain available to address any further concerns or questions.

Please let us know. Thanks.

Questions

Do self-distillation really work in small model continual training ?

Fig. 4 in your paper and Fig. 7 in [3] indicated the performance decreased, when the n of iteration became large. From a theoretical perspective on synthetic data, the data variance decreases with multiple generations n [4]. To prevent variance reduction, i.e., to enhance data diversity, this paper incorporates data synthesized by LLMs throughout the iterative process. This operation is very important. From data perspective, authors should analyze the data distribution shifting across n.

From a theoretical perspective on synthetic data, the data variance decreases with multiple generations as discussed in [4]. However [4] uses a model that self-distills itself using a fixed strategy. On the contrary, SIKeD uses multiple strategies which allows the model to switch strategies over iterations for a given query, potentially increasing the data variance. Moreover, we also mix data from LLM in 2 variations: a) All b) Sparse which also changes over iterations. We analyzed the change in strategy distribution in Fig. 2 which shows that diversity may or may not increase but mixing data reduces overall KL divergence between the training data distribution and the distribution generated by the smaller model leading to better overall performance.

Is there any new theoretical insights ?

The authors observed shifts in generation strategies over iterations n. What causes this phenomenon? Additionally, as more data is generated, the overall dataset size increases.

• We observed that smaller models can learn to pick the suitable strategy for a given task over iterations but this requires careful mixing of data between the LLM and self generated data controlled by the parameter ‘alpha’. Simply using all the strategy data combined from the LLM wont work as shown in the Results table (Combined baseline in Table 1).
• In terms of overall data size, SIKeD increases the data size for ‘All’ variation, but not for ‘Sparse’. The dataset used for training ‘Sparse’ models utilizes correctly generated data by the smaller model, and for questions where the smaller model is unable to generate correct answers, corresponding data from LLM is added. Over iterations, as the smaller model improves, additional training data from LLM decreases, essentially keeping the training dataset approximately constant. However, we observe that ‘Sparse’ performs as well as the ‘All’ variation, if not better.

We again thank the reviewer for all the important points raised. Please let us know if all the concerns are addressed. We are happy to discuss more. Thank you.

We thank the reviewer for providing detailed and easy to follow feedback on our paper. We specially thank the reviewer for citing relevant literature along with the comments. We address the weaknesses and questions highlighted by the reviewer below:

Weaknesses

W1: The mathematical notation is overly verbose.

While some of the notation may look complex, it is necessary to formulate the iterative self-guided framework in detail, and Algorithm 1 is required to do so. Furthermore, we use the mathematical framework to define and show the importance of the data mixing rate "alpha", without which it would have been difficult to discuss the two settings: "All" and "Sparse". Finally, the original knowledge distillation paper by Hinton et al. uses KL divergence to discuss distillation effects, and it was important to draw a comparison between SIKeD and KL divergence.

W2: Limited generalization

The approach only enhances the GSM8K dataset, but reasoning tests should be conducted on more realistic datasets, such as MATH, Arc-Challenge and so on. And more reasoning tasks also need to be evaluated, such as commonsense reasoning and symbolic reasoning.

We wanted to demonstrate the self-iterative knowledge distillation approach where a smaller model can learn to pick the right strategy to solve a given task following LLM. We used Chain of Thought, subquestion decomposition and Program of Thought as our strategies to pick from. Essentially we were looking for a task that possesses the qualities of intermediate reasoning, can be decomposed into smaller tasks and can be represented in the form of a program. We found mathematical reasoning as the perfect fit for our use case and other reasoning tasks were hard to fit into the chosen strategies. This is why we limited ourselves to mathematical reasoning and also added that to the title to not oversell our idea. We are collecting a dataset in a different domain that can work across multiple strategies but that is for future work.

Also, to compare our work with the previous works of knowledge distillation, we limited ourselves to four datasets that were commonly used in the past work on Magister et al. (https://arxiv.org/abs/2212.08410), Shridhar et al. (https://arxiv.org/abs/2212.00193) and Zhu et al (https://arxiv.org/abs/2401.11864) which also form our baselines. Finally, we did initial analysis on the MATH dataset and found PoT and L2M to be a weaker strategy compared to CoT which biased the model to always pick CoT.

W3: The absence of important references. The self-distillation in small models is already studied in [1,2,3].

Self-distillation has been studied in the past but our work demonstrates that direct distillation and self-distillation are on the two extremes. A mixture of the two (controlled by alpha in our work) performs better as shown in Fig. 6. In addition, our work also focuses on teaching multiple strategies to smaller models so that they can learn to pick the right strategy for a given task, which is missing from all previous work on self-distillation. We will add the highlighted reference in our paper. Thanks for that.

We want to sincerely thank the reviewers for their time and effort in evaluating our paper. We would appreciate it if you could kindly confirm that the rebuttal was received and let us know if any additional steps or clarifications are required from our side. Your feedback is highly important to us, and we remain available to address any further concerns or questions.

Please let us know. Thanks.

Thank you very much for your clarification. I summarize my concerns as follows:

• The main contribution overlaps with previous works. [6] demonstrated that keeping the source real data in iteratively generated data will avoid model collapse. They illustrate this conclusion in both language, vision, and molecular data. They also explain why this happens theoretically. If you keep the source data in your iteration, there is an upper bound for test error. In other words, these conclusions in [6] overlap a lot with your main contribution. Put above in your paper is iterative data mixing. I acknowledge your implementation is important. However, based on previous works, you didn’t provide new conclusions for self-distillation or synthetic data research areas.
• Validation is limited compared to previous works. I know math is suitable, and everyone knows it. However, compared with the classic paper in this line, Chain-of-thought [7] conducted experiments on math, commonsense reasoning, and symbolic reasoning, a total 11 datasets. Your paper only trains on 1 dataset, test on the other four. More recently, [3] also conducted experiments on 7 datasets. I think the validation in this paper is unacceptable, it lacks persuasiveness.
• Needing more concise and direct presentation. I acknowledge that rigorous formalization is important, but what you've done is actually quite straightforward and direct—why not just state it plainly? As I summarized earlier, and I'll reiterate here, ”during the training process using self-distilled data, it is necessary to incorporate the original real data. “ However, as mentioned above, these things have been well-stuied both emperically and theoretically in [6].

In conclusion, based on overlapping contributions, limited validation, and no theoretical insights, I will keep my rating “3, reject “. I think this paper needs much improvement to link with the recent literature. Hope my opinions can help you. If I am wrong, please correct me.

[6] Gerstgrasser M, Schaeffer R, Dey A, et al. Is model collapse inevitable? breaking the curse of recursion by accumulating real and synthetic data[J]. arXiv preprint arXiv:2404.01413, 2024.

[7] Wei J, Wang X, Schuurmans D, et al. Chain-of-thought prompting elicits reasoning in large language models[J]. Advances in neural information processing systems, 2022, 35: 24824-24837.

Below is a discussion of all the issues the reviewers wanted to discuss.

Q1: Have the authors considered applying SIKeD to tasks outside of mathematical reasoning to test the generalizability of the strategy selection mechanism?

Please refer to the first point under weakness above.

Q2: Does SIKeD require the ground-truth answers of the training data? If not, how does it handle the tasks where the ground-truth answers are not available?

Since the knowledge distillation works with the LLM generating the training data for smaller models, we need to verify that the data generated by the LLM is correct.. If the ground truth data is unavailable, we can either use the LLM to verify the generations of the smaller model or do a self-consistency over the generated samples (since we have K=10) and use the most consistent answer. However, we did not explore these in our work due to the availability of the ground truth data.

Q3: Is there any knowledge distillation baseline that could be used to compare the performance of SIKeD on mathematical reasoning tasks? Current experiments only compare SIKeD against CoT, L2M, PoT and Combined

Comparison with the baselines of CoT, PoT, L2M and combined are based on the past knowledge distillation works. CoT based distillation (CoT in Table1) was proposed by Magister et al. (https://arxiv.org/abs/2212.08410), L2M based distillation (L2M in Table 1) was proposed by Shridhar et al. (https://arxiv.org/abs/2212.00193) and the combined one is inspired from Zhu et al (https://arxiv.org/abs/2401.11864). All these comparisons are presented in Table 1 and are considered as baselines against which we compare our proposed methodology SIKeD. We will clarify this in the paper that our baselines come from previous distillation works.

Q4: Can the authors explain why does the improvement on Qwen 1.5B model is less significant compared to the other base models? What are the potential reasons for this discrepancy?

A recent work “Careful Examination of LLM performance on GSM” (https://arxiv.org/abs/2405.00332) has suggested that some models are over optimized for GSM8K style tasks and we suspect that could be one of the reasons for minor improvements.

Q5: The small models are tuned with LoRA, what if the parameters of the small models are fully tuned? Would the performance of SIKeD be further improved?

Experimenting with fully fine-tuning small models showed that the performance worsened, possibly due to over-fitting. We again refer to “Careful Examination of LLM performance on GSM”” (https://arxiv.org/abs/2405.00332) which suggests that some models are over-optimized on GSM8K style tasks

Q6: "The iterative training is stopped when accuracy shows only marginal improvements or declines." What specific criteria is used to determine the optimal number of iterations?

We trained the Gemma 2B model further till 5th iteration and the overall performance went down. We noticed that training for one or two iterations after the accuracy starts to plateau or goes down is a good stopping criteria. For our experiments, training up to 3 iterations seems to be a good stopping point.

Please let us know if there are any more concerns. We are happy to discuss more.

We thank the reviewer for providing positive feedback for our experimental process, and also for a very clear, detailed and easy to follow list of questionnaires. We address the weaknesses and questions raised by the reviewer below:

Weakness1: The proposed method is only evaluated on mathematical reasoning tasks. It’s unclear how well SIKeD would generalize to other domains that require more nuanced strategy selection.

We have addressed this concern in details in the general comment section. In short, we wanted to demonstrate the self-iterative knowledge distillation approach where a smaller model can learn to pick the right strategy to solve a given task following LLM. We used Chain of Thought, subquestion decomposition and Program of Thought as our strategies to pick from based on the past works. Essentially we were looking for a task that possesses the qualities of intermediate reasoning, can be decomposed into smaller tasks and can be represented in the form of a program. We found mathematical reasoning to be the perfect fit for our use case. The other reasoning tasks were difficult to fit into the chosen strategies. Therefore, we limited ourselves to mathematical reasoning and also added that to the title to not oversell our idea. We are collecting a dataset in a different domain that can work across multiple strategies but that is for future work.

Weakness2: The paper lacks comparison with knowledge distillation methods.

We compared SIKeD with past knowledge distillation works like CoT based distillation (CoT in Table 1) which was proposed by Magister et al. (https://arxiv.org/abs/2212.08410), L2M based distillation (L2M in Table 1) which was proposed by Shridhar et al. (https://arxiv.org/abs/2212.00193) and combined is inspired from Zhu et al (https://arxiv.org/abs/2401.11864). All these comparisons are presented in Table 1 and are considered as baselines against which we compare our proposed methodology SIKeD. We will clarify this in the paper that our baselines come from previous distillation works.

The code is not available for reproduction.

We are currently in the process of cleaning our code and we’ll add a Github link soon.

We want to sincerely thank the reviewers for their time and effort in evaluating our paper. We would appreciate it if you could kindly confirm that the rebuttal was received and let us know if any additional steps or clarifications are required from our side. Your feedback is highly important to us, and we remain available to address any further concerns or questions.

Please let us know. Thanks.

Concern3: Computational efficiency of SIKeD and number of iterations required for convergence

We performed a full Self-guided iterative training using SIKeD and with the “Sparse” version of our approach, we only take the problems from the LLM that the smaller model could not solve. For all queries with incorrect solutions generated by the smaller model, we take the query output pairs from the LLM.

For example, Gemma 7B has a baseline accuracy of ~70%. Assuming the same accuracy on the training dataset (which is true), only 30% of the data from LLM is required for the next iteration. Since all baseline models and each iteration of SIKeD are trained for 3 epochs, training SIKeD with 30% data for 3 epochs (1 iteration of SIKeD) corresponds approximately to 1 epoch of baseline models. Thus, 3 iterations of SIKeD correspond to 3 epochs of baseline models, which is equivalent to 1 additional training round of baseline models. In other words, the training cost of SIKeD is 2X of the baseline model with no additional test time cost. Note that we trained the baseline model for 3 more epochs and it led to a worse performance possibly due to overfitting.

Convergence - We trained the Gemma 2B model further till 5th iteration and the overall performance went down. We noticed that training for one or two iterations after the accuracy starts to plateau or goes down is a good stopping criteria. For our experiments, training up to 3 iterations seems to be a good stopping point.

Resources Needed - All the models have been fine-tuned using LoRA using a single RTX GPU with 24 GB RAM. In terms of the compute time, here are the training time per model:

Per iteration

• Gemma 2B - 1 hour
• Qwen 1.5B - 1 hour
• Qwen 0.5B - 30 mins
• Gemma 7B - 2 hours
• SmolLm - 1 hour

Each training iteration was done under 1 hour (except for Gemma 7B that takes 2 hours per iteration). For inference with VLLM, all models except Gemma 7B were run on a single RTX GPU with 24 GB RAM. For inference with Gemma 7B, a single A100 GPU was needed. The inference was completed within 2 hours.

Concern4: Scaling of SIKeD with model size

Since all the models used in the paper were initially trained on different datasets (which in most cases are not public), it is hard to judge the scalability of SIKeD with model sizes. Also models of different sizes are often trained on different overall tokens which makes the comparison hard for us to do. Nevertheless, we found that SIKeD works well across models of all sizes with no significant impact due to scaling the model sizes.

Please let us know if we addressed your concerns and we are happy to discuss more. Thank you.

We thank the reviewer for recognizing and valuing our work, as well as for their positive feedback on our paper writing, presented methodology and experimentation choice.

We address the weaknesses highlighted by the reviewers, along with the questions raised, in detail below:

Concern1: SIKeD is contingent upon the quality of the initial LLM data

This is an important point raised and knowledge distillation as a concept depends on the quality of the teacher model data. However, SIKeD is not limited by the initial LLM data as SIKeD is a self-guided iterative learning framework. This means that with a weaker LLM, the quality of distillation data will be relatively poor compared to a stronger LLM, leading to a weaker distilled smaller model as baseline. But our approach improves the baseline performance irrespective of the teacher model used.

We replace Llama3 70B as LLM with Llama3 8B and report our results for Gemma 2B and 7B models for GSM8K below:

Gemma-2B

Gemma-7B

Overall SIKeD improved the performance over baselines even when a weaker LLM was used. We observed similar results with the Qwen models.

Concern2: Study primarily focused on mathematical reasoning tasks

We have addressed this concern in details in the general comment section. In short, we used Chain of Thought, subquestion decomposition and Program of Thought as our strategies to pick from based on the past works. Essentially we were looking for a task that possesses the qualities of intermediate reasoning, can be decomposed into smaller tasks and can be represented in the form of a program. We found mathematical reasoning to be the perfect fit for our use case. The other reasoning tasks were difficult to fit into the chosen strategies. Therefore, we limited ourselves to mathematical reasoning and also added that to the title to not oversell our idea. We are collecting a dataset in a different domain that can work across multiple strategies but that is for future work.

We want to sincerely thank the reviewers for their time and effort in evaluating our paper. We would appreciate it if you could kindly confirm that the rebuttal was received and let us know if any additional steps or clarifications are required from our side. Your feedback is highly important to us, and we remain available to address any further concerns or questions.

Please let us know. Thanks.

We thank the reviewer for the positive feedback on our experimental analysis and framework. We discuss the one concern highlighted by the reviewer below:

Concern: I’m curious to see that among the three methods, 1) training with the proposed method, 2) pure distilling and 3) pure self-generating, which method can make the model generate the most diverse trajectories and whether the diversity is aligned with model’s performance on OOD tasks

We refer to Figure 6 which shows that the best performing model is neither a pure distillation (alpha=1) nor pure self-generation (alpha=0), but somewhere in the middle. SIKeD improves the performance of smaller models by tuning them to select the most accurate reasoning strategy often, but in addition it also enables the smaller model to switch strategies over iterations. Figure 5 and the qualitative analysis (Figure 7) shows this.

From a diversity perspective, pure distillation incorporating all the reasoning strategies (combined in our baseline) has the highest diversity, as a larger model is able to generate more diverse reasoning chains. The diversity of purely self-generating distillation is the lowest, as a smaller model is always biased towards one strategy (Please check Figure 1 for reference).

Hence, Diversity(Pure Distillation) > Diversity(SIKeD) > Diversity(Pure Self-Generating)

However, Performance(SIKeD) > Performance(Pure Distillation) and Performance(SIKeD) > Performance(Purely Self-generating). This is mostly because only using LLM data has a distributional gap with the smaller model while purely using smaller model data has limited correct samples to improve. A mix of both worlds yields the best results as presented in our work.

We are happy to discuss more. Please let us know if any more concerns. Thank you.

We want to sincerely thank the reviewers for their time and effort in evaluating our paper. We would appreciate it if you could kindly confirm that the rebuttal was received and let us know if any additional steps or clarifications are required from our side. Your feedback is highly important to us, and we remain available to address any further concerns or questions.

Please let us know. Thanks.