PaperlayerView exact revision on arXiv

Source-linked AI summary

Model Fusion via Optimal Transport

Sidak Pal Singh, Martin Jaggi

arXiv:1910.05653v6cs.LGstat.ML

TL;DR

The paper asks how differently trained neural networks can be fused into one model when ensembles are too costly and retraining or data sharing is undesirable. It aligns neurons layer by layer with optimal transport before averaging parameters. The resulting OT fusion outperforms vanilla averaging, supports one-shot transfer and different-sized models, and can replace ensembles after mild fine-tuning.

  • Problem

    Ensembles require maintaining and evaluating multiple models, while direct weight averaging lacks reliable neuron correspondences; the paper studies one-shot fusion without retraining or training-data sharing.

  • Method

    The method uses layer-wise optimal transport to softly align neurons and weights before averaging parameters, interpreted as a Wasserstein barycenter.

  • Results

    OT fusion successfully merges differently weighted and differently tasked networks, significantly outperforms vanilla averaging, and can serve as an efficient ensemble proxy with mild fine-tuning.

  • Takeaways & Limitations

    The approach supports one-shot knowledge transfer, data-free structured-pruning post-processing, and fusion of models with different sizes.

  • Takeaways & Limitations

    The ideal joint search over permutation matrices across all layers is computationally intractable, so the method uses a layer-wise greedy procedure.

Abstract

from arXiv · show

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, CIFAR100, 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. The code is available at the following link, https://github.com/sidak/otfusion.

1 Introduction

The paper addresses how to fuse differently trained neural networks into one model without retraining or sharing training data. It proposes optimal-transport alignment before parameter averaging and reports improvements over vanilla averaging, with applications to model exchange and compression.

  • Method: The method targets one-shot fusion using only network weights, avoiding the need for retraining and training-data sharing.The setting includes networks trained differently but having the same number of layers.
  • Motivation: Ensembles improve prediction performance and robustness but require storing and evaluating all K trained models.This creates memory and computation demands that grow with the number of models.
  • Motivation: Vanilla averaging directly averages corresponding parameters, but differing neuron roles make those correspondences unreliable and often produce poor performance.A neuron in one model may match a differently indexed neuron in another model.
  • Method: The proposed layer-wise method uses optimal transport to softly align neurons and weights before averaging parameters.The alignment cost is based on activation or incoming-weight similarity, and the average can be interpreted as a Wasserstein barycenter.
  • Results: The authors report successful merging across differing weights and slightly different tasks, outperforming vanilla averaging in all scenarios.The method can inherit abilities unique to one parent and outperform that parent on the other network's task.
  • Applications: The approach is presented for model exchange, different-sized model fusion, federated learning, and privacy- or communication-sensitive settings.These applications seek to encourage model exchange instead of data exchange.

2 Related Work

Related work combines models through ensembles, distillation, assignment-based fusion, and neuron alignment. The paper positions OT fusion as a broader parameter-aggregation approach whose Wasserstein-barycenter formulation extends beyond federated learning.

  • Ensembling: Ensembling combines model outputs but conflicts with the goal of maintaining and running only one model.The related-work discussion treats multiple-model inference as infeasible for this setting.
  • Distillation: Distillation transfers a teacher's abilities to a student using teacher output probabilities as an additional training signal.The student is typically smaller than the pre-trained teacher.
  • Distillation: Distillation is computationally costly, requires hyper-parameter tuning, and still requires sharing training data with the teacher.The paper contrasts this with its data-sharing-avoiding approach.
  • Network fusion: Prior assignment-based fusion methods can be limited to parent models whose weights are already close.The cited Elastic Weight Consolidation approach formulates an assignment problem using diagonal Hessian approximations.
  • Alignment-based methods: The paper develops OT-based alignment with Wasserstein barycenters for general model fusion, including different-sized models and structured pruning.It contrasts this scope with concurrent alignment work focused on federated learning.
  • Novelty: The authors state that applying Wasserstein barycenters to neural-network weight averaging had not previously been considered, to their knowledge.This is presented as the method's novelty claim.

3 Background on Optimal Transport (OT)

Optimal transport compares probability distributions by minimizing the cost of moving mass between points in a shared ground space. Wasserstein barycenters extend this idea to define a principled notion of averaging across multiple measures.

  • OT formulation: Optimal transport compares two probability distributions using a ground-space distance or transport cost.The paper introduces this discrete formulation as notation for the remainder of the work.
  • OT formulation: Empirical measures are represented as weighted sums of Dirac distributions over points in the ground space.The weights lie in a probability simplex.
  • OT formulation: The ground-cost matrix C specifies the cost of moving each source point x(i) to each target point y(j).The transport optimization uses these pairwise costs.
  • OT formulation: The transport map T records how much probability mass moves from each source point to each target point.The optimal map is obtained from the linear-program formulation.
  • Wasserstein distance: For S = R^d with metric-based costs, optimal transport induces the p-Wasserstein distance between probability distributions.The distance is defined through the optimal transport objective.
  • Wasserstein barycenters: A Wasserstein barycenter is a measure minimizing the weighted sum of p-th-power Wasserstein distances to K input measures.The weights η belong to the K-dimensional probability simplex.

