SIKeD: Self-guided Iterative Knowledge Distillation for Mathematical Reasoning

## Abstract

Abstract Large Language Models (LLMs) can transfer their reasoning skills to smaller models by teaching them to generate the intermediate reasoning process required to solve multistep reasoning tasks. While LLMs can accurately solve reasoning tasks through a variety of strategies, even without fine-tuning, smaller models are not expressive enough to fit the LLMs distribution on all strategies when distilled and tend to prioritize one strategy over the others.
This reliance on one strategy poses a challenge for smaller models when attempting to solve reasoning tasks that may be difficult with their preferred strategy.
To address this, we propose a distillation method SIKeD : S elf-guided I terative K nowledge D istillation 1 1 1 SIKeD is pronounced as “psyched” , where the LLM teaches the smaller model to approach a task using different strategies and the smaller model uses its self-generated on-policy outputs to choose the most suitable strategy for the given task. The training continues in a self-guided iterative manner, where for each training iteration, a decision is made on how to combine the LLM data with the self-generated outputs. Unlike traditional distillation methods, SIKeD allows the smaller model to learn which strategy is suitable for a given task while continuously learning to solve a task using different strategies.
Our experiments on various mathematical reasoning datasets show that SIKeD significantly outperforms traditional distillation techniques across smaller models of different sizes. https://github.com/kumar-shridhar/SIKeD

## 1 Introduction

Figure: Figure 1: Histogram of strategy choices for the LLM and the smaller model. LLM tends to select several reasoning strategies, but the smaller model is biased towards one strategy. The comparison was done on 1000 data points randomly sampled from the GSM8K train set.
Refer to caption: x1.png

Large language models (LLMs), with tens to hundreds of billions of parameters, generally outperform smaller models (with billions of parameters or fewer) in a variety of reasoning tasks [, ]. One notable strength of large models is their ability to reason and perform multistep reasoning tasks, often considered an important aspect of intelligence []. However, the significant size and computational demands of these large models present several challenges. For example, LLaMA3 models [] are trained using clusters of 24,000 GPUs, limiting their accessibility to many researchers and practitioners.

To bridge this gap, a key approach involves teaching smaller models to replicate the knowledge of a larger model, often referred to as *knowledge distillation* []. Typically, smaller models can be taught to replicate the multistep reasoning capabilities of larger models by incorporating a set of intermediate sequences [, ].
However, these intermediate steps can be derived from a number of different strategies, such as Chain of Thought (CoT) [], Subquestion Decomposition [, ], and Program of Thoughts (PoT) [], among others.
A viable solution is to distill these reasoning capabilities into smaller models either by distilling individual strategies [, , ] or by incorporating multiple strategies simultaneously [, ].
Although smaller models have demonstrated impressive performance when distilled with a single strategy, they often struggle to master multiple strategies equally well. An example is presented in Figure 1 where a larger model can use multiple strategies to generate data but upon distilling, a smaller model tends to favor one over the others. This is because reasoning through a variety of strategies tends to emerge as a result of scaling language models, making it difficult for smaller models to replicate this behavior [].

On the other hand, learning to solve a task using multiple strategies can help smaller models overcome the limitations of relying on a single approach. However, a key challenge arises when, despite being trained on a fixed dataset containing various strategies, a distribution mismatch occurs between the data generated by the LLM and the outputs produced by the smaller model during inference. This mismatch can hinder the ability of the smaller model to generalize across different reasoning strategies. This issue, often discussed in imitation learning [, ], results in the student model consistently choosing one strategy, even when a different approach would be more appropriate. As a result, the student generates outputs with strategy choices that are highly unlikely to match those produced by the teacher.

To address this challenge, we introduce our distillation methodology, SIKeD: Self-guided Iterative Knowledge Distillation. The process begins with the LLM teaching the smaller model to approach tasks using a variety of reasoning strategies, providing a strong foundation for the smaller model to understand different problem-solving approaches. However, due to inherent biases and its limited capacity, the smaller model may still struggle to match the LLM’s distribution of strategy choices effectively.
To resolve this, we take inspiration from constructivist learning theory [], where the learner builds knowledge during the “assimilation phase” and refines their understanding during the “accommodation phase” to incorporate new insights.
We propose generating outputs using the smaller model in an on-policy setup and selecting the best strategies for the task. By mixing the LLM-generated data with self-generated outputs, we leverage the strengths of both datasets. We iteratively fine-tune the smaller model allowing it to recognize strategies that it learned from the LLM but did not initially apply. With this approach, we align the smaller model with its own learned knowledge rather than forcing its distribution to mirror that of the LLM’s.

Our proposed method extends beyond traditional one-step distillation, as each iteration of SIKeD leads to an updated policy that better grasps new information. We repeat multiple iterations of SIKeD based on the accuracy-cost tradeoff (does the improvement justify the cost of another iteration), allowing for continuous refinement and improvement of the model’s reasoning capabilities.
We demonstrate the effectiveness of SIKeD on several mathematical reasoning tasks using models with fewer than 7 billion parameters.(^2^22We acknowledge that “smaller model” is a relative term, and we consider models with fewer than 7 billion parameters to be smaller models.) On four mathematical datasets—GSM8K [], SVAMP [], ASDiv [], and MultiArith []—our approach achieves improvements of up to +5 points over traditional distillation strategies. Additionally, we show that multiple rounds of SIKeD allow the model to select the appropriate strategy for a given problem, while traditional distillation using LLM’s data tends to leave it biased.

