Branch-Merge Distillation Methods

• Branch-merge distillation is a method that consolidates specialized submodels into a single compact model via selective output, weight, or feature merging.
• It employs diverse mechanisms such as output-level distillation, weight-space averaging, and layer-wise alignment to harness expert strengths while avoiding common pitfalls of naive merging.
• This approach improves model compression, multi-domain integration, and robustness against catastrophic forgetting, making it ideal for scalable and efficient AI systems.

Branch-merge distillation is a class of methodologies designed to integrate knowledge from multiple specialized or pre-trained submodels (“branches”) into a single, compact model (“merged” or “student” model) via distillation-based or parameter-merging schemes. This paradigm has become central to model merging, large-scale knowledge transfer, robust compression of LLMs, and fusion of domain-specific or multi-modal capabilities. Branch-merge methods exploit the diversity and strengths of individual expert models, overcoming the pitfalls of joint multitask training, naive weight averaging, or single-teacher distillation.

1. Fundamental Principles and Motivation

Branch-merge distillation addresses the challenge of consolidating diverse capabilities from multiple teacher models, each possibly specialized for different domains, tasks, or data distributions, into one student of fixed or reduced capacity. The motivation stems from empirical findings that joint multitask pretraining or data-mixed distillation is often computationally prohibitive and may induce conflicting optimization signals, while naive merging (e.g., task arithmetic, weight averaging) can suffer catastrophic forgetting or sharp drops in accuracy on some tasks.

In its archetypal realization, branch-merge distillation proceeds in two conceptual phases:

• Branching: Each branch corresponds to a teacher model fine-tuned or distilled on a specific target (language, domain, modality, or reasoning style), independently of the others.
• Merging: The parameters, representations, or outputs of these branches are selectively fused—either via logits-level distillation, weight-space arithmetic, progressive feature alignment, or selective integration—to create a multi-skilled merged model that preserves (and even enhances) the constituent expertise.

Prominent motivations for this framework include multilingual language representation (Khanuja et al., 2021), data- and compute-efficient LLM compression (Sun et al., 6 Mar 2025), robust multi-task model integration (Yoshida et al., 2 Aug 2025, Dalili et al., 24 Dec 2025), and distillation of reasoning skills from diverse LLMs (Shen et al., 10 Sep 2025).

2. General Branch-Merge Frameworks

A unifying taxonomy of branch-merge distillation methods encompasses several algorithmic regimes:

• Output-level Distillation: Each branch produces predictions (logits, soft targets, or rationales) on unlabeled samples. The student is trained to minimize a combined loss that aligns its outputs with all branches, sometimes with an additional supervised component if labeled data is available. MergeDistill (Khanuja et al., 2021) exemplifies this approach, where knowledge from multiple language-specific LMs is unified via a masked language modeling (MLM) loss plus average Kullback-Leibler (KL) divergence to all relevant teacher logits.
• Weight-space Merging: After independent fine-tuning or SFT (Supervised Fine-Tuning) on each branch, the corresponding sets of model parameters are merged using arithmetic (e.g., convex or element-wise averaging), importance-based masking, or consensus filtering. Merge-of-Thought distillation (MoT) (Shen et al., 10 Sep 2025) iteratively alternates between branch-wise training and weight-averaged merging, yielding a single model that encapsulates consensus reasoning from competing teachers.
• Feature- or Layer-wise Alignment: Progressive or layer-wise distillation matches intermediate or final representations between student and experts, often proceeding sequentially through layers or blocks. ProDistill (Xu et al., 18 Feb 2025) employs a staged MSE alignment at each layer of large vision or NLP models across tasks, producing scalable yet accurate merges.
• Parameter Selection and Fusion: Selective integration schemes, such as Arcee Fusion (Sun et al., 6 Mar 2025), identify and incorporate only the most salient weight updates from each branch, using parameter-wise importance scores derived from KL divergence of output distributions and thresholds based on the distribution of these scores.
• Multi-branch to Single-branch Feature Migration: In vision applications, merging dual-branch (e.g., spatial and frequency) teachers into single-branch students is enabled by dedicated feature projectors and gradient modulation, as in Two-in-One Knowledge Distillation (TOKD) (Zhou et al., 2023).