4 Proposed Algorithm

The proposed fusion method uses optimal transport to softly align neurons layer by layer before averaging model parameters. It supports one-shot fusion, unequal widths, and data-free alignment, while acknowledging computational and data-related trade-offs.

  • Motivation: Optimal transport provides soft neuron correspondences when direct parameter matching is unavailable, including between layers with different neuron counts.The transport map replaces an ideal permutation by optimally transporting neurons from one model to another.
  • General procedure: The algorithm defines neuron measures, initializes uniform masses, aligns incoming weights using the previous layer’s transport map, and computes transport maps sequentially.After alignment, corresponding layer weights are averaged to form the fused model.
  • General procedure: The procedure starts after the input layer and treats the final output-layer transport as identity because input and output neuron orders are shared.This preserves the fixed ordering at the network boundaries.
  • Multiple models: For multiple models, the method estimates a fused model, aligns every model to that estimate, and averages the aligned weights.The two-model procedure is recovered when the initial fused estimate is one parent model.
  • Alignment strategies: Alignment can use neuron activations from unlabeled samples or incoming weight vectors, with Euclidean distance serving as the optimal-transport ground metric.Activation-based alignment uses pre-activations, whereas weight-based alignment represents each neuron by its incoming weights.
  • Alignment strategies: Weight-based alignment requires no dataset samples, while activation-based alignment can target domains using unlabeled data; their fusion results are reported as similar.Activation-based alignment experiments typically use approximately 100–400 samples, although 25 samples still outperform vanilla averaging.
  • Limitations: The ideal joint permutation search is computationally intractable for deep networks, motivating a layer-wise greedy fusion procedure.Entropy-regularized optimal transport can improve runtime but slightly reduces test accuracy relative to exact optimal transport.

5 Experiments

The experiments evaluate OT-based fusion across heterogeneous tasks, multiple models, unequal widths, pruning, and teacher–student transfer. Across these settings, OT alignment improves fusion outcomes over vanilla averaging and supports efficient, data-free or one-shot transfer.

  • 5.1 Fusion in the setting of heterogeneous data and tasks: OT fusion significantly outperforms vanilla parameter averaging and improves over individual models in heterogeneous skill-transfer settings without retraining.Models A and B combine specialist and generalist abilities through activation-based alignment, including across different or shared initializations.
  • 5.1 Fusion in the setting of heterogeneous data and tasks: 93.11% overall test accuracy is achieved by OT averaging, versus 89.78% and 87.35% for the individual models in one heterogeneous-data experiment.The result confirms skill transfer from both parent models in the reported setting.
  • 5.2 Fusing different sized models: Unequal-width fusion maps weights from a bigger model to a smaller one before averaging, enabling model compression and federated-learning use cases.The same procedure can also combine parameters in the opposite direction, from smaller to larger networks.
  • 5.2 Fusing different sized models: ≈10% or more test-accuracy gains occur in the high-sparsity regime when OT fusion post-processes structured pruning of VGG11 on CIFAR10.The approach fuses the original dense network into the pruned model and remains independent of the pruning algorithm used.
  • 5.2 Fusing different sized models: 89.4% accuracy is recovered after CONV_8 pruning from 90.3% to 81.5%, while CONV_7 pruning improves from 87.6% to 90.1% at 41% overall sparsity.These examples show performance recovery after structured pruning using OT fusion.
  • 5.3 Efficient alternative to ensembling: For multiple VGG11 models on CIFAR100, OT averaging plus fine-tuning gains approximately 1.4%, 1.7%, and 2% over the best individual models for 4, 6, and 8 models.Vanilla averaging approaches random performance and fails to retrain, whereas fine-tuned OT fusion remains more efficient than ensembling.
  • 5.4 Teacher-Student Fusion: OT averaging plus fine-tuning transfers teacher knowledge into smaller students, outperforming both the original student and fine-tuning the student alone.Using the OT-fused model as the distillation student also outperforms random or pretrained-student initialization, without requiring a temperature sweep.

6 Conclusion