## 2 Preliminaries: LLM based Distillation

#### Problem Setup

We consider the standard setup of LLM-based distillation (also referred to as *knowledge distillation*), where the data is sampled from the larger model (LLM) with intermediate reasoning and the smaller model is fine-tuned (distilled) over it [, ]. Two auto-regressive sequence models are involved in the process: a larger model or the LLM denoted as $p_{L}$ and a smaller model to be distilled as $p_{\texttt{sm}}^{\theta}$ (with learnable parameters $\theta$). In this work, we consider a reasoning dataset $\mathcal{D}$ consisting of a question $q_{i}$ and a numerical answer $a_{i}$ pairs for $n$ data points, i.e. $i\in\{1,\ldots,n\}$. Since our work focuses on improving reasoning in smaller models by teaching them to solve a variety of reasoning strategies, we consider three reasoning strategies in this work: Chain-of-Thought (CoT), Least-to-Most (L2M), and Program-of-Thought (PoT).
For a specific *reasoning strategy*, denoted as $s\in S$, we generate the reasoning chain or *rationale*, denoted as $r_{i}$ leading to the final answer as: $r_{i}\sim p_{L}(.\mid\texttt{pr}_{s},q_{i})$,
where, $\texttt{pr}_{s}$ represents the strategy-specific prompt, and $s\in\{\texttt{CoT},\texttt{L2M},\texttt{PoT}\}$. Prompts used are provided in section 8.

### 2.1 LLM based Distillation

We add the generated rationales to the the dataset $\mathcal{D}$ to create an initial training dataset $\mathcal{D}_{\texttt{LLM}}$ consisting of a quadruple of $\{q_{i},a_{i},s,r_{i}\}$ for each data point.
We perform a data filtering by extracting the final answer $\hat{a}_{i}$ from the generated rationale $r_{i}$ and comparing it with the ground truth answer $a_{i}$. We discard all samples that do not match, i.e., we keep samples where $\hat{a}_{i}=a_{i}$. This filtering process eliminates incorrect rationales, ensuring that only high-quality data is used for distilling the smaller models.

We start the distillation process by training the smaller model with the created dataset $\mathcal{D}_{\texttt{LLM}}$. The question $q_{i}$ is provided as input, and the smaller model $p_{\texttt{sm}}^{\theta}$ (with learnable parameters $\theta$) is first instructed to generate the strategy $s$, followed by the rationale $r_{i}$ that leads to the final answer $a_{i}$. The loss $\mathcal{L}_{\texttt{L}}(\theta)$ is defined as:

$$ $\mathcal{L}_{\texttt{L}}(\theta)=-\mathbb{E}_{(q_{i},s,r_{i})\sim\mathcal{D}_{ \texttt{LLM}}}\left[\log\ p_{\texttt{sm}}^{\theta}\ (s\mid q_{i},I)+\sum_{t=1} ^{M}\log p_{\texttt{sm}}^{\theta}\ (r_{i,t}\mid r_{i,<t},s,q_{i},I)\right],$ (2.1) $$

where $M$ represents the number of tokens decoded over time $t$ in an autoregressive manner, and $I$ is the instruction used during fine-tuning. Note that this is analogous to traditional knowledge distillation from LLMs except that we make a strategy choice before generating rationales.

#### Limitations

Training solely on LLM-generated data $\mathcal{D}_{\texttt{LLM}}$ can lead to a distribution mismatch between the training data and the smaller model’s own output distribution. Specifically, the larger model due to its larger capacity, may produce correct reasoning across multiple strategies that the smaller model can find difficult to replicate directly []. A comparison of the strategy selected by the LLM and the smaller model on 1K samples is presented in Figure 1. The smaller model performs poorly when generating outputs on its own, as the training data distribution $P_{\text{train}}(x)$ is different from the model’s output distribution $P_{\texttt{sm}}^{\theta}(x)$ as $P_{\text{train}}^{(1)}(x)=P_{\text{LLM}}(x),$ where $x$ represents the samples $(q_{i},s,r_{i})$, and $P_{\text{LLM}}(x)$ is the distribution of data generated by the LLM $p_{L}$.

#### Proposed Solution

To mitigate the distributional shift in strategy choice between the LLM and the smaller model, we propose to incorporate the smaller model’s own correct outputs into the training data. This *self-guided* training with data mixing aligns the training data distribution more closely with the smaller model’s output distribution, making learning more effective. A visualization of the data mixing approach is presented in Figure 2 that demonstrates that data mixing reduces the distribution shift, bringing the LLM and the smaller model’s output distribution closer. This allows the smaller model to choose the right strategy for a given task, much like the LLM.

## 3 SIKeD : S elf-guided I terative K nowledge D istillation

We propose SIKeD, an *iterative* training where smaller models can take advantage of their own generations to refine their strategy choices for a given task. In a nutshell, we generate data from the smaller model, filter out the correct samples based on whether the generated solutions are correct, and mix this data with the LLM-generated data to adjust its strategy preferences. The smaller distilled model is used to iteratively generate data in an on-policy setting where it updates itself by leveraging both the LLM data and its own generations. This iterative process allows the smaller model to improve its reasoning abilities and strategy selection over time by leveraging the LLM’s knowledge and its own prior learning. The following paragraphs discuss the steps involved in our proposed iterative distillation methodology and the training objective.

