Key-Point-Driven Mathematical Reasoning Distillation of Large Language Model

Mathematical reasoning tasks, exemplified by benchmarks such as GSM8K (Cobbe et al., 2021) and SVAMP (Patel et al., 2021), present a substantial challenge for LLMs. To enhance LLMs’ performance in this domain, researchers have identified two primary strategies.

Chain-of-Thought Reasoning LLMs’ reasoning ability can be enhanced by prompting them to articulate intermediate steps towards a solution, as demonstrated by Wei et al. Wei et al. (2022). This insight has spurred various advancements (Chen et al., 2023; Wang et al., 2023b, a; Liu et al., 2023) that refine reasoning paths: Chen et al. Chen et al. (2023) prompt LLMs to generate executable code; Wang et al. Wang et al. (2023b) use multiple reasoning paths with a voting mechanism; Wang et al. Wang et al. (2023a) have LLMs create a plan before reasoning; Liu et al. Liu et al. (2023) employ diverse reasoning prompts for problem-solving; Zhong et al. Zhong et al. (2024) encourage LLMs to deeply understand problems and leverage key information for better reasoning. Building on these methods, our work introduces Key-Point-Driven Mathematical Reasoning Distillation (KPDD) to further enhance SLMs’ mathematical reasoning.

Finetuning-based Reasoning refines LLMs such as Llama2 (Touvron et al., 2023b), Qwen (Bai et al., 2023), and Baichuan2 (Yang et al., 2023) by integrating techniques from advanced models like GPT-4 (OpenAI, 2023) and PaLM-2 (Anil et al., 2023). Notably, Yuan et al.Yuan et al. (2023) utilize Rejection Sampling Fine-Tuning (RFT) to enhance LLMs’ mathematical reasoning, while WizardMath(Luo et al., 2023) employs Reinforcement Learning from Evolved Instructions Feedback (RLEIF) to improve LLaMA-2’s reasoning abilities. MAmmoTH (Yue et al., 2023) combines Chain-of-Thought (CoT) and Program-of-Thought (PoT) rationales for more effective instruction-tuning of LLMs in math problem-solving. Despite their effectiveness, the large model sizes of these LLMs continue to limit their deployment efficiency.

Knowledge Distillation optimizes LLMs for practical use by transferring knowledge from larger models to smaller, efficient ones (Zhu et al., 2023b). Recent research (Magister et al., 2023; Shridhar et al., 2023; Ho et al., 2023; Fu et al., 2023) focuses on endowing compact models ($\leq$ 1B parameters) like T5 (Raffel et al., 2020) and GPT-2 (Radford et al., 2019) with advanced reasoning capabilities from LLMs such as GPT-4 (OpenAI, 2023) and PaLM-2 (Anil et al., 2023). For instance, Ho et al. Ho et al. (2023) fine-tune student models using accurate reasoning paths from LLMs, Shridhar et al. Shridhar et al. (2023) train dual-model systems on sub-questions and solutions, and Fu et al. Fu et al. (2023) propose scaling down general competencies of smaller models to enhance task-specific performance. Our work introduces a novel distillation approach where two SLMs independently extract the core question and key problem-solving information from an original question. These key points are then utilized to guide another SLM in solving the original question effectively.

Our KPDD method begins by creating a mathematical reasoning dataset from LLMs, which is then used for SLM fine-tuning. In our paper, we use in-context learning (Dong et al., 2023; Min et al., 2022; Rubin et al., 2022) to prompt LLMs for constructing the reasoning dataset. Furthermore, in stage 3, our KPDD method employs two distillation approaches: one distills the SLM to generate CoT rationales for problem-solving, and the other distills the SLM to generate PoT rationales for problem-solving. In other words, our KPDD method can be divided into two approaches: KPDD-CoT and KPDD-PoT.

After constructing these reasoning datasets, we use them to fine-tune the SLMs. In the KPDD, we fine-tune three SLMs: the first SLM, called KPDD-CoT/PoT-core, is used to extract the core question from the original problem, the second SLM, called KPDD-CoT/PoT-info, extracts the problem-solving information, and the third SLM, called KPDD-CoT/PoT-solve, uses both the core question and problem-solving information to solve the original question.

Figure 4 illustrates the inference process of KPDD. After fine-tuning, the process for solving a given question involves three main steps:

1. 1.

   Core Question Extraction: First, we use the KPDD-CoT/PoT-core model to extract the core question from the original problem. This step isolates the essential part of the problem that needs to be addressed.

