TinyR1-32B-Preview: Boosting Accuracy with Branch-Merge Distillation

📄 arXiv: 2503.04872v2 📥 PDF

作者: Lin Sun, Guangxiang Zhao, Xiaoqi Jian, Yuhan Wu, Weihong Lin, Yongfu Zhu, Change Jia, Linglin Zhang, Jinzhu Wu, Junfeng Ran, Sai-er Hu, Zihan Jiang, Junting Zhou, Wenrui Liu, Bin Cui, Tong Yang, Xiangzheng Zhang

分类: cs.CL, cs.AI

发布日期: 2025-03-06 (更新: 2025-03-17)

备注: Preprint

💡 一句话要点

提出Branch-Merge蒸馏方法，提升小模型在数学、编程和科学等领域的性能。

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

关键词: 大型语言模型 模型压缩 知识蒸馏 领域特定学习 模型融合

📋 核心要点

  1. 现有模型蒸馏和迁移学习方法在压缩大型语言模型时，往往难以保持较高的准确率，这是一个挑战。
  2. Branch-Merge蒸馏方法通过分支阶段的知识选择性提炼和合并阶段的跨领域知识转移，提升模型压缩效果。
  3. 实验表明，TinyR1-32B-Preview在数学、编程和科学等领域超越了基线模型，并在AIME 2024上接近教师模型性能。

📝 摘要（中文）

本文提出了一种Branch-Merge蒸馏方法，旨在解决大型语言模型（LLM）小型化过程中性能下降的问题。该方法通过两个阶段增强模型压缩：首先是分支阶段，通过领域特定的监督微调（SFT），将大型教师模型的知识选择性地提炼到专门的学生模型中；然后是合并阶段，将这些学生模型合并，以实现跨领域知识转移并提高泛化能力。实验结果表明，使用DeepSeek-R1作为教师模型，DeepSeek-R1-Distill-Qwen-32B作为学生模型，经过Branch-Merge蒸馏得到的TinyR1-32B-Preview在多个基准测试中优于DeepSeek-R1-Distill-Qwen-32B，包括数学（+5.5分）、编程（+4.4分）和科学（+2.9分），并在AIME 2024上取得了与DeepSeek-R1几乎相同的性能。Branch-Merge蒸馏方法为创建计算成本和时间更低的小型、高性能LLM提供了一种可扩展的解决方案。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）压缩过程中，现有蒸馏方法难以在高精度下有效缩小模型规模的问题。现有方法通常在知识迁移过程中存在信息损失，导致学生模型性能下降，尤其是在需要复杂推理的领域，例如数学和编程。

核心思路：论文的核心思路是利用领域特定的知识蒸馏，将大型教师模型的知识分解并选择性地迁移到多个专门的学生模型中，然后在合并阶段将这些学生模型融合，从而实现跨领域知识的迁移和泛化能力的提升。这种“分而治之”的策略旨在更有效地利用教师模型的知识，并避免单一蒸馏过程中的信息瓶颈。

技术框架：Branch-Merge蒸馏方法包含两个主要阶段： 1. 分支阶段（Branch Phase）：使用领域特定的监督微调（SFT）训练多个学生模型。每个学生模型专注于学习教师模型在特定领域的知识。选择合适的领域和数据集是关键。 2. 合并阶段（Merge Phase）：将训练好的学生模型合并成一个统一的模型。合并策略可能包括简单的权重平均、更复杂的模型融合技术或微调。目标是使合并后的模型能够综合利用各个学生模型的知识，从而提高整体性能。

关键创新：该方法最重要的创新点在于将知识蒸馏过程分解为“分支”和“合并”两个阶段。分支阶段允许针对不同领域进行专门的知识提炼，避免了传统蒸馏方法中知识的混淆和稀释。合并阶段则通过融合不同领域的知识，增强了模型的泛化能力。与传统的单阶段蒸馏相比，Branch-Merge蒸馏能够更有效地利用教师模型的知识，并生成性能更优的学生模型。

关键设计：论文中使用的关键设计包括： 1. 领域选择：根据任务需求选择合适的领域进行分支训练。例如，对于数学能力，可以选择数学相关的训练数据。 2. 监督微调（SFT）：使用领域特定的数据集对学生模型进行微调，以使其更好地学习教师模型在该领域的知识。 3. 模型合并策略：论文可能采用了简单的权重平均或其他模型融合技术来合并学生模型。具体的合并策略可能需要根据实验结果进行调整。 4. 损失函数：在分支阶段，可能使用了标准的知识蒸馏损失函数，例如KL散度损失，来衡量学生模型和教师模型输出之间的差异。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，TinyR1-32B-Preview在多个基准测试中显著优于DeepSeek-R1-Distill-Qwen-32B，具体提升包括：数学（+5.5分）、编程（+4.4分）和科学（+2.9分）。此外，TinyR1-32B-Preview在AIME 2024上取得了与教师模型DeepSeek-R1几乎相同的性能，证明了Branch-Merge蒸馏方法在保持模型性能方面的有效性。

🎯 应用场景

该研究成果可应用于各种需要小型化、高性能语言模型的场景，例如移动设备上的本地部署、资源受限环境下的应用、以及需要快速响应的边缘计算等。通过Branch-Merge蒸馏，可以降低模型部署的计算成本和延迟，同时保持较高的模型性能，从而扩展LLM的应用范围。

📄 摘要（原文）

The challenge of reducing the size of Large Language Models (LLMs) while maintaining their performance has gained significant attention. However, existing methods, such as model distillation and transfer learning, often fail to achieve high accuracy. To address this limitation, we introduce the Branch-Merge distillation approach, which enhances model compression through two phases: (1) the Branch Phase, where knowledge from a large teacher model is \textit{selectively distilled} into specialized student models via domain-specific supervised fine-tuning (SFT); And (2) the Merge Phase, where these student models are merged to enable cross-domain knowledge transfer and improve generalization. We validate our distillation approach using DeepSeek-R1 as the teacher and DeepSeek-R1-Distill-Qwen-32B as the student. The resulting merged model, TinyR1-32B-Preview, outperforms its counterpart DeepSeek-R1-Distill-Qwen-32B across multiple benchmarks, including Mathematics (+5.5 points), Coding (+4.4 points) and Science (+2.9 points), while achieving near-equal performance to DeepSeek-R1 on AIME 2024. The Branch-Merge distillation approach provides a scalable solution for creating smaller, high-performing LLMs with reduced computational cost and time.