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

1. Problem Statement

The paper addresses the challenge of reducing the size of Large Language Models (LLMs) while maintaining high reasoning performance. While distilling knowledge from large "teacher" models (like DeepSeek-R1) into smaller "student" models is a common strategy, existing methods face significant limitations:

• Gradient Interference: Naive mixed-data distillation (training a single model on a combined dataset of Math, Coding, and Science) often leads to conflicting gradients. This causes a "seesaw effect," where improving performance on one domain degrades performance on another.
• Suboptimal Cold Starts: Small models require a strong Supervised Fine-Tuning (SFT) cold start to unlock reasoning capabilities before Reinforcement Learning (RL) can be effective. Traditional distillation often fails to provide a robust starting point for specialized tasks.
• Computational Cost: Retraining models on massive mixed datasets to resolve domain conflicts is computationally expensive and time-consuming.

2. Methodology: Branch-Merge Distillation

The authors propose a novel two-phase distillation strategy called Branch-Merge, which decouples domain-specific learning and then selectively integrates the knowledge.

Phase 1: The Branch Phase (Specialization)

Instead of training one model on mixed data, the authors train multiple specialized "expert" models from a single base student model (DeepSeek-R1-Distill-Qwen-32B).

• Process: The base model is fine-tuned separately on three distinct, high-quality datasets:
  • Math: 58k samples from NuminaMath 1.5.
  • Coding: 20k trajectories from OpenThoughts.
  • Science: 8.6k Chain-of-Thought (CoT) trajectories generated by DeepSeek-R1.
• Outcome: Three specialized models (Math Expert, Coding Expert, Science Expert) are created, each optimized for its specific domain without interference from other tasks.

Phase 2: The Merge Phase (Integration)

The specialized experts are merged into a single unified model using Arcee Fusion, a parameter-space merging technique.

• Mechanism: Arcee Fusion compares parameters between a "Left Model" (existing) and a "Right Model" (new expert).
• Importance Scoring: It calculates an importance score ($SIS$) for each parameter update based on the divergence (KL-divergence) between the parameter distributions of the two models.
• Dynamic Thresholding: A threshold ($S_{THR}$) is calculated dynamically using the median and interquartile range of the importance scores.
• Selective Integration:
  • If a parameter update in the new model is highly significant (above the threshold), it is retained.
  • If the update is minor or noisy, the original parameter is kept.
  • Formula: $\theta^i_M = \theta^i_L + (\theta^i_R - \theta^i_L) \cdot \lceil \max(0, S^i_{IS} - S_{THR}) \rceil$
• Result: This acts as a "high-pass filter," preserving crucial domain-specific logic while suppressing noise, effectively approximating the union of expert capabilities rather than a compromised average.

3. Key Contributions

• TinyR1-32B-Preview: The release of a 32B parameter model that outperforms its backbone and approaches the performance of the much larger DeepSeek-R1 (671B) on reasoning benchmarks.
• Novel Distillation Strategy: Introduction of the Branch-Merge approach, which solves the gradient interference problem inherent in multi-task distillation by decoupling training and using saliency-guided merging.
• Efficiency & Cost Reduction:
  • Time: The merging process takes only 4 hours on 4 H800 GPUs, compared to 23 hours on 32 H800 GPUs required for traditional mixed-data retraining.
  • Cost: The estimated reproduction cost is 744 H800 GPU hours ($1,500), representing a 90% reduction in merging-stage computational overhead.
• Open Source: Full release of the model, training data, scripts, and evaluation code on Hugging Face and GitHub.

4. Experimental Results

The model was evaluated on three primary benchmarks: AIME 2024 (Math), LiveCodeBench (Coding), and GPQA-Diamond (Science).

• Performance Gains: TinyR1-32B-Preview improves upon the backbone by +5.5 (Math), +4.4 (Coding), and +2.9 (Science).
• Comparison to Teacher: It achieves math performance (78.1) very close to the 671B DeepSeek-R1 teacher (79.8), a feat traditional distillation has not achieved for 32B models.
• Ablation Study: The paper confirms that the "Data Mixture" baseline suffers from the seesaw effect (improving Math/Science but hurting Coding relative to experts), whereas the Branch-Merge approach retains coding strength while boosting other domains.
• Generalization: The model also shows strong generalization on instruction-following (IFEval) and general reasoning (BBH, HumanEval), outperforming individual experts in these areas.

5. Significance

• Democratization of Reasoning: By achieving near-teacher performance with a 32B model, this work makes advanced reasoning capabilities accessible for local deployment and smaller organizations, reducing reliance on massive, expensive APIs.
• Paradigm Shift in Distillation: The paper challenges the standard "mixed-data SFT" approach, demonstrating that decoupled training + selective merging is a more effective and efficient path for multi-domain capability.
• Foundation for RL: The authors posit that this high-quality SFT model provides a robust "cold start" for future Reinforcement Learning (RL) optimization, potentially accelerating the development of even more capable small models.
• Reproducibility: The complete open-sourcing of the pipeline sets a new standard for transparency in LLM research, allowing the community to verify and build upon the Branch-Merge methodology.