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)의 크기를 줄이면서도 성능을 유지하는 과제는 상당한 주목을 받고 있습니다. 그러나 모델 증류(model distillation) 및 전이 학습(transfer learning)과 같은 기존 방법들은 종종 높은 정확도를 달성하지 못합니다. 이러한 한계를 해결하기 위해, 우리는 두 단계를 통해 모델 압축을 강화하는 Branch-Merge 증류 접근법을 소개합니다: (1) Branch 단계에서는 대형 교사 모델의 지식을 도메인별 지도 미세 조정(SFT)을 통해 특화된 학생 모델들로 선택적으로 증류하고, (2) Merge 단계에서는 이러한 학생 모델들을 병합하여 도메인 간 지식 전달을 가능하게 하고 일반화를 개선합니다. 우리는 교사 모델로 DeepSeek-R1을, 학생 모델로 DeepSeek-R1-Distill-Qwen-32B를 사용하여 이 증류 접근법을 검증했습니다. 그 결과 병합된 모델인 TinyR1-32B-Preview는 수학(+5.5점), 코딩(+4.4점), 과학(+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.