Merging Models with Fisher-Weighted Averaging

In plain words

Teams often want one model that knows what several separately trained models know. A cheap trick is to average their weights, which works when the models start from the same point. This paper asks why averaging works and whether a better rule exists.

Its answer is that a plain average treats every weight as equally important. The authors weight each weight by how much the model's predictions depend on it, measured with a quantity called the Fisher information. Weights that matter a lot for a model keep their value in the merge, and weights that barely matter give way to the other models. The merge is one weighted average, with no extra training run.

The authors use this for ensembles, for making fine-tuned vision models more robust, and for moving skills from one NLP task to another. On the last one they report results close to ordinary gradient-based training at a lower cost. They also merge models in ways that are hard to do with gradient training.

TL;DR

Fisher merging replaces the plain average used in model merging with a per-parameter weighted average. The weight of each parameter in each model is that model's diagonal Fisher information §2.3 eq. 4. Gains over plain averaging are reported for ensembling, WiSE-FT robust fine-tuning and merging with donor tasks §3.1 §3.2 §3.3. The official code covers only the GLUE setting with HuggingFace BERT and RoBERTa, and it has no license file.

Problem

Averaging the parameters of models trained from the same initialization is the common way to combine them: federated learning, Polyak averaging, model ensembling by averaging, and WiSE-FT robust fine-tuning all do it §2.1. Such an average is the closed-form solution of the joint-likelihood problem when each model's posterior is an isotropic Gaussian centred on its parameters §2.1. That posterior is a poor match for real networks, so the authors ask whether a better posterior estimate yields a better merge §1.

Gradient-based transfer (fine-tuning, intermediate-task training, domain adaptation) has two further limits: an improvement to an ancestor model is not passed to its descendants, and combining several sources either needs many saved checkpoints or suffers catastrophic forgetting §1.

Assumptions: the models share an architecture and an initialization §2; the Fisher is a local property of each parameter, so the method may work worse when parameters are far apart in parameter space §2.4; and the trained parameters are assumed to be a local maximum of the posterior §2.3.

Method

Setting. Given M trained networks with parameters θ_1 … θ_M and weights λ_i ≥ 0 with Σ λ_i = 1, merging seeks the parameters that maximise the λ-weighted joint log-likelihood of the per-model Gaussian posteriors §2.2. With identity precision, the solution is the λ-weighted average of the parameters, which the paper calls isotropic merging §2.1 §2.2.

Posterior precision. The Laplace approximation gives each posterior a mean at the trained parameters and a precision equal to the Hessian at that point. At a mode, the Hessian equals the Fisher information, so the paper uses the Fisher as the precision §2.3. The Fisher is the expected outer product of the gradient of the log-likelihood of the model's output eq. 1. The full matrix is too large to store, so only its diagonal is used, the same approximation as in elastic weight consolidation §2.3 §4.

Estimate. The diagonal is estimated on N training examples. For each example, the expectation over output labels is computed exactly when there are few classes, or estimated by sampling §2.3 eq. 2.

Merge. Each merged parameter is a weighted average of the models' parameters, where each model's weight is its λ times its Fisher value for that parameter §2.4 eq. 4. Setting every Fisher to one recovers isotropic merging §2.4.

Head parameters. Only the shared body is merged. Task-specific classification heads stay as they are in the target model. The paper notes this can shift the distribution of head inputs, and reports that it worked in its experiments §2.4.

Use cases. Ensembling by averaging parameters of models fine-tuned on the same task §3.1; robust fine-tuning by averaging a pre-trained model with a fine-tuned one §3.2 fig. 1; merging a fine-tuned model with a donor-task model, as an alternative to intermediate-task training §3.3 fig. 1; and merging a model that was already trained on an intermediate task with another donor §3.3 fig. 1. Domain-adaptive pre-training checkpoints are merged with task fine-tuning checkpoints in the same way §3.4.

# Diagonal Fisher of one trained model  [eq. 2]
# Input:  model; examples (N training inputs)
# Output: one Fisher tensor per body parameter, same shape as the parameter

def diagonal_fisher(model, examples):
    # Start every accumulator at zero  [eq. 2]
    fisher = zeros_like(model.body_parameters)
    for example in examples:
        # Sum over each label, weighted by its probability  [eq. 2]
        for label in classes:
            # Gradient of log p(label | example)  [eq. 2]
            grad = gradient(log_prob(model, example, label))
            fisher += prob(model, example, label) * grad ** 2
    # Average over the N examples  [eq. 2]
    return fisher / len(examples)