#### Data generation

For each question $q_{i}$ and its associated reasoning strategy $s$, we first generate $K$ rationales using the current smaller model $p_{\texttt{sm}}^{\theta}$ as: $r_{i}^{(k)}\sim p_{\texttt{sm}}^{\theta}(\cdot\mid s,q_{i},I),\quad\text{for }
k=1,\ldots,K.$
Note that we generate multiple samples $K$
as the likelihood of a correct answer being present in one of the rationales increases significantly with additional generations for smaller models [, ].

#### Data Filtering

Next, we extract the predicted answer $\hat{a}_{i}^{(k)}$ from each rationale $r_{i}^{(k)}$ and compare it with the ground truth $a_{i}$. We collect the correct samples, where $\hat{a}_{i}^{(k)}=a_{i}$, into a new dataset $\mathcal{D}_{\texttt{self}}=\left\{(q_{i},s,r_{i}^{(k)})\ \middle|\ \hat{a}_{i
}^{(k)}=a_{i}\right\}.$

Figure: Figure 2: Alignment of the smaller model’s strategy distribution with the LLM over iterations. Each subplot represents an iteration in the training process, showing the probability distributions over reasoning strategies: PoT, L2M, and CoT. The blue bars depict the LLM’s distribution $P_{L}$, while the orange bars represent the smaller model’s distribution $P_{SM}$, which is biased towards CoT. The green bars show the training data distribution $P_{\text{train}}^{(t)}$, a mixture of $P_{L}$ and $P_{SM}$ weighted by the mixing rate $\alpha$. As $\alpha$ decreases over iterations (from 0.90 to 0.20), $P_{\text{train}}^{(t)}$ shifts from being similar to the LLM’s distribution towards the smaller model’s distribution. The KL divergence between the training data and the smaller model distributions decreases accordingly, indicating increased similarity.
Refer to caption: x2.png

#### Data mixing

We combine the LLM-generated dataset $\mathcal{D}_{\texttt{LLM}}$ with the self-generated dataset $\mathcal{D}_{\texttt{self}}$ to form the mixed dataset $\mathcal{D}_{\texttt{mix}}=\mathcal{D}_{\texttt{LLM}}\cup\mathcal{D}_{\texttt{
self}}$.

Note that, we do not always use all the data from LLM in $\mathcal{D}_{\texttt{mix}}$, and study two variations:
*All* when all LLM data is used in $\mathcal{D}_{\texttt{mix}}$, and *Adaptive* when only queries that have no correct generations in $\mathcal{D}_{\texttt{self}}$ are taken from $\mathcal{D}_{\texttt{LLM}}$. *Adaptive* uses less generated data from the LLM, resulting in more computationally efficient training.

The corresponding training data distribution changes to a mixture of the LLM data distribution and the model’s own output distribution:

$$ $P_{\text{train}}^{(2)}(x)=\alpha P_{\text{LLM}}(x)+(1-\alpha)P_{\texttt{sm}}^{ \theta}(x),$ $$

where $\alpha=\frac{|\mathcal{D}_{\texttt{LLM}}|}{|\mathcal{D}_{\texttt{LLM}}|+|
\mathcal{D}_{\texttt{self}}|}$ serves as a normalized mixing rate between the two datasets.

#### Training objective

By including $\mathcal{D}_{\texttt{self}}$ in the training data, we reduce the divergence between $P_{\text{train}}^{(2)}(x)$ and the model’s own output distribution $P_{\texttt{sm}}^{\theta}(x)$, thus minimizing the distribution shift and improving training effectiveness of choosing the right strategy for a given task.

We continue training the smaller model on $\mathcal{D}_{\texttt{mix}}$ using the following loss function:

$$ $\mathcal{L}_{\texttt{mix}}(\theta)=-\mathbb{E}_{(q_{i},s,r_{i})\sim\mathcal{D} _{\texttt{mix}}}\left[\log p_{\texttt{sm}}^{\theta}\ (s\mid q_{i},I)+\sum_{t=1 }^{M}\log p_{\texttt{sm}}^{\theta}\ (r_{i,t}\mid r_{i,<t},s,q_{i},I)\right]$ (3.1) $$

The expected loss over the training data is:

$$ $\mathcal{L}_{\texttt{mix}}(\theta)=-\mathbb{E}_{x\sim P_{\text{train}}^{(2)}(x )}\left[\log p_{\texttt{sm}}^{\theta}(x)\right]$ $$

where $x=(q_{i},s,r_{i})$, and $p_{\texttt{sm}}^{\theta}(x)$ denotes the probability assigned by the model to the sample $x$.

#### Analogous to minimizing the KL divergence

Mixing the data is analogous to minimizing the Kullback-Leibler (KL) divergence [] between the training data distribution $P_{\text{train}}^{(2)}(x)$ and the model’s output distribution $P_{\texttt{sm}}^{\theta}(x)$:

$$ $D_{\text{KL}}(P_{\text{train}}^{(2)}(x)\parallel P_{\texttt{sm}}^{\theta}(x))= \sum_{x}P_{\text{train}}^{(2)}(x)\log\frac{P_{\text{train}}^{(2)}(x)}{P_{ \texttt{sm}}^{\theta}(x)}.$ $$