The paper presents optimal-transport-based neuron alignment followed by weight averaging as a versatile approach to fusing neural networks. It supports one-shot knowledge transfer, pruning, different model sizes, and replacing ensembles after fine-tuning.

  • Layer-wise soft neuron alignment via optimal transport precedes weight averaging for fusing models in varied settings.
  • The approach enables one-shot knowledge transfer without sharing training data.
  • The method provides data-free, algorithm-independent post-processing for structured pruning.
  • The framework combines parameters from models with different sizes and, after fine-tuning, can replace complete ensembles with one model at inference.
  • The authors identify distributed optimization, continual learning, and generative-model fusion as future use cases.

Broader Impact

Model fusion is framed as direct knowledge transfer that can support privacy-preserving collaborative learning. Reliable fusion could help multiple agents exchange knowledge without requiring centralized sharing of sensitive data.

  • Model fusion is a building block for collaborative learning because it transfers knowledge directly between trained neural networks.
  • In federated learning, fusion can help agents participate in joint training and knowledge exchange while preserving data ownership and privacy.
  • The authors position reliable fusion as a step toward privacy-preserving and efficient collaborative machine learning.

S1.1 Experimental Details

The experiments evaluate VGG11, ResNet18, and MLPNET under specified SGD training settings, using best checkpoints for selected convolutional models. Biases are fixed to zero, and batch normalization is omitted in the ResNet18 experiments.

  • Training configurations: VGG11 is trained for 300 epochs with SGD, a 0.05 initial learning rate, momentum 0.9, weight decay 0.0005, and batch size 128.The learning rate is halved every 30 epochs, and the best test-accuracy checkpoint is selected.
  • Training configurations: MLPNET uses SGD with a constant learning rate of 0.01, momentum 0.5, and batch size 64.
  • Training configurations: ResNet18 is trained for 300 epochs with SGD, learning-rate drops at epochs 150 and 250, momentum 0.9, weight decay 0.0001, and batch size 256.Batch normalization is skipped in the current experiments.
  • Alignment settings: The activation-based alignment results use pre-activation values, which generally perform slightly better than postactivation values.
  • Alignment settings: The regularization constant for activation-based alignment in Table S2 is 0.05.
  • Implementation boundary: Neuron biases are set to zero in all experiments rather than treated as regular weights.Handling biases as regular weights is left for future work.

S1.2 Combining weights and activations for alignment

The activation-based alignment variant incorporates neuron activation statistics into optimal-transport probability masses. Mean and standard-deviation information are combined into neuron importance weights before the remaining algorithm proceeds unchanged.

  • Activation outputs provide feature-presence signals that can be combined with weight-based alignment through probability mass values.
  • Mean activation measures neuron significance, while standard deviation captures consistently high absolute activations across inputs.
  • Neuron importance is computed as the elementwise product of mean and standard-deviation activation statistics.The paper denotes this product with ⊙.
  • The resulting importance values define the probability masses b^(l)_k, while the rest of the alignment algorithm remains unchanged.

S1.3 Optimal Transport

The implementation uses Python Optimal Transport on CPUs for Wasserstein distances and barycenters, while GPU implementations could further improve efficiency.

  • The method uses the Python Optimal Transport library to compute Wasserstein distances and barycenters on CPU.
  • GPU implementations are possible and could further boost computational efficiency.
  • CPU execution is currently sufficient for the reported timings.

S1.4 Timing information

The experiments show that OT averaging is computationally practical and consistently supports stronger fusion, retraining, heterogeneous model combination, unequal-width transfer, pruning, and distillation initialization than simpler baselines.

  • Timing information: ≈3 seconds for MLPNET on MNIST, ≈5 seconds for VGG11 on CIFAR10, and ≈7 seconds for ResNet18 on CIFAR10 quantify activation-based alignment costs.These timings include computing activations over the mini-batch.
  • Retraining performance: OT averaging outperforms vanilla averaging by a large margin during MLPNet retraining, including early training stages where vanilla averaging performs worse.
  • Multi-model retraining: OT averaging successfully retrains across settings except at the original learning rate of 0.05, whereas vanilla averaging fails to retrain effectively.For four and six models, vanilla averaging reaches best accuracies of 12.40 and 11.01, respectively.
  • Heterogeneous data splits: 95.3 mean test accuracy versus 95.1 for vanilla averaging shows OT fusion outperforming vanilla averaging across tested data proportions.The reported OT result has standard deviation 0.1 at proportions 0.1 and 0.9.
  • Extensions and applications: OT fusion transfers performance across unequal model widths, improves post-processing for CIFAR100 pruning, and provides stronger initialization than smaller-model or random alternatives.Finetuning OT averages outperforms finetuning the smaller model across multiple runs and most tested hyperparameter settings.
  • One-shot skill transfer: 93.6% test accuracy gives OT averaging a ≈6% gain over the best individual model while remaining 4× more efficient than ensembling.In the same setting, ensembling reaches 95.0% and vanilla averaging 80.6%.
Loading 1910.05653v6…