Improving Mathematical Reasoning Capabilities of Small Language Models via Feedback-Driven Distillation

TL;DR

A Feedback-Driven Distillation (FDD) framework to enhance SLMs' mathematical reasoning capabilities and a multi-round distillation paradigm to iteratively enrich the distillation datasets, thereby progressively improving the mathematical reasoning abilities of SLMs.

Abstract

Large Language Models (LLMs) demonstrate exceptional reasoning capabilities, often achieving state-of-the-art performance in various tasks. However, their substantial computational and memory demands, due to billions of parameters, hinder deployment in resource-constrained environments. A promising solution is knowledge distillation, where LLMs transfer reasoning capabilities to Small Language Models (SLMs, $\le$ 1B parameters), enabling wider deployment on low-resource devices. Existing methods primarily focus on generating high-quality reasoning rationales for distillation datasets but often neglect the critical role of data quantity and quality. To address these challenges, we propose a Feedback-Driven Distillation (FDD) framework to enhance SLMs' mathematical reasoning capabilities. In the initialization stage, a distillation dataset is constructed by prompting LLMs to pair mathematical problems with corresponding reasoning rationales. We classify problems into easy and hard categories based on SLM performance. For easy problems, LLMs generate more complex variations, while for hard problems, new questions of similar complexity are synthesized. In addition, we propose a multi-round distillation paradigm to iteratively enrich the distillation datasets, thereby progressively improving the mathematical reasoning abilities of SLMs. Experimental results demonstrate that our method can make SLMs achieve SOTA mathematical reasoning performance.

Figures (4)

• Figure 1: The overview of our distillation method. In the initialization stage, we prompt the LLM to create an initial distillation mathematical dataset based on the original mathematical dataset. Next, in the question generation stage, we prompt the LLM to generate more diverse questions based on hard questions and more complex questions from easy ones. These generated questions are used to expand the distillation dataset, enhancing its diversity and complexity. This expanded distillation dataset is then used to fine-tune SLMs from scratch to enhance mathematical reasoning abilities and mitigate catastrophic forgetting. Additionally, we propose a multi-round distillation paradigm to iteratively expand the dataset and further improve the mathematical reasoning capabilities of the SLMs.
• Figure 2: Effect of the Number of Rounds. The experimental results demonstrate the effectiveness of the multi-round distillation paradigm, while additional rounds further enhance the mathematical reasoning performance of SLMs.
• Figure 3: Effect of the Number of Reasoning Paths. The experimental results show that more reasoning paths can make SLMs have better mathematical reasoning performance.
• Figure 4: Analysis for Data Leakage. The experimental results demonstrate that in our approach, the data generated by the LLM exhibits minimal correlation with the data in the test datasets, effectively eliminating the impact of data leakage on the mathematical reasoning performance of the SLM.