🤖 machine learning

Model Fusion via Retrofitting

This paper introduces a neuron-centric model fusion framework that treats fusion as a representation-matching problem by grouping intermediate neurons and using attribution scores to align salient features, achieving superior performance over existing methods—particularly in zero-shot and non-IID scenarios—across diverse architectures like VGGs, ResNets, and ViTs without retraining.

Original authors: Phoomraphee Luenam, Andreas Spanopoulos, Amit Sant, Thomas Hofmann, Sotiris Anagnostidis, Sidak Pal Singh

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Phoomraphee Luenam, Andreas Spanopoulos, Amit Sant, Thomas Hofmann, Sotiris Anagnostidis, Sidak Pal Singh

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 have two expert chefs who have spent years perfecting their own unique recipes. Chef A specializes in Italian cuisine, and Chef B specializes in Japanese cuisine. They have never worked together, they use different tools, and they might even measure ingredients in different units.

Now, imagine you want to create a single "Master Chef" who can cook both Italian and Japanese dishes perfectly, but you cannot ask them to go back to culinary school to relearn everything from scratch. You just want to combine their existing knowledge into one person immediately.

This is the problem of Model Fusion. In the world of Artificial Intelligence (AI), we often have different neural networks (the "chefs") trained separately. We want to merge them into one powerful model without the expensive and time-consuming process of retraining them on new data.

The Problem: The "Language Barrier"

The paper explains that simply averaging the weights of these two models (like mixing their recipe books page by page) usually fails. Why? Because even if they are trained on similar data, they learn things in different ways.

The "Permutation" Issue: Think of Chef A's "spice rack" where the "cumin" is in the top drawer, while Chef B's "cumin" is in the bottom drawer. If you just average their drawers, you end up with a messy mix that doesn't make sense. The AI neurons (the "ingredients") are in different orders and locations.
The "Zero-Shot" Challenge: Most existing methods work okay if the chefs have similar backgrounds. But if they trained on completely different data (like one on Italian food, one on Japanese food), existing methods fail miserably. The combined model becomes confused and performs worse than the individual chefs.

The Solution: "Retrofitting" with a "Target"

The authors propose a new method called Model Fusion via Retrofitting. Instead of just averaging the chefs, they use a clever two-step process that acts like a translator and a coach.

Step 1: The "Grouping" (Finding the Twins)

First, the algorithm looks at what the neurons (the neurons are like the individual neurons in a human brain, or the specific "skills" of the chefs) are actually doing when they see data.

It groups similar neurons from Chef A and Chef B together.
The Creative Twist: The paper introduces Neuron Attribution Scores. Imagine that not all ingredients are equally important. Some spices are critical for the flavor; others are just garnish. The algorithm identifies which neurons are the "star players" (high importance) and which are "bench warmers" (low importance).
It then creates a "Target" for each group. Think of this as a "Gold Standard" recipe card that represents the perfect average of the grouped neurons, weighted by how important they are.

Step 2: The "Retrofit" (The Coach)

Now, the algorithm builds the new "Master Chef" (the fused model). It doesn't just copy the old chefs; it trains the new model's layers to match those "Gold Standard" targets it just created.

It's like taking a new apprentice and saying, "Don't just copy Chef A or Chef B. Look at this specific target dish we designed. Adjust your cooking until your output matches this target perfectly."
This happens layer by layer, from the bottom of the network to the top.

Why This is Better (The Results)

The paper tested this on various AI architectures (like VGG, ResNet, and Vision Transformers) and found:

It Works When Others Fail: In "Zero-Shot" scenarios (where the models are merged without any extra training on new data), existing methods often collapse into random guessing. The authors' method maintains high accuracy, even when the models were trained on completely different, non-overlapping data (like the "Sharded" setup where one model sees only red cars and the other sees only blue cars).
It Respects Importance: By using those "Neuron Attribution Scores," the method ensures that the most critical features of the original models are preserved, rather than getting lost in the noise of less important features.
It's Flexible: It works on different types of AI architectures, not just one specific kind.

The Trade-off

The paper admits that this method is a bit slower and computationally heavier than the "quick and dirty" averaging methods. However, they argue this is worth it because:

It produces a usable model immediately (Zero-Shot), whereas other methods often require hours or days of extra fine-tuning to work at all.
In the long run, the time saved on fixing broken models outweighs the extra time spent on the initial fusion.

Summary Analogy

Old Way: You take two different languages, translate them word-for-word into a third language, and hope the meaning stays. It usually results in gibberish.
New Way (Retrofitting): You identify the core concepts (the "targets") that both chefs understand, assign them a value based on importance, and then teach a new chef to speak that specific, unified language perfectly.

The paper concludes that this approach allows us to combine independent AI models effectively, even when they are very different, without needing to retrain them from scratch.

Technical Summary: Model Fusion via Retrofitting

Problem Statement

Model fusion aims to combine independently trained neural networks into a single model without retraining, offering a solution for scenarios where retraining is prohibitively expensive or infeasible due to data privacy constraints (e.g., Federated Learning). However, this task is complicated by representational divergence caused by permutation invariance, random initialization, and heterogeneous training data distributions.

