PreprintCited by 1Domain Adaptation and Few-Shot LearningPrivacy-Preserving Technologies in DataAdvanced Neural Network Applications

Improving Model Fusion by Training-time Neuron Alignment with Fixed Neuron Anchors

Open access
Li, Zexi, Li, Zhiqi, Lin, Jie, Shen, Tao, Xiao, Jun, Guo, Yike, Lin, Tao, Wu, Chao·2024·arXiv

Abstract

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.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2402.01342

Cites

32 of 32 references verified0 citation errors · 0 not checkable

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • The random mask of gradients is superior to pruning because it does not harm network performance while regularizing optimization in a subspace.

    “It is superior to pruning because it does not harm the network performance while regularizing optimization in a subspace [44].”

    From WISE: Rethinking the Knowledge Memory for Lifelong Model Editing of Large Language Models (2024) · cites this paper as [44]

    The paper describes its gradient-mask method as lossless in terms of generalization compared with pruning (Hypothesis and Preliminary Finding, p. 5), and says the shared mask makes models learn in a parameter subspace (Training-time Neuron Alignment, p. 5). But the claim that it does not harm performance is too absolute: the ColD Fusion discussion (p. 9) says an overly large mask ratio can slightly reduce single-model accuracy. (LLM, conf=0.85)

    Experimental — accuracy not yet measured

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

Part of this paper’s reference list could not be split into single entries.

Flagged first, then by citations

