PaperlayerView exact revision on arXiv

Source-linked AI summary

Merging Models with Fisher-Weighted Averaging

Michael Matena, Colin Raffel

arXiv:2111.09832v2cs.LG

TL;DR

The paper asks how knowledge and capabilities can be transferred across trained models without relying exclusively on repeated gradient-based training. It introduces Fisher merging, which uses Fisher-weighted parameter averaging under a Laplace posterior approximation, and finds comparable or better performance than traditional transfer learning at significantly lower cost.

  • Problem

    Existing gradient-based transfer methods require iterative retraining, discard prior downstream work when checkpoints improve, and generally transfer information from only one model.

  • Method

    Fisher merging combines model parameters by treating their posteriors with a Laplace approximation whose precision is the diagonal Fisher information.

  • Results

    Fisher merging outperforms unweighted parameter averaging and achieves comparable or sometimes better performance than gradient-based transfer learning at significantly lower cost.

  • Takeaways & Limitations

    Merging provides a fast way to combine capabilities across models and supports transfer strategies that are onerous with traditional gradient-based training.

  • Takeaways & Limitations

    The full Fisher matrix is impractical for all but the smallest models, so the method uses a diagonal approximation in practice.

Abstract

from arXiv · show

Averaging the parameters of models that have the same architecture and initialization can provide a means of combining their respective capabilities. In this paper, we take the perspective that this "merging" operation can be seen as choosing parameters that approximately maximize the joint likelihood of the posteriors of the models' parameters. Computing a simple average of the models' parameters therefore corresponds to making an isotropic Gaussian approximation to their posteriors. We develop an alternative merging procedure based on the Laplace approximation where we approximate each model's posterior as a Gaussian distribution whose precision matrix corresponds to its Fisher information. We first show that our "Fisher merging" technique provides a performance boost in settings where simple parameter averaging is currently used -- specifically, robust fine-tuning and model ensembling. Then, we compare merging to standard gradient-based transfer learning and demonstrate that merging enables a fundamentally different method for transferring capabilities across models. Specifically, we show that Fisher merging is competitive with gradient-based transfer learning approaches (while being significantly cheaper) in intermediate-task training and domain-adaptive pre-training. We also show that our merging procedure makes it possible to combine models in previously unexplored ways. We release our code to facilitate future research into methods for merging models.

1 Introduction

The paper reframes model merging as a way to transfer and combine capabilities without repeated gradient-based training. It introduces Fisher merging and evaluates it across ensembling, robust fine-tuning, and transfer-learning settings.

  • Motivation: Gradient-based transfer learning requires iterative fine-tuning and can force researchers to repeat downstream training when improved or intermediate-task checkpoints become available.Existing approaches also primarily transfer information from a single model, while sequential multi-task training risks checkpoint growth or catastrophic forgetting.
  • Existing merging: Parameter merging averages models sharing an architecture and initialization, and is used for federated learning, robust fine-tuning, and parameter-based ensembling.Compared with prediction ensembling, parameter averaging produces one model rather than requiring inference through all ensemble members.
  • Method: Fisher merging interprets merging as approximately maximizing joint posterior likelihood and weights each parameter using the diagonal Fisher information.Simple parameter averaging corresponds to an isotropic Gaussian posterior approximation, whereas Fisher merging uses a Laplace approximation.
  • Results: Fisher merging outperforms isotropic merging across model ensembling and robust fine-tuning, and is competitive with gradient-based transfer learning at significantly lower cost.The paper evaluates intermediate-task transfer and domain-adaptive pre-training in addition to established merging applications.
  • Implications: The experiments show that merging can combine models through transfer strategies that are onerous or impossible with existing gradient-based methods.The paper presents these strategies as a way to explore previously unexplored combinations of trained models.

2 Weighted Parameter Averaging for Model Merging

