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

2025-03-06

AI Summary

To address the significant accuracy degradation in large language model (LLM) compression, this paper proposes a novel branch-and-fusion distillation paradigm. First, multiple lightweight student models are domain-specifically fine-tuned via supervised fine-tuning (SFT) to specialize across distinct domains. Second, parameter fusion is performed to enable cross-domain knowledge transfer, overcoming the limitations of conventional unidirectional knowledge distillation. The method innovatively adapts the DeepSeek-R1 teacher model to Qwen-architecture students and introduces a two-stage co-optimization mechanism. Evaluated on mathematical, programming, and scientific benchmarks, TinyR1-32B-Preview achieves gains of +5.5, +4.4, and +2.9 points, respectively, and approaches DeepSeek-R1’s performance on AIME 2024—while substantially reducing inference cost.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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

Problem

• Reducing LLM size while maintaining performance
• Enhancing model compression via Branch-Merge distillation
• Improving accuracy across multiple benchmarks

Innovation

• Branch-Merge distillation for model compression
• Selective knowledge transfer via domain-specific fine-tuning
• Merging student models for cross-domain generalization