Branch-and-Merge (BaM) Technique

• Branch-and-Merge (BaM) is a machine learning technique that splits data into independent branches and merges specialized models to reduce catastrophic forgetting.
• It employs a two-phase approach—with branch training followed by parameter-space fusion using methods like linear interpolation or Slerp—to enhance generalization and compression.
• BaM has demonstrated robust performance in language adaptation and distillation tasks, yielding lower forgetting and superior retention compared to conventional methods.

Branch-and-Merge (BaM) refers to a family of machine learning adaptation and distillation techniques that operate by (1) splitting data or tasks into independent "branches," training specialized model variants in parallel or sequence on these branches, and then (2) merging these specialized models via parameter-space fusion. The overarching aim is to optimize for both rapid adaptation and effective knowledge retention across domains, especially under severe catastrophic forgetting constraints or cross-task conflicts. Contemporary BaM methods have been shown to outperform or strictly dominate standard continued pretraining, conventional instruction finetuning, and naïve mixed-domain distillation, achieving stronger retention of base knowledge, more scalable model compression, and improved generalization on diverse benchmarks (Alexandrov et al., 2024, Sun et al., 6 Mar 2025).

1. Foundations and Motivation

Catastrophic forgetting—where models adapted to new domains or languages lose previously acquired capabilities—poses a practical barrier in both continual language adaptation and multi-domain model distillation. Standard adaptation paradigms (continued pretraining, mixed-data instruction finetuning) exacerbate this issue by entangling gradient updates from diverse tasks, resulting in large-magnitude parameter shifts and negative transfer. BaM mitigates this by isolating initial learning to smaller, homogeneous slices of data or domains, then merging the independently adapted model checkpoints. The approach is formalized via the following setup (Alexandrov et al., 2024):

• Pretrained model: $f_{\theta}$ with $\theta\in\mathbb{R}^d$
• Target data: $\mathcal{X}_{\mathrm{train}}$, typically partitioned into $N$ slices $\{\mathcal{X}_i\}_{i=1}^N$
• Objective: minimize $$L_{\mathrm{target}}(f_{\theta'};\mathcal{X}_{\mathrm{train}})$$ subject to minimal increase in $L_{\mathrm{source}}(f_{\theta'};\text{held-out})$

BaM’s two-phase framework has also been applied to LLM distillation, as in Branch-Merge distillation for TinyR1-32B-Preview, where each expert is fine-tuned on a domain-specific dataset and merged via an importance-driven mask, achieving performance comparable to much larger teacher models (Sun et al., 6 Mar 2025).

2. Algorithmic Structure and Merging Strategies

2.1 Adaptation BaM (Language Transfer)

The canonical BaM algorithm for language adaptation proceeds as follows:

1. Data Partition: Split $\mathcal{X}_{\mathrm{train}}$ into $N$ non-overlapping slices.
2. Branch Training: For each slice, initialize from prior checkpoint and train to convergence: $$\theta_i\gets\text{TrainUntilConvergence}(\theta^{i-1}, \mathcal{X}_i)$$.
3. Merging: After every $K$ branches or at end of epoch, merge checkpoints using either linear interpolation or Slerp:
   • Linear: $\theta'=(1-c)\theta_1 + c\theta_2$
   • Slerp: $$\theta'=\theta+\frac{\sin((1-c)\phi)}{\sin\phi}\tau_1+\frac{\sin(c\phi)}{\sin\phi}\tau_2$$, where $\tau_j=\theta_j-\theta$, $$\phi=\arccos\left(\frac{\tau_1\cdot\tau_2}{\|\tau_1\|\|\tau_2\|}\right)$$ and $c$ typically $0.5$.
4. Iteration: Repeat for all slices and branches.

2.2 BaM for Distillation

BaM distillation (Branch-Merge) divides training into two phases (Sun et al., 6 Mar 2025):

1. Branch Phase: For each domain $d$, fine-tune from a backbone on supervised data $D_d$ with standard cross-entropy or hybrid SFT/KD losses: $$L^{d}_{\mathrm{branch}}(\theta) = \mathbb{E}_{(x,y)\in D_d}\left[-\sum_{t=1}^{|y|}\log p_\theta(y_t|y_{<t},x)\right]$$.
2. Merge Phase: Merge expert checkpoints $\{\theta_{d}\}$ using the Arcee fusion procedure:
   • Compute $\Delta = \theta_T - \theta_S$
   • Compute per-parameter importance score: $IS_i = \text{KL}(T\|S)\cdot\Delta_i$
   • Form mask $M_{i}=1$ if $$IS_i>\text{median}(IS)+\alpha_{\text{merge}}\cdot IQR(IS)$$, otherwise $0$ (typically $\alpha_{\text{merge}}=1.5$)
   • Merge: $$\theta_{\mathrm{merge}} = \theta_S + M\odot(\theta_T-\theta_S)$$