Research works printed but not linked: 41

  1. [1][1] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever et al. , 'Language models are unsupervised multitask learners,' OpenAI blog , vol. 1, no. 8, p. 9, 2019. -
  2. [2][2] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell et al. , 'Language models are few-shot learners,' Advances in neural information processing systems , vol. 33, pp. 1877-1901, 2020. -
  3. [6][6] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, 'Highresolution image synthesis with latent diffusion models,' in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2022, pp. 10 684-10 695. -
  4. [12][12] M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S. Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith et al. , 'Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time,' in International Conference on Machine Learning . PMLR, 2022, pp. 23 965-23 998. -
  5. [13][13] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, 'Communication-efficient learning of deep networks from decentralized data,' in Artificial intelligence and statistics . PMLR, 2017, pp. 12731282. -
  6. [15][15] Z. Li, T. Lin, X. Shang, and C. Wu, 'Revisiting weighted aggregation in federated learning with neural networks,' in Proceedings of the 40th International Conference on Machine Learning , ser. Proceedings of Machine Learning Research, vol. 202. PMLR, 23-29 Jul 2023, pp. 19 767-19 788. -
  7. [16][16] J. Frankle, G. K. Dziugaite, D. Roy, and M. Carbin, 'Linear mode connectivity and the lottery ticket hypothesis,' in International Conference on Machine Learning . PMLR, 2020, pp. 3259-3269. -
  8. [17][17] F. Draxler, K. Veschgini, M. Salmhofer, and F. Hamprecht, 'Essentially no barriers in neural network energy landscape,' in International conference on machine learning . PMLR, 2018, pp. 1309-1318. -
  9. [18][18] R. Entezari, H. Sedghi, O. Saukh, and B. Neyshabur, 'The role of permutation invariance in linear mode connectivity of neural networks,' in International Conference on Learning Representations , 2022. -
  10. [19][19] S. Ainsworth, J. Hayase, and S. Srinivasa, 'Git re-basin: Merging models modulo permutation symmetries,' in The Eleventh International Conference on Learning Representations , 2022. -
  11. [21][21] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly et al. , 'An image is worth 16x16 words: Transformers for image recognition at scale,' in International Conference on Learning Representations , 2021. -
  12. [23][23] S. P. Singh and M. Jaggi, 'Model fusion via optimal transport,' Advances in Neural Information Processing Systems , vol. 33, pp. 22 045-22 055, 2020. -
  13. [29][29] D. A. E. Acar, Y. Zhao, R. Matas, M. Mattina, P. Whatmough, and V. Saligrama, 'Federated learning based on dynamic regularization,' in International Conference on Learning Representations , 2020. -
  14. [30][30] T. Garipov, P. Izmailov, D. Podoprikhin, D. P. Vetrov, and A. G. Wilson, 'Loss surfaces, mode connectivity, and fast ensembling of dnns,' Advances in neural information processing systems , vol. 31, 2018. -
  15. [31][31] G. Benton, W. Maddox, S. Lotfi, and A. G. G. Wilson, 'Loss surface simplexes for mode connecting volumes and fast ensembling,' in International Conference on Machine Learning . PMLR, 2021, pp. 769-779. -
  16. [32][32] G. Ilharco, M. T. Ribeiro, M. Wortsman, L. Schmidt, H. Hajishirzi, and A. Farhadi, 'Editing models with task arithmetic,' in The Eleventh International Conference on Learning Representations , 2023. -
  17. [34][34] T. J. Vlaar and J. Frankle, 'What can linear interpolation of neural network loss landscapes tell us?' in International Conference on Machine Learning . PMLR, 2022, pp. 22 325-22 341. -
  18. [36][36] P. Zhao, P.-Y. Chen, P. Das, K. N. Ramamurthy, and X. Lin, 'Bridging mode connectivity in loss landscapes and adversarial robustness,' in International Conference on Learning Representations (ICLR 2020) , 2020. -
  19. [37][37] E. S. Lubana, E. J. Bigelow, R. P. Dick, D. Krueger, and H. Tanaka, 'Mechanistic mode connectivity,' in International Conference on Machine Learning . PMLR, 2023, pp. 22 965-23 004. -
  20. [38][38] B. Simsek, F. Ged, A. Jacot, F. Spadaro, C. Hongler, W. Gerstner, and J. Brea, 'Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances,' in International Conference on Machine Learning . PMLR, 2021, pp. 9722-9732. -
  21. [40][40] M. Yurochkin, M. Agarwal, S. Ghosh, K. Greenewald, N. Hoang, and Y. Khazaeni, 'Bayesian nonparametric federated learning of neural networks,' in International conference on machine learning . PMLR, 2019, pp. 7252-7261. -
  22. [41][41] H. Wang, M. Yurochkin, Y. Sun, D. Papailiopoulos, and Y. Khazaeni, 'Federated learning with matched averaging,' in International Conference on Learning Representations , 2020. -
  23. [42][42] N. Tatro, P.-Y. Chen, P. Das, I. Melnyk, P. Sattigeri, and R. Lai, 'Optimizing mode connectivity via neuron alignment,' Advances in Neural Information Processing Systems , vol. 33, pp. 15 300-15 311, 2020. -
  24. [45][45] R. Ye, M. Xu, J. Wang, C. Xu, S. Chen, and Y. Wang, 'Feddisco: Federated learning with discrepancy-aware collaboration,' in International Conference on Machine Learning . PMLR, 2023, pp. 39 879-39 902. -
  25. [46][46] T. Lin, L. Kong, S. U. Stich, and M. Jaggi, 'Ensemble distillation for robust model fusion in federated learning,' Advances in Neural Information Processing Systems , vol. 33, pp. 2351-2363, 2020. -
  26. [49][49] M. Choi, H. Lee, G. Nam, and J. Lee, 'Sparse weight averaging with multiple parti-cles for iterative magnitude pruning,' ICLR 2024 . -
  27. [50][50] Y. Lin, S. Han, H. Mao, Y. Wang, and B. Dally, 'Deep gradient compression: Reducing the communication bandwidth for distributed training,' in International Conference on Learning Representations , 2018. -
  28. [51][51] T. Vogels, S. P. Karimireddy, and M. Jaggi, 'Powersgd: Practical lowrank gradient compression for distributed optimization,' Advances in Neural Information Processing Systems , vol. 32, 2019. -
  29. [52][52] Y. LeCun and C. Cortes, 'MNIST handwritten digit database,' arxiv , 2010. [Online]. Available: http://yann.lecun.com/exdb/mnist/ -
  30. [55][55] J. von Oswald, C. Henning, B. F. Grewe, and J. Sacramento, 'Continual learning with hypernetworks,' in International Conference on Learning Representations , 2019. -
  31. [56][56] Z. Liu, E. J. Michaud, and M. Tegmark, 'Omnigrok: Grokking beyond algorithmic data,' in The Eleventh International Conference on Learning Representations , 2022. -
  32. [58][58] A. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts, 'Learning word vectors for sentiment analysis,' in Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies , 2011, pp. 142-150. -
  33. [62][62] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever, 'Language models are unsupervised multitask learners,' OpenAI Blog , vol. 1, no. 8, 2019. -
  34. [63][63] S. Merity, C. Xiong, J. Bradbury, and R. Socher, 'Pointer sentinel mixture models,' in Proceedings of the 5th International Conference on Learning Representations (ICLR) , 2017. -
  35. [64][64] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, and I. Sutskever, 'Learning transferable visual models from natural language supervision,' in Proceedings of the 38th International Conference on Machine Learning , ser. Proceedings of Machine Learning Research, M. Meila and T. Zhang, Eds., vol. 139. PMLR, 18-24 Jul 2021, pp. 8748-8763. [Online]. Available: https://proceedings.mlr.press/v139/radford21a.html -
  36. [65][65] S. Han, J. Pool, J. Tran, and W. Dally, 'Learning both weights and connections for efficient neural network,' Advances in neural information processing systems , vol. 28, 2015. -
  37. [68][68] J. Wang, Q. Liu, H. Liang, G. Joshi, and H. V. Poor, 'Tackling the objective inconsistency problem in heterogeneous federated optimization,' Advances in neural information processing systems , vol. 33, pp. 76117623, 2020. -
  38. [69][69] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, 'Scaffold: Stochastic controlled averaging for federated learning,' in International Conference on Machine Learning . PMLR, 2020, pp. 5132-5143. -
  39. [71][71] H.-Y. Chen and W.-L. Chao, 'On bridging generic and personalized federated learning for image classification,' in International Conference on Learning Representations , 2022. -
  40. [73][73] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, 'Federated optimization in heterogeneous networks,' Proceedings of Machine Learning and Systems , vol. 2, pp. 429-450, 2020. -
  41. [75][75] H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein, 'Visualizing the loss landscape of neural nets,' Advances in neural information processing systems , vol. 31, 2018. -