As we include more self-generated data, $(1-\alpha)$ increases, and $P_{\text{train}}^{(2)}(x)$ becomes closer to $P_{\texttt{sm}}^{\theta}(x)$. This reduces the KL divergence and aligns the training data distribution with the model output distribution, leading to more effective learning. Figure 2 demonstrates that as the value of $\alpha$ decreases, the distribution of the training data strategy choices aligns with the smaller model, with a reduction in their KL value over iterations. This allows the smaller model to better capture the strategy distribution of the larger model.

### 3.1 Iterative Self-Training of SIKeD

We repeat the data generation, filtering, mixing, and training steps iteratively. In each iteration $t$, the smaller model potentially generates new correct rationales that are added to the training data. The training data distribution at iteration $t$ becomes:

$$ $P_{\text{train}}^{(t)}(x)=\alpha^{(t)}P_{\text{LLM}}(x)+(1-\alpha^{(t)})P_{ \texttt{sm}}^{\theta^{(t-1)}}(x),$ $$

where $\theta^{(t-1)}$ are the model parameters from the previous iteration, and $\alpha^{(t)}$ is updated based on the sizes of $\mathcal{D}_{\texttt{LLM}}$ and $\mathcal{D}_{\texttt{self}}^{(t)}$ at iteration $t$. Note that the generated samples from the smaller model automatically govern the value of $\alpha^{(t)}$ based on the size of $\mathcal{D}_{\texttt{self}}^{(t)}$.

This iterative process continues until the model’s performance converges or a predefined stopping criterion is met. Over multiple iterations, the model’s own output distribution $P_{\texttt{sm}}^{\theta^{(t)}}(x)$ gradually improves, and the training data distribution becomes increasingly aligned with it. We present an end-to-end training methodology in Algorithm 1.

Figure: Algorithm 1 SIKeD: Self-guided Iterative Knowledge Distillation

## 4 Experimental Details

#### Dataset

Our work demonstrates the effectiveness of selecting an appropriate strategy for a given task. We consider multi-step mathematical reasoning datasets in our work, as various strategies can solve the task fairly well.
We trained SIKeD on the GSM8K training set [], which includes 7,473 samples, and tested it on the corresponding test set of 1,319 samples. To assess the domain transferability of our distillation method, we also evaluated it on three additional mathematical datasets: SVAMP [] with 1,000 samples, ASDiv [] with 2,300 test samples, and MultiArith [] with 180 samples. As the GSM8K training set was used to train the smaller model, we classify it as an *in-distribution* dataset. In contrast, no training data from SVAMP, ASDiv, or MultiArith was used, as they were exclusively employed for testing purposes and thus considered *out-of-distribution*.

#### Implementation Details

We used the Llama3 70B model [] as the large language model (LLM) to generate the rationales. We performed distillation on different smaller models ranging from 0.5B to 7B parameters, including Qwen2 0.5B [], Qwen2 1.5B [], SmolLM 1.7B [], Gemma 2B [], and Gemma 7B []. All smaller models were fine-tuned using LoRA [] with a rank of 16, and alpha of 32. We used a learning rate of 3e-4 for Qwen models with a cyclic scheduler, while we set 2e-4 as the learning rate for other models and used a linear scheduler. We train all models for 3 epochs. We implemented all our experiments using the Unsloth FastLanguageModel [] and used the VLLM library [] for inference. We set the temperature $t=0$ for data generation from the LLM while $t=0.7$ was used for generating samples from the smaller model at each iteration. We set the number of generated samples or $K$ to 10. We report Top-1 accuracy (maj@1).

Our proposed approach was compared against a set of single-strategy distillation methods. In this work, we employed three reasoning strategies: Chain-of-Thought (CoT) [], Program-of-Thought (PoT) [], and Least-to-Most (L2M) [].

## 5 Results and Discussion