Each paradigm admits variants governed by data regime (e.g., unlabeled vs. labeled, few-shot vs. large-scale), merging granularity (layerwise, parameterwise, or whole-network), and reconciliation strategy for teacher conflicts or capacity mismatches.

3. Distillation and Merging Objectives

The objective functions and operational mechanisms of branch-merge distillation are tightly coupled to the nature and abundance of available data:

• Task-agnostic Distillation: Objective is typically a convex combination of supervised loss (cross-entropy with true labels) and knowledge-distillation loss (e.g., KL divergence) between student and teacher soft predictions. For multiple branches, averaging or summing over losses from all relevant teachers is standard:

$$L_{\text{ALL}}(x) = \alpha \cdot \frac{1}{|T_k|}\sum_{t \in T_k} L_{\text{KD}}^t(x) + (1-\alpha) L_{\text{MLM}}(x)$$

as realized in MergeDistill (Khanuja et al., 2021), with $\alpha$ annealed during training.

• Weight-space Averaging: Branch variants are merged by simple arithmetic, e.g., $\theta^{*} = \sum_k \alpha_k \theta^{(k)}$, or by more advanced rules using element-wise masks or importance scores. MoT (Shen et al., 10 Sep 2025) and Arcee Fusion (Sun et al., 6 Mar 2025) exemplify strictly weight-space perspectives, with consensus filtering emerging from averaging parameters across divergent branches.
• Progressive, Layer-wise Surrogates: Layer-specific merging coefficients $\{\lambda^{(l)}_i\}$ are optimized to minimize feature-wise MSE at each stage, using dual-input schemes to align merged-layer and expert-layer representations (Xu et al., 18 Feb 2025).
• Heterogeneity and Confidence Conditioning: Distillation may involve harmonizing task-vector norms (as in DisTaC (Yoshida et al., 2 Aug 2025)) or pre-conditioning teacher branches via secondary distillation to mitigate confidence disparities and reduce merging-induced interference.
• Flatness and Generalization Guarantees: Recent innovations incorporate PAC-Bayes generalization bounds with explicit flatness and heterogeneity penalties to guide both merging weights and the use of flat minima (via Sharpness-Aware Minimization) (Dalili et al., 24 Dec 2025).

These objectives are further modified by architectural constraints, data availability, and the degree of overlap in teacher coverage across domains.

4. Key Algorithms and Methodological Variants

The practical realization of branch-merge distillation has led to a diverse set of algorithms, each optimized for different scenarios.

Algorithm / Paper	Merge Modality	Distillation/Objective	Granularity
MergeDistill (Khanuja et al., 2021)	Logits-level, MLM	Cross-entropy + KL sum	Token, Mask
MoT (Shen et al., 10 Sep 2025)	Weight averaging	NLL on SFT data	Full-param, round-wise
ProDistill (Xu et al., 18 Feb 2025)	Progressive layer-wise	MSE (activations)	Per-layer, elementwise
TinyR1-32B (Sun et al., 6 Mar 2025)	Arcee Fusion (weights)	KL-importance masking	Parameterwise
DisTaC (Yoshida et al., 2 Aug 2025)	Task-vector conditioning	Soft-target KD + $\ell_2$	Task vector
SAMerging (Dalili et al., 24 Dec 2025)	SAM-based merging	Multi-teacher KD (KL)	Task vector, flatness
TOKD (Zhou et al., 2023)	Feature projection	Cosine/MSE + rotation	Head/proj, branchwise

• MergeDistill orchestrates multilingual LM merging by offline teacher logits storage, token vocabulary union, and shuffled multi-language MLM+KD training.
• MoT alternates branch-specific SFT using filtered teacher reasoning traces with simple weight-space averaging, repeating for several rounds to maximize consensus features and minimize catastrophic forgetting.
• ProDistill iterates over layers of deep models, at each step optimizing merging coefficients to best match per-task activation statistics using only few-shot domain samples.
• TinyR1-32B’s branch-merge distillation leverages domain-specific specialists, then fuses them via Arcee Fusion’s parameterwise salience estimation and masking, scaling to billion-parameter models within minimal additional compute budget.
• DisTaC preconditions or sharpens task vectors via distilled soft targets with KL divergence and norm rescaling before plug-and-play merging by established rules.
• SAMerging integrates flatness-aware PAC-Bayes theory, leveraging SAM to penalize sharp regions and heterogeneity, explicitly linking generalization error to the KD objective.

