Training-time Neuron Alignment through Permutation Subspace for Improving Linear Mode Connectivity and Model Fusion

2024-02-02

arXiv.org

AI Summary

Deep neural network model fusion is hindered by linear mode connectivity (LMC) barriers arising from weight-space solution dispersion, primarily due to inconsistent neuron permutations across distinct training configurations. To address this, we propose Training-time Neuron Alignment (TNA), a mechanism that enhances fusion performance without increasing inference overhead. TNA introduces permutation subspaces into the training phase—enabling lossless neuron alignment for the first time—and instantiates the TNA-PFN algorithm, which we theoretically prove reduces LMC barriers and supports federated fusion under heterogeneous data. Experiments demonstrate substantial improvements: enhanced generalization of Vision Transformers (ViTs) in Model Soup, superior performance of large language models (LLMs) in ColD fusion, and high-accuracy, low-communication wide-model fusion in federated learning settings.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

In deep learning, stochastic gradient descent often yields functionally similar yet widely scattered solutions in the weight space even under the same initialization, causing barriers in the Linear Mode Connectivity (LMC) landscape. Overcoming these barriers is crucial for understanding deep learning dynamics and enhancing model-fusion algorithms. Previous studies highlight the role of permutation symmetry in reducing post-training barriers through network permutation. However, these post-hoc methods, demanding extra computations, are less effective for larger, complex models (e.g., ViT, LLM) due to numerous permutation matrices. Thus, in this paper, we study training-time neuron alignment. Our hypothesis suggests that training-time permutation subspace can reduce LMC barriers for free. We find that pruning at initialization supports this. Beyond pruning, we introduce TNA-PFN, a simple yet lossless algorithm using a partial gradient mask during training. TNA-PFN is theoretically and empirically validated for reducing LMC barriers. It excels in wide model fusion applications, especially in federated learning, two algorithms based on TNA-FPN that are proposed to show its prospects even under heterogeneous datasets. Moreover, TNA-PFN can enhance the generalization of model soup for vision transformers and ColD fusion for pretrained language models.

Problem

• Aligning neuron permutations during training to improve model fusion
• Reducing loss barriers between models for better parameter integration
• Enabling efficient federated learning through training-time neuron alignment

Innovation

• Training-time neuron alignment with fixed anchors
• Partially fixed neuron weights reduce permutations
• Improves model fusion without post-training matching