The paper frames model merging as maximizing joint posterior likelihood and derives Fisher merging as a Fisher-weighted alternative to isotropic parameter averaging. The method combines shared-architecture, same-initialization models while retaining practical approximations and safeguards for large or incompatible parameter sets.

  • Isotropic merging: Simple parameter averaging is equivalent to an isotropic Gaussian posterior approximation, while model-specific coefficients control each model’s importance.Equal coefficients suit ensembling; higher recipient weight can mimic intermediate-task transfer.
  • Merging formulation: Model merging seeks parameters that maximize the joint likelihood of posterior distributions from models sharing architecture and initialization.The merged model is intended to inherit capabilities from all input models.
  • Practical approximation: The diagonal Fisher approximation is used because storing the full Fisher requires O(|θ|^2) memory and is impractical for all but small models.Estimating the diagonal Fisher requires per-example gradients and has roughly the cost of training on the sampled examples.
  • Fisher merging: Fisher merging replaces isotropic precision with each model’s diagonal Fisher information, producing a parameter-wise weighted average.The Laplace approximation motivates using Fisher information as the Gaussian posterior precision.
  • Scope and safeguards: Fisher merging may underperform when models are far apart in parameter space, so the paper focuses on models trained from the same initialization.The Fisher is a local property of an individual parameter value.
  • Scope and safeguards: For parameters with near-zero Fisher values across models, the method defaults to a privileged target model; task-specific heads remain unchanged when parameters are unshared.These choices address numerical instability and incompatible parameter sets.

3 Experiments

Experiments evaluate Fisher merging for ensembling, robust fine-tuning, intermediate-task transfer, and domain adaptation. Across these settings, Fisher merging improves on isotropic merging and can match gradient-based transfer or output ensembling at substantially lower cost.

  • Ensembling: Fisher merging significantly outperforms isotropic merging for BERT ensembling while matching prediction ensembling with a single model.The evaluation uses five fine-tuned BERT checkpoints on RTE, MRPC, and SST-2.
  • Ensembling: 5× cheaper: Fisher merging performs comparably to output ensembling on validation accuracy for RTE, MRPC, and SST-2.The comparison concerns ensembles of five fine-tuned BERT models.
  • Robust Fine-Tuning: Fisher merging produces a better IID–OOD accuracy trade-off than isotropic merging in WiSE-FT across five OOD datasets.At the λ1 value giving the best average OOD accuracy, Fisher merging yields about 1% higher IID accuracy.
  • Intermediate-task training: Intermediate-task transfer usually benefits from merging, but Fisher merging remains worse than standard gradient-based training when MNLI is the donor task.Gradient-based intermediate-task training hurts on some datasets, whereas merging always helps in the reported RTE experiments.
  • Intermediate-task training: Fisher merging an intermediate-task-trained RTE model with donor-task models boosts performance over gradient-based intermediate-task training for all tested GLUE tasks.The additional boost when merging with an MNLI-trained model indicates a complementary transfer path.
  • Costs: 6× lower total cost: Fisher merging is substantially cheaper than fine-tuning an RTE checkpoint, while isotropic merging is 275× cheaper under the reported FLOP estimates.The Fisher estimate includes merged-checkpoint computation, validation evaluation, and one-time Fisher estimation; the Fisher matrix can be reused across merges.

4 Related Work

The paper situates merging alongside posterior-based continual-learning methods, output-based knowledge transfer, and checkpoint averaging. It distinguishes its cross-task use of merging from prior averaging settings.

  • Posterior-based continual learning: EWC uses a Laplace-approximated posterior to regularize continual learning and prevent catastrophic forgetting.Unlike merging, EWC restricts incorporated models to previous checkpoints from a continually trained model.
  • Output-based combination: Distillation and ensembling combine knowledge through model outputs, whereas merging directly combines parameters.The paper notes that parameter representations may contain more knowledge than model outputs.
  • Checkpoint averaging: Federated learning and Polyak averaging use isotropic checkpoint averaging, typically across checkpoints from one model’s training run.The paper presents cross-task transfer through checkpoint averaging as a distinct application.
  • Fisher-based methods: Natural gradient descent uses the Fisher information matrix as a metric or preconditioner, providing related motivation for Fisher-based model merging.Both approaches incorporate Fisher information, but the passage describes natural gradient descent as an optimization procedure.

5 Conclusion