Existing methods, such as OTFusion (Singh & Jaggi, 2020) and Git Re-Basin (Ainsworth et al., 2023), rely on the Linear Mode Connectivity (LMC) conjecture, which posits that independently trained networks lie in the same low-loss basin. While LMC holds for models trained on identical data, prior studies indicate that even models trained on the same data may learn divergent internal features. Consequently, existing fusion methods struggle significantly in zero-shot settings (fusion without further training) and under non-IID (non-Independent and Identically Distributed) data regimes, often collapsing to near-random performance. Furthermore, current approaches often lack generality across architectures (e.g., struggling with Transformers) and rely on base models that may not be fully optimized, obscuring the true potential of fusion algorithms.

Methodology

The authors propose a neuron-centric family of fusion algorithms that reframes model fusion as a principled representation-matching problem. Instead of simply averaging weights or finding permutation matrices, the method treats fusion as a two-stage process: grouping neurons from parent models into target representations and then training the fused model's sub-networks to approximate these targets.

Core Framework

Representation Cost: The objective is defined by the squared $L_2$ distance between the activations of base model neurons and the closest neurons in the fused model, weighted by neuron importance scores (attribution scores). This penalizes the misrepresentation of salient features.
Decomposition (Theorem 1): The authors introduce an auxiliary target vector $T$ $T$ to decompose the representation cost into two interpretable components:
- Grouping Error: Measures how well base model neurons cluster around importance-weighted centroids ( $T_k$ ).
- Approximation Error: Measures how well the fused model's output neurons approximate these centroids.
Two-Stage Algorithm:
- Step 1 (Grouping): Neurons from the base models are clustered to minimize grouping error. This yields importance-weighted centroids that serve as targets.
- Step 2 (Retrofitting): The weights of the fused model at the current level are optimized (via closed-form solution for linear layers or SGD for non-linear layers) to minimize the approximation error against these targets. This process is performed bottom-up through the network levels.

Algorithm Variants

The paper introduces two primary algorithms based on the grouping strategy:

Hungarian Fusion (HF): Designed for equal-sized models with level-wise one-to-one matching. It uses the Hungarian algorithm to find the optimal assignment.
K-means Fusion (KF): Designed for the general case, including models of arbitrary widths and depths. It uses heuristic K-means clustering to group neurons.

Both algorithms have Linear (closed-form solution for affine layers) and Gradient-based (SGD/Adam for non-linear layers) variants.

Key Innovations

Neuron Attribution: Unlike prior work, this approach incorporates neuron importance scores (e.g., Conductance, DeepLIFT) to bias alignment toward salient features, improving performance in heterogeneous settings.
Architectural Flexibility: The method is applicable to any architecture modularizable as a Directed Acyclic Graph (DAG) of levels, empirically validated on VGGs, ResNets, and Vision Transformers (ViTs).
Zero-Shot Capability: By explicitly optimizing the approximation of clustered representations, the method achieves strong performance without the extensive fine-tuning required by previous methods.

Experimental Results

The authors evaluated their methods across three training regimes: Sharded (disjoint class sets), Non-IID (skewed class distributions), and Full Dataset.

Sharded/Non-IID Regimes: In extreme non-IID settings (e.g., 6-way class splits on CIFAR-100), existing methods like OTFusion and Git Re-Basin collapsed to near-random accuracy (often <5%). In contrast, KF-Gradient and HF-Gradient maintained robust zero-shot accuracy (e.g., 34.7% on CIFAR-100 for 6-way splits vs. <1.5% for baselines).
Full Dataset Regimes: Even when models are trained on the full dataset, the proposed methods outperform baselines in zero-shot fusion. For example, on CIFAR-100, KF-Gradient achieved 63.0% zero-shot accuracy compared to 4.3% for Transformer OTFusion. After fine-tuning, the proposed methods consistently approached ensemble performance while maintaining single-model inference costs.
Compression and Heterogeneity: The method successfully fused models of different widths and compressed larger models (ResNet-50) into smaller ones (ResNet-18) with significantly higher accuracy than standard Knowledge Distillation (KD) when data was limited.
Robustness: The approach demonstrated robustness to varying fusion dataset sizes, with data augmentation effectively mitigating the need for large fusion sets.

Significance and Claims

The paper claims to address three key gaps in prior research:

Reproducibility: Providing a flexible, open-source reimplementation that facilitates systematic benchmarking across architectures.
Base Model Quality: Demonstrating that fusion methods can yield significant improvements even when base models are well-trained, challenging the notion that poor fusion results are solely due to weak baselines.
Heterogeneous Data: Showing that existing methods fail in zero-shot fusion under non-IID conditions, while the proposed neuron-centric approach successfully integrates models trained on disjoint or skewed data distributions.

The authors emphasize that their approach is the first to successfully incorporate neuron importance scores into model fusion, leading to consistent improvements over existing baselines, particularly in zero-shot and non-IID scenarios where other methods fail. They position the method as a practical solution for Federated Learning and privacy-preserving model aggregation, where retraining is not an option. While acknowledging that gradient-based variants are computationally more expensive than lightweight baselines, the authors argue this overhead is justified by the ability to produce viable fused models in regimes where competitors yield unusable results.