Proving Linear Mode Connectivity of Neural Networks via Optimal Transport

Damien Ferbach,Baptiste Goujaud,Gauthier Gidel,Aymeric Dieuleveut
2024-03-02
Abstract:The energy landscape of high-dimensional non-convex optimization problems is crucial to understanding the effectiveness of modern deep neural network architectures. Recent works have experimentally shown that two different solutions found after two runs of a stochastic training are often connected by very simple continuous paths (e.g., linear) modulo a permutation of the weights. In this paper, we provide a framework theoretically explaining this empirical observation. Based on convergence rates in Wasserstein distance of empirical measures, we show that, with high probability, two wide enough two-layer neural networks trained with stochastic gradient descent are linearly connected. Additionally, we express upper and lower bounds on the width of each layer of two deep neural networks with independent neuron weights to be linearly connected. Finally, we empirically demonstrate the validity of our approach by showing how the dimension of the support of the weight distribution of neurons, which dictates Wasserstein convergence rates is correlated with linear mode connectivity.
Machine Learning
What problem does this paper attempt to address?
### Problems the Paper Aims to Solve This paper aims to address the phenomenon of Linear Mode Connectivity (LMC) in neural networks and provide a theoretical foundation to explain this phenomenon. Specifically: 1. **Research Background**: - Training deep neural networks on complex tasks is a high-dimensional non-convex optimization problem. - Previous studies have found that different solutions obtained through methods like Stochastic Gradient Descent (SGD) can often be connected via simple continuous paths (e.g., linear paths), as long as the weights are appropriately permuted. 2. **Main Contributions**: - Theoretically proving that two sufficiently wide two-layer neural networks, after being trained with SGD, can be linearly connected if the weights are appropriately permuted. - Providing upper and lower bounds on the width of each layer in deep neural networks to achieve linear mode connectivity. - Experimentally validating the relationship between the support dimension of weight distribution and linear mode connectivity. 3. **Core Theory**: - Utilizing Optimal Transport (OT) theory to find the best match between two sets of trained neural networks. - Within the mean-field theory framework, proving the linear mode connectivity of two-layer neural networks after SGD training and providing boundary conditions for the width of hidden layers. - For multilayer perceptrons (MLP), proposing a general strategy to demonstrate their linear mode connectivity, considering the independence assumption of weight distributions. 4. **Experimental Validation**: - Validating the effectiveness of the proposed theoretical framework, demonstrating the relationship between the support dimension of weight distribution and linear mode connectivity, and proposing new weight matching methods. In summary, this paper aims to reveal the linear mode connectivity between different solutions in neural networks through theoretical analysis and experimental validation, providing practical guidance for real-world applications.