5. Theoretical Analysis and Empirical Findings

Theoretical advances have established the limits and guarantees of different branch-merge strategies:

• Necessity of Data for Merging: ProDistill (Xu et al., 18 Feb 2025) proves that any purely weight-based, data-agnostic merging algorithm can be worst-case suboptimal, motivating the use of even a small validation set for scalable merging.
• Flatness and Heterogeneity Bounds: SAMerging (Dalili et al., 24 Dec 2025) shows that generalization for merged models is bounded by sharpness (via a flatness proxy), cross-task heterogeneity, and multi-teacher KD loss, formally linking the practical minimization of KL divergence to improvement in multi-task excess risk.
• Mitigation of Overfitting and Forgetting: MoT (Shen et al., 10 Sep 2025) demonstrates that averaging across branch SFT variants yields flatter minima, better consensus, higher robustness to distribution shift, and less catastrophic task forgetting versus standard or naive union merges.
• Gradient Homogenization: In multi-branch-to-single-branch scenarios, as in TOKD (Zhou et al., 2023), trainable rotation modules resolve conflicting gradients between differently specialized branches, enabling effective fusion without degradation.

Empirical results consistently indicate that branch-merge distillation either matches or surpasses baseline models in aggregate task accuracy, even when the student operates at a fraction of the parameter count or data footprint. For example, MergeDistill’s student LMs outperform or reach within 1–3% of larger monolingual/multilingual teachers on NER and QA benchmarks (Khanuja et al., 2021), while TinyR1-32B achieves gains of +5.5, +4.4, and +2.9 points in math, coding, and science relative to its baseline (Sun et al., 6 Mar 2025).

6. Practical Considerations and Application Scenarios

Branch-merge distillation is applicable in a range of settings, each posing specific operational constraints and incentives:

• Model Compression: Reducing the footprint of LLMs or vision transformers for deployment without sacrificing accuracy, particularly by exploiting cross-domain knowledge synergy (Sun et al., 6 Mar 2025).
• Multi-lingual/Domain LM Construction: Combining monolingual, multilingual, or domain-specific LMs into a unified student with balanced, task-agnostic coverage (Khanuja et al., 2021).
• Multi-teacher Transfer in Few-shot Regimes: Leveraging only a handful of validation examples per task/domain (even as low as 16 shots) to enable robust model fusion when large calibration sets are unavailable (Xu et al., 18 Feb 2025, Dalili et al., 24 Dec 2025).
• Fusion of Modal, Reasoning, or Feature Complementarities: Merging spatial and frequency information in forensic detectors (Zhou et al., 2023), or integrating chain-of-thought reasoning styles in LLMs from heterogeneous sources (Shen et al., 10 Sep 2025).
• Robustness to Source Model Pathologies: Pre-conditioning of task vectors for norm or confidence disparities (DisTaC), or mitigating training set heterogeneity and label smoothing effects.

Pragmatically, these methods require that the student model architecture is compatible (or at least mappable) with the collection of teacher branches, and that adequate calibration or validation data is available for the layerwise or KD objectives.

7. Current Limitations and Future Directions

Despite considerable advances, branch-merge distillation faces open technical challenges:

• Extensions Beyond Homogeneous Architectures: Most current algorithms assume identical or easily-aligned architectures; robust cross-architecture or cross-modality merges remain more challenging.
• Automated or Differentiable Hyperparameter Selection: Setting of KD loss balance coefficients, merge thresholds, flatness/SAM radii, and temperature.
• Scaling to Larger and More Diverse Branch Pools: Handling an increasing number of teacher models with mutually incompatible or vastly different priors, languages, or modalities introduces new heterogeneity and generalization issues.
• Adaptive and Continual Merging: Ensuring that models can continuously incorporate new expert knowledge via incremental branch-merging without catastrophic drift.
• Benchmarking Under Distribution Shift: While some methods are robust to peer-level and distribution-shifted teachers (Shen et al., 10 Sep 2025), formal guarantees and expanded empirical evaluation are ongoing research.

A plausible implication is that future frameworks will integrate explicit regularization for feature diversity and flatness, adopt differentiable or learning-to-merge paradigms, and exploit cross-modal or hybrid architectures for increasingly general knowledge integration. Branch-merge distillation thus occupies a central position in scalable, efficient, and robust neural model construction across contemporary ML domains.