2. 2.

   Problem-Solving Information Extraction: Next, the KPDD-CoT/PoT-info model extracts the relevant problem-solving information. This model identifies and lists the necessary context and data required to solve the core question.

3. 3.

   Solution Generation: Finally, based on the original question, the core question, and the problem-solving information, the KPDD-CoT/PoT-solve model generates rationales in either CoT or PoT format to solve the original question. For KPDD-PoT, this involves generating Python code that can compute the answer.

This structured approach ensures that each model focuses on a specific aspect of the problem-solving process, leading to more accurate and reliable solutions.

In our paper, the training dataset is derived from the GSM8K training set, which comprises diverse grade school math word problems (Cobbe et al., 2021). Additionally, the mathematical reasoning capabilities of the SLMs are evaluated using the GSM8K test set, along with other datasets including ASDiv, which contains diverse math word problems (Miao et al., 2020), SVAMP, which features math word problems with varying structures (Patel et al., 2021), and MultiArith, which consists of arithmetic word problems (Roy and Roth, 2015). The statistics of these datasets are summarized in Table 1. This comprehensive evaluation approach ensures that the SLMs’ mathematical reasoning capabilities are thoroughly tested across a variety of problem types and structures, providing a robust assessment of their performance.

We employ GPT-4 as the teacher LLM to construct our training dataset and utilize FlanT5 models—Small (60M), Base (250M), and Large (760M) (Chung et al., 2022)—as student SLMs. We manually create 8 demonstrations to guide GPT-4 in generating 4 reasoning paths for each dataset (KPDD-CoT and KPDD-PoT). Fine-tuning of all student SLMs is conducted using the Huggingface library (Wolf et al., 2020) on an NVIDIA 3090 GPU with 24 GB RAM. The learning rate for fine-tuning is set to 5e-4, with a total of 10 fine-tuning epochs.

Proprietary Large Language Models We present CoT prompting results from an array of SoTA LLMs, such as OpenAI’s GPT-4, ChatGPT (gpt-3.5-turbo), Google’s PaLM-2, and Anthropic’s Claude-2.

Open-Source Large Language Models We present mathematical reasoning performance of Llama-2-7B, CodeLLaMA-7B, and their fine-tuned versions, such as Platpus-2, WizardMath, TORA.

Fine-tuned Small Language Models We present some works that try to fine-tune SLMs under 1B, such as Ho et al. Ho et al. (2023) fine-tune GPT-3-ada, Fu et al. Fu et al. (2023) fine-tune FlanT5, and Shridhar et al. Shridhar et al. (2023) fine-tune GPT-2.

Table 2 showcases our method’s performance on four mathematical datasets, revealing key insights:

1. 1.

   KPDD-CoT Enhances Mathematical Reasoning: KPDD-CoT significantly improves the mathematical reasoning capabilities of SLMs, with absolute improvements ranging from 5.01% to 15.51% across tasks. Traditional baselines typically rely on CoTD, which involves generating numerous steps and performing extensive calculations. However, CoTD often encounters semantic misunderstanding errors that hinder the improvement of SLMs’ mathematical reasoning abilities. In contrast, KPDD-CoT employs extra SLMs to extract key points (including the core question and problem-solving information) of the question and uses these key points to guide the SLMs’ reasoning. This approach significantly reduces the semantic misunderstanding errors of CoTD, making KPDD-CoT better suited for improving the mathematical reasoning ability of SLMs.

2. 2.

   KPDD-PoT Outperforms State-of-the-Art: KPDD-PoT surpasses previous state-of-the-art fine-tuned SLMs at all scales, with absolute improvements between 32.18% and 54.63% across tasks. Furthermore, KPDD-PoT’s accuracy is higher than that of KPDD-CoT, highlighting the advantage of rationales in PoT format in enhancing SLMs’ reasoning capabilities. Our analysis finds that the mathematical reasoning performance of CoTD is limited not only by semantic misunderstanding errors but also by calculation errors. PoTD converts rationales from CoT format into PoT format, formulating the reasoning process into a Python program and sending it to an extra Python interpreter to generate the final answer. This method transfers numerical computation from SLMs to a Python interpreter, avoiding calculation errors. Additionally, by extracting key points of the question, KPDD-PoT implicitly enhances the SLMs’ understanding of the question, thereby improving their overall mathematical reasoning capabilities.