# Fisher merge of one body parameter  [eq. 3] [eq. 4]
# Input:  params[m], fishers[m] for each model m; coeffs[m] sum to 1
# Output: merged parameter tensor

def fisher_merge(params, fishers, coeffs):
    # Weight of model m at each element: coeff times Fisher  [eq. 4]
    weights = [c * f for c, f in zip(coeffs, fishers)]
    # Numerator: sum of weight times parameter  [eq. 4]
    numerator = sum(w * p for w, p in zip(weights, params))
    # Denominator: sum of weights  [eq. 4]
    denominator = sum(weights)
    # Element-wise division gives the merged parameter  [eq. 4]
    return numerator / denominator

The same computation with the released code's defaults, which the paper does not describe. Each model's coefficient is divided by the global L2 norm of its Fisher, and the target model's Fisher is floored at 1e-6 merging.py merging.py. See Paper vs code.

# Fisher merge as coded by default  [code: model_merging/merging.py#L36-L71]
# Input:  params, fishers, coeffs as above; index 0 is the target model
# Output: merged parameter tensor (not in the paper)

def fisher_merge_as_coded(params, fishers, coeffs, floor=1e-6):
    # One scalar per model: L2 norm over all of its Fisher tensors
    # [code: model_merging/merging.py#L74-L76]
    norms = [global_l2_norm(fisher) for fisher in fishers]
    weights = []
    for m, fisher in enumerate(fishers):
        # Only the target model's Fisher is floored
        # [code: model_merging/merging.py#L62-L64]
        if m == 0:
            fisher = maximum(fisher, floor)
        # Coefficient divided by that model's Fisher norm
        # [code: model_merging/merging.py#L53-L55]
        weights.append(coeffs[m] / norms[m] * fisher)
    numerator = sum(w * p for w, p in zip(weights, params))
    denominator = sum(weights)
    return numerator / denominator

Hyperparameters

Results

Author-reported. Setup: BERT-base on GLUE and ensembles of five fine-tuned checkpoints from HuggingFace §3.1 app. A; ViT-B/16 for WiSE-FT §3.2; BERT-base and RoBERTa-large on GLUE §3.3; RoBERTa-base for domain adaptation §3.4.

• Ensembling (RTE, MRPC, SST-2): Fisher merging beats isotropic merging in every case and is comparable to averaging the predictions of the five models, at one model's inference cost instead of five §3.1 fig. 2. Numbers are only in the figure.
• WiSE-FT (ImageNet to five OOD sets): Fisher merging gives a better IID and OOD trade-off than isotropic merging. At the λ_1 with the best average OOD accuracy, IID accuracy is about 1 point higher §3.2 fig. 3 app. B fig. 7.
• Intermediate-task training (BERT-base, RTE target): gradient-based intermediate training hurts on some donors, while merging never hurt in the experiments. Fisher is comparable to or better than isotropic merging, with the largest gap for MNLI §3.3 fig. 4.
• Merging with an MNLI-trained RTE model (BERT-base, RTE target): Fisher merging with an MNLI model is worse than gradient-based training, which reaches 77.4 ± 1.6 against 73.2 ± 5.1 for Fisher §3.3 tab. A3 tab. A1. Fisher merging still improves on gradient-based intermediate training for all donors, per fig. 5.
• RoBERTa-large (RTE after MNLI): the largest boost is 2.2 points, from Fisher merging with another RTE checkpoint §3.3 fig. 6.
• Domain adaptation (RoBERTa-base, five trials):

• Fewer Fisher examples (BERT-base, MNLI donor, RTE target): 73.4% with the full training set, 72.7% with 256 examples, and 72.2% for isotropic merging §3.3 tab. A4.
• Cost (estimated FLOPs for one RTE checkpoint, BERT-base): fine-tuning for 10 epochs is about 5.5e14 FLOPs. Merging takes about 4.0e8 FLOPs, the evaluation about 2.0e12, and the Fisher about 9.1e13. The authors put the total at about 6 times less than fine-tuning for Fisher merging and 275 times less for isotropic merging §3.3 [Costs paragraph].

Table A1 (Fisher merging on GLUE with BERT-base, rows = donor, columns = target) and Table A2 (isotropic) give every GLUE pair app. E tab. A1 tab. A2. Table A3 gives sequential gradient fine-tuning tab. A3.

Paper vs code

The official repository is https://github.com/mmatena/model_merging, linked from the paper's abstract. It is owned by the first author, at commit 46802f98319fea0a606b2881eac5cdc01dbb48e7. The repository has no LICENSE file, so this card paraphrases it and quotes no code.

