Saved in:
Bibliographic Details
Main Authors: Li, Zexi, Li, Zhiqi, Lin, Jie, Shen, Tao, Xiao, Jun, Guo, Yike, Lin, Tao, Wu, Chao
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2402.01342
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866918171560116224
author Li, Zexi
Li, Zhiqi
Lin, Jie
Shen, Tao
Xiao, Jun
Guo, Yike
Lin, Tao
Wu, Chao
author_facet Li, Zexi
Li, Zhiqi
Lin, Jie
Shen, Tao
Xiao, Jun
Guo, Yike
Lin, Tao
Wu, Chao
contents Model fusion aims to integrate several deep neural network (DNN) models' knowledge into one by fusing parameters, and it has promising applications, such as improving the generalization of foundation models and parameter averaging in federated learning. However, models under different settings (data, hyperparameter, etc.) have diverse neuron permutations; in other words, from the perspective of loss landscape, they reside in different loss basins, thus hindering model fusion performances. To alleviate this issue, previous studies highlighted the role of permutation invariance and have developed methods to find correct network permutations for neuron alignment after training. Orthogonal to previous attempts, this paper studies training-time neuron alignment, improving model fusion without the need for post-matching. Training-time alignment is cheaper than post-alignment and is applicable in various model fusion scenarios. Starting from fundamental hypotheses and theorems, a simple yet lossless algorithm called TNA-PFN is introduced. TNA-PFN utilizes partially fixed neuron weights as anchors to reduce the potential of training-time permutations, and it is empirically validated in reducing the barriers of linear mode connectivity and multi-model fusion. It is also validated that TNA-PFN can improve the fusion of pretrained models under the setting of model soup (vision transformers) and ColD fusion (pretrained language models). Based on TNA-PFN, two federated learning methods, FedPFN and FedPNU, are proposed, showing the prospects of training-time neuron alignment. FedPFN and FedPNU reach state-of-the-art performances in federated learning under heterogeneous settings and can be compatible with the server-side algorithm.
format Preprint
id arxiv_https___arxiv_org_abs_2402_01342
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Improving Model Fusion by Training-time Neuron Alignment with Fixed Neuron Anchors
Li, Zexi
Li, Zhiqi
Lin, Jie
Shen, Tao
Xiao, Jun
Guo, Yike
Lin, Tao
Wu, Chao
Machine Learning
Model fusion aims to integrate several deep neural network (DNN) models' knowledge into one by fusing parameters, and it has promising applications, such as improving the generalization of foundation models and parameter averaging in federated learning. However, models under different settings (data, hyperparameter, etc.) have diverse neuron permutations; in other words, from the perspective of loss landscape, they reside in different loss basins, thus hindering model fusion performances. To alleviate this issue, previous studies highlighted the role of permutation invariance and have developed methods to find correct network permutations for neuron alignment after training. Orthogonal to previous attempts, this paper studies training-time neuron alignment, improving model fusion without the need for post-matching. Training-time alignment is cheaper than post-alignment and is applicable in various model fusion scenarios. Starting from fundamental hypotheses and theorems, a simple yet lossless algorithm called TNA-PFN is introduced. TNA-PFN utilizes partially fixed neuron weights as anchors to reduce the potential of training-time permutations, and it is empirically validated in reducing the barriers of linear mode connectivity and multi-model fusion. It is also validated that TNA-PFN can improve the fusion of pretrained models under the setting of model soup (vision transformers) and ColD fusion (pretrained language models). Based on TNA-PFN, two federated learning methods, FedPFN and FedPNU, are proposed, showing the prospects of training-time neuron alignment. FedPFN and FedPNU reach state-of-the-art performances in federated learning under heterogeneous settings and can be compatible with the server-side algorithm.
title Improving Model Fusion by Training-time Neuron Alignment with Fixed Neuron Anchors
topic Machine Learning
url https://arxiv.org/abs/2402.01342