3. 3.

   Importance of Model Size: The efficacy of mathematical reasoning distillation in SLMs is highly dependent on model size; larger models assimilate more reasoning knowledge, leading to superior performance. For instance, under KPDD-PoT, FlanT5-Small achieves 20.77% accuracy on GSM8K, FlanT5-Base reaches 34.57%, and FlanT5-Large attains 46.32%.

4. 4.

   Strong Transferability of KPDD: KPDD exhibits strong transferability. The distillation dataset of KPDD is constructed based on the GSM8K training dataset, and we evaluate our SLMs on several mathematical reasoning datasets, including the GSM8K test dataset, ASDiv dataset, SVAMP dataset, and MultiArith dataset. Our experimental results show that KPDD not only achieves good reasoning performance on the GSM8K test dataset but also performs well on the ASDiv, SVAMP, and MultiArith datasets. These results demonstrate that KPDD has strong transferability and further corroborate that SLMs do not improve their reasoning performance through data leakage.

In this subsection, we delve into the impact of various components within KPDD. We have considered five distinct categories, which include: 1. Original SLMs without any fine-tuning; 2. SLMs with original CoT/PoT distillation; 3. SLMs with core distillation combined with CoT/PoT distillation; 4. SLMs with problem-solving information distillation combined with CoT/PoT distillation; 5. SLMs with KPDD. For each of the latter four categories, we have constructed corresponding reasoning datasets, each containing a single reasoning path per question. Following this, we have utilized FlanT5-base as our foundation for SLMs, and we have fine-tuned these models using the aforementioned reasoning datasets. To evaluate the reasoning capabilities of these SLMs, we have tested them on the GSM8K test dataset, as well as on the ASDiv, SVAMP, and MultiArith datasets.

Tables 3 and 4 present the results of our experiments, from which we make several observations: (1) We observe a significant performance improvement in Category 2 compared to original SLMs. Specifically, under CoT reasoning, Category 2 achieves an average accuracy gain of 4.61% across multiple datasets, while under PoT reasoning, it achieves a substantial average accuracy improvement of 32.91%. These experimental results indicate that CoTD and PoTD can markedly enhance the mathematical reasoning ability of SLMs. (2) We find that Categories 3 and 4 exhibit a further performance increase relative to Category 2. Specifically, in the context of CoT reasoning, Categories 3 and 4 achieve average accuracy gains of 0.55% and 0.45% respectively over Category 2 across multiple datasets. Under PoT reasoning, the gains are more pronounced with Categories 3 and 4 achieving average accuracy improvements of 4.21% and 6.13% respectively. This suggests that SLMs can deepen their understanding of questions by focusing on key points, thereby further enhancing their mathematical reasoning ability. (3) In Category 5, we combine the core questions with the problem-solving information to guide SLMs in addressing the questions. The results are promising: Category 5 achieves an average accuracy of 9.98% under CoT reasoning and a remarkable 45.20% under PoT reasoning across multiple datasets. This indicates that key points in questions play a crucial role in boosting the reasoning capabilities of SLMs, and that combining several key points provides richer information, leading to further improvements in their reasoning abilities.

In this subsection, we investigate the impact of SLM Quantity in KPDD. We consider five distinct categories: $\mathrm{I}$. Using one SLM to simultaneously extract the core question and problem-solving information, and solve the original question; $\mathrm{II}$. Using one SLM to extract the core question and problem-solving information, and another SLM to solve the original question; $\mathrm{III}$. Using one SLM to extract the core question, another SLM to extract the problem-solving information, and a third SLM to solve the original question; $\mathrm{IV}$. Using one SLM to extract the problem-solving information, another SLM to extract the core question, and both to solve the original question; $\mathrm{V}$. Using one SLM to extract the core question, another SLM to extract the problem-solving information, and a third SLM to solve the original question. For each category, we create corresponding reasoning datasets, each containing a single reasoning path per question. We utilize FlanT5-base as our base SLMs, fine-tuning them on these reasoning datasets. To assess their reasoning capabilities, we evaluate these SLMs on the GSM8K test dataset, as well as on the ASDiv, SVAMP, and MultiArith datasets.

