Harmony in Diversity: Merging Neural Networks with Canonical Correlation Analysis

TL;DR

This paper introduces CCA Merge, a novel model fusion method based on Canonical Correlation Analysis, which improves the merging of neural networks by better aligning features, leading to enhanced performance over previous methods.

Contribution

The paper presents a new model merging algorithm using CCA that outperforms existing permutation-based methods, especially when merging multiple models trained on different data splits.

Findings

CCA Merge achieves higher accuracy than past methods.

The method effectively merges more than two models.

Performance improvements are consistent across different data splits.

Abstract

Combining the predictions of multiple trained models through ensembling is generally a good way to improve accuracy by leveraging the different learned features of the models, however it comes with high computational and storage costs. Model fusion, the act of merging multiple models into one by combining their parameters reduces these costs but doesn't work as well in practice. Indeed, neural network loss landscapes are high-dimensional and non-convex and the minima found through learning are typically separated by high loss barriers. Numerous recent works have been focused on finding permutations matching one network features to the features of a second one, lowering the loss barrier on the linear path between them in parameter space. However, permutations are restrictive since they assume a one-to-one mapping between the different models' neurons exists. We propose a new model merging algorithm, CCA Merge, which is based on Canonical Correlation Analysis and aims to maximize the correlations between linear combinations of the model features. We show that our alignment method leads to better performances than past methods when averaging models trained on the same, or differing data splits. We also extend this analysis into the harder setting where more than 2 models are merged, and we find that CCA Merge works significantly better than past methods. Our code is publicly available at https://github.com/shoroi/align-n-merge