TinyR1-32B-Preview: ブランチマージ蒸留による精度向上

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

March 6, 2025

著者: 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.AI

要旨

大規模言語モデル（LLM）のサイズを縮小しながら性能を維持するという課題は、大きな注目を集めています。しかし、モデル蒸留や転移学習といった既存の手法では、高い精度を達成できないことが多いです。この制約を克服するため、我々はBranch-Merge蒸留アプローチを提案します。この手法は、2つのフェーズを通じてモデル圧縮を強化します：（1）Branch Phaseでは、大規模な教師モデルから得られた知識を、ドメイン固有の教師ありファインチューニング（SFT）を介して専門化された学生モデルに選択的に蒸留します。（2）Merge Phaseでは、これらの学生モデルを統合し、ドメイン間の知識転移を可能にし、汎化性能を向上させます。我々は、教師モデルとしてDeepSeek-R1、学生モデルとしてDeepSeek-R1-Distill-Qwen-32Bを使用して、この蒸留アプローチを検証しました。その結果得られた統合モデル、TinyR1-32B-Previewは、Mathematics（+5.5ポイント）、Coding（+4.4ポイント）、Science（+2.9ポイント）など複数のベンチマークで、対応するDeepSeek-R1-Distill-Qwen-32Bを上回り、AIME 2024ではDeepSeek-R1とほぼ同等の性能を達成しました。Branch-Merge蒸留アプローチは、計算コストと時間を削減しつつ、小型で高性能なLLMを作成するためのスケーラブルなソリューションを提供します。

English

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 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.