📊 statistics

Variational Model Merging for Pareto Front Estimation in Multitask Finetuning

This paper proposes Variational Model Merging, a novel Bayesian framework that leverages flexible non-Gaussian posteriors to provably improve the quality and estimation of Pareto fronts in multitask finetuning, outperforming existing methods that rely on simpler Gaussian assumptions.

Original authors: Hugo Monzón Maldonado, Nico Daheim, Thomas Möllenhoff, Iryna Gurevych, Mohammad Emtiyaz Khan

Published 2026-06-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Hugo Monzón Maldonado, Nico Daheim, Thomas Möllenhoff, Iryna Gurevych, Mohammad Emtiyaz Khan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef trying to create the perfect "fusion" dish that balances three distinct flavors: spicy, sweet, and sour. You want a recipe that doesn't just taste good overall, but offers the best possible combinations of these flavors without one overpowering the others. In the world of Artificial Intelligence, this is called Multitask Finetuning. You have a smart AI model, and you want it to be good at three different jobs (like recognizing cats, translating French, and solving math problems) all at once.

The problem is that finding the perfect "recipe" (the right mix of weights for each task) is incredibly expensive and time-consuming. It's like having to cook the dish from scratch thousands of times, changing the ingredients slightly each time, just to see which version tastes best.

This paper introduces a clever shortcut called Variational Model Merging. Here is how it works, broken down into simple concepts:

1. The Old Way: "The Average Soup"

Previously, researchers tried to guess the best mix by simply taking the "average" of three separate models. Imagine you have three chefs: one who makes great spicy food, one great at sweet, and one at sour. The old method was to take a spoonful of each chef's soup, mix them in a bowl, and hope the result is delicious.

The paper argues this is like Simple Averaging. It's cheap and fast, but often the result is bland or misses the mark. It's a "lazy" guess that doesn't understand the complex chemistry between the flavors.

2. The New Idea: "The Bayesian Recipe Book"

The authors propose a smarter way using a concept from statistics called Bayesian Inference. Instead of just averaging the soups, they treat each chef's work as a "probability map" or a recipe book that describes not just what the flavor is, but how confident the chef is about it.

The "Posterior" (The Recipe Book): When a model learns a task, it creates a "posterior distribution." Think of this as a map showing all the possible ways to tweak the model to get that specific task right.
Merging the Maps: Instead of mixing the final soups, the new method mixes these maps. It multiplies the recipe books together to create a new, combined map that shows the best possible compromises between the tasks.

3. The Secret Sauce: "Flexible Maps"

The paper's biggest discovery is that the shape of these maps matters.

Rigid Maps (Isotropic Gaussians): Imagine a map that assumes the flavor is a perfect, round circle. This is simple but often wrong. This leads to the old "Simple Averaging" method.
Flexible Maps (Full Gaussians & Mixtures): Imagine a map that can stretch, squish, and even split into multiple blobs to capture complex shapes. This is what the authors call Variational Model Merging.

The Analogy:
If you are trying to find the best spot to park a car in a crowded lot:

Simple Averaging is like guessing the middle of the lot. You might hit a car.
Hessian-Weighted Merging is like drawing a circle around the empty spots. Better, but still might miss a tight corner.
Variational Model Merging (Mixture of Gaussians) is like drawing a flexible, stretchy net that wraps perfectly around every single empty parking spot, no matter how weird the shape is.

4. The Result: A Better "Pareto Front"

In math, a Pareto Front is a list of all the "perfect" trade-offs. It's the set of solutions where you can't get more "spicy" without losing some "sweetness."

The paper claims that by using these flexible, stretchy maps (specifically, mixtures of Gaussian distributions), they can estimate this list of perfect trade-offs much more accurately than before.

They tested this on image recognition (teaching AI to see) and language models (teaching AI to translate).
The Finding: Their new method found better "recipes" (task mixes) than simple averaging or even more complex previous methods. It got closer to the results you would get if you had the time and money to cook the dish from scratch thousands of times, but it did it in a fraction of the time.

Summary

The paper doesn't invent a new way to cook the dish (train the model from scratch). Instead, it invents a smarter way to mix the leftovers.

By treating the models not as fixed points but as flexible, probabilistic maps, and then merging those maps using advanced math, they can predict the best way to combine different AI skills without needing to do the expensive, repetitive work of training the AI over and over again. It's like having a super-intelligent sous-chef who can taste three different soups and instantly tell you the exact ratio needed to make the perfect fusion dish, saving you hours of cooking.

Technical Summary: Variational Model Merging for Pareto Front Estimation in Multitask Finetuning

Problem Statement

Multitask finetuning requires practitioners to determine optimal task-mixing strategies (weighting vectors $\alpha$ ) to balance performance across multiple objectives. While Pareto fronts provide a rigorous framework for understanding task relationships and identifying non-dominated solutions, computing them is prohibitively expensive. It typically requires numerous independent finetuning runs to evaluate different weighting configurations, a cost that scales poorly with large models (e.g., LLMs) and high-dimensional parameter spaces.