• Fisher normalisation (code default on, not in paper). Each model's coefficient is divided by the L2 norm of that model's Fisher merging.py merging.py merging.py merge_and_evaluate.py. Paper eq. 4 has no such term eq. 4. Whether the paper's results used it is not stated.
• Fisher floor. By default the code floors only the target model's Fisher, at 1e-6 merging.py merge_and_evaluate.py. The paper says the target value is used when the Fisher is near zero for all models [§2.4 Numerical Issues]. The code's rule is different: a target Fisher of zero still gets weight 1e-6 in the denominator.
• Coefficient grid. The code's default is 51 points including both endpoints merge_and_evaluate.py merging.py. The paper uses 50 points for GLUE §3.3.
• Selection split. The code selects the best merge by average score on the same split it reports, with defaults of validation and 4096 examples merge_and_evaluate.py merge_and_evaluate.py evaluation.py. The paper selects on the first 2048 validation examples for GLUE §3.3 and on the first 2048 test examples for DAPT app. D.
• Fisher expectation. The paper allows sampled expectations and used one sample per example for DAPT §2.3 app. D. The code asserts the exact expectation fisher.py.
• Dropout. The code runs the Fisher forward pass with training=False fisher.py. The paper does not say.
• Multi-model merges. The grid mode accepts exactly two models merge_and_evaluate.py. Random Dirichlet coefficients handle more models merging.py. The paper leaves merges of more than two models to future work §3.3.
• STS-B range. The code bins [0, 5] into 25 buckets data.py data.py. A comment in the code says the paper states the range starts at 1. This version of the paper says 0 to 5 §3.3.
• Scope of the code. The repository covers GLUE with HuggingFace BERT and RoBERTa [README.md]. The ensembling, WiSE-FT and DAPT experiments are not in it, and their hyperparameters are not released in the repository.
• Consistent. Only the body is merged and the head comes from the target model, as in the paper hf_util.py merging.py.

Pitfalls

• Normalisation changes the weights. Dividing by each model's Fisher norm gives a model with a large-norm Fisher less weight. Implementers should decide which variant to use and say so. The paper's equations describe the un-normalised form.
• Floor semantics. The floor applies to the target only, so donor parameters with Fisher near zero get almost no weight, while a near-zero target Fisher gets weight 1e-6 times its coefficient. The paper says the choice of fallback value had little effect §2.4.
• Fisher cost. Estimating the diagonal needs roughly the same compute as training on N examples, because it needs a per-example gradient for each class §2.3. The Fisher is the most expensive step, but it can be reused across merges §3.3.
• Cost numbers are estimates. The FLOP figures use heuristics from a scaling-law paper, not measured wall-clock times §3.3. The Fisher term dominates the total.
• Selection on the reporting split. λ is tuned on validation data that the results report, which inflates the reported scores. The DAPT section says test examples were used for selection app. D. Reproduce with a held-out split.
• Variance. Several differences are within one standard deviation: ChemProt (Fisher 83.1 ± 0.4 against isotropic 82.8 ± 0.4) and SciERC, where isotropic merging is the best row tab. 1. On RTE with an MNLI donor, the Fisher standard deviation is 5.1 tab. A1.
• Donor-task merges can lose to gradient training. For an MNLI donor on BERT-base, sequential fine-tuning beats Fisher merging on RTE tab. A1 tab. A3. The paper's claim of a cheaper alternative holds for some donors, not all.
• Ensembling Fisher not described. The paper does not state how the Fisher was computed for the ensembling experiments §3.1.
• Laplace assumptions. The diagonal ignores correlations between parameters, and treats the trained parameters as a local maximum of the posterior §2.3.
• Only one multi-stage order is reported. The paper merges an MNLI-then-RTE model with donors §3.3 fig. 5. It does not report other orders, and its merges use two models at a time §3.3.

Lineage

• Builds on: Laplace approximation and Fisher information (EWC uses the same diagonal approximation for regularisation) §2.3 §4; federated averaging §1; Polyak averaging and Model Soups for ensembling §3.1; WiSE-FT for robust fine-tuning §3.2; intermediate-task training §3.3; domain-adaptive pre-training §3.4. The paper cites these by name only; no arXiv ids are given in the text.
• Compared against: isotropic merging (plain averaging) §2.1; gradient-based intermediate-task training §3.3; sequential gradient fine-tuning on the same donor chain app. E tab. A3; prediction ensembling §3.1.
• Superseded by: not stated in this wiki.