TinyR1-32B-Preview: Boosting Accuracy with Branch-Merge Distillation (2503.04872v2)

Abstract: The challenge of reducing the size of 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 \textit{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.

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

The paper presents a novel approach to reducing the size of LLMs without compromising performance, introducing a method termed Branch-Merge Distillation. This research addresses the inherent challenges in model distillation and transfer learning, which often struggle to maintain high accuracy. The authors propose a two-phased strategy to enhance model compression by leveraging domain-specific knowledge, which is then consolidated into a unified model to improve cross-domain understanding.

Branch-Merge Distillation Method

The Branch-Merge Distillation approach is divided into two primary phases:

1. Branch Phase: This phase involves selectively distilling knowledge from a large "teacher" model into several smaller, domain-specific "student" models. The DeepSeek-R1 model acts as the teacher model, and the knowledge distillation occurs through supervised fine-tuning (SFT) on specialized datasets in domains such as mathematics, coding, and science. The aim is to create specialized models that excel in specific areas by utilizing tailored data sets for training.
2. Merge Phase: The domain-specific student models are subsequently merged to facilitate cross-domain knowledge transfer, leading to a generalized model. This phase employs the Arcee Fusion technique, which integrates significant parameter updates from the teacher model to enhance the student's capabilities.

Experimental Results and Implications

The research substantiates the efficacy of the Branch-Merge approach by demonstrating a marked improvement in model performance across various domains. The consolidated TinyR1-32B-Preview model exhibited notable performance gains over its counterparts, including DeepSeek-R1-Distill-Qwen-32B, showcasing improvements across benchmarks such as Mathematics (+5.5 points), Coding (+4.4 points), and Science (+2.9 points). Furthermore, the computation efficiency was also a substantial outcome, reducing the model merging phase's computational cost by 90% compared to traditional methodologies.

Contributions and Future Potential

The paper makes significant contributions in both methodology and practical application domains:

• Scalability and Accuracy: The Branch-Merge methodology demonstrates that it is feasible to scale down the computations while maintaining, if not improving, model accuracy across specialized domains.
• Cost Efficiency: The approach effectively reduces the computational requirements, thereby enabling smaller-scale high-performing models with reduced training costs.
• Open Source Contribution: The authors commit to releasing their models, data, and training/evaluation codes, promoting reproducibility and further research opportunities within the community.

Speculation on Future Developments

The proposed methodology and findings could pave the way for several future research directions, including:

• Generalization Across Additional Domains: Potential to apply this branch-merge approach to other domains and integrate additional modality data, potentially influencing how models are trained across various data types.
• Broader Adoption: As computationally efficient yet accurate LLM models are highly desirable, such an approach could influence the deployment of LLMs in environments with limited resources.
• Advanced Optimization Techniques: Future exploration into refining the Arcee Fusion rules and integrating them with other advanced optimization techniques to further enhance accuracy and reduction efficiency could drive innovation in AI model development.

The work detailed in this paper represents a significant step toward more efficient, domain-generalizing LLMs, offering valuable insights and methodologies for future research in model distillation and optimization.