**Table 1: Top-1 (maj@1) accuracy comparison across different sized models (Gemma 2B and 7B, SmolLM 1.7B, Qwen 0.5B and 1.5B) on four mathematical datasets: GSM8K, ASDiv, MultiArith, and SVAMP. “Combined” refers to the scenario where data from all three reasoning strategies are merged and then used for distillation. We report the two variants of SIKeD: “Adaptive” and “All”, and the best performance is highlighted.**
|  |  | Dataset |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Model | Method | In Distribution | Out of Distribution |  |  |
|  |  | GSM8K | ASDiv | MultiArith | SVAMP |
|  | CoT | 67.40 | 68.76 | 98.33 | 66.80 |
|  | L2M | 69.29 | 64.69 | 96.11 | 64.80 |
|  | PoT | 71.34 | 67.85 | 98.89 | 75.00 |
| Gemma 7B | Combined | 70.74 | 69.11 | 99.44 | 69.40 |
|  | SIKeD |  |  |  |  |
|  | Adaptive | 73.84<br>($\uparrow$ +2.5) | 70.59<br>($\uparrow$ +1.5) | 99.44<br>(-) | 72.90<br>($\downarrow$ -2.1) |
|  | All | 71.42<br>($\uparrow$ +0.1) | 70.76<br>($\uparrow$ +1.6) | 100.0<br>($\uparrow$ +0.6) | 76.40<br>($\uparrow$ +1.4) |
|  | CoT | 36.54 | 54.01 | 87.22 | 41.90 |
|  | L2M | 36.92 | 43.47 | 81.67 | 31.60 |
|  | PoT | 44.05 | 58.13 | 90.56 | 56.80 |
| Gemma 2B | Combined | 44.05 | 57.96 | 84.44 | 56.20 |
|  | SIKeD |  |  |  |  |
|  | Adaptive | 47.23<br>($\uparrow$ +3.2) | 59.05<br>($\uparrow$ +0.8) | 91.11<br>($\uparrow$ +0.6) | 58.60<br>($\uparrow$ +1.8) |
|  | All | 45.26<br>($\uparrow$ +1.2) | 58.00<br>($\downarrow$ -0.1) | 88.89<br>($\downarrow$ -0.6) | 56.70<br>($\downarrow$ -0.1) |
|  | CoT | 16.38 | 30.37 | 58.89 | 22.60 |
|  | L2M | 18.73 | 22.13 | 53.89 | 17.90 |
|  | PoT | 23.73 | 43.77 | 61.11 | 34.50 |
| SmolLM 1.7B | Combined | 24.56 | 46.77 | 67.22 | 35.90 |
|  | SIKeD |  |  |  |  |
|  | Adaptive | 27.98<br>($\uparrow$ +3.4) | 47.20<br>($\uparrow$ +0.4) | 72.22<br>($\uparrow$ +5.0) | 37.80<br>($\uparrow$ +1.9) |
|  | All | 27.75<br>($\uparrow$ +3.2) | 48.16<br>($\uparrow$ +1.4) | 72.22<br>($\uparrow$ +5.0) | 39.50<br>($\uparrow$ +3.6) |
|  | CoT | 55.57 | 68.76 | 99.44 | 66.30 |
|  | L2M | 54.59 | 63.69 | 96.67 | 62.30 |
|  | PoT | 64.22 | 66.94 | 95.56 | 74.30 |
| Qwen 1.5B | Combined | 64.44 | 67.64 | 98.89 | 73.20 |
|  | SIKeD |  |  |  |  |
|  | Adaptive | 64.97<br>($\uparrow$ +0.5) | 68.98<br>($\uparrow$ +1.3) | 99.44<br>(-) | 75.40<br>($\uparrow$ +1.1) |
|  | All | 64.14<br>($\downarrow$ -0.3 ) | 68.55<br>($\uparrow$ +0.9) | 98.89<br>($\downarrow$ -0.5) | 73.50<br>($\downarrow$ -0.8) |
|  | CoT | 36.47 | 54.66 | 83.89 | 43.00 |
|  | L2M | 33.59 | 49.76 | 76.67 | 44.60 |
|  | PoT | 41.62 | 56.83 | 92.22 | 51.40 |
| Qwen 0.5B | Combined | 42.38 | 57.79 | 90.56 | 51.40 |
|  | SIKeD |  |  |  |  |
|  | Adaptive | 43.14<br>($\uparrow$ +0.8) | 58.44<br>($\uparrow$ +0.7) | 93.33<br>($\uparrow$ +1.1) | 51.70<br>($\uparrow$ +0.3) |
|  | All | 43.06<br>($\uparrow$ +0.8) | 58.61<br>($\uparrow$ +0.8) | 92.22<br>(-) | 52.10<br>($\uparrow$ +0.7) |

#### LLM Based Distillation

We start by distilling smaller models using the reasoning dataset generated using the LLM in two variations: using data from a single strategy (CoT, PoT, or L2M), and a combination of all three strategies (referred to as “Combined”).
Table 1 compares the accuracies of the approaches across four mathematical datasets. The “Combined” approach benefited smaller models, yielding slight improvements for the Qwen 0.5B, Qwen 1.5B, and SmolLM 1.7B models. However, it showed little to no improvement, and sometimes even worse performance, for the larger Gemma 2B and 7B models. This indicates that simply merging the distillation data for each strategy is not sufficient for effective multi-strategy distillation.

#### Consistent improvement across in-distribution dataset

Compared to the traditional LLM-based distillation approaches, we observe consistent improvements with SIKeD across all models, ranging from 0.5B to 7B parameters as highlighted in Table 1. On the in-distribution GSM8K dataset, both Gemma 2B and 7B show significant gains of +3.2 points and +2.5 points respectively (44.05 $\rightarrow$ 47.23 and 71.34 $\rightarrow$ 73.84, respectively). Similarly, SmolLM showed the largest improvement of +3.4 points (24.56 $\rightarrow$ 27.98). Similarly, the smaller Qwen models see gains of +0.5 points for the larger variant (1.5B) and +0.8 points for the smaller variant (0.5B).

#### SIKeD performs well on out-of-distribution datasets

For the out-of-distribution datasets, there is a steady improvement on the ASDiv dataset, with Gemma 7B gaining +1.6 points (69.11 $\rightarrow$ 70.76), +0.8 points for Gemma 2B (58.13 $\rightarrow$ 59.05), +1.4 points for SmolLM (46.77 $\rightarrow$ 48.16), +1.3 points for Qwen 1.5B (67.64 $\rightarrow$ 68.98), and +0.8 points for Qwen 0.5B (57.79 $\rightarrow$ 58.61). A similar trend is seen for the MultiArith dataset, where SmolLM shows the largest gain of +5 points. It is followed by Qwen 0.5B with +1.1 points, while other models outperform the baseline. In particular, Gemma 7B achieves a perfect score of 100. The results are similar for the SVAMP dataset, with Qwen 0.5B, Qwen 1.5B, SmolLM 1.7B, Gemma 2B, and Gemma 7B gaining +0.7, +1.1, +3.6, +1.8 and +1.4 points, respectively.

