Model Merging by Uncertainty-Based Gradient Matching

Nico Daheim,Thomas Möllenhoff,Edoardo Maria Ponti,Iryna Gurevych,Mohammad Emtiyaz Khan
2024-08-24
Abstract:Models trained on different datasets can be merged by a weighted-averaging of their parameters, but why does it work and when can it fail? Here, we connect the inaccuracy of weighted-averaging to mismatches in the gradients and propose a new uncertainty-based scheme to improve the performance by reducing the mismatch. The connection also reveals implicit assumptions in other schemes such as averaging, task arithmetic, and Fisher-weighted averaging. Our new method gives consistent improvements for large language models and vision transformers, both in terms of performance and robustness to hyperparameters. Code available here.
Machine Learning,Artificial Intelligence,Computation and Language
What problem does this paper attempt to address?