Cited by

Sorted by citations

Authors
Li, Z., Li, Z., Lin, J., Shen, T., Xiao, J., Guo, Y., Lin, T., & Wu, C.
Affiliations
  • Georgia Institute of Technology (US)
  • Zhejiang University (CN)
  • Hong Kong University of Science and Technology (HK)
  • Westlake University (CN)
Venue
arXiv
Publisher
arXiv
Published
2024
Volume
48
Issue
3
Pages
2186-2199
ISSN
0162-8828 (ISSN-L)
OpenAlex
W4415445268
Cited by
0
References
16
FWCI
0.00
Type
journal article
Language
en
Open access
Locations
  • IEEE Transactions on Pattern Analysis and Machine Intelligence (publishedVersion)
  • PubMed (publishedVersion)
  • arXiv (Cornell University) (submittedVersion)
  • Rare & Special e-Zone (The Hong Kong University of Science and Technology) (submittedVersion)
  • arXiv (Cornell University)
Funders
  • Westlake University
  • National Natural Science Foundation of China (62441617)
  • National Science and Technology Major Project (2022ZD0115101)
  • Fundamental Research Funds for the Central Universities (226-2025-00057)
  • Westlake Education Foundation
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Domain Adaptation and Few-Shot Learning
  • Physical Sciences › Computer Science › Artificial Intelligence › Privacy-Preserving Technologies in Data
  • Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Advanced Neural Network Applications
Keywords
model fusion, permutation invariance, loss landscape, multi-model fusion, vision transformer, cold fusion, federated learning, heterogeneous federated learning, parameter averaging
MeSH terms
Deep Learning, Algorithms, Humans, Neurons, Neural Networks, Computer

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review