#### Biasing SIKeD in favor of our strategy of choice

Figure: Figure 3: Accuracy comparison between single distillation strategies of CoT, PoT, and L2M with SIKeD biased training using the same strategy using the Gemma 7B model.
Refer to caption: x3.png

For some tasks, one strategy might be better than the others due to its performance, lower cost, or better suitability for some use cases (for example, PoT is significantly better for SVAMP compared to other strategies). In such cases, it would be beneficial to bias the selection towards that strategy (^3^33Note that this is different from the already biased selection of the smaller model, as our biased strategy may not be the default biased choice of the smaller model.). This can be done by taking only the sample from our choice of biased strategy when more than one strategy is correct from the model-generated samples. For example, if for a given data point, a smaller model samples both CoT and PoT correctly, and our biased strategy choice is PoT, we will ignore the CoT output and take only the PoT. Figure 3 compares the individual distillation strategy with the biased SIKeD. Using Gemma 7B as a smaller model across all datasets, SIKeD outperforms individual distillation strategies by a margin of 2-4 points, highlighting the effectiveness of SIKeD over other distillation approaches.

#### How many iterations to run for SIKeD

Figure: Figure 4: Iterative accuracy comparison for the Gemma 2B model across all datasets. The process is stopped when the gains diminish or when it is no longer cost effective to continue.
Refer to caption: x4.png

With each iteration of SIKeD, the model learns to solve a task using different strategies and adjusts its strategy choice for a given task. This allows for continuous training of SIKeD.
Figure 4 illustrates the accuracy improvements across iterations for the Gemma 2B model on various datasets. The iterative training is stopped when accuracy shows only marginal improvements or declines.
In our experiments, three iterations have consistently proven to be the optimal balance across different models and datasets.

#### How the strategy distribution changes over iterations

Figure 5 illustrates the strategy distribution across different iterations for the GSM8K dataset using the SmolLM 1.7B model. Iteration 0 represents the baseline “combined” training from Table 1, and as expected, the smaller model is initially biased towards one strategy (PoT in this case). Iterations 1, 2, and 3 show the model’s progression using SIKeD, where it learns to diversify and select the suitable strategy for the given problem. Notably, while PoT remains the dominant strategy, the model improves its usage of the other two strategies—CoT and L2M—which were absent in the biased baseline. This diversification of strategies results in an overall gain of +3 points over the baseline.

Figure: Figure 5: Strategy distribution over iterations for GSM8K dataset using SmolLM 1.7B model.
Refer to caption: x5.png

#### Training from the last checkpoint vs training from pre-trained checkpoint

In our work, we iteratively train from the last checkpoint using on-policy training as we expect continuous improvements in the model performance with a newly learned strategy. However, an alternative approach uses off-policy training (training the pre-trained model at each iteration) to achieve strong performance [].
We compared on-policy training (our proposed approach) with off-policy training (as in Gulcehre et al. []), utilizing both LLM-generated and self-generated data, and observed a notable decrease in the overall accuracy with off-policy training. Note that we used all of the LLM data at each iteration for off-policy training as the training is done on the pre-trained model. On the GSM8K dataset, our on-policy approach outperformed off-policy training by +6 points (45.26 vs 38.90) using the Gemma 2B model. A similar trend was observed on out-of-distribution datasets, where SIKeD outperformed with a gain of +4-7 points on both the SVAMP and ASDiv datasets (58.6 vs 51.3 for SVAMP and 59.05 vs 55.44 for ASDiv) and a gain of +2 points on the MultiArith dataset (91.11 vs 88.33).

#### Is data mixing better than LLM-based distillation or self-distillation

Figure: Figure 6: Accuracy comparison with different data mix controlled by $\alpha$. $\alpha=1$ means all data is from the LLM only, while $\alpha=0$ means only smaller model data is used.
Refer to caption: x6.png

Mixing data from the smaller model with the LLM helps in bridging the distributional gap between the LLM and the smaller model.
We explore the role of the data mixing rate alpha ($\alpha$) to validate our hypothesis regarding data mixing. When $\alpha$ is set to 1, only LLM-generated data is used, while at the other extreme, $\alpha=0$ means only data generated by the smaller model is used. As shown in Figure 6, an $\alpha$ value between 0 and 1 performs better than either of the extreme cases.

#### Qualitative analysis demonstrating that smaller models learn to choose the right strategy over iterations

The increase in the performance of smaller models can be attributed to the change in the distribution of reasoning strategies as the model becomes more aligned with its strategy choices. Figure 7 shows that a model can solve a given problem in the third iteration by switching the reasoning strategy which it initially struggled with. Furthermore, Figure 8 shows an example of a model correcting its reasoning chain over time due to its improved reasoning capabilities.

Figure: Figure 7: A qualitative example showing how Gemma 2B shifts its strategy selection over iterations to solve a given task.

Figure: Figure 8: Over iteration, Gemma 2B model learnt to fix its mistake using the same strategy

## 6 Related Work

#### Knowledge Distillation for Reasoning Tasks

