Model Merging via Multi-Teacher Knowledge Distillation

TL;DR

The paper addresses the challenge of merging multiple fine-tuned models without access to their training data by deriving a flatness-aware PAC-Bayes generalization bound that introduces a cross-task heterogeneity term. It then reframes model merging as multi-teacher knowledge distillation on unlabeled data, showing that minimizing the student–teacher KL divergence tightens the bound on the merged model's excess risk. Building on this, SAMerging combines the KL-based KD objective with Sharpness-Aware Minimization to seek flat minima, yielding strong, data-efficient generalization across vision and NLP benchmarks. Empirically, SAMerging achieves state-of-the-art results on several task suites with minimal unlabeled calibration data, while remaining post-hoc and inference-efficient, underscoring its practical impact for privacy-preserving, multi-task deployment.

Abstract

Model merging has emerged as a lightweight alternative to joint multi-task learning (MTL), yet the generalization properties of merged models remain largely unexplored. Establishing such theoretical guarantees is non-trivial, as the merging process typically forbids access to the original training data and involves combining fine-tuned models trained on fundamentally heterogeneous data distributions. Without a principled understanding of these dynamics, current methods often rely on heuristics to approximate the optimal combination of parameters. This dependence is most critical in coefficient scaling, the weighting factors that modulate the magnitude of each fine-tuned model's contribution to the shared parameter. However, without a principled objective to guide their selection, these methods lead to brittle performance and are highly sensitive to scaling initialization. We address this gap by (i) establishing a novel flatness-aware PAC-Bayes generalization bound specifically for the model merging setting. This analysis introduces a "cross-task heterogeneity" term that formally captures the mismatch between diverse fine-tuned model priors and the target multi-task distributions. Guided by this theoretical insight, (ii) we frame model merging as multi-teacher knowledge distillation on scarce, unlabeled data. We formally demonstrate that minimizing the student-teacher Kullback-Leibler divergence directly tightens the upper bound on the merged model's excess risk. Guided by the flatness-aware bound derived, (iii) we operationalize this objective via SAMerging, a method that employs Sharpness-Aware Minimization (SAM) to find flat minima. Empirically, SAMerging establishes a new state of the art across vision and NLP benchmarks, achieving remarkable performance. The code is available at https://github.com/arshandalili/SAMerging.

Figures (8)

• Figure 1: Sensitivity to merge scaling and initialization. On TA-8, we compare merge scaling (TA/TIES) and initialization (AdaMerging/SAMerging). While designed to learn coefficients, AdaMerging's performance is sensitive to initialization, suggesting its objective/optimizer is a bottleneck. In contrast, SAMerging attains higher and more stable accuracy across the range.
• Figure 2: The loss landscape around the merged model of SAMerging and AdaMerging on TA-8 with perturbing along EuroSAT and SUN397.
• Figure 3: The loss behavior around the merged model of SAMerging and AdaMerging on TA-8 by perturbing along task vectors of EuroSAT and SUN397.
• Figure 4: Data-dependent methods gain in performance with increasing number of calibration data on TA-8 using ViT-B/32.
• Figure 5: Loss surface for MNIST and DTD on TA-8.

Theorems & Definitions (31)

• Lemma 1
• Proposition 1
• Theorem 1
• Lemma 2
• Lemma 3
• Lemma 4
• Theorem 2
• Remark 1
• Definition 1
• Lemma 5: Single-Task Excess Risk Bound