Table 5 and 6 present the results of our experiments, from which we make several observations: (1) Compared to other categories, Category $\mathrm{I}$ performed worse. For KPDD-CoT, Category $\mathrm{I}$ achieved an average accuracy of 7.16% across multiple datasets, while for KPDD-PoT, it achieved an average accuracy of 39.61%. This suggests that the limited model size of a single SLM hinders its performance across multiple tasks. (2) Category $\mathrm{II}$ outperformed Categories $\mathrm{III}$ and $\mathrm{IV}$ in reasoning performance. For KPDD-CoT, Category $\mathrm{II}$ achieved an average accuracy of 9.51% across multiple datasets, while for KPDD-PoT, it achieved an average accuracy of 41.80%. We attribute this result to the importance of the KPDD-CoT/PoT-solve component, where using a single SLM for this phase yields the best reasoning performance. (3) For KPDD-CoT, Category $\mathrm{V}$ achieved an average accuracy of 9.98% across multiple datasets, while for KPDD-PoT, it achieved an average accuracy of 45.20%. This is the highest reasoning performance among all categories, indicating that our approach of using a separate SLM for each component maximizes the performance of each component, thereby maximizing the reasoning performance of KPDD.

In this subsection, we fine-tune CodeT5-Base on our reasoning datasets, which are differentiated by the number of reasoning paths they contain, to analyze the effect of reasoning path multiplicity on the reasoning performance of SLMs. This examination aims to discern how the quantity of reasoning paths in training data influence the model’s ability to perform reasoning tasks.

Figure 5 presents the results of our experiments, which demonstrate that a variety of reasoning paths can bolster the reasoning performance of SLMs. For instance, CodeT5-Base, when trained on an KPDD-PoT dataset featuring four reasoning paths, attains a 34.57% accuracy on the GSM8K test dataset and a 52.29% accuracy on ASDiv. In contrast, CodeT5-Base trained on an KPDD-PoT dataset with only one reasoning path achieves 49.33% accuracy on GSM8K test dataset and 46.1% accuracy on ASDiv. This suggests that the inclusion of multiple reasoning paths in training data can significantly enhance the model’s performance, particularly in tasks requiring explanation generation.

In this subsection, our aim is to verify whether KPDD can indeed reduce semantic misunderstanding errors. KPDD-PoT implicitly includes the reasoning process within its rationales, making it challenging to conduct error analysis on rationales in PoT format. Conversely, rationales in CoT format explicitly contain the reasoning steps, allowing us to clearly understand how the SLM solves the questions step by step, thus facilitating error analysis. Therefore, in this part, we focus on error analysis for rationales in CoT format. To achieve our goal, we randomly sample 100 examples from GSM8K/SVAMP and perform error analysis on the questions with incorrect answers. For a better understanding of KPDD’s effect, we also consider three other scenarios: (1) vanilla CoTD, (2) reasoning that combines vanilla CoTD and core question extraction, and (3) reasoning that combines vanilla CoTD and problem-solving information extraction. Furthermore, to simplify our analysis, we use flanT5-base as our SLMs, and the corresponding reasoning datasets still contain a single reasoning path per question.

The detailed quantitative results are illustrated in Figure 6. By analyzing the experimental results, we found that: (1) Combination of Multiple Errors in SLMs: SLMs tend to exhibit combinations of multiple errors, with calculation errors having the most significant impact on reasoning performance. Specifically, vanilla CoTD on the GSM8K dataset showed 51 understanding errors, 79 calculation errors, and 34 step missing errors, resulting in a total of 164 errors. This number far exceeds the original number of problems, with calculation errors outnumbering other types of errors. Similar results were observed in the SVAMP dataset. This explains why PoTD achieves better reasoning performance than CoTD: PoTD converts vanilla rationales into Python programs, delegating the calculation process to an external Python interpreter to avoid calculation errors. (2) Reduction of Understanding Errors with Key Points: Introducing key points of the original questions effectively reduces understanding errors. Specifically, when core questions were introduced in vanilla CoTD, the number of understanding errors on the GSM8K dataset decreased to 50, and on the SVAMP dataset, it decreased to 53. When problem-solving information was introduced in vanilla CoTD, the number of understanding errors decreased to 48 on GSM8K and to 51 on SVAMP. These results indicate that key points of the original questions help SLMs better understand the questions, thereby reducing understanding errors and improving reasoning performance. (3) Further Reduction of Understanding Errors with Multiple Key Points: Combining multiple key points can further reduce understanding errors. Specifically, KPDD reduced the number of understanding errors to 46 on GSM8K and to 50 on SVAMP. This suggests that KPDD’s method of integrating multiple key points can deepen SLMs’ understanding of the original questions, further reducing understanding errors and enhancing reasoning performance.