This structured parameter-wise merge preserves salient, high-importance information from each branch while preventing destructive averaging.

3. Theoretical Insights and Empirical Behavior

The BaM mechanism achieves dual goals: limiting the total $L_2$ parameter drift (curbing forgetting) and increasing gradient signal quality via noise averaging. Theoretically, each “task vector” $\tau_i$ found by training on a single branch can be decomposed as $\tau_i=\tau^*+\epsilon_i$ with $\epsilon_i$ zero-mean estimation noise. Merging $K$ task vectors induces variance reduction: $\text{Var}(\epsilon')=\text{Var}(\epsilon_i)/K$, meaning the denoised update better approximates the ideal adaptation direction. Empirical results illustrate that for any given $L_2$ parameter change, BaM achieves higher target performance and substantially lower source-domain degradation compared to monolithic or reduced-LR adaptation (Alexandrov et al., 2024).

In distillation, Arcee fusion selectively propagates only "important" changes, as measured by local KL divergence, mitigating over-regularization and preserving domain-specific expertise (Sun et al., 6 Mar 2025). This yields models that, while significantly compressed, approach or even match teacher-level performance across multiple domains.

4. Application Domains and Experimental Results

BaM methodology has been validated on several fronts:

• Language Transfer (Alexandrov et al., 2024):
  • Adaptation from English LLMs (Meta-Llama 3-8B, Mistral 7B) to Bulgarian and German, with/without approximate experience replay.
  • BaM with $K=2, N=8$ (Bulgarian) or $N=4$ (German) achieves:
  • Bulgarian CPT: Avg BG = 53.40, Avg EN = 66.24 (vs. CPT: 53.11/64.84)
  • German CPT: Avg DE = 57.68, Avg EN = 65.79 (vs. CPT: 53.51/60.79)
  • Instruction Finetuning: BaM i.i.d. and split ordering further boost both target and English performance over standard mix or EN-only IFT.
• Distillation (Sun et al., 6 Mar 2025):
  • Compression of DeepSeek-R1 into TinyR1-32B: BaM achieves Math/Code/Science pass@1 of 78.1/61.6/65.0 (up to +5.5 points over DeepSeek-R1-Distill-Qwen-32B baseline), approaching DeepSeek-R1 teacher accuracy within 1.7-6.5 points across domains.
  • Resource cost: Merge phase executes in $<4$ H800-GPU hours, over 90% faster than mixed-data retraining.

Task/Domain	Standard Adaptation	BaM Performance
BG CPT Avg	53.11	53.40
DE CPT Avg	53.51	57.68
Math	72.6	78.1
Coding	57.2	61.6
Science	62.1	65.0

5. Ablations, Sensitivity, and Methodological Variants

Systematic ablations confirm that BaM’s trade-offs are controlled by the number of branches ($N$), parallelism factor ($K$), and the merge coefficient ($c$):

• Increasing $N$ or $K$ further reduces both adaptation magnitude and forgetting, but may diminish target learning.
• Merge method (Linear interpolation vs. Slerp) has marginal effect.
• Quality of experience replay is critical for source-domain retention.
• Ordered merging in distillation produces robust results across different combine orders.
• Extensions to alternative backbone models and inclusion of additional domains or tasks yield further generalization potential (Alexandrov et al., 2024, Sun et al., 6 Mar 2025).

6. Scalability, Limitations, and Extensions

BaM introduces no extra compute requirements versus standard approaches during total training, and its merge operation is scalable to large parameter counts and commodity hardware. However, reliance on (approximate) experience replay in language adaptation and careful tuning of structural hyperparameters ($K, N, c$, mask thresholds) remain necessary for optimal results. To date, validation is limited to $7$–$8$B parameter LLMs and two languages.

Potential extensions include deeper theoretical characterization of the loss surface under repeated merging, application to vision-language or other multimodal models, integration with continual-learning regularizers (EWC, Mixout), and finer-grained parameter importance weighting in merging. Expansion to more domains, additional languages, and even larger-scale models is a plausible direction.

7. Broader Impact and Generalization

BaM exemplifies a general paradigm for mitigating catastrophic forgetting, negative transfer, and capacity loss in large foundation models. It supplies a recipe for scalable, resource-efficient, and effective adaptation and distillation—demonstrated in both language transfer and supervised compression scenarios. Its capacity to balance low-magnitude parameter drift with high-fidelity adaptation positions it as a robust baseline for robust, continual, or multi-domain learning in modern neural architectures (Alexandrov et al., 2024, Sun et al., 6 Mar 2025).