Technical Summary: Bayesian Model Merging (BMM)

Problem Definition

Model merging aims to combine multiple task-specific expert models into a single architecture without joint retraining or access to original training data. This paradigm addresses the storage and deployment overhead of maintaining separate experts and the prohibitive costs of multi-task learning under data silos or privacy constraints.

Existing merging methods face two primary limitations:

1. Underutilization of Inductive Bias: Most methods estimate merged weights from scratch, failing to leverage the valuable prior knowledge embedded in strong "anchor" models (e.g., existing merging solutions like Task Arithmetic or TSV).
2. Lack of Global Optimization: Advanced methods often operate module-wise in isolation or rely on shared hyperparameters across heterogeneous network modules, lacking a strategy to globally coordinate regularization strengths to maximize end-to-end performance.

Furthermore, many methods require auxiliary calibration data (data-assisted), which is often unavailable due to privacy or storage constraints, while data-free methods often lack the theoretical grounding to approximate activation statistics effectively.

Methodology

The paper introduces Bayesian Model Merging (BMM), a plug-and-play bi-level optimization framework.

1. Activation-Based Bayesian Regression (Inner Loop)

BMM formulates module-wise merging as a Bayesian linear regression problem.

• Prior Induction: Instead of estimating weights from scratch, BMM uses a strong anchor model (e.g., TA, TIES, TSV) to define a prior. The merged task vector $U$ is regularized by an anchor task vector $U^{(0)}$ via an element-wise independent Gaussian prior.
• Likelihood: In the data-assisted setting, the likelihood is derived from empirical activation statistics ($X$) and residual outputs ($Y$) collected from a small calibration set.
• Closed-Form Solution: The Maximum A Posteriori (MAP) estimate of the merged task vector is derived analytically:
  $$U_{MAP} = (YX^\top + \lambda U^{(0)})(XX^\top + \lambda I)^{-1}$$
  where $\lambda$ is a regularization weight balancing empirical fit and the anchor prior. This provides an efficient closed-form solution for module-wise merging.

2. Bayesian Optimization for Global Coordination (Outer Loop)

To address the heterogeneity of network modules, BMM does not use a single shared hyperparameter. Instead, it employs Bayesian Optimization (BO) to search for module-specific regularization strengths ($\lambda$) globally.

• Objective: The outer loop maximizes the predictive evidence (approximated by validation set performance) over the hyperparameters $\lambda$.
• Efficiency: A Gaussian Process (GP) surrogate model guides the search, balancing exploration and exploitation. To manage computational complexity in large models (e.g., ViT-L/14), the authors propose a block-wise parameter tying strategy, grouping modules by function (e.g., attention-in, MLP-out) to reduce the search space dimensionality.

3. Data-Free Variant via Theoretical Alignment

To enable merging without auxiliary calibration data, the authors derive a data-free variant based on a theoretical alignment between activation statistics and task vectors.

• Theorem 1: Under mild assumptions regarding the terminal phase of fine-tuning (neural collapse and representation geometry), the Gram matrix of input activations ($E[xx^\top]$) is positively correlated with the Gram matrix of task vectors ($U^\top U$).
• Approximation: This alignment allows the replacement of empirical activation statistics with a surrogate derived solely from expert weights: $XX^\top \approx c \sum (U^{(t)})^\top U^{(t)}$.
• Result: This yields a data-free MAP estimate that retains the closed-form solution, requiring no calibration data, only the task vectors of the experts and the anchor.

Key Contributions

1. Bi-Level Optimization Framework: BMM combines a closed-form inner-loop MAP estimator (leveraging anchor priors) with an outer-loop Bayesian optimization for global hyperparameter coordination.
2. Theoretical Data-Free Derivation: The paper establishes a theoretical link between activation statistics and task vectors (Theorem 1), enabling a data-free variant of Bayesian merging that preserves the closed-form solution.
3. Plug-and-Play Flexibility: The framework is modular, allowing any existing merging solution to serve as an anchor prior, and supports both data-assisted and data-free settings.

Experimental Results

The authors evaluated BMM on vision (ViT-B/32, ViT-L/14) and language (Llama-3.2-3B, Llama-3.1-8B) benchmarks across 8, 14, and 20-task merging scenarios.

• Performance Gains: BMM consistently outperformed all plug-and-play anchor baselines (TA, TIES, RegMean, TSV, WUDI-Merging, ISO-CTS) in both data-assisted and data-free settings.
  • On the ViT-L/14 8-task benchmark, a single BMM-merged model achieved 95.1% accuracy (data-assisted) and 95.0% (data-free), closely matching the average performance of the eight individual task-specific experts (95.8%).
  • Relative improvements over weaker anchors (e.g., TA, TIES) reached up to 27%.
• Scalability: BMM maintained effectiveness as the number of merged tasks increased (up to 20 tasks), successfully alleviating task interference.
• Efficiency: Despite the BO search, BMM remained computationally competitive. For instance, on ViT-B/32, BMM with only 20 BO trials outperformed baselines requiring exhaustive grid searches (225 configurations) while reducing wall-clock time significantly.
• Ablation Studies:
  • Global Coordination: Replacing BO with shared-$\lambda$ or random search resulted in lower performance, confirming the necessity of module-specific global optimization.
  • Data-Free Validity: The data-free variant performed comparably to the data-assisted variant, validating the theoretical alignment between activation statistics and task vectors.
  • Uncertainty Calibration: A sampling-based variant of BMM (using the full Gaussian posterior) demonstrated improved Expected Calibration Error (ECE) compared to the point estimate, without sacrificing accuracy.

Significance and Claims

The paper claims that BMM provides a practical and scalable path toward model merging when auxiliary calibration data is limited or unavailable. By combining informative anchor priors with global hyperparameter coordination, BMM bridges the gap between the performance of individual task-specific experts and unified multi-task models.

The authors emphasize that their approach does not require joint retraining or access to original training data, making it suitable for environments with data silos or privacy constraints. The derivation of the data-free variant is highlighted as a key theoretical contribution, demonstrating that the geometric structure of task vectors can effectively substitute for empirical activation statistics. The results suggest that leveraging existing merging solutions as priors, rather than treating merging as a zero-shot combination, is a critical factor in achieving state-of-the-art performance.