Symmetries in Overparametrized Neural Networks: A Mean-Field View (2405.19995v2)

Abstract: We develop a Mean-Field (MF) view of the learning dynamics of overparametrized Artificial Neural Networks (NN) under data symmetric in law wrt the action of a general compact group $G$. We consider for this a class of generalized shallow NNs given by an ensemble of $N$ multi-layer units, jointly trained using stochastic gradient descent (SGD) and possibly symmetry-leveraging (SL) techniques, such as Data Augmentation (DA), Feature Averaging (FA) or Equivariant Architectures (EA). We introduce the notions of weakly and strongly invariant laws (WI and SI) on the parameter space of each single unit, corresponding, respectively, to $G$-invariant distributions, and to distributions supported on parameters fixed by the group action (which encode EA). This allows us to define symmetric models compatible with taking $N\to\infty$ and give an interpretation of the asymptotic dynamics of DA, FA and EA in terms of Wasserstein Gradient Flows describing their MF limits. When activations respect the group action, we show that, for symmetric data, DA, FA and freely-trained models obey the exact same MF dynamic, which stays in the space of WI laws and minimizes therein the population risk. We also give a counterexample to the general attainability of an optimum over SI laws. Despite this, quite remarkably, we show that the set of SI laws is also preserved by the MF dynamics even when freely trained. This sharply contrasts the finite-$N$ setting, in which EAs are generally not preserved by unconstrained SGD. We illustrate the validity of our findings as $N$ gets larger in a teacher-student experimental setting, training a student NN to learn from a WI, SI or arbitrary teacher model through various SL schemes. We last deduce a data-driven heuristic to discover the largest subspace of parameters supporting SI distributions for a problem, that could be used for designing EA with minimal generalization error.

• The paper demonstrates that exploiting weakly invariant parameter distributions yields equivalent optimization outcomes through data augmentation and feature averaging.
• It reveals that mean-field dynamics preserve strong invariant properties during training in the large-N limit, even without explicit symmetry constraints.
• The study proposes heuristic design strategies for constructing equivariant architectures, aiming to enhance generalization and reduce overfitting.

Analysis of "Symmetries in Overparametrized Neural Networks: A Mean-Field View"

The paper, titled "Symmetries in Overparametrized Neural Networks: A Mean-Field View" by Javier Maass Martínez and Joaquín Fontbona, offers a rigorous exploration of how symmetries can be incorporated into the training dynamics of overparametrized neural networks (NNs) through a Mean-Field (MF) framework. The authors provide a comprehensive analysis of learning dynamics within the context of distributional symmetries established by a compact group action and explore several symmetry-leveraging techniques (SL). The exploration addresses the implications for population risk minimization and proposes potential heuristics for architectural design in machine learning.

The authors embark on this endeavor by introducing a generalized class of shallow NNs under a MF view, emphasizing the role of symmetries imposed by a group action. They classify parameter distribution symmetries into weakly and strongly invariant types and present an analysis of SL techniques, such as Data Augmentation (DA), Feature Averaging (FA), and Equivariant Architectures (EA). In particular, they offer a thorough investigation of how these methods relate to each other and the equilibrium solutions they lead to regarding population risk minimization.

Key Theoretical Contributions

1. Symmetry in Parameter Distributions: The paper distinguishes between weakly invariant (WI) and strongly invariant (SI) parameter distributions, which correspond to different levels of symmetry in model parameters. It argues that WI distributions play a crucial role in representing invariant shallow models, highlighting their significance in constructing equivariant functions.
2. Equivalence and Optimization: A central result is that DA and FA yield equivalent optimizations of the population risk exclusively over WI measures. Moreover, owing to the invariance of the risk function under symmetric data, these procedures are equivalent to freely trained models in terms of their MF dynamics. However, the paper provides a counterexample illustrating that optimization over SI measures can be restrictive.
3. Preservation of Symmetry in MF Training: One notable finding is that when initialized in a SI space, MF dynamics preserve this space throughout training, even without explicit constraints to do so. This phenomenon, observed in the large-N limit, starkly contrasts with the finite-N scenario where free training dynamics may not retain SI configurations.
4. Implications for Architecture Design: The paper speculates on a heuristic for discovering the largest subspace of parameters supporting SI distributions, guiding the construction of EAs with minimal generalization error. This heuristic is built on empirical observations that MF dynamics tend to remain within SI subspaces, suggesting a path to more efficient architecture search strategies.

Practical and Theoretical Implications

Practically, this paper's insights into symmetry utilization in NNs can lead to more efficient model training regimes that naturally respect the symmetrical nature of input data, potentially reducing overfitting and improving generalization. The techniques discussed also optimize computational resources by leveraging the inherent symmetrical properties of the data.

Theoretically, the paper extends the foundational understanding of how distributional symmetries impact learning in overparametrized regimes. By integrating symmetry considerations into MF theory, the authors open new avenues for research into the dynamics of neural networks, symmetry transformations, and their applications in real-world datasets structured on underlying group symmetries.

Moving forward, one might explore the scalability of these theoretical guarantees to deep networks with complex inner configurations and extend this understanding to broader classes of symmetries or non-compact groups. Additionally, exploring the convergence rates of different training schemes could offer valuable insights into the comparative performance of SL techniques.

In summary, this work significantly enhances the comprehension of neural network training dynamics through a symmetry-focused lens within the MF framework, proposing both theoretical advancements and practical methods that could refine current machine learning practices.