Summary
This paper presents a mean-field analysis on a class of overparameterized neural networks that are expressed as an ensemble of $N$ units, trained with SGD under symmetries in data distribution and possible use of symmetry-leveraging techniques (data augmentation, feature averaging, and equivariant architectures). While the analysis is based on prior work on mean-field analysis on the same class of neural networks e.g. based on Wasserstein gradient flows [1], the contribution of this work is in extending to learning under symmetries. The setup is as follows:
- In considered neural networks, a unit computes $\sigma_*:(\mathbf{x},\theta_i)\mapsto\hat{\mathbf{y}}$ and the network computes $\Phi_{\theta}^N:(\mathbf{x},\theta)\mapsto \frac{1}{N}\sum_{i=1}^N\sigma_*(\mathbf{x},\theta_i)$.
- Overparameterization means the width $N\to\infty$, yielding a limiting ensemble $\Phi_\mu$ where $\mu$ is understood a probability measure on the space of parameters $\theta_i$.
- The parameters $\theta$ are optimized using noisy SGD.
The mean-field theory of the neural networks considered above [1] aims to analyze the training dynamics of the limiting ensemble $\Phi_\mu$, given that its optimization using a convex loss function becomes a convex optimization problem (unlike $\Phi_{\theta}^N$ which induces non-convex optimization) and that its global optimum may be approximated by training $\Phi_{\theta}^N$ with a large width $N$. Specifically, it is known that SGD on $\Phi_{\theta}^N$ approximates Wasserstein gradient flow on $\Phi_\mu$ in the scaling limit $N\to\infty$ under certain assumptions, which then, under certain assumptions, converges to the global optima of the convex optimization problem.
Given the background, the paper presents a mean-field analysis given that the data generating distribution is possibly under a symmetry described by a compact group $G$, and symmetry-leveraging techniques are possibly used. The authors consider data augmentation, feature averaging (i.e. symmetrization; also called the Reynolds operator), and equivariant architectures, in accordance to literature [2]. The idea for the analysis is to distinguish between two types of symmetries for the measure $\mu$ on parameters: weakly-invariant and strongly-invariant, where weakly-invariant refers to an invariant measure on the subspace containing possibly non-invariant elements, and strongly-invariant refers to a measure on invariant subspace. The authors first show in Section 3.1 that, given that the unit $\sigma_*(\cdot, \cdot)$ is jointly equivariant, $\Phi_\mu$ is equivariant if and only if $\mu$ is weakly-invariant. Then, feature averaging (and data augmentation) can be understood as being related to weakly-invariant distributions obtained through symmetrization, and equivariant architecture is related to strongly-invariant distributions.
The authors then, in Section 3.2, begin mean-field analysis on optimality under a set of assumptions including that the unit $\sigma_*(\cdot, \cdot)$ is jointly equivariant (Assumption 1), and show that:
- (Proposition 3) Optimizing with respect to an invariant loss function yields a weakly-invariant optimum if it has a unique minimizer,
- (Corollary 1) If input data distribution is invariant, optimizing under data augmentation and feature averaging are equivalent and corresponds to approximating the symmetrized version of mean of the labeling function,
- (Corollary 2) If data distribution is equivariant, optimizing under data augmentation and feature averaging provide no advantage,
- (Proposition 4 and 5) But even when data distribution is equivariant, equivariant architectures (i.e. strongly invariant measures $\mu$) may not achieve the optimum with respect to a given invariant loss function, unless universality properties that are "good enough" for the given data distribution are assumed.
Then, in Section 3.3, the authors provide mean-field analysis on training dynamics, and show that:
- (Theorem 3, Corollary 3, and Theorem 4) Under an invariant loss function and Assumption 1 (among others), a training dynamics $\mu_{0:t}$ following the Wasserstein gradient flows yields a unique weakly-invariant solution $\mu_t$. If the data distribution is equivariant, from Corollary 2 this result applies to freely-trained neural networks without symmetry leveraging techniques, and data augmentation and feature augmentation again provide no advantages.
- (Theorem 5 and 6) Under an invariant loss function and Assumption 1 (among others), a training dynamics $\mu_{0:t}$ starting at a strongly-invariant initial condition $\mu_0$ and following the Wasserstein gradient flows yields strongly-invariant solutions $\mu_t$. Importantly, this means that if the data distribution is equivariant, $\mu_t$ stays strongly-invariant throughout training even though it is not explicitly forced or constrained to be invariant. Furthermore, equivariance of data distribution can be dropped if one of the symmetry leveraging techniques are used, which all yield coinciding solutions.
In Section 4, the authors provide synthetic experiments focusing on validating Theorem 3-6, and then propose and demonstrate an empirical method for finding good equivariant architecture for the task at hand (i.e. the most expressive among invariant subspace) by training according to Theorem 5 and 6 until convergence.
[1] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.
Strengths
- S1. The paper targets an important and challenging problem of understanding training dynamics of overparameterized neural networks under symmetries and their leveraging. The technical approach that applies mean-field analysis of generalized shallow neural networks [1] to learning under symmetries is original as far as I am aware.
- S2. The paper considers three types of symmetry-leveraging techniques: data augmentation, feature averaging, and equivariant architectures. These three are sufficiently general to represent currently used techniques [2]. The descriptions as well as strengths and weaknesses of each technique discussed in Section 2.3 are correct.
- S3. The authors' approach of imposing distributional symmetries on the measure $\mu$ to describe equivariant architectures as strongly-invariant (and feature averaging as weakly-invariant given that the unit $\sigma_*(\cdot, \cdot)$ is jointly equivariant) is interesting and could be potentially useful in future work.
- S4. The results presented in Section 3 and 4 are interesting, in particular Theorems 5 and 6 which shows that overparameterization and equivariant units, combined with equivariant data distribution or symmetry leveraging, leads to the property that strongly-invariant parameter initialization stays strongly-invariant throughout the training even without explicit constraints to do so. (But also see W1.)
[1] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.
Weaknesses
- W1. I have concerns regarding Assumption 1, which is a key assumption underlying Section 3 and assumes that the unit $\sigma_*:(\mathbf{x},\theta_i)\mapsto\hat{\mathbf{y}}$ is jointly equivariant with respect to the input $\mathbf{x}$, the parameters $\theta_i$, and the output $\hat{\mathbf{y}}$. This assumption states that we have already (partially) baked in the symmetry constraint to our model class. This may be natural for analyzing equivariant architectures, but does not coincide with the practical uses of data augmentation and feature averaging (e.g. please see [2]), as they are generally applied to an arbitrary, non-equivariant function, while in the setup of this work they are applied on already equivariant units. This limits the usefulness of the results regarding data augmentation and feature averaging. Perhaps the setup of this work could be relevant to symmetry breaking in already equivariant architectures [3], but I am not entirely sure.
- W2. In Line 92-93, the authors claim that "Note that requiring an infinite and i.i.d. sample from $\pi$ is key when letting later $N\to\infty$. This is not truly a restriction, since one can always sample from the empirical distribution of finite data points". This seems problematic since (1) if we can infinitely sample from $\pi$, then we have access to it, which violates basic assumptions in machine learning (Line 82), and (2) if we infinitely sample from an empirical distribution supported on finite set of data points, the resulting distribution would not be $\pi$, and consequently the SGD scheme in Equation (1) would not be optimizing the objective in Equation (2). (Please correct me if I am wrong.)
- W3. The writing of the paper could be in general improved. For example, the paper uses an excessive number of abbreviations (MF, MFL, NN, SL, DA, FA, EA, WI, SI, WGF, RMD, V), which is hurting readability. It would be also helpful to provide an example of a measure and its (non-equivalent) invariant symmetrization and projection in Definition 5; these are important concepts that are used in all later sections, but I was not able to directly grasp them from reading the text. For example, considering a uniform distribution on $[0, 1]\subset\mathbb{R}$, and a multiplicative group {1, -1} acting on $\mathbb{R}$ by multiplication, may be sufficiently informative.
- W4. The architecture discovery algorithm proposed in Section 4.2 was a bit questionable to me, since the architecture of the ensemble is already specified (at least up to a significant portion) by the jointly equivariant unit(s) $\sigma_*(\cdot,\cdot)$, and the algorithm searches for their parameters within a strongly-invariant regime. I am not sure why this algorithm can be understood as doing architecture discovery, instead of doing the usual parameter optimization.
- W5. Currently, it is not immediate how the study of the ensemble considered in the paper would benefit the development of practical neural networks that leverage symmetry. Showing how some of the currently used symmetry-leveraging deep neural networks can be viewed as such ensembles could be informative in this directions.
[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.
[3] Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, and Tess E. Smidt, Discovering symmetry breaking in physical systems with
relaxed group convolution, arXiv preprint arXiv:2310.02299, 2023.
Questions
Please see the Weaknesses section.
Limitations
The authors have partially addressed limitations of their work, for example in Section 5. Since the discussions are currently scattered across the main text, I encourage the authors to gather them in a dedicated limitations section.