SIKeD: Self-guided Iterative Knowledge Distillation for mathematical reasoning

📄 arXiv: 2410.18574v1 📥 PDF

作者: Shivam Adarsh, Kumar Shridhar, Caglar Gulcehre, Nicholas Monath, Mrinmaya Sachan

分类: cs.AI

发布日期: 2024-10-24

🔗 代码/项目: GITHUB

💡 一句话要点

提出SIKeD，通过自引导迭代知识蒸馏提升小模型在数学推理任务上的性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 数学推理 大型语言模型 自引导学习 迭代训练

📋 核心要点

  1. 现有知识蒸馏方法在数学推理任务中，小模型难以学习LLM的多样化解题策略，倾向于依赖单一策略。
  2. SIKeD方法通过LLM指导小模型学习多种解题策略，并利用小模型自身生成的结果选择最佳策略。
  3. 实验表明，SIKeD在多个数学推理数据集上显著优于传统蒸馏方法，提升了小模型的性能。

📝 摘要（中文）

大型语言模型(LLM)可以通过教授小模型生成解决多步推理任务所需的中间推理过程，从而将其推理能力迁移到小模型。尽管LLM可以通过多种策略准确地解决推理任务，即使没有经过微调，但小模型在蒸馏时，其表达能力不足以拟合LLM在所有策略上的分布，并且倾向于优先考虑一种策略。这种对单一策略的依赖给小模型带来了挑战，当尝试使用其首选策略难以解决的推理任务时，问题尤为突出。为了解决这个问题，我们提出了一种蒸馏方法SIKeD（用于数学推理的自引导迭代知识蒸馏），其中LLM教导小模型使用不同的策略来处理任务，而小模型使用其自身生成的on-policy输出来选择最适合给定任务的策略。训练以自引导迭代的方式持续进行，在每个训练迭代中，都会决定如何将LLM数据与自生成输出相结合。与传统的蒸馏方法不同，SIKeD允许小模型学习哪种策略适合给定的任务，同时不断学习使用不同的策略来解决任务。在各种数学推理数据集上的实验表明，SIKeD在不同规模的小模型上显著优于传统的蒸馏技术。

🔬 方法详解

问题定义：论文旨在解决小模型在数学推理任务中，通过知识蒸馏学习LLM的推理能力时，无法有效学习LLM提供的多种解题策略，从而导致模型泛化能力不足的问题。现有方法通常让小模型拟合LLM的输出分布，但忽略了LLM可能采用多种不同的推理路径，而小模型的容量有限，难以同时学习所有策略。

核心思路：论文的核心思路是让LLM指导小模型学习多种解题策略，并让小模型根据自身生成的输出，选择最适合当前问题的策略。通过这种自引导的方式，小模型可以逐步掌握不同策略的适用范围，从而提高解决复杂问题的能力。

技术框架：SIKeD的整体框架是一个迭代的知识蒸馏过程。首先，LLM作为教师模型，为小模型提供多种解题策略的示例。然后，小模型利用这些示例进行学习，并生成自己的解题过程。接着，小模型根据自身生成的输出，判断哪种策略更适合当前问题。最后，根据判断结果，调整LLM数据和自生成数据的组合方式，进行下一轮迭代训练。

关键创新：SIKeD的关键创新在于引入了自引导机制，让小模型能够根据自身生成的输出，选择最合适的解题策略。这与传统的知识蒸馏方法不同，传统方法通常只是简单地让小模型拟合LLM的输出，而忽略了策略选择的重要性。

关键设计：SIKeD的关键设计包括：1) LLM提供的多种解题策略的生成方式；2) 小模型选择策略的判断依据；3) LLM数据和自生成数据的组合方式。具体的参数设置、损失函数和网络结构等细节在论文中进行了详细描述，但摘要中未提供具体信息，因此未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在多个数学推理数据集上进行了实验，结果表明SIKeD方法显著优于传统的知识蒸馏技术。具体的性能数据和提升幅度在摘要中未提供，因此未知。但实验结果表明，SIKeD能够有效提升小模型在数学推理任务上的性能，使其能够更好地解决复杂问题。

🎯 应用场景

SIKeD方法可应用于各种需要复杂推理能力的场景，例如数学问题求解、代码生成、逻辑推理等。通过将大型模型的知识迁移到小型模型，可以在资源受限的环境中部署高性能的推理系统，例如移动设备、嵌入式系统等。该方法还有助于提高模型的鲁棒性和泛化能力，使其能够更好地应对各种复杂问题。

📄 摘要（原文）

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 (Self-guided Iterative Knowledge Distillation for mathematical reasoning), 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. Our code is available at: https://github.com/kumar-shridhar/SIKeD