Knowledge distillation [, ] is a widely-used technique for transferring knowledge from a large language model (LLM) to a smaller model. Previous research has focused on transferring intermediate reasoning steps to smaller models, either step-by-step [, , ] or by combining multiple strategies simultaneously [, ]. These approaches can be viewed as aggregating diverse data sources for distillation, similar to the LLM data approach in our work. In contrast, [] and [] focus on *self-distillation*, where a model improves its performance without external data or knowledge. Specifically, [] leverages word embeddings, while [] uses temporal model checkpoints as a proxy for ground truth. However, both approaches rely solely on data generated by the smaller model and exclude LLM data. Our method strikes a balance between these two extremes by using LLM data to learn multiple strategies and self-generated data to optimize for the right strategy choice.

#### Self-learning

Previous studies, such as [, , , ], have shown the effectiveness of the *self-training paradigm* in NLP tasks but are limited by the choice of tasks. While ReST [] uses off-policy training, we find on-policy training more suitable for our case both in terms of data efficiency and performance. On-policy training also allows a better choice of learning strategies, since the model can use its most recent learning. Agarwal et al. [] introduces Generalized Knowledge Distillation (GKD), an on-policy training method that aligns the distributions of large language models (LLMs) and smaller models by incorporating output sequences sampled from the student during training. However, the task was limited to the distribution alignment and not to aligning the strategy choices in a multi-strategy distillation. Simply applying GKD would not address this issue, as it would force the smaller model to learn all strategies, which is impractical given its limited capacity.

Finally, we compare our distillation strategies with LLM-based distillation using both individual strategies [, , ] and a combination of several strategies at once [, ].

## 7 Conclusion

We propose SIKeD: Self-guided Iterative Knowledge Distillation, that addresses the challenge of distilling multistep reasoning skills from large language models (LLMs) to smaller models. Unlike traditional methods, which often leave smaller models biased towards a single strategy, SIKeD uses iterative *self-guided* training, combining LLM and self-generated data to improve overall reasoning in smaller models. We demonstrate our approach across various mathematical reasoning datasets and demonstrate that SIKeD improves the ability of smaller models to handle complex reasoning, achieving significant performance gains.

## 8 Appendix

Figure: Figure 9: 8-shot Chain of Thought demonstration.

Figure: Figure 10: 8-shot Least to Most demonstration.

Figure: Figure 11: 8-shot Program of Thought demonstration.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can the authors comment on the potential of SIKeD to be effective in domains outside of mathematical reasoning? Have there been any preliminary experiments or considerations in this direction?

2. Could the authors elaborate on the computational efficiency of SIKeD, especially in terms of the number of iterations required for convergence and the resources needed for each iteration?

3. The paper discusses various smaller models, but does not discuss how the size of the smaller model affects the distillation process and the final performance. Are there any insights on how SIKeD scales with different model sizes?
soundness: 3
strengths: 1. The paper presents a unique approach to knowledge distillation by introducing the concept of self-guided iterative training. This method allows smaller models to dynamically adjust their strategy preferences, which is a creative solution to the challenge of strategy distribution mismatch in traditional distillation.

2. The experiments are well-designed and conducted across various mathematical reasoning datasets, providing a thorough evaluation of SIKeD's effectiveness. The improvements in performance metrics are substantial and clearly demonstrated.

3. The paper is well-organized, with a clear problem statement and a logical flow of ideas. The methodology is explained in detail, making it easy for readers to understand the proposed approach and its implications.
summary: The paper presents SIKeD, a knowledge distillation approach to enhance smaller models with reasoning skills from Large Language Models (LLMs). SIKeD employs an iterative process that enables smaller models to learn multiple strategies and select the most appropriate for a given task, addressing the issue of strategy bias found in conventional distillation methods. The paper reports that SIKeD outperforms traditional techniques on mathematical reasoning tasks across various smaller model sizes.
weaknesses: 1. The paper's methodology, SIKeD, is contingent upon the quality of the initial LLM data. There is a need for further exploration on how fluctuations in LLM data quality might influence the distillation process and the performance of the resulting smaller models.

2. The study primarily focuses on mathematical reasoning tasks, with less clarity on the transferability of SIKeD to other reasoning domains such as commonsense or symbolic reasoning. Additional investigation into the broader applicability of SIKeD could be valuable.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: - Q1: Have the authors considered applying SIKeD to tasks outside of mathematical reasoning to test the generalizability of the strategy selection mechanism?
- 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?
- 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.
- 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?
- 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?
- 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?
soundness: 3
strengths: - The idea is novel. SIKeD leverages the idea of constructivist learning theory and uses an iterative self-guided approach for multi-strategy distillation, compared to previous single-step distillation.
- SIKeD shows generalization across in- and out-of-distribution mathematical reasoning datasets, demonstrating its effectiveness in diverse contexts.
- SIKeD is evaluated using various small model types, showing consistent improvements over the baselines across different model types.
- The paper is well-written and easy to follow.
summary: This paper proposes SIKeD, a novel knowledge distillation approach for transferring multistep reasoning skills from large language models (LLMs) to smaller models, particularly for mathematical reasoning tasks. Traditional distillation methods struggle with strategy selection, often resulting in smaller models biased towards a single strategy. SIKeD addresses this by allowing the smaller model to iteratively learn and apply various strategies, combining LLM-generated data with its self-generated outputs to refine strategy selection. The method demonstrates improvements over single-strategy and combined-strategy distillations, achieving superior results across in-distribution and out-of-distribution mathematical reasoning datasets.
weaknesses: - W1: 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.
- W2: The paper lacks comparison with knowledge distillation methods.
- W3: The code is not available for reproduction.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see weakness
soundness: 3
strengths: 1. intuitive and effective method

