SIK e D : Self-guided Iterative Knowledge Distillation for Mathematical Reasoning (2025)

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**nowledg**e** **D**istillation, 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.

## introduction 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 @xcite . one notable strength of large models is their ability to reason and perform multistep reasoning tasks, often considered an important aspect of intelligence @xcite . however, the significant size and computational demands of these large models present several challenges. for example, llama3 models @xcite 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 @xcite . typically, smaller models can be taught to replicate the multistep reasoning capabilities of larger models by incorporating a set of intermediate sequences @xcite . however, these intermediate steps can be derived from several different strategies, such as chain of thought (cot) @xcite , subquestion decomposition @xcite , and program of thoughts (pot) @xcite , among others. a viable solution is to distill these reasoning capabilities into smaller models either by distilling individual strategies @xcite @xcite or by incorporating multiple strategies simultaneously @xcite . 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 (llama-3 70b) can use multiple strategies to generate data but upon distilling, a smaller model (gemma 2b) tends to favor one over the others. we used three-shot prompting on both the larger and smaller model, showing the models one example of each strategy. the goal for the model was then to 1) select a strategy and 2) solve the reasoning task using the selected strategy for a given problem. the results displayed in figure 1 confirm our hypothesis. this happens 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 @xcite . 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 @xcite , 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 @xcite , where the learner builds knowledge during the "assimilation phase" and refines their under-standing during the "accommodation phase" to incorporate new insights. we propose generating outputs using the smaller model in an onpolicy 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 7 billion parameters or less. 2 on four mathematical datasets-gsm8k @xcite , svamp @xcite , asdiv @xcite , and multiarith (roy and roth, 2015)-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. ## preliminaries: llm based distillation problem setup we consider the standard setup for llm-based distillation (also referred to as knowledge distillation), where distillation data is sampled from a larger model with intermediate reasoning steps and a smaller distilled model is fine-tuned on the data @xcite . 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 θ sm (with learnable parameters θ). in this work, we consider a reasoning distillation setup where the distillation dataset d consists of pairs of math questions q i and their numerical answers a i , @xmath0, . . . , n}. our work focuses on improving reasoning in smaller models by teaching them to utilize 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 each question q i and a specific reasoning strategy, denoted as @xmath1, we generate the reasoning chain (or rationale), denoted as r i leading to the final answer: r i ∼ p l (. | pr s , q i ), where, pr s represents the strategy-specific prompt. the prompts used for the generation of reasoning chains are provided in appendix a. ## siked: self-guided iterative knowledge distillation we propose siked, an iterative training procedure 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 prior learning. the following paragraphs discuss the steps involved in our proposed iterative distillation methodology and the training objective. ## 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 @xcite , 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 @xcite with 1,000 samples, asdiv (miao et al., 2020) with 2,300 test samples, and multiarith (roy and roth, 2015) 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. ## conclusion we propose siked: self-guided iterative knowledge distillation, which 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 evaluate 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.