Model Fusion via Optimal Transport

By Sidak Pal Singh, Martin Jaggi. 2019

318 citations recorded by the source corpus.

Combining different models is a widely used paradigm in machine learning applications. While the most common approach is to form an ensemble of models and average their individual predictions, this approach is often rendered infeasible by given resource constraints in terms of memory and computation, which grow linearly with the number of models. We present a layer-wise model fusion algorithm for neural networks that utilizes optimal transport to (soft-) align neurons across the models before averaging their associated parameters. We show that this can successfully yield "one-shot" knowledge transfer (i.e, without requiring any retraining) between neural networks trained on heterogeneous non-i.i.d. data. In both i.i.d. and non-i.i.d. settings, we illustrate that our approach significantly outperforms vanilla averaging, as well as how it can serve as an efficient replacement for the ensemble with moderate fine-tuning, for standard convolutional networks (like VGG11), residual networks (like ResNet18), and multi-layer perceptrons on CIFAR10 and MNIST. Finally, our approach also provides a principled way to combine the parameters of neural networks with different widths, and we explore its application for model compression. Code is available under the following link https://github.com/modelfusion.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Mathematical Optimization

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Mathematical Optimization

Citations in the corpus

Attack of the Tails: Yes, You Really Can Backdoor Federated Learning (2020, methodology): “Alternatives to simple weighted averaging have also been proposed: PFNM [47] and FedMA [48] proposes using neural matching, and [49] uses optimal transport to achieve model aggregation.”
Federated Learning with Matched Averaging (2020, background): “Concurrent work of Singh & Jaggi (2019) explores the Wasserstein barycenter variant of equation 2.”
Co-Transport for Class-Incremental Learning (2021, neutral): “model fusion [37], domain adaptation [7], object detection [10], model reuse [56, 57], and generative models [2, 4].”
Dataless Knowledge Fusion by Merging Weights of Language Models (2022, unknown): “We apply weight-based matching in OTFusion (Singh & Jaggi, 2020).”
BioMistral: A Collection of Open-Source Pretrained Large Language Models for Medical Domains (2024, methodology): “Methods include averaging model weights or considering permutation invariance (Il-harco et al., 2022; Choshen et al., 2022; Singh and Jaggi, 2020; Ainsworth et al., 2023).”
Git Re-Basin: Merging Models modulo Permutation Symmetries (2022, methodology): “• Singh & Jaggi (2020) introduces a method for merging multiple models simultaneously, but only demonstrates results on at most 8 models at a time and performs continued training after merging.”
FjORD: Fair and Accurate Federated Learning under heterogeneous targets with Ordered Dropout (2021, methodology): “Recent techniques tackle this problem by matching clients’ neurons before averaging [71, 60, 65].”
The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks (2021, background): “Singh & Jaggi (2020) proposed a layer-wise model fusion algorithm for making ensembles.”
Parameter-Efficient Fine-Tuning Methods for Pretrained Language Models: A Critical Review and Assessment (2023, methodology): “It employs the optimal transport method [78], [79] to align the parameters of adapters based on weights and activations, which gives better performance with fewer trainable parameters compared to AdapterFusion.”
WoodFisher: Efficient second-order approximations for model compression (2020, methodology): “(2015), or aligning and fusing intermediate or output layers via optimal transport (Singh and Jaggi, 2019) or by weight sharing (Chou et al.”
Fine-tuning Global Model via Data-Free Knowledge Distillation for Non-IID Federated Learning (2022, unknown): “In addition, [37] a r X i v : 2203 .”
Meta-learning without data via Wasserstein distributionally-robust model fusion (2022, methodology): “Similar to FedMA, [Singh and Jaggi, 2020] utilizes opti-2046 mal transport to align neurons across models in a layer-wise fashion.”
TIES-Merging: Resolving Interference When Merging Models (2023, methodology): “Previous works merge models by summing the individual model weights with different weighting schemes, either via a simple average [9, 28, 83], via more sophisticated means that incorporate parameter importance [45] or account for permutation invariances [1, 31, 70, 74, 42].”