2. thorough experimental analysis
summary: It is hard for a small LLM to learn multiple correct trajectories towards the same ground truth answer, which might affect its reasoning capability, especially the generalization to OOD reasoning tasks. The paper has tested distilling multi-trajectory training data  for post-training smaller LLM from large LLMs and observed unsatisfactory results. To alleviate this problem, the paper proposes a mix of self-generated multi-trajectory training data together with the distilled data for SFT. Surprisingly it achieves better reasoning performance compared to existing methods, especially on OOD tasks.
weaknesses: The only concern I  have on the experiment part is that only testing the small model’s preference on COT, POT and L2M is a bit constrained. 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. Because in each strategy, for example in COT, a model can also generate multiple cot trajectories that lead to the correct answer. I’m curious to see which of the methods can improve the general diversity of the model’s output trajectories the most and whether this diversity is aligned with model’s OOD performance.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
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. 

**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.

------

[1] Ho N, Schmid L, Yun S Y. Large language models are reasoning teachers[J]. arXiv preprint arXiv:2212.10071, 2022.

[2] Fu Y, Peng H, Ou L, et al. Specializing smaller language models towards multi-step reasoning[C]//International Conference on Machine Learning. PMLR, 2023: 10421-10430.

[3] Zhu X, Qi B, Zhang K, et al. PaD: Program-aided Distillation Can Teach Small Models Reasoning Better than Chain-of-thought Fine-tuning[J]. arXiv preprint arXiv:2305.13888, 2023.

[4] Mobahi H, Farajtabar M, Bartlett P. Self-distillation amplifies regularization in hilbert space[J]. Advances in Neural Information Processing Systems, 2020, 33: 3351-3361.

[5] Dohmatob E, Feng Y, Yang P, et al. A tale of tails: Model collapse as a change of scaling laws[J]. arXiv preprint arXiv:2402.07043, 2024.
soundness: 2
strengths: This paper carefully presents the relevant method analysis and explores how generation strategies evolve across multiple iterations in small models, which is a very intriguing phenomenon.
summary: This paper proposes **S**elf-guided **I**terative **K**nowledge **D**istillation (SIKeD), which utilizes outputs of LLMs (teacher) and small models (student) to Iterative train the student model. The main process can be summarized as follows:

source data (GSM8K training set) → LLMs generate Cot, PoT and other format reasoning → training small models → small models augment data → filtering → training small models → small models augment data → filtering → … … n …

Source data: GSM8K training set (7,473 samples)

Model: Qwen2 0.5B, 1.5B, SmolLM 1.7B, Gemma 2B / 7B,
weaknesses: 1. The mathematical notation is overly verbose. There is no proof or any other theoretical contributions. However, this paper employs overly complex notation. And there is also no revision on loss function. The overly complex notation will prevent quick understanding. The main process of proposed method can be summarized in one or two sentences. And all operations performed at the dataset level.

2. 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.

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

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: 1. Can the authors elaborate on the underlying reasons for smaller models' bias towards specific strategies? Does the preferred strategy vary?
2. 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.
3. Please explain the mechanism by which mixing the data promotes diverse strategy selection in smaller models.
4. How to judge if the generated r_i is correct? (L294)
5. Is it a typo in L211 & L215 that the same notation appears in different contexts?
6. 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
soundness: 3
strengths: The proposed problem is worth exploring.
The paper is well-written and easy to follow.
summary: This paper addresses the challenge of knowledge transfer from large language models (LLMs) to smaller, more efficient models. The authors identify that smaller models typically demonstrate constrained capabilities and tend to favor single reasoning strategies. To address this limitation, the authors introduce SIKED, an approach that leverages self-generated outputs to create iteratively mixed training datasets. The method promotes the development of diverse reasoning strategies during the knowledge transfer process. The effectiveness of SIKED is demonstrated through empirical evaluations on several mathematical reasoning datasets, showing improvements over baseline approaches.
weaknesses: 1. 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.
2. 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.

Meta-review: **Summary:** This paper proposes SIKeD, a Self-guided Iterative Knowledge Distillation approach aimed at improving reasoning capabilities of smaller models by leveraging both large language models (LLMs) and self-generated outputs. The method iteratively adjusts the proportion of teacher and student-generated data, guided by a mixing parameter. The evaluation, primarily conducted on mathematical reasoning datasets, shows improvement over traditional knowledge distillation baselines in both in-distribution and out-of-distribution tasks.

**Decision:** While SIKeD presents a novel iteration-based distillation framework for improving reasoning in smaller models, it falls short in several areas. The evaluation is narrowly focused on mathematical reasoning tasks, lacking broader validation across diverse reasoning domains (JLQW, YiHw, ytP5). Furthermore, the contributions overlap with existing works on iterative data mixing and self-distillation, offering limited methodological novelty (YiHw). The lack of theoretical innovations also somewhat undermines the significance of the proposed approach (YiHw). Lastly, the method’s dependence on the quality of LLM data introduces practical limitations (JLQW).

Decision: Reject