The paper concludes that Fisher merging combines model capabilities through Fisher-weighted parameter averaging and can rival gradient-based transfer at lower cost. It also identifies broader combinations of models as a direction for future work.

  • Contribution: Fisher merging combines different models’ capabilities by computing a weighted average of their parameters.The method is motivated by viewing parameter averaging as maximizing the joint likelihood of model posteriors.
  • Results: Using Fisher information to weight parameters outperforms an unweighted average across the reported experiments.
  • Results: Fisher merging achieves comparable and sometimes better performance than gradient-based transfer-learning methods at significantly lower cost.
  • Implications: The experiments demonstrate merging strategies that would be onerous with traditional gradient-based training.The authors propose investigating alternative approximations of Fisher information and model posteriors.

A Checkpoints used for Ensembling

The ensembling experiments use fine-tuned BERT-base checkpoints from the Hugging Face model hub, spanning MRPC, SST-2, and RTE.

  • Checkpoint source: The experiments use fine-tuned BERT-base checkpoints obtained from the Hugging Face model hub.
  • RTE checkpoints: The RTE checkpoints include textattack, yoshitomo-matsubara, Ruizhou, howey, and anirudh21 model-hub entries.
  • MRPC checkpoints: The MRPC checkpoints include textattack, yoshitomo-matsubara, and Maelstrom77 model-hub entries.
  • SST-2 checkpoints: The SST-2 checkpoints include aviator-neural, howey, yoshitomo-matsubara, ikevin98, TehranNLP-org, and other BERT-base entries.

B Individual dataset results for robust fine-tuning

The paper reports individual robust fine-tuning results for five out-of-domain datasets evaluated with isotropic or Fisher merging.

  • Dataset-level results: Individual results are reported for WiSE-FT on five out-of-domain datasets using either isotropic or Fisher merging.

C GLUE Fine-tuning Details

Low-resource GLUE tasks were fine-tuned for 10 epochs under a fixed optimizer, batch size, learning rate, and five-run evaluation protocol.

  • Low-resource tasks: CoLA, MRPC, RTE, and STS-B were fine-tuned for 10 epochs using batch size 16.The optimizer was Adam with a learning rate of 1e-5.
  • Optimization: The low-resource-task fine-tuning used the Adam optimizer with a learning rate of 1e-5.
  • Evaluation: The study ran five independent fine-tuning runs for the low-resource tasks.

D Domain-Adaptive Pre-training Details

Domain-adaptive pre-training used RoBERTa-base on BIOMED and CS S2ORC data, then merged fine-tuning checkpoints with domain-specific DAPT checkpoints.

  • Domain-adaptive pre-training: RoBERTa-base received 32,768 additional domain-adaptive pre-training steps on the BIOMED and CS S2ORC splits.Training used batch size 32, Adam, and learning rate 1e-5.
  • OOD evaluation: Figure 7 reports individual OOD dataset results for WiSE-FT with isotropic or Fisher merging, using color to indicate increasing λ1 from 0 to 1.
  • Checkpoint merging: Each fine-tuning checkpoint was merged with the DAPT checkpoint from the task’s domain.The procedure saved checkpoints after each of 10 fine-tuning epochs.

E Full results for intermediate-task training

Appendix tables A1–A3 report intermediate-task training results across all possible GLUE datasets considered as target tasks.

  • GLUE coverage: Tables A1–A3 report results for intermediate-task training on GLUE.
  • Target-task evaluation: The appendix considers all possible datasets in GLUE as target tasks.
  • Experimental scope: The reported results evaluate intermediate-task training across GLUE target-task choices.

F Using fewer examples to estimate the Fisher

The appendix examines intermediate-task merging across GLUE, sequential fine-tuning, and how the number of Fisher examples affects performance.

  • Fisher estimation: Table A4 varies the examples used for MNLI and RTE Fisher computation and reports RTE validation accuracy.Original RTE checkpoints averaged 63.7 accuracy, while isotropic merging averaged 72.2.
  • Appendix comparisons: Tables A1–A3 compare intermediate-task Fisher merging, isotropic merging, and sequential fine-tuning on GLUE with BERT-base.Columns represent target tasks and rows represent intermediate tasks.
  • Result reporting: The appendix reports standard deviations across runs and marks unmerged or directly fine-tuned target-checkpoint scores in italics.
Loading 2111.09832v2…