Existing approaches attempt to mitigate this by using model merging to create cheap surrogate models for Pareto front estimation. However, current methods often rely on simple averaging or fixed quadratic surrogates without a theoretical guarantee that these merging strategies yield high-quality Pareto estimates. There is a lack of research explicitly designing new model-merging methods to directly and provably improve the quality of Pareto-front estimates.

Methodology: Variational Model Merging

The authors propose Variational Model Merging, a Bayesian framework that unifies model merging with the estimation of Pareto fronts. The core insight is to view model merging as distributed Bayesian computation where the joint posterior is formed by multiplying task-specific distributions raised to power $\alpha_t$ .

Theoretical Framework

Posterior Merging: The paper establishes that the minimizer of a weighted multitask loss (Equation 1) corresponds to the Maximum A Posteriori (MAP) estimate of a merged posterior $p_\alpha(\theta) \propto p_0(\theta)^\gamma \prod p_t(\theta)^{\alpha_t}$ .
Variational Approximation: Since exact posteriors are intractable, the authors replace task-specific posteriors $p_t$ and the prior $p_0$ with tractable variational approximations $q_t$ and $q_0$ from an exponential family.
Surrogate Losses: The merging operation minimizes a weighted sum of surrogate losses $\hat{\ell}_t(\theta) = -\log q_t(\theta)$ $\hat{ℓ}_{t} (θ) = - lo g q_{t} (θ)$ . The choice of the variational family $Q$ $Q$ dictates the expressiveness of the surrogate:
- Isotropic Gaussian: Leads to simple averaging (linear surrogates), equivalent to standard model soups or task arithmetic.
- Full Gaussian: Incorporates Hessian information, leading to quadratic surrogates (e.g., Fisher-weighted averaging).
- Mixture of Gaussians (MoG): Provides highly flexible, non-convex surrogates capable of capturing complex loss landscape geometries.

Key Algorithmic Contributions

The authors derive a practical merging strategy using Mixtures of Exponential Families (specifically Mixtures of Gaussians).

Algorithm 1: An Expectation-Maximization (EM) procedure is used to find the mode of the merged mixture posterior. This involves iteratively updating the mixture weights $\hat{\pi}_{tk}$ and the merged parameters $\theta_\alpha$ .
Efficiency: The method allows for the generation of Pareto front estimates for many $\alpha$ values with negligible computational overhead after the initial single-task finetuning and posterior approximation steps.
Implementation Variants: The paper details four specific implementations ranging from simple averaging (AdamW-SG) to Hessian-weighted merging (IVON-Hess) and Mixture-weighted merging (MultiIVON-Hess and SoupIVON-Hess).

Key Theoretical Result

The paper proves that more flexible posteriors necessarily yield better merging methods.

By framing the problem in terms of Kullback-Leibler (KL) divergence between the true optimal posterior $q^*_\alpha$ and the approximated merged posterior $q_\alpha$ , the authors show that the approximation error is minimized when the variational family is more expressive.
Specifically, if variational family $Q_1$ is a strict subset of $Q_2$ ( $Q_1 \subset Q_2$ ), the merged solution using $Q_2$ is guaranteed to be at least as good as, and typically better than, that using $Q_1$ . This implies that Mixture-weighted merging should theoretically outperform Hessian-weighted and Simple merging in estimating the Pareto front.

Experimental Results

The authors validate their theory across vision and language transformers (ResNet-20, ViT-B/32, RoBERTa, GEMMA-2B) on tasks including image classification, sentiment analysis, and machine translation.

Pareto Front Quality: Across all benchmarks, Mixture-Weighted Merging consistently produces Pareto front estimates that are closer to the true multitask finetuning front than Hessian-Weighted or Simple Merging.
Coverage: Flexible posteriors (Mixtures) are able to discover trade-off regions and solutions that simple scalarization or less expressive merging methods miss. For instance, in ResNet-20 experiments, Mixture merging improved accuracy from 76.9% (Simple) to 87.2%, approaching the 90.4% achieved by full multitask finetuning.
Cost Efficiency: While generating the initial surrogates requires single-task training, estimating the Pareto front for hundreds of $\alpha$ combinations takes seconds via merging, compared to hours or days required for full multitask finetuning sweeps.
Scalability: The method scales to billion-parameter models (GEMMA-2B) using parameter-efficient finetuning (LoRA), where merging takes approximately 1 minute compared to 17 hours for joint training.

Significance and Claims

The paper claims to be a rare instance where Bayesian ideas are explicitly used to improve Pareto analysis in deep learning.

Theoretical Gap Filled: It provides the first theoretical proof that increasing the flexibility of the posterior approximation in model merging directly improves the quality of Pareto front estimates.
Practical Utility: It offers a practical recipe for practitioners to quickly explore task trade-offs and identify high-performing weighting strategies without the prohibitive cost of exhaustive multitask finetuning.
Novelty: Unlike previous works that use merging merely as a post-hoc tool or rely on fixed quadratic surrogates, this work proposes a new class of merging methods (Variational Model Merging) derived from first principles to optimize the estimation process itself.

The authors acknowledge limitations, noting that exact Bayesian posteriors are computationally intractable for large networks, necessitating approximations (e.g., diagonal Hessians) which introduce some performance gap compared to exact solutions. However, the empirical results suggest that the proposed flexible approximations significantly narrow this gap compared to existing baselines.