TinyR1-32B-Preview : Amélioration de la précision grâce à la distillation par fusion de branches

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

March 6, 2025

Auteurs: 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

Résumé

Le défi de réduire la taille des grands modèles de langage (LLM) tout en maintenant leurs performances a suscité une attention considérable. Cependant, les méthodes existantes, telles que la distillation de modèles et l'apprentissage par transfert, échouent souvent à atteindre une haute précision. Pour pallier cette limitation, nous introduisons l'approche de distillation Branch-Merge, qui améliore la compression des modèles à travers deux phases : (1) la phase Branch, où les connaissances d'un grand modèle enseignant sont distillées de manière sélective dans des modèles étudiants spécialisés via un fine-tuning supervisé spécifique à un domaine (SFT) ; et (2) la phase Merge, où ces modèles étudiants sont fusionnés pour permettre un transfert de connaissances interdomaines et améliorer la généralisation. Nous validons notre approche de distillation en utilisant DeepSeek-R1 comme modèle enseignant et DeepSeek-R1-Distill-Qwen-32B comme modèle étudiant. Le modèle fusionné résultant, TinyR1-32B-Preview, surpasse son homologue DeepSeek-R1-Distill-Qwen-32B sur plusieurs benchmarks, notamment en Mathématiques (+5,5 points), en Codage (+4,4 points) et en Sciences (+2,9 points), tout en atteignant une performance quasi équivalente à celle de DeepSeek-R1 sur l'AIME 2024. L'approche de distillation Branch-Merge offre une solution évolutive pour créer des LLM plus petits et performants, avec un coût et un temps de calcul réduits.

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.