TinyR1-32B-Preview: Migliorare la Precisione con la Distillazione Branch-Merge

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

March 6, 2025

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

Abstract

La sfida di ridurre le dimensioni dei Large Language Models (LLM) mantenendo le loro prestazioni ha attirato una significativa attenzione. Tuttavia, i metodi esistenti, come la distillazione del modello e il transfer learning, spesso non riescono a raggiungere un'elevata accuratezza. Per affrontare questa limitazione, introduciamo l'approccio di distillazione Branch-Merge, che migliora la compressione del modello attraverso due fasi: (1) la Fase Branch, in cui la conoscenza di un grande modello insegnante viene distillata selettivamente in modelli studente specializzati tramite fine-tuning supervisionato specifico per dominio (SFT); e (2) la Fase Merge, in cui questi modelli studente vengono fusi per consentire il trasferimento di conoscenza tra domini e migliorare la generalizzazione. Validiamo il nostro approccio di distillazione utilizzando DeepSeek-R1 come insegnante e DeepSeek-R1-Distill-Qwen-32B come studente. Il modello risultante, TinyR1-32B-Preview, supera la sua controparte DeepSeek-R1-Distill-Qwen-32B in diversi benchmark, tra cui Matematica (+5,5 punti), Codifica (+4,4 punti) e Scienza (+2,9 punti), raggiungendo prestazioni quasi pari a DeepSeek-R1 su AIME 2024. L'approccio di distillazione Branch-Merge fornisce una soluzione scalabile per creare LLM più piccoli e ad alte prestazioni con costi e tempi computazionali ridotti.

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.