Quantifying water-energy nexus for urban water systems: A case study of Addis Ababa city

1.   Introduction

The ability of neural networks to learn rich representations—or features—of their input data is commonly observed in state-of-the art models ^[1,2] and often thought to be the reason behind their good practical performance [3,Chap. 1]. Yet, our theoretical understanding of how feature learning arises from simple gradient-based training algorithms remains limited. Much progress (discussed in Section 1.3) has been made recently to understand the power and limitations of gradient-based learning with neural networks, showing in particular their superiority over fixed-feature methods on some difficult tasks. However, positive results are often obtained for algorithms that differ in substantial ways from plain (stochastic) gradient descent (e.g., the layers trained separately, or the algorithm makes just one truly non-linear step, etc.).

In this work, we take the algorithm as a given and instead adopt a descriptive approach. Our goal is to improve our understanding of how neural networks behave in the presence of symmetries in the data with plain gradient descent (GD) on two-layer fully-connected ReLU neural networks. To this end, we investigate situations with strong symmetries on the data, the target function and on the initial parameters, and study the properties of the training dynamics and the learned predictor in this context.

1.1.   Problem setting

We denote by $d$ the input dimension, $\rho$ the input data distribution which we assume to be uniform over the unit sphere $\mathbb{S}^{d-1}$ of $\mathbb{R}^d$, and by $\mathcal{P}_2(\Omega)$ the space of probability measures with finite second moments over a measurable space $\Omega$. We call $\sigma$ the activation function, which we take to be $\text{ReLU}$, that is $\sigma(z) = \max(0, z)$, $\ell : \mathbb{R} \times \mathbb{R} \to \mathbb{R}$ the loss function, which we assume to be continuous in both arguments and continuously differentiable w.r.t. its second argument and we denote by $\partial_2 \ell$ this derivative.

Mean-field limit of two-layer networks. In this work, we consider the infinite-width limit in the mean-field regime of the training dynamics of two-layer networks without intercept with a $\text{ReLU}$ activation function. Given a measure $\mu \in \mathcal{P}_2(\mathbb{R} \times \mathbb{R}^d)$, we consider the infinitely wide two-layer network parameterized by $\mu$, defined, for any input $x \in \mathbb{R}^d$, by

where, for any $c = (a, b) \in \mathbb{R} \times \mathbb{R}^d$, $\phi(c; x) = a \sigma \left(b^\top x \right)$. Note that width-$m$ two-layer networks with input weights $(b_j)_{j \in [1, m]} \in (\mathbb{R}^d)^{m}$ and output weights $(a_j)_{j \in [1, m]} \in \mathbb{R}^m$ can be recovered by a measure $\mu_m = (1/m) \sum_{j = 1}^m \delta_{(m a_j, b_j)}$ with $m$ atoms.

Objective and Wasserstein gradient flow. We consider the problem of minimizing the population loss objective for a given target function $f^*:\mathbb{R}^d \rightarrow \mathbb{R}$, which we assume to be bounded on the unit sphere, that is

The Fréchet derivative of the objective function $F$ at $\mu$ is given by the function $F^{\prime}_{\mu}(c) = \mathbb{E}_{x \sim \rho} \left[\partial_2 \ell \left(f^*(x), f(\mu; x) \right) \phi(c; x) \right]$ for any $c = (a, b) \in \mathbb{R} \times \mathbb{R}^d$ (for more details, see Appendix B.1). Starting from a given measure $\mu_0 \in \mathcal{P}_2(\mathbb{R} \times \mathbb{R}^d)$, we study the Wasserstein gradient flow (GF) of the objective (1.2) which is a path $(\mu_t)_{t \geq 0}$ in the space of probability measures satisfying, in the sense of distributions, the partial differential equation (PDE) known as the continuity equation:

Initialization. We make the following assumption on the initial measure $\mu_0 \in \mathcal{P}_2(\mathbb{R} \times \mathbb{R}^d)$: $\mu_0$ decomposes as $\mu_0 = \mu_0^1 \otimes \mu_0^2$ where $\mu_0^1, \mu_0^2 \in \mathcal{P}_2(\mathbb{R}) \times \mathcal{P}_2(\mathbb{R}^d)$. This follows the standard initialization procedure at finite width. Because no direction should a priori be favored, we assume $\mu_0^2$ to have spherical symmetry, i.e., , it is invariant under any orthogonal transformation, and we additionally assume that $|a| = ||b||$ almost surely at initialization. It is shown in [4,Lemma 26], and [5,Section 2.5], that with this assumption, $\mu_t$ stays supported on the set $\{|a| = ||b||\}$ for any $t \geq 0$.

Comment on the assumptions. The assumption that $\mu_0$ decomposes as a product of two measures is to stay as close as possible to what is done in practice (independent initialization for different layers). The assumption that $|a| = ||b||$ is of a technical nature, and, along with the regularity conditions on the loss $\ell$ and the input data distribution $\rho$, ensures that the Wasserstein GF (1.3) is well-defined [5,Lemmas 3.1 and 3.9] when using ReLU as a activation function (which bears technical difficulties because of its non-smoothness). The results of Section 2 hold for others activation function which potentially require less restrictive assumptions on $\mu_0$ and $\rho$ but still requires $\mu_0$ to decompose as a product of measures. In contrast, the results of Sections 3 and 4 are specific to $\sigma = \text{ReLU}$ and thus require the assumptions above on $\mu_0$ and $\rho$. Since our work focuses mostly on ReLU, we choose to state the results of all sections with the (more restrictive) assumptions stated above on $\mu_0$ and $\rho$.

Relationship with finite-width GD. If $\mu_0 = (1/m) \sum_{j = 1}^m \delta_{(a_j(0), b_j(0))}$ is discrete, the Wasserstein GF (1.3) is exactly continuous-time GD on the parameters of a standard finite-width neural network, and discretization errors (w.r.t. the number of neurons) can be provided ^[6,7].

1.2.   Summary of contributions

Our main object of study is the gradient flow of the population risk of infinitely wide two-layer ReLU neural networks without intercept. Our motivation to consider this idealistic setting—infinite data and infinite width—is that it allows, under suitable choices for $\rho$ and $\mu_0$, the emergence of exact symmetries which are only approximate in the non-asymptotic setting^*.

^*In contrast, our focus on GF is only for theoretical convenience and most of our results could be adapted to the case of GD.

Symmetries, structure, and convergence. In this work, we are interested in the structures learned by the predictor $f(\mu_t; \cdot)$ under GF as $t$ grows large. Specifically, we make the following contributions:

● In Section 2, we prove that if $f^*$ is invariant under some orthogonal linear map $T$, then $f(\mu_t; \cdot)$ inherits this invariance under GF (Theorem 2.1).

● In Section 3, we study the case when $f^*$ is an odd function and show that the network converges to the best linear approximator of $f^*$ at an exponential rate (Theorem 3.2). Linear predictors are optimal over the hypothesis class in that case, in particular because there is no intercept in our model.

● In Section 4, we consider the multi-index model where $f^*$ depends only on the orthogonal projection of its input onto some sub-space $H$ of dimension $d_H$. We prove that the dynamics can be reduced to a PDE in dimension $d_H$. If in addition, $f^*$ is the Euclidean norm of the projection of the input, we show that the dynamics reduce to a one-dimensional PDE (Theorem 4.3). In the latter case, we were not able to prove theoretically the convergence of the neurons of the first layer towards $H$, and leave this as an open problem but we provide numerical evidence in favor of this result.

The code to reproduce the results of the numerical experiments can be found at: https://github.com/karl-hajjar/learning-structure.

1.3.   Related work

Infinite-width dynamics. It has been shown rigourously that for infinitely wide networks there is a clear distinction between a feature-learning regime and a kernel regime ^[8,9]. For shallow networks, this difference stems from a different scale (w.r.t. width) of the initialization where a large initialization leads to the Neural Tangent Kernel (NTK) (a.k.a. the "lazy regime") which is equivalent to a kernel method with random features ^[10] whereas a small initialization leads to the so-called mean-field (MF) limit where features are learned from the first layer ^[8,9]. However, it is unclear in this setting exactly what those features are and what underlying structures are learned by the network. The aim of the present work is to study this phenomenon from a theoretical perspective for infinitely wide networks and to understand the relationship between the ability of networks to learn specific structures and the symmetries of a given task.

A flurry of works study the dynamics of infinitely wide two-layer neural networks. ^{[5,6,11,12,13]} study the gradient flow dynamics of the MF limit and show that they are well-defined in general settings and lead to convergence results (local or global depending on the assumptions). On the other hand, Jacot et al. ^[10] study the dynamic of the NTK parameterization in the infinite-width limit and show that it amounts to learning a linear predictor on top of random features (fixed kernel), so that there is no feature learning.

Convergence rates. In the MF limit, convergence rates are in general difficult to obtain in a standard setting. For instance, ^[5,11] show the convergence of the GF to a global optimum in a general setting but this does not allow convergence rates to be provided. To illustrate the convergence of the parameterizing measure to a global optimum in the MF limit, Ma et al. ^[14] prove local convergence (see Section 7) for one-dimensional inputs and a specific choice of target function in $O(t^{-1})$ where $t$ is the time step. At finite-width, Daneshmand and Bach ^[15] also prove convergence of the parameters to a global optimum in $O(t^{-1})$ using an algebraic idea which is specific to the ad-hoc structure they consider (inputs in two dimensions and target functions with finite number of atoms).

In Section 3, we show convergence of the MF limit at an exponential rate when the target function is odd. In the setting of this section, the training dynamics are degenerate and although input neurons move, the symmetries of the problem imply that the predictor is linear.

Low-dimensional structure. Studying how neural networks can adapt to hidden low-dimensional structures is a way of approaching theoretically the feature-learning abilities of neural networks. Bach ^[16] studies the statistical properties of infinitely wide two-layer networks, and shows that when the target function only depends on the projection on a low-dimensional sub-space, these networks circumvent the curse of dimensionality with generalization bounds which only depend on the dimension of the sub-space. In a slightly different context, Chizat and Bach ^[4] show that for a binary classification task, when there is a low-dimensional sub-space for which the projection of the data has sufficiently large inter-class distance, only the dimension of the sub-space (and not that of the ambient space) appears in the upper bound on the probability of misclassification. Whether or not such a low-dimensional sub-space is actually learned by GD is not addressed in these works.

Similarly, ^[17,18] focus on learning functions which have a hidden low-dimensional structure with neural networks. They consider a single step of GD on the input layer weights and show that the approximation / generalization error adapts to the structure of the problem: they provide bounds on the number of data points / parameters needed to achieve negligible error, which depend on the reduced dimension and not the dimension of the ambient space. In a similar context, Mousavi-Hosseini et al. ^[19] consider (S)GD on the first layer only of a finite-width two-layer network and show that with sufficient $L_2$-regularization and with a standard normal distribution on the input data the first layer weights align with the lower-dimensional sub-space when trained for long enough. They then use this property to then provide statistical results on networks trained with SGD.

In a setting close to ours but on a classification task with finite-data and at finite-width, Paccolat et al. ^[20] compare the feature learning regime with the NTK regime in the presence of hidden low-dimensional structure and quantify for each regime the scaling law of the test error w.r.t. the number of training samples, mostly focusing on the case $d_H = 1$.

In a similar setting to that of ^[16], Abbe et al. ^[21] study how GF for infinitely wide two-layer networks can learn specific classes of functions which have a hidden low-dimensional structure when the inputs are Rademacher variables. This strong symmetry assumption ensures that the learned predictor shares the same low-dimensional structure at any time step (from the $t = 0$) and this allows them to characterize precisely what classes of target functions can or cannot be learned by GF in this setting. In contrast, we are interested in how infinitely wide networks learn those low-dimensional structures during training, and in the role of symmetries in enabling such a behaviour after initialization.

Learning representations. An existing line of work ^{[18,22,23,24,25]} studies in depth the representations learned by neural networks trained with (S)GD at finite-width from a different perspective focusing on the advantages of feature-learning in terms of performance comparatively to using random features. In contrast, our aim is to describe the representations themselves in relationship with the symmetries of the problem.

Symmetries. We stress that the line of work around symmetries of neural networks dealing with finding network architectures for which the output is invariant (w.r.t. to its input or parameters) by some group of transformations (see ^[26,27,28], and references therein) is entirely different from what we are concerned with in the present work. In contrast, the setting of Mei et al. ^[6] is much closer to ours as they study how the invariances of the target function / input data can lead to simplifications in the dynamics of infinitely wide two-layer networks in the mean-field regime which allows them to prove global convergence results.

1.4.   Notations

We denote by $\mathcal{M}_+(\Omega)$ the space of non-negative measures over a measurable space $\Omega$. For any measure $\mu$ and measurable map $T$, $T_{\#} \mu$ denotes the pushforward measure of $\mu$ by $T$. We denote by $\mathcal{O}(p)$ and $\text{id}_{\mathbb{R}^p}$ respectively the orthogonal group and the identity map of $\mathbb{R}^p$ for any $p \in \mathbb{N}$. Finally, $\left < {\cdot}, \, {\cdot} \right >$ is the Euclidean inner product and $||\, \cdot \, ||$ the corresponding norm.

2.   Invariance under orthogonal symmetries

In this section, we demonstrate that if the target function $f^*$ is invariant under some orthogonal transformation $T$, since the input data distribution is also invariant under $T$, then $f(\mu_t; \cdot)$ is invariant under $T$ as well for any $t \geq 0$. This invariance property of the dynamics w.r.t. orthogonal symmetries is possible with an infinite number of neurons but is only approximate at finite-width. It is noteworthy that the results of this section hold for any activation function $\sigma$ and input data distribution $\rho$ which has the same symmetries as $f^*$, provided that the Wasserstein GF (1.3) is unique. We start with a couple of definitions:

Definition 2.1 (Function invariance). Let $T$ be a map from $\mathbb{R}^d$ to $\mathbb{R}^d$, and $f: \mathbb{R}^d \rightarrow \mathbb{R}$. Then, $f$ is said to be invariant (resp. anti-invariant) under $T$ if for any $x \in \mathbb{R}^d$, $f(T(x)) = f(x)$ (resp. $f(T(x)) = -f(x)$).

Definition 2.2 (Measure invariance). Let $\Omega \subset \mathbb{R}^d$, $T$ be a measurable map from $\Omega$ to $\Omega$, and $\mu$ be a measure on $\Omega$. Then, $\mu$ is said to be invariant under $T$ if $T_{\#} \mu = \mu$, or equivalently, if for any continuous and compactly supported $\varphi: \Omega \to \mathbb{R}$, $\int \varphi(x) \mathrm{d} \mu(x) = \int \varphi(T(x)) \mathrm{d} \mu(x)$.

We are now ready to state the two main results of this section.

Proposition 2.1 (Learning invariance). Let $T \in \mathcal{O}(d)$, and assume that $f^*$ is invariant under $T$. Then, for any $t \geq 0$, the Wasserstein GF $\mu_t$ of Eq (1.3) is invariant under $\tilde{T}: (a, b) \in \mathbb{R} \times \mathbb{R}^d \mapsto (a, T(b))$, and the corresponding predictor $f(\mu_t; \cdot)$ is invariant under $T$.

Proposition 2.2 (Learning anti-invariance). Under the same assumptions as in Proposition 2.1 except now we assume $f^*$ is anti-invariant under $T$, and assuming further that $\partial_2 \ell(-y, -\hat{y}) = -\partial_2 \ell(y, \hat{y})$ for any $y, \hat{y} \in \mathbb{R}$, and that $\mu_0^1$ is symmetric around $0$ (i.e., , invariant under $:a \in \mathbb{R} \mapsto - a$), we then have that for any $t \geq 0$, the Wasserstein GF $\mu_t$ in Eq (1.3) is invariant under $\tilde{T}: (a, b) \in \mathbb{R} \times \mathbb{R}^d \mapsto (-a, T(b))$, and the corresponding predictor $f(\mu_t; \cdot)$ is anti-invariant under $T$.

Remark. The results above also hold for networks with intercepts at both layers. The conditions of Proposition 2.2 are satisfied by both the squared loss and the logistic loss (a.k.a. the cross-entropy loss).

Essentially, those results show that training with GF preserves the orthogonal symmetries of the problem: the invariance of the target function under an orthogonal transformation leads to the same invariance for $\mu_t$ and $f(\mu_t; \cdot)$. The proof, presented in Appendix C, relies crucially on the fact that $T$ is an orthogonal map which combines well with the structure of $\phi(c; x)$ involving an inner product. The idea is essentially that the orthogonality of $T$ allows us to relate the gradient of $\phi$ (and consequently of $F^\prime_{\mu_t}$) w.r.t. $c$ at $(T(c); x)$ to the same gradient at $(c; T^{-1}(x))$ and then to use the invariance of $f^*$ and $\rho$ to conclude.

In the following sections we discuss the particular cases where functions are (anti-)invariant under $-\text{id}_{\mathbb{R}^d}$ (i.e., , even or odd functions) or some sub-group of $\mathcal{O}(d)$.

3.   Exponential convergence for odd target functions

We consider here an odd target, function, i.e., , for any $x \in \mathbb{R}^d$, $f^*(-x) = -f^*(x)$.

Linearity of odd predictors. Proposition 2.2 ensures that the predictor $f(\mu_t; \cdot)$ associated with the Wasserstein GF of Eq (1.3) is also odd at any time $t \geq 0$, and we can thus write, for any $x$, $f(\mu_t; x) = \frac{1}{2} \left(f(\mu_t; x) - f(\mu_t; -x) \right)$, which yields

where the last equality stems from the fact that for ReLU, $\sigma(x) - \sigma(-x) = x$. Put differently, the predictor is linear: it is the same as replacing $\sigma$ by $\frac{1}{2} \text{id}_{\mathbb{R}^d}$, and $f(\mu_t; x) = {w(t)}^\top x$, where

This degeneracy is not surprising as in fact, a linear predictor is the best one can hope for in this setting. Indeed, consider the following assumption and the next lemma:

Assumption 1 (Squared loss function). The loss function $\ell$ is the squared loss, i.e., , $\ell(y, \hat{y}) = \frac{1}{2} (y - \hat{y})^2$, and thus satisfies the condition of Proposition 2.2.

We make this assumption in order to provide an explicit convergence rate in Theorem 3.2 below.

Lemma 3.1 (Optimality of odd predictors). Let $f$ be a predictor in the hypothesis class $\mathcal{F} : = \left\{ :x \mapsto \int a\sigma(b^\top x) \mathrm{d} \mu(a, b); \mu \in \mathcal{P}_2(\mathbb{R} \ \times \mathbb{R}^d) \right)\}$. Then, denoting $f_{\mathit{\text{odd}}}(x) : = \frac{1}{2}(f(x) - f(-x))$ (resp. $f_{\mathit{\text{even}}} : = \frac{1}{2}(f(x) + f(-x))$) the odd (resp. even) part of $f$, one has:

Proof. The result readily follows from the decomposition $f = f_{\text{odd}} + f_{\text{even}}$ which leads to

We then get that $L(f) \geq L(f_{\text{odd}})$ with equality if and only if $\mathbb{E}_{x \sim \rho} \left[\left(f_{\text{even}}(x)\right)^2 \right] = 0$, i.e., , $f_{\text{even}}(x) = 0$ for $\rho$-almost every $x$. Finally, if $\mu \in \mathcal{P}_2(\mathbb{R}^{d+1})$, then $\nu : = \frac{1}{2}(\mu + S_{\#}\mu) \in \mathcal{P}_2(\mathbb{R}^{d+1})$, where $S: (a, b) \in \mathbb{R}^{d+1} \mapsto (-a, -b)$, and $f(\nu; \cdot) = f_{\text{odd}}(\mu; \cdot)$, which shows $f_{\text{odd}}(\mu; \cdot) \in \mathcal{F}$.

Since, as shown above, any odd predictor turns out to be linear because of the symmetries of ReLU, in this context, the best one can expect is thus to learn the best linear predictor.

Exponential convergence for linear networks. We are thus reduced to studying the dynamics of linear networks (which in our case are infinitely wide), which is an interesting object of study its own right (Ji and Telegarsky ^[29] show a result similar to our result below in the finite-width case with the logistic loss on a binary classification task). In this case, the Wasserstein GF (1.3) (with ReLU replaced by $\frac{1}{2} \text{id}_{\mathbb{R}^d}$) is defined for more general input distributions $\mathbb{P} \in \mathcal{P}_2(\mathbb{R}^d)$ (e.g., empirical measures) and target functions $f^*$. The objective in this context is thus to learn:

with the dynamics of linear infinitely wide two-layer networks described by the Wasserstein GF (1.3) where the activation function $\sigma$ is replaced by $\frac{1}{2} \text{id}_{\mathbb{R}^d}$. Theorem 3.2 below shows exponential convergence to a global minimum of $Q$ as soon as the problem is strongly convex. Note that although in this case both $\phi(\cdot; \cdot)$ (see Eq (1.1)) and the predictor in the objective $Q$ are linear w.r.t. the input, only the predictor in $Q$ is linear in the parameters (ordinary least squares).

Theorem 3.2. Assume that the smallest eigenvalue $\lambda_{\mathit{\text{min}}}$ of $\mathbb{E}_{x \sim \mathbb{P}}[x x^\top]$ is positive. Let $(\mu_t)_{t \geq 0}$ be the Wasserstein GF associated to (1.3) with activation function $\frac{1}{2} {\rm{id}}_{\mathbb{R}^d}$ instead of $\sigma = {\rm{{ ReLU}}}$, and call $w(t) = \frac{1}{2} \int a b \, \mathrm{d} \mu_t(a, b) \in \mathbb{R}^d$. Then, there exits $\eta > 0$ and $t_0 > 0$ such that, for any $t \geq t_0$,

Remark. Note that as soon as $\mathbb{P}$ has spherical symmetry, the problem becomes strongly convex by Lemma B.3. Note that although $F(\mu_t) = Q(w(t))$, $(w(t))_{t \geq 0}$ is not a gradient flow for the (strongly) convex objective $Q$ (which would immediately guarantee exponential convergence to the global minimum).

The proof, provided in Appendix E, proceeds in two steps: first it is shown that $w^\prime(t) = - H(t) \nabla Q(w(t))$ for some positive definite matrix $H(t)$ whose smallest eigenvalue is always lower-bounded by a positive quantity, then we prove that this leads to exponential convergence. Figure 1 illustrates that the dynamics of GF on $F$ remain non-linear in that they do not reduce to GF on $Q$ (although the paths are close). To simulate GF on $F$ we use a large (but finite) number of neurons $m = 1024$ and a small (but positive) step-size $10^{-2}$ and simply proceed to do GD on the corresponding finite-dimensional objective (see comment in Section 1.1 on relationship between the Wasserstein GF and finite-width GD). The dashed curves are the trajectories of $w(t)$ (see Theorem 3.2) for $10$ random initializations, which corresponds to GD on $F$, whereas the blue curve is the trajectory of GD on $Q$ starting from $0$. $\mathbb{P} = \frac{1}{n} \sum_{i = 1}^n \delta_{x_i}$ is the empirical distribution over $n = 100$ samples, with $x_i \sim \mathcal{U}([-1, 1]^d)$ and $y_i = f^*(x_i)$ drawn i.i.d. from $\mathcal{N}(y^*, 2)$ where $y^* \sim \mathcal{N}(0, 1)$.

4.   Learning the low-dimensional structure of the problem

Consider a linear sub-space $H$ of dimension $d_H < d$ (potentially much smaller than the ambient dimension), and assume $f^*$ has the following structure: $f^*(x) = f_H(p_H(x))$ where $p_H$ is the orthogonal projection onto $H$ (which we also write $x^H$ for simplicity, and we reserve sub-scripts for denoting entries of vectors) and $f_H:H \rightarrow \mathbb{R}$ is a given function.

In this context it is natural to study whether the learned function shares the same structure as $f^*$. As observed in Figure 2 this is not the case in finite time, but it is reasonable however to think that the learned predictor $f(\mu_t; \cdot)$ shares the same structure as $f^*$ as $t \to \infty$, and we give numerical evidence in this direction. On the other hand, we prove rigorously that the structure of the problem allows to reduce the dynamics to a lower-dimensional PDE. In this section, we consider for simplicity that $\mu_0^1$ is the uniform distribution over $\{-1, +1\}$ and that $\mu_0^2$ is the uniform distribution over $\mathbb{S}^{d-1}$.

Comment on the assumptions for this section. The assumptions that $|a| = ||b||$ on the support of $\mu_0$ is crucial here. This ensures that the Wasserstein GF (1.3) is well-defined and that $\mu_t$ stays supported on the set $\{|a| = ||b||\}$ for any $t \geq 0$, a fact which is used in the proofs. The assumption that $\rho$ is the uniform distribution over the unit sphere bears some importance but could likely be replaced by other measures with spherical symmetry provided that the dynamics would still be well-defined and at the cost of more technical proofs.

4.1.   Symmetries and invariance

The structure of $f^*$ implies that it is invariant by any $T \in \mathcal{O}(d)$ which preserves $H$, i.e., , such that its restrictions to $H$ and $H^\perp$ are $T_{|H} = \text{id}_H$ and $T_{|H^\perp} \in \mathcal{O}(d_\perp)$, where $\mathcal{O}(d_\perp)$ is the orthogonal group of $H^\perp$ whose dimension is $d_\perp = d - d_H$. By Proposition 2.1, such transformations also leave the predictor $f(\mu_t; \cdot)$ invariant for any $t \geq 0$ since $\rho$ is spherically symmetric. Lemma 4.1 below then ensures that $f(\mu_t; x)$ depends on the projection $x^\perp$ onto $H^\perp$ only through its norm, that is $f(\mu_t; x) = \tilde{f}_t(x^H, ||x^\perp||)$ for some $\tilde{f}_t : H \times \mathbb{R}_+ \rightarrow \mathbb{R}$.

Lemma 4.1 (Invariance by a sub-group of $\mathcal{O}(d)$). Let $f:\mathbb{R}^d \rightarrow \mathbb{R}$ be invariant under any $T \in \mathcal{O}(d)$ such that $T_{|H} = {\rm{id}}_H$ and $T_{|H^\perp} \in \mathcal{O}(d_\perp)$. Then, there exists some $\tilde{f}: H \times \mathbb{R}_+ \rightarrow \mathbb{R}$ such that for any $x \in \mathbb{R}^d$, $f(x) = \tilde{f}(x^H, ||x^\perp||)$.

Proof. Consider $\tilde{f}:(x^H, r) \in H \times \mathbb{R}_+ \mapsto f(x^H + r e^\perp_1)$ where $e^\perp_1$ is the first vector of an orthonormal basis of $H^\perp$, and let $x \in \mathbb{R}^d$. If $x^\perp = 0$, the result is obvious. Otherwise, consider an orthogonal linear map $T_x$ such that ${T_x}_{|H} = {\text{id}}_H$ and $T_x$ sends $x^\perp / ||x^\perp||$ on $e^{\perp}_{1}$. The invariance of $f$ under $T_x$ implies $f(x) = f(T_x(x)) = f(x^H + ||x^\perp|| e^\perp_1) = \tilde{f}(x^H, ||x^\perp||)$.

Figure 2 shows that the dependence in $||x^\perp||$ cannot be removed in finite time: $f(\mu_t; u_H + r e^{\perp}_{1})$ does depend on the distance $r \in \mathbb{R}_+$ to $H$, but this dependence tends to vanish as $t \to \infty$. The plots of Figure 2 are obtained by discretizing the initial measure $\mu_{0, m} = \frac{1}{m} \sum_{j = 1}^m \delta_{(a_j(0), b_j(0))}$ with $m = 1024$ atoms, and sampling $a_j(0) \sim \mathcal{U}(\{-1, +1\})$ and $b_j(0) \sim \mathcal{U}(\mathbb{S}^{d-1})$. We perform GD with a finite step-size $\eta =$ and a finite number $n = 256$ of fresh i.i.d. samples from the data distribution per step with $f^*(x) = ||x^H||$, $d = 20$ and $d_H = 5$.

Dynamics over the sphere $\mathbb{S}^{d-1}$. Using the positive $1$-homogeneity of ReLU, and with the assumptions on $\mu_0$, the dynamics on $\mu_t \in \mathcal{P}_2(\mathbb{R}^{d+1})$ can be reduced to dynamics on the space $\mathcal{M}_+(\mathbb{S}^{d-1})$ of non-negative measures over $\mathbb{S}^{d-1}$: only the direction of neurons matter and their norm only affects the total mass. From this point of view, neurons with positive and negative output weights behave differently and have separate dynamics. Indeed, consider the pair of measures $(\nu_t^+, \nu_t^-) \in \mathcal{M}_+(\mathbb{S}^{d-1})^2$ characterized by the property that for any continuous test function $\varphi: \mathbb{S}^{d-1} \rightarrow \mathbb{R}$,

where we have used the superscript $^\pm$ to denote either or $\nu_t^+$ or $\nu_t^-$ and the right-hand side is changed accordingly (the integration domain) depending on the sign $+$ or $-$. Because ReLU is positively $1$-homogeneous, we have $f(\mu_t; x) = \int \sigma(u^\top x) \mathrm{d}(\nu_t^+ - \nu_t^-)(x)$. It is shown in Appendix E.1.1 that $\nu_t^\pm$ satisfies, in the sense of distributions, the equation

where, for any $u \in \mathbb{S}^{d-1}$,

Equation (4.2) can be interpreted as a Wasserstein-Fisher-Rao GF ^[30] on the sphere since $\tilde{v}_t(u) = \text{proj}_{\{u\}^\perp}(\nabla g_t(u))$.

Closed dynamics over $[0, \pi/2] \times \mathbb{S}^{d_H-1}$. The dynamics on the pair $(\nu_t^+, \nu_t^-)$ can be further reduced to dynamics over $[0, \pi/2] \times \mathbb{S}^{d_H-1}$. Indeed, by positive 1-homogeneity of $f(\mu_t; \cdot)$ we may restrict ourselves to inputs $u \in \mathbb{S}^{d-1}$, and $f(\mu_t; u)$ depends only on $u^H$ and $||u^\perp||$. However, because $||u^H||^2 + ||u^\perp||^2 = 1$, this dependence translates into a dependence on the direction $u^H / ||u^H||$ of the projection onto $H$ and the norm $||u^H||$. The former is an element of $\mathbb{S}^{d_H-1}$ while the latter is given by the angle $\theta$ between $u$ and $H$, that is $\theta : = \arccos(u^\top u^H/||u^H||) = \arccos(||u^H||)$. This simplification leads to the following lemma:

Lemma 4.2. Define the measures $\tau_t^+, \tau_t^-$ by $\tau_t^\pm = P_{\#} \nu_t^\pm \in \mathcal{M}_{+}([0, \pi/2] \times \mathbb{S}^{d_H-1})$ via $P: u \in \mathbb{S}^{d-1} \backslash H^\perp \mapsto (\arccos(||u_H||), u_H/||u_H||) \in [0, \pi/2] \times \mathbb{S}^{d_H-1}$. Then, the measures $\tau_t^+, \tau_t^-$ satisfy the equation

where $G_t: [0, \pi/2] \times \mathbb{S}^{d_H-1} \rightarrow \mathbb{R}$, and $V_t: [0, \pi/2] \times \mathbb{S}^{d_H-1} \rightarrow \mathbb{R}^{d_H+1}$ are functions depending only on $(\tau_t^+, \tau_t^-)$, and furthermore, $f(\mu_t; \cdot)$ can be expressed solely using $\tau_t^+, \tau_t^-$ (exact formulas are provided in Appendix E.1.2).

Abbe et al. ^[21] show a similar result with a lower-dimensional dynamics in the context of infinitely wide two-layer networks when the input data have i.i.d coordinates distributed uniformly over $\{-1, +1\}$ (i.e., , Rademacher variables), except that they do not have the added dimension due to the angle $\theta$ as we do thanks to their choice of input data distribution.

Lemma 4.2 above illustrates how the GF dynamics of infinitely wide two-layer networks adapts to the lower-dimensional structure of the problem: the learned predictor and the dynamics can described only in terms of the angle $\theta$ between the input neurons and $H$ and their projection on the unit sphere of $H$.

4.2.   One dimensional reduction

Since the predictors we consider are positively homogeneous, one cannot hope to do better than learn a positively homogeneous function. A natural choice of such a target function to learn is the Euclidean norm. With the additional structure that the target only depends on the projection onto $H$, this leads to considering $f^*(x) = ||x^H||$ which has additional symmetries compared to the general case presented above: it is invariant by any linear map $T$ such that $T_{|H} \in \mathcal{O}(d_H)$ and $T_{|H^\perp} \in \mathcal{O}(d_\perp)$. By Proposition 2.1 those symmetries are shared by $\mu_t$ and $f(\mu_t; \cdot)$, and we show that in this case the dynamic reduces to a one-dimensional dynamic over the angle $\theta$ between input neurons and $H$.

We prove a general disintegration result for the uniform measure on the sphere in the Appendix (see Lemma B.4) which allows, along with some spherical harmonics analysis, to describe the reduced dynamics and characterize the objective that they optimize. This leads to the following result:

Theorem 4.3 (1d dynamics over the angle $\theta$). Assume that $f^*(x) = ||x^H||$, and define the measures $(\tau_t^+, \tau_t^-) \in {\mathcal{M}_+([0, \pi/2])}^2$ from $(\nu_t^+, \nu_t^-)$ via $P: u \in \mathbb{S}^{d-1} \mapsto \arccos(||u_H||) \in [0, \pi/2]$: $\tau_t^\pm = P_{\#} \nu_t^\pm$. Then, the pair $(\tau_t^+, \tau_t^-)$ follows the Wasserstein-Fisher-Rao GF for the objective $A(\tau^+, \tau^-) : = \mathbb{E} \left[\ell \left(f(\tau^+, \tau^-; x), f^*(x)\right) \right]$ over the space $\mathcal{M}_+([0, \pi/2]) \times \mathcal{M}_+([0, \pi/2])$, where $f(\tau^+, \tau^-; x)$ is the expression (with a slight overloading of notations) of $f(\mu; x)$ in function of $(\tau^+, \tau^-)$ (see Appendix E.2 for more details):

where $B$ is the Beta function, and

Additionally, $f(\mu_t; \cdot)$, $G_t$, and $V_t$ only depend on the pair $(\tau_t^+, \tau_t^-)$, and for any $t \geq 0$, it holds that $F(\mu_t) = A(\tau_t^+, \tau_t^-)$.

Remark. The result should still hold for general $\rho$ which are spherically symmetric as long as the Wasserstein GF (1.3) is well-defined but the proof is more technical. In addition, this result shows that even with more structure than in Lemma 4.2, the dynamics of infinitely wide two-layer networks are still able to adapt to this setting: these dynamics, as well as the learned predictor, can be fully characterized solely by the one-dimensional dynamics over the angle $\theta$ between input neurons and $H$. This is noteworthy since this angle determines the alignment of the neurons with $H$, and thus measures how much the representations learned by the network have adapted to the structure of the problem. Furthermore, as discussed below, this reduction with exact formulas enables efficient numerical simulation in one dimension.

Daneshmand and Bach ^[15] prove the global convergence of a reduced one-dimensional dynamics in a context similar to ours but their original problem is two-dimensional and with a choice of activation function that leads to specific algebraic properties.

Expression of $f(\mu_t; \cdot)$. Because of the symmetries of $f(\mu_t; \cdot)$, which result from that of $f^*$, $f(\mu_t; x)$ depends only on $||x^H||$ and $||x^\perp||$. What is more, since $f(\mu_t; \cdot)$ is positively $1$-homogeneous (because ReLU is) it actually holds that $f(\mu_t; x) = ||x|| \tilde{f}_t(\varphi_x)$ where $\varphi_x = \arccos(||x^H|| / ||x||)$ is the angle between $x$ and $H$, and $\tilde{f}_t(\varphi) : = \int_{\theta} \tilde{\phi}(\theta; \varphi) \mathrm{d}(\tau_t^+ - \tau_t^-)(\theta)$, $\tilde{\phi}$ depending only on $\sigma$ and fixed probability measures (see Appendix E.1.2 for an exact formula).

Learning the low-dimensional structure as $t \to \infty$. Although, as shown in Figure 2, $f(\mu_t; \cdot)$ does not learn the low-dimensional structure in finite-time, it is reasonable to expect that as $t \to \infty$, the measures $\tau_t^\pm$ put mass only on $\theta = 0$, indicating that the only part of the space that the predictor is concerned with for large $t$ is the sub-space $H$. Since we assume here that the target function $f^*$ is non-negative, the most natural limits for $\tau_t^+$ and $\tau_t^-$ are $\tau_t^+ \to \alpha \delta_{0}$ with $\alpha > 0$, and $\tau_t^- \rightarrow 0$ (in the sense that $\tau_t^-([0, \pi/2]) \rightarrow 0$) as $t \rightarrow \infty$, because then the "negative" output weights do not participate in the prediction in the large $t$ limit.

The global convergence result of Chizat and Bach ^[5,11] still holds but is not quantitative and moreover does not guarantee that the limit is the one described above. We leave the proof of this result as an open problem, but we provide numerical evidence supporting this conjecture. Indeed, we take advantage of the one-dimensional reduction from Theorem 4.3, and numerically simulate the resulting dynamics by parameterizing $\tau_t^\pm$ via weight and position ^[31] as $\mu_{m, t} = (1/m) \sum_{j = 1}^m c_j^\pm(t) \delta_{\theta_j^\pm(t)}$, and simulating the corresponding dynamics for $c_j^\pm(t)$ and $\theta_j^\pm(t)$. The corresponding results are depicted in Figure 3 which are again obtained by discretizing the initial measures $\tau_0^+, \tau_0^-$ and performing GD with finite step-size (see more details in Appendix 63141592631415927). Figure 3(a), (b) show that the mass of $\tau_t^+$ tends to concentrate around $0$ while that of $\tau_t^-$ tends to concentrate around $\pi/2$, indicating that $\tau_t^+$ adapts to the part of the space relevant to learning $f^*$ while $\tau_t^-$ puts mass close to the orthogonal to that space.

Total mass of particles at convergence. If $\tau_\infty^- = 0$ and $\tau_\infty^+ = \alpha \delta_{0}$ as described above, we have $f(\mu_\infty; x) = \alpha ||x|| \tilde{\phi}(0; \varphi_x) = \alpha \frac{\Gamma(d_H/2)}{2 \sqrt{\pi} \Gamma((d_H+1)/2)} ||x|| \cos(\varphi_x) = \frac{\alpha \Gamma(d_H/2)}{2 \sqrt{\pi} \Gamma((d_H+1)/2)} ||x^H||$. To recover exactly $f^*$, it must hold that $\alpha = \tau_\infty^+([0, \pi/2]) = \frac{2 \sqrt{\pi} \Gamma((d_H+1)/2)}{\Gamma(d_H/2)}$. Defining the normalized probability measure $\tilde{\tau}_t^\pm = \tau_t^\pm / \tau_t^\pm([0, \pi/2])$, we thus expect $\tilde{\tau}_t^+$ to grow close to $\delta_{0}$ and $\tilde{\tau}_t^-$ to $\delta_{\pi/2}$. In terms of total mass, we expect that $\tau_t^+([0, \pi/2])$ gets closer to $\alpha$ while $\tau_t^-([0, \pi/2])$ gets closer to $0$.

The numerical behaviour depicted in Figure 3(c) seems to follow our intuitive description, at least until a critical time $t^*$ in the numerical simulation which corresponds to the first time $t$ where $\tau_t^+([0, \pi/2]) > \alpha$. While the total mass of $\tau_t^\pm$ (dashed lines) seems to approach its limit rapidly before $t^*$ it slowly moves further away from it for $t \geq t^*$. On the other hand, while the angles only slowly change before $t^*$, they start converging fast towards the corresponding Dirac measures after $t^*$. It is unclear whether this slight difference in behaviour (around the critical time $t^*$) between what we intuitively expected and the numerical simulation is an artefact of the finite width and finite step size or if it actually corresponds to some phenomenon present in the limiting model. For more details concerning the numerical experiments, see Appendix 63141592631415927. Note that there is a priori not a unique global optimum: $\tau_\infty^+$ and $\tau_\infty^-$ (if they exist) can compensate on parts of the space $[0, \pi/2]$ and lead to the same optimal predictor for different choices of measures. Our numerical experiments suggest that the GF dynamics select a "simple" solution where $\tau_\infty^+$ is concentrated on $\{\theta = 0\}$ and $\tau_\infty^-$ vanishes (puts $0$ mass everywhere), which is a form of implicit bias.

5.   Conclusions

We have explored the symmetries of infinitely wide two-layer ReLU networks and we have seen that: $(i)$ they adapt to the orthogonal symmetries of the problem, $(ii)$ they reduce to the dynamics of a linear network in the case of an odd target function and lead to exponential convergence, and $(iii)$ when the target function depends only on the orthogonal projection onto a lower-dimensional sub-space $H$, the dynamics can be reduced to a lower-dimensional PDE. In particular, when $f^*$ is the Euclidean norm, this PDE is over a one-dimensional space corresponding to the angle $\theta$ between the particles and $H$. We have presented numerical experiment indicating that the positive particles converge to the subspace $H$ in this case and leave the proof of this result as an open problem. We also leave as an open question whether the results of Section 2 extend to deeper networks.

Acknowledgments

Karl Hajjar received full support from the Agence Nationale de la Recherche (ANR), reference ANR-19-CHIA-0021-01 "BiSCottE". The authors thank Christophe Giraud for useful discussions and comments.

Conflict of interest

The authors declare there is no conflict of interest.

Appendix

A.   Additional notations and preliminary results

A.1.   Notations for the appendix

We introduce in this section additional notation that we use throughout the Appendix.

Residual: we call $R_t(y) : = - \partial_2 \ell(f^*(y), f(\mu_t; y))$, the "residual", which is equal to the difference $f^*(y) - f(\mu_t; y)$ when $\ell$ is the squared loss.

Identity matrix: we denote by $I_p$ the identity matrix in $\mathbb{R}^{p \times p}$ for any $p \in \mathbb{N}$.

Indicator functions: we denote by $\mathbf{1}_{A}$ the indicator of a set $A$, that is $\mathbf{1}_{A}(z) = 1 \iff z \in A$, and $\mathbf{1}_{A}(z) = 0$ otherwise.

Total variation: for any measure $\nu$, we denote by $|\nu|$ its total variation, which should cause no confusion with the absolute value given the context.

Beta / Gamma function and distribution: for $\alpha, \beta > 0$, we denote by $B(\alpha, \beta)$ the Beta function equal to $\Gamma(\alpha) \Gamma(\beta) / \Gamma(\alpha + \beta)$ where $\Gamma$ is the Gamma function, and by $\text{Beta}(\alpha, \beta)$ the beta law with density equal to $u^{\alpha-1} u^{\beta-1} / B(\alpha, \beta)$ on $[0, 1]$.

Gaussian / spherical measures: we call $\rho_p$ the standard Gaussian measure in $\mathbb{R}^p$ (corresponding to $\mathcal{N}(0, I_p)$) for any $p \in \mathbb{N}$.

Whenever $\tau \in \mathcal{M}_+(\Omega)$ has finite and non-zero total variation, we denote by $\tilde{\tau} \in \mathcal{P}_2(\Omega)$ its normalized counterpart (which is a probability measure), that is $\tilde{\tau} = \tau / \tau(\Omega) = \tau / |\tau|$.

For any $p \in \mathbb{N}$, we call $\omega_p$ the Lebesgue (spherical) measure over the unit sphere $\mathbb{S}^{p-1}$ of $\mathbb{R}^p$, that is the measure such that $\tilde{\omega}_p$ is the uniform measure on $\mathbb{S}^{p-1}$. We then denote by $|\mathbb{S}^{p-1}|$ the surface area of $\mathbb{S}^{p-1}$, that is $|\mathbb{S}^{p-1}| : = |\omega_p| = \omega_p(\mathbb{S}^{p-1}) = 2\pi^{p/2} / \Gamma(p/2)$.

Smooth functions: we denote by $\mathcal{C}(\Omega)$ (resp. $\mathcal{C}^1_c(\Omega)$) the set of continuous (resp. continuously differentiable and compactly supported) functions from a set $\Omega$ to $\mathbb{R}$.

A.2.   General results on invariance for measures and functions

In this section, we list a number of lemmas related to symmetries of measures and functions which will prove helpful in the proofs presented in the Appendix.

Lemma A.1 (Invariance under invertible maps). Let $\mu$ be a measure invariant under some measurable and invertible map $T$. Then, assuming $T^{-1}$ is also measurable, one has that $\mu$ is also invariant under $T^{-1}$.

Remark. A similar result holds for a function $f$ invariant under an invertible map.

Proof. Because $\mu$ is invariant under $T$, we have for any measurable set $A$, $\mu(A) = \mu(T^{-1}(A))$. Since $T^{-1}$ is assumed to be measurable, for any measurable set $A$, $T(A)$ is also measurable ($T(A) = (T^{-1})^{-1}(A)$) and thus $\mu(T(A)) = \mu(T^{-1}(T(A)) = \mu(A)$ which shows $\mu$ is invariant under $T^{-1}$.

Lemma A.2 (Invariance of the density). Let $\nu$ be a measure with density $p$ w.r.t. some measure $\mu$, and assume both $\nu$ and $\mu$ are $\sigma$-finite and invariant under some measurable and invertible map $T$, whose inverse $T^{-1}$ is also measurable. Then $p$ is also invariant under $T$ $\mu$-almost everywhere, i.e., , $p(T(x)) = p(x)$ for $\mu$-almost every $x$.

Proof. For any measurable $\varphi$ (w.r.t. $\mu$, and thus w.r.t. $\nu$ as well), $\varphi \circ T^{-1}$ is also measurable, and we have, on the one hand

and on the other hand

which shows that $\int \varphi \left(p \circ T \right) \mathrm{d} \mu = \int \varphi \, p \, \mathrm{d} \mu$, and thus that $p \circ T = p$ $\mu$-almost everywhere.

Lemma A.3 (Projected variance with spherical symmetry). Let $\zeta$ be a spherically symmetric measure on $\mathbb{R}^p$ (i.e., , such that for any orthogonal linear map $T \in \mathcal{O}(p)$, $T_{\#} \zeta = \zeta$), with finite second moment. Then we have the following matrix identity:

Proof. The $(i, j)$-th entry of the matrix on the left-hand-side is $\int_{z} z_i z_j \mathrm{d} \zeta(z)$, and it is readily seen that the terms outside the diagonal are $0$. Indeed, let $(i, j) \in [1, p]^2$ with $i \neq j$, and consider the orthogonal map $T_j: z \in \mathbb{R}^p \mapsto (z_1, \ldots, z_{j-1}, -z_j, z_{j+1}, \ldots, z_p)^\top$. The spherical symmetry of $\rho$ implies that it is invariant under $T_j$, which yields $\int_z z_i z_j \mathrm{d} \rho(z) = - \int_z z_i z_j \mathrm{d} \rho(z)$, thereby showing that the latter is $0$. To see that the diagonal terms are all equal, it suffices to consider the orthogonal map $S_{i}$ which swaps the 1st and $i$-th coordinates of a vector $z$. The invariance of $\rho$ under $S_i$ yields $\int_z (z_1)^2 \mathrm{d} \rho(z) = \int_z (z_i)^2 \mathrm{d} \rho(z)$, which concludes the proof.

A.3.   A disintegration result on the unit sphere $\mathbb{S}^{d-1}$

Consider a $u \in \mathbb{S}^{d-1}$. $u$ is determined by: $(i)$ its angle $\theta : = \arccos(||u^H||) \in [0, \pi/2]$ with $H$ (i.e., , its angle with its projection $u^H$ onto $H$), $(ii)$ the direction $z^H = u^H / ||u^H|| \in \mathbb{S}^{d_H-1}$ of its projection $u^H$ onto $H$, and finally $(iii)$ the direction $z^\perp = u^\perp / ||u^\perp|| \in \mathbb{S}^{d_\perp-1}$ of its projection $u^\perp$ onto $H^\perp$. Since $||u^H||^2 + ||u^\perp||^2 = 1$, the angle $\theta$ gives both the norms of the projections onto $H$ and $H^\perp$: $||u^H|| = \cos(\theta)$ and $||u^\perp|| = \sin(\theta)$.

When $z$ ranges over the unit sphere $\mathbb{S}^{d-1}$, the angle $\theta$ and the directions $z^H, z^\perp$ range over $[0, \pi/2]$, $\mathbb{S}^{d_H-1}$, and $\mathbb{S}^{d_\perp - 1}$ respectively. We wish to understand what measures we obtain on these three sets when $z$ is distributed on the sphere according to the Lebesgue measure $\omega_d$. We show below below that after the change of coordinates described above (from $u \in \mathbb{S}^{d-1}$ to $(\theta, z^H, z^\perp) \in [0, \pi/2] \times \mathbb{S}^{d_H-1} \times \mathbb{S}^{d_\perp - 1}$), the corresponding measures over $\mathbb{S}^{d_H-1}$ and $\mathbb{S}^{d_\perp-1}$ are uniform measures and the measure over $\theta$ is given by a push-forward of a Beta distribution as defined below:

Definition A.1 (Distribution $\gamma$ of the angle $\theta$). We define the measure $\gamma$ on $[0, \pi/2]$ with the following density w.r.t. the Lebesgue measure on $[0, \pi/2]$:

Remark. $\gamma$ is in fact simply given by $(\arccos \circ \sqrt{\cdot})_{\#} \text{Beta}(d_H/2, d_\perp/2)$. Note that the total variation of gamma is $|\gamma| = \gamma([0, \pi/2]) = \frac{1}{2} B(\frac{d_H}{2}, \frac{d_\perp}{2})$, and the corresponding normalized (probability) measure is $\mathrm{d} \tilde{\gamma}(\theta) = \mathrm{d} \gamma(\theta) / |\gamma| = \frac{2}{B(\frac{d_H}{2}, \frac{d_\perp}{2})} \cos(\theta)^{d_H-1} \sin(\theta)^{d_\perp-1} \mathrm{d} \theta$.

We now state the disintegration theorem and give its proof:

Theorem A.4 (Disintegration of the Lebesgue measure on the sphere). Let $\omega_d$ denote the Lebesgue measure on the sphere measure on the sphere of $\mathbb{R}^d$, and let $\gamma$ be the measure of Definition 3. Then, one has

where

Proof. Denoting $\tilde{\omega}_d$ the uniform measure on the sphere, $|\mathbb{S}^{d-1}| : = \frac{2 \pi^{d/2}}{\Gamma(d/2)}$ the surface are of the sphere in dimension $d$, and $\rho_p$ the standard Gaussian distribution in $\mathbb{R}^p$ for any $p$. Using the well-known fact that $\tilde{\omega}_d = \Pi_{\#} \rho_d$ with $\Pi: x \in \mathbb{R}^d \backslash \{0\} \mapsto x/||x|| \in \mathbb{S}^{d-1}$, we have, for any measurable test function $\varphi : \mathbb{S}^{d-1} \rightarrow \mathbb{R}$,

with

Doing the polar change of variables $(r_H, r_\perp) \in \mathbb{R}_{+}^2 \rightarrow (R, \theta) \in \mathbb{R}_{+} \times [0, \pi/2]$, we get:

where

which concludes the proof.

Remark. A similar disintegration result holds for the uniform measure $\tilde{\omega}_d$ on the sphere. The corresponding measures which are then pushed-forward by the same $\Phi$ are the normalized counterparts of the measures in the theorem above: $\tilde{\omega}_d = \Phi_{\#} (\tilde{\omega}_{d_H} \otimes \tilde{\omega}_{d_\perp} \otimes \tilde{\gamma})$. This readily comes from noting that a simple calculation yields $|\omega_d| = |\omega_{d_H}| \, |\omega_{d_\perp}|\, |\gamma|$.

B.   Gradient flows on the space of probability measures

B.1.   First variation of a functional over measures

Given a functional $F: \mathcal{P}_2(\mathbb{R}^p) \to \mathbb{R}$, its first variation or Fréchet derivative at $\mu \in \mathcal{P}_2(\mathbb{R}^p)$ is defined as a measurable function, denoted $\frac{\delta F }{\delta \mu}(\mu) : \mathbb{R}^p \to \mathbb{R}$, such that, for any $\nu \in \mathcal{P}_2(\mathbb{R}^p)$ for which $\mu + t \nu \in \mathcal{P}_2(\mathbb{R}^p)$ in a neighborhood (in $t$) of $t = 0$,

See Santambrogio [32,Definition 7.12], or [33,p.29] for more details on the first variation.

In the case of the functional defined in Eq (1.2) corresponding to the population loss objective, using the differentiability of the loss $\ell$ w.r.t. its second argument, one readily has that

since

B.2.   Wasserstein gradient flows in the space $\mathcal{P}_2(\mathbb{R}^{d+1})$

A Wasserstein gradient flow for the objective $F$ defined in Eq (1.2) is a path $(\mu_t)_{t \geq 0}$ in the space of probability measures $\mathcal{P}_2(\mathbb{R}^{d+1})$ which satisfies the continuity equation with a vector field $v_t$ which is equal to the opposite of the gradient of the first variation of the functional $F$. This means that we have, in the sense of distributions,

That a pair $((\mu_t)_{t \geq 0}, v_t)$ consisting of a path in $\mathcal{P}_2(\mathbb{R}^p)$ and a (time-dependent) vector field in $\mathbb{R}^p$ satisfies the continuity equation $\partial_t \mu_t = - \text{div} (v_t \mu_t)$ in the sense of the distributions simply means that for any test function $\varphi \in \mathcal{C}^1_c(\mathbb{R}^{p})$,

where $\partial_t$ stands for the time derivative $\frac{d}{dt}$. Similarly, when we say that the advection-reaction equation $\partial_t \mu_t = - \text{div} \left(v_t \mu_t \right) + g_t \mu_t$ is satisfied for some function $g_t: \mathbb{R}^{p} \to \mathbb{R}$, we mean that it is in the sense of distributions: for any test function $\varphi \in \mathcal{C}^1_c(\mathbb{R}^{p})$,

An alternative description of the Wasserstein gradient flow of the objective $F$ is to consider a flow $X_{\bullet}(\cdot)$ in $\mathbb{R}_+ \times \mathbb{R}^{d+1}$ such that, for any $c \in \mathbb{R}^{d+1}$,

and to define $\mu_t = (X_t)_{\#} \mu_0$.

For more details on Wasserstein gradient flows in the space of probability measures see Santambrogio [32,Section 5.3], and [33,Section 4], and for more details on the equivalence between the continuity equation and the flow-based representation of the solution see Santambrogio [32,Theorem 4.4].

C.   Proofs of the symmetry results of Section 2

There are two main ideas behind the proof. Call $\tilde{T}: (a, b) \in \mathbb{R} \times \mathbb{R}^d \mapsto (\pm a, T(b))$ (depending on whether $f^*$ is invariant or anti-invariant under $T$) and consider the following two facts:

Structure of $\phi((a, b);x)$. Since $T$ is orthogonal, so is $\tilde{T}$, and the structure of $\phi((a, b);x) = a\sigma(b^\top x)$ is such that $\phi(\tilde{T}(a, b); x) = \pm \phi((a, b); T^{-1}(x))$ because $T$ is orthogonal (its adjoint is thus its inverse).

Conjugate gradients. Computing the gradient of a function whose input has been transformed by $\tilde{T}^{-1}$ is the same as the conjugate action of $\tilde{T}$ on the gradient: $\nabla (\varphi \circ \tilde{T}^{-1}) = \tilde{T} \circ (\nabla \varphi) \circ \tilde{T}^{-1}$ (this is due to the fact that the adjoint of $\tilde{T}^{-1}$ is $\tilde{T}$ because $\tilde{T}$ is orthogonal). Note that we similarly get $\nabla (\varphi \circ \tilde{T}) = \tilde{T}^{-1} \circ (\nabla \varphi) \circ \tilde{T}$.

C.1.   Preliminaries

We present here arguments that are present in both the proofs of Propositions 2.1 and 2.2. Let $T$ be a linear orthogonal map such that $f^*(T(x)) = \pm f^*(x)$, where the $\pm$ is because we deal with both cases at the same time since the logic is the same. Let $t \geq 0$, and define $\nu_t : = \tilde{T}^{-1}_\# \mu_t$. We aim to show that $(\nu_t)_{t \geq 0}$ is also a Wasserstein gradient flow for the same objective as $(\mu_t)_{t \geq 0}$.

Prediction function. Let $x \in \mathbb{R}^d$. We have, using the fact that $T$ is orthogonal (and thus that $\left < {T(x)}, \, {y} \right > = \left < {x}, \, {T^{-1}(y)} \right >$),

Time derivative. Let $\varphi \in \mathcal{C}^1_c(\mathbb{R}^d)$. Because $\mu_t$ satisfies the continuity Eq (1.3) in the sense of distributions, and using the remark above on conjugate gradients as well as the orthogonality of $\tilde{T}$, we have:

Conjugate velocity field. The equality above actually shows that $\nu_t$ satisfies the continuity equation with the conjugate velocity field $\tilde{T}^{-1} \circ v_t \circ \tilde{T}$ instead of $v_t$. We show below that the former is closely related to the latter (and is in fact equal to $-\nabla F^{\prime}_{\nu_t}$ with sufficient assumptions on $\partial_2 \ell$, which is the step proven in Appendices C.2 and C.3). Indeed, because $v_t$ is a gradient: $v_t = - \nabla F^{\prime}_{\mu_t}$, we have using again the remark above on conjugate gradients:

Computing the function on the right-hand-side, for any $(a, b) \in \mathbb{R} \times \mathbb{R}^d$, we get, using the remark above on the structure of $\phi$,

$\rho$ is invariant under $T$ since it spherically symmetric by assumption (and thus invariant under any orthogonal map) and we can therefore replace $y$ by $T(y)$ in the integral above, which yields

and thus we get

One can already notice that if $f^*$ is invariant under $T$ (as opposed to anti-invariant), that is if we keep the "$+$" in $\pm$, we get $\tilde{T}^{-1} \circ v_t \circ \tilde{T} = -\nabla F^{\prime}_{\nu_t}$.

C.2.   Proof of Proposition 2.1

Proof. We first prove $\nu_0 = \mu_0$ and then prove that both $(\mu_t)_{t \geq 0}$ and $(\nu_t)_{t \geq 0}$ are Wasserstein gradient flows of the objective $F$ defined in Eq (1.2), starting from the initial condition $\mu_0$ at $t = 0$. The unicity of such a gradient flow then guarantees that $\mu_t = \nu_t$ and thus $f(\mu_t; T(x)) = f(\mu_t; x)$ by the preliminaries above on the prediction function (see Appendix C.1).

Initialization: $\nu_0 = \mu_0$. By definition, $\tilde{T}(a, b) = (a, T(b))$. Since $\mu_0 = \mu_0^1 \otimes \mu_0^2$ by assumption, and $\mu_0^2$ is invariant under $T$ since it has spherical symmetry, it is clear that $\mu_0$ is invariant under $\tilde{T}$, and thus under $\tilde{T}^{-1}$ by Lemma A.1, which gives $\nu_0 = \mu_0$ because $\nu_t = \tilde{T}^{-1}_\# \mu_t$ for any $t$ by definition.

Time derivative. From the preliminary results above (see Appendix C.1) we have

which shows that $(\nu_t)_{t \geq 0}$ is also a Wasserstein gradient flow of the objective $F$. By unicity of the latter (starting from the initial condition $\mu_0$), it must hold that $\mu_t = \nu_t$ for any $t \geq 0$ which concludes the proof.

C.3.   Proof of Proposition 2.2

The proof follows the exact same pattern as that of Proposition 2.1 (see Appendix C.2). We now have by definition, $\tilde{T}(a, b) = (-a, T(b))$ and the added symmetry assumption on $\mu_0^1$ ensures that $\nu_0 = \mu_0$ still holds in this case. As for the time derivative, the preliminaries above (see Appendix C.1) ensure that

where we have used the extra assumption that $\partial_2 \ell(-y, -\hat{y}) = - \partial_2 \ell(y, \hat{y})$. This yields

and the conclusion follows from the same logic as for Proposition 2.1.

D.   Proof of the exponential convergence for linear networks: Theorem 3.2

Proof. The proof is divided in three steps: $(i)$ we derive the dynamics in time of the vector $w(t) = \frac{1}{2} \int a b \, \mathrm{d} \mu_t(a, b)$, $(ii)$ we show that the positive definite matrix $H(t)$ appearing in these dynamics has its smallest eigenvalue lower-bounded by some positive constant after some $t_0 > 0$, and $(iii)$ we show that this implies the exponential convergence to the global minimum.

Generalities on the objective $Q$. Expanding the square in the definition of $Q$ (3.2), we have

If $C \neq 0$, $Q(w) \to \infty$ as $||w|| \to \infty$ and since $Q$ is lower-bounded by $0$, it thus admits at least one global minimum. This minimizer $w^\star$ is unique as soon as $Q$ is strongly convex, i.e., , $C$ is definite positive, which holds in this case as we have assumed the smallest eigenvalue $\lambda_{\text{ min}}$ of $C$ to be $> 0$. Note that $\nabla Q(w) = Cw - \beta = \int_x \left((x^\top w) - f^*(x) \right) x \mathrm{d} \mathbb{P}(x) \in \mathbb{R}^d$.

First step: dynamics of $w(t)$. Let $k \in \{1, \ldots, d\}$, the $k$-th coordinate $w_k(t)$ of $w(t)$ is given by $w_k(t) = \int a b_k \, \mathrm{d} \mu_t(a, b)$, and its time derivative is given by

where $v_t$ is given by Eq (1.3) except we replace $\sigma$ by $\frac{1}{2} \text{id}_{\mathbb{R}^d}$ and $\rho$ by $\mathbb{P}$ in $F_{\mu_t}$, that is

On the other hand, $\nabla_{(a, b)} (a b_k) = \begin{pmatrix} b_k \\ a e_k \end{pmatrix} \in \mathbb{R}^{1 + d}$ where $e_k$ is the $k$-th element of the canonical orthonormal basis of $\mathbb{R}^d$. Note that here, $R_t(y) = f^*(y) - \langle w(t), y \rangle$. We thus get

Note that the term on the right in the inner products is in fact equal to $-\nabla Q(w(t))$, which yields the following dynamics for the vector $w(t)$:

Second step: lower bound on the smallest eigenvalue of $H(t)$. At initialization, by symmetry one has $w(0) = 0$, and using Lemma A.3, one has that $H(0) = \frac{1}{4} \left(\frac{1}{d} + 1\right) I_d$, so that

If $\beta = 0$, then $\nabla Q(0) = 0$ and since $w(0) = 0$, $w(t)$ starts at the global optimum and thus stays constant equal to $0$. Otherwise, if $||\beta|| > 0$, one has $\frac{d}{dt} Q(w(t)) \Big|_{t = 0} < 0$, which ensures that there is a $t_0 > 0$ such that $Q(w(t)) < Q(w(0)) = Q(0)$ for any $t \in (0, t_0]$. Call $\varepsilon : = \left[Q(0)- Q(w(t_0)) \right]/2 > 0$. The continuity of $Q$ at $0$ guarantees that there is a $\delta > 0$ such that for any $w \in \mathbb{R}^d$, if $||w|| < \delta$, then $|Q(w) - Q(0)| \leq \varepsilon$.

Now assume that there exists $t_1 \geq t_0$ such that $\int a^2 \mathrm{d} \mu_{t_1}(a, b) \leq \delta$. Then, one has

where we have used in the penultimate inequality that $\mu_{t_1}$ is supported on the set $\{ |a| = ||b|| \}$ because of the assumptions on the initialization $\mu_0$ (see Section 1.1). This ensures that $|Q(w(t_1)) - Q(0)| \leq \varepsilon$. Since $:t \mapsto Q(w(t))$ is decreasing ($Q(w(t)) = F(\mu_t)$ and it is classical that the objective is decreasing along the gradient flow path, see third step below) and $t_1 \geq t_0$, this means that

which is a contradiction. Therefore, for any $t \geq t_0$, $\int a^2 \mathrm{d} \mu_t(a, b) \geq \delta$. Calling $\eta : = \delta/4 > 0$, we thus have that for any $t \geq t_0$, the smallest eigenvalue of $H(t)$ is larger than $\eta$ because $H(t)$ is the sum of the positive semi-definite matrix $\frac{1}{4} \int b b^\top \mathrm{d} \mu_t(a, b)$ and of the positive definite matrix $\frac{1}{4} \int a^2 \mathrm{d} \mu_t(a, b) I_d$ whose smallest eigenvalue is at least $\eta$ for $t \geq t_0$.

Third step: exponential convergence. We have:

which shows that because $H(t)$ is positive definite, the objective $Q$ is decreasing along the path $(w(t))_{t \geq 0}$. Since after $t_0 > 0$, the smallest eigenvalue of $H(t)$ is lower bounded by a constant $\eta > 0$, we have that, for any $t \geq t_0$:

Because $Q$ is $\lambda_{\text{ min}}$-strongly convex (as the smallest eigenvalue of $C$ is $\lambda_{\text{ min}} > 0$), one has the classical inequality

Plugging this into Eq (D.1) gives

which by Gronwall's lemma in turn yields for any $t \geq t_0$

thereby proving exponential convergence.

Exponential convergence in distance. Given that $\nabla Q(w^\star) = 0$ because $w^\star$ is and optimum, it holds $C w^\star = \beta$. Using this fact, it easily follows that

and the right-hand-side is lower bounded by $\frac{ \lambda_{\text{ min}}}{2} ||w - w^\star||^2$, from which we conclude that

and the exponential decrease of the right-hand-side allows to conclude.

E.   Proofs of Section 4: $f^*$ depends only on the projection on a sub-space $H$

E.1.   The general case

E.1.1.   Closed dynamics on the sphere $\mathbb{S}^{d-1}$

We wish to show here that the pair of measures $(\nu_t^+, \nu_t^-)$ defined through Eq (4.1) satisfy Eq (4.2) and that the corresponding dynamic is closed is the sense that it can be expressed solely using $(\nu_t^+, \nu_t^-)$ (without requiring to express quantities in function of $\mu_t$). Below, we use $\kappa(z) = \max(0, z)$. We do this do differentiate it from the activation function $\sigma$ (which is also equal to ReLU) so as avoid confusion because the $\kappa$ which appears below has nothing to do with the activation function of the network and simply comes from the integration domain in the calculations.

Equations of the dynamics on the sphere. Let $\varphi \in \mathcal{C}^1_c(\mathbb{S}^{d-1})$. One has

Let us compute the components of the gradient above. We have

The Jacobian of the map $:b \in \mathbb{R}^d \mapsto b / ||b||$ is equal to $\frac{1}{||b||}(I_d - b b^\top / ||b||^2)$ which is a symmetric (or self-adjoint) matrix, so that the gradient w.r.t. $b$ is

On the other hand, the first component of $v_t(a, b)$ (corresponding to the gradient w.r.t. $a$) is

and the last $d$ components (corresponding to the gradient w.r.t. $b$) are

When computing the inner product $\nabla_{(a, b)} \left(\kappa(\pm a) ||b|| \varphi \left(\frac{b}{||b||} \right) \right)^\top v_t(a, b)$, we can re-arrange the terms to keep one term where $\varphi$ appears and the other where $\nabla \varphi$ appears. Using the facts that the Jacobian computed above is symmetric, that $\kappa(z) = \kappa^\prime(z) z$ for any $z \in \mathbb{R}$, and that $\mathbf{1}_{\{ \pm a \geq 0 \}} a = \pm \mathbf{1}_{\{ \pm a \geq 0 \}} |a| = \pm \kappa(\pm a)$, we get,

where, for $u \in \mathbb{S}^{d-1}$

Finally, because $\mu_t$ stays on the cone $\{(a, b) \in \mathbb{R}^{d+1}; |a| \ = ||b||\}$ for any $t$ (see Chizat and Bach [4,Lemma 26], Wojtowytsch [5,Section 2.5]), when integrating against $\mu_t$, we can replace $||b||$ by $|a|$ and vice-versa. We thus get that the time derivative we initially computed is the sum of two terms:

which shows that $\nu^\pm_t$ satisfies Eq (4.2) in the sense of distributions.

Closed dynamics. We want to show that $g_t$ and $\tilde{v}_t$ can be expressed using only $\nu_t^+$ and $\nu_t^-$. Both these quantities depend on $t$ only through the residual $R_t$, which itself only depends on $t$ through $f(\mu_t; \cdot)$. We thus show that the latter can be expressed using only $\nu_t^+$ and $\nu_t^-$, which easily follows from writing, for any $y \in \mathbb{R}^d$,

E.1.2.   Closed dynamics in $d_H + 1$ dimensions

Proof. We first prove that the Eq (4.4) for $\tau_t^\pm$ holds in the sense of distributions, and then show that the corresponding dynamics are closed because the $V_t$ and $g_t$ appearing in Eq (4.4) can be expressed with $(\tau_t^+, \tau_t^-)$ (and not only with $(\nu_t^+, \nu_t^-)$ for instance). We show this by expressing $f(\mu_t; \cdot)$ only in function of the pair $(\tau_t^+, \tau_t^-)$.

The pair $(\tau_t^+, \tau_t^-)$ satisfy Eq (4.4). First, we show that $g_t$ and $\tilde{v}_t$ defined in Eq (4.3) admit modified expressions that match the structure of the pushforward transforming $\nu_t^\pm$ into $\tau_t^\pm$. Indeed, since $\rho$ is assumed to be spherically symmetric, it is invariant by any orthogonal transformation. In particular, for a fixed $u \in \mathbb{S}^{d-1}$ such that $u^\perp \neq 0$, we consider the orthogonal map $T^u: \mathbb{R}^d \rightarrow \mathbb{R}^d$ such that $T^u_{|H} = \text{id}_H$ and $T^u_{|H^\perp}$ sends the canonical orthonormal basis $(e^\perp_1, \ldots, e^\perp_{d_\perp})$ of $H^\perp$ on $(u^\perp/||u^\perp||, u_2, \ldots, u_{d_\perp})$ where $(u_2, \ldots, u_{d_\perp}) \in (H^\perp)^{d_\perp - 1}$ is an orthonormal family, orthogonal to $u^\perp$, so that for any $y^\perp \in H^\perp$ with coordinates $y^\perp_1, \ldots, y^\perp_{d_\perp}$ in the basis $(e^\perp_1, \ldots, e^\perp_{d_\perp})$, $T^u_{|H^\perp}(y^\perp) = y^\perp_1 u^\perp/||u^\perp|| + h_u(y_\perp)$ with $h_u(y_\perp) \perp u^\perp$.

Note that since $f^*(y) = f_H(y^H)$ and $f(\mu_t; y) = \tilde f_t(y^H, ||y^\perp||)$, the residual $R_t(y) = f^*(y) - f(\mu_t; y)$ is invariant by any orthogonal transformation which preserves $H$ (and in particular by $T^u$). We thus have

Now consider, for any $(\theta, z^H) \in [0, \pi/2] \times \mathbb{S}^{d_H-1}$,

We show below that $(\tau_t^+, \tau_t^-)$ satisfy Eq (4.4) with the $G_t$ and $V_t$ defined above. Let $\varphi \in \mathcal{C}^1_c([0, \pi/2] \times \mathbb{S}^{d_H-1})$. Since $\tau_t^\pm$ is defined as a push-forward measure obtained from $\nu_t^\pm$ we have:

By definition of the pushforward, and since $u^H = \cos(\arccos(||u^H||)) u^H/||u^H||$ and $||u^\perp|| = \sin(\arccos(||u^H||))$ for $u \in \mathbb{S}^{d-1}$, the first integral is equal $\int \varphi(\theta, z^H) G_t(\theta, z^H) \mathrm{d} \tau_t^\pm(u)$. For the second integral, let us first compute the gradient. One has

We observe that the gradient above belongs to $H$ which implies that when computing its inner product with $\tilde{v}_t(u)$ we can consider only the component of the latter along $H$. Additionally, we note that $I_{d_H} - u^H (u^H)^\top / ||u^H||^2$ is actually the orthogonal projection onto $\{u^H\}^\perp$, so that it yields $0$ when applied to $u$. Using that $||u^\perp|| = \sqrt{1 - ||u^H||^2}$ for $u \in \mathbb{S}^{d-1}$, we then get:

where $\nabla \varphi (\theta, z^H) = \begin{pmatrix} \partial_\theta \varphi (\theta, z^H) \\ \nabla_{z^H} \varphi (\theta, z^H) \end{pmatrix}$. This shows that

which proves that $\tau_t^\pm$ indeed satisfies Eq (4.3) in the sense of distributions.

The dynamics are closed in the pair $(\tau_t^+, \tau_t^-)$. The only thing left to prove to show that the dynamics are closed for the pair $(\tau_t^+, \tau_t^-)$ is that $G_t$ and $V_t$ can be expressed using only the pair $(\tau_t^+, \tau_t^-)$. The only dependence of these quantities on $t$ is through the residual $R_t$ which itself depends on $t$ only through $f(\mu_t; \cdot)$. Let $y \in \mathbb{R}^d$. We have already shown at the end of the previous Section E.1.1 that by definition of $\nu_t^+$ and $\nu_t^-$, we have

On the other hand, we show below that the integral of any measurable function $\varphi: \mathbb{S}^{d-1} \rightarrow \mathbb{R}$ against $\nu_t^\pm$ can be expressed as an integral against $\tau_t^\pm$ in the case where $\nu_t^\pm$ admits a density w.r.t. the uniform measure on $\mathbb{S}^{d-1}$ (which is the case for $\nu_0^\pm$), the case of a general measure $\nu_t^\pm$ being a simple extension via a weak convergence argument. Thus call $p_t^\pm$ the density of $\nu_t^\pm$ w.r.t. $\tilde{\omega}_d$. Since $\nu_t^\pm$ is invariant by any linear map $T$ such that $T_{|H} = \text{id}_{H}$ $T_{|H^\perp} \in \mathcal{O}(d_\perp)$ (because of the symmetries on $\mu_t$ given by Proposition 2.1), and since this is also the case for $\tilde{\omega}_d$ because $\tilde{\omega}_d$ has spherical symmetry and $T$ is orthogonal, we have by Lemma A.2 that $p_t^\pm$ is invariant by any such $T$, which then leads to $p_t$ having the form $p_t(u) = \tilde{p}_t^\pm(u^H, ||u^\perp||)$ by Lemma 4.1.

First step. We show that $\tau_t^\pm$ has the density $q_t^\pm(\theta, z^H) = |\mathbb{S}^{d_\perp-1}| \tilde{p}_t^\pm(\cos(\theta) z^H, \sin(\theta))$ w.r.t. $\tilde{\gamma} \otimes \tilde{\omega}_{d_H}$ where the measure $\tilde{\gamma}$ is the normalized counterpart of the measure in Definition A.1. Indeed, let $\varphi: [0, \pi/2] \times \mathbb{S}^{d-1} \rightarrow \mathbb{R}$ be any measurable function w.r.t. $\tau_t^\pm$. Using the disintegration Lemma A.4 on $\tilde{\omega}_d$, one has that

which proves the desired density for $\tau_t^\pm$.

Second step. Consider a measurable $\varphi: \mathbb{S}^{d-1} \rightarrow \mathbb{R}$ w.r.t. $\nu_t^\pm$. One has with similar calculations as above

Applying this to $f(\mu_t; y)$ shows that the latter quantity can be expressed solely using $(\tau_t^+, \tau_t^-)$, which proves that the dynamics is indeed closed and therefore concludes the proof when $\nu_t^\pm$ has a density.

Third step: extending to any measure. It is known that for any measure $\nu$ over $\mathbb{S}^{d-1}$, there exists a sequence of measure $(\nu_n)_{n \in \mathbb{N}}$ such that: $(i)$ $\nu_n$ has a density $p_n$ w.r.t. the uniform measure $\tilde{\omega}_d$ over $\mathbb{S}^{d-1}$, and $(ii)$ the sequence $(\nu_n)_{n \in \mathbb{N}}$ converges weakly to $\nu$, that is, for any continuous (and thus automatically bounded because the unit sphere is compact) $\varphi$, $\int \varphi \mathrm{d} \nu_n \xrightarrow[n \to \infty]{} \int \varphi \mathrm{d} \nu$. Let thus $\nu \in \mathcal{M}_+(\mathbb{S}^{d-1})$, and consider a sequence $(\nu_n)_{n \in \mathbb{N}}$ with density converging weakly towards $\nu$. Let $\tau$ (resp. $\tau_n$) be defined from $\nu$ (resp. $\nu_n$) as $\tau_t^\pm$ is defined from $\nu_t^\pm$, that is for any measurable $\varphi : [0, \pi/2] \times \mathbb{S}^{d_H-1} \to \mathbb{R}$,

Let thus $\varphi$ be a continuous map from $\mathbb{S}^{d-1} \to \mathbb{R}$ (having in mind the example of $:u \mapsto \sigma(u^\top y)$ for a fixed $y$). By the result of Step $2$, since $\nu_n$ has a density for every $n$, we have that

and taking the limit $n \to \infty$, the left-hand-side of Eq (E.1) converges to $\int \varphi \mathrm{d} \nu$ by assumption. Now let us look at the right-hand-side of (E.1). Calling $\psi(\theta, z^H) = \int_{z^\perp} \varphi(\cos(\theta) z^H + \sin(\theta) z^\perp) \mathrm{d} \tilde{\omega}_{d_\perp}(z^\perp)$ and $\Phi(u) = \int_{z^\perp} \varphi(u^H + ||u^\perp|| z^\perp) \mathrm{d} \tilde{\omega}_{d_\perp}(z^\perp)$, the right-hand-side is in fact $\int \psi \mathrm{d} \tau_n$ and, for any $n \in \mathbb{N}$, is equal to:

and a similar result holds for $\tau$ and $\nu$. Now, the continuity of $\Phi$ is readily obtained from that of $\varphi$, and thus the right-hand-side in the last equality above converges to $\int \Phi \mathrm{d} \nu$ which is also equal to $\int \psi \mathrm{d} \tau$ by the same calculations as above. The right-hand-side in (E.1) therefore converges to $\int \psi \mathrm{d} \tau$, and since the limits of both sides are equal, we get $\int \varphi \mathrm{d} \nu = \int \psi \mathrm{d} \tau$, which is the claim of Step 2 for a general measure $\nu$ which does not necessarily admit a density, thereby concluding the proof.

E.2.   Case when $f^*$ is the euclidean norm: Theorem 4.3

Here, we give the proof of Theorem 4.3 which shows that when $f^*(x) = ||x^H||$ the dynamics can be reduced to a single variable: the angle $\theta \in [0, \pi/2]$ between particles and the subs-space $H$.

We decompose the proof in three steps: first we show that the pair of measures $(\tau_t^+, \tau_t^-) \in \mathcal{M}_+([0, \pi/2])$ as defined in Section 4.2 indeed follows Eq (4.5); then we show that the dynamics are indeed closed by proving that the terms $V_t$ and $G_t$ appearing in the GF depend only on $(\tau_t^+, \tau_t^-)$; and finally, we show that Eq (4.5) indeed corresponds to a Wasserstein-Fisher-Rao GF on a given objective functional over $\mathcal{M}_+([0, \pi/2])^2$.

E.2.1.   Proof of the GF equation

Proof. We first use the added symmetry to simplify the terms $g_t$ and $\tilde{v}_t$ which appear in the GF with $(\nu_t^+, \nu_t^-)$ (see Section 4.1) and express them only with $||u^H||$ and $||u^\perp||$. Then we use the equations satisfied by $(\nu_t^+, \nu_t^-)$ to obtain equations for $(\tau_t^+, \tau_t^-)$.

Equations for $(\tau_t^+, \tau_t^-)$. Let $\varphi \in \mathcal{C}^1_c([0, \pi/2])$. We have

One has that

which belongs to $H$. We recall here the expressions of $\tilde{v}_t$ and $g_t$: for any $u \in \mathbb{S}^{d-1}$, we have

Since, $f^*(x) = ||x^H||$, $f^*$ is now invariant under any orthogonal map $T$ preserving $H$ and $H^\perp$, that is such that the restrictions $T_{|H} \in \mathcal{O}(d_H)$ and $T_{|H^\perp} \in \mathcal{O}(d_\perp)$. Proposition 2.1 then ensures that so is $f(\mu_t, \cdot)$, which in turn implies that the residual $R_t(\cdot) = \partial_2 \ell(f(\mu_t; \cdot), f^*(\cdot))$ also shares that invariance property. Using a similar change of variable as in Appendix E.1.2, and because $\rho$ is spherically symmetric, one gets that $g_t$ can be re-written

Calling

one has $g_t(u) = G_t(\arccos(||u^H||))$ because $u \in \mathbb{S}^{d-1}$, so that $||u^\perp|| = \sqrt{1 - ||u^H||^2}$. Then, by definition of $\tau_t^\pm$, the second integral in the time derivative above is equal to $\int \varphi(\theta) G_t(\theta) \mathrm{d} \tau_t^\pm$. For the first integral appearing in that time derivative, we get

Expanding the inner product inside the integral, we have

Calling

and using again the spherical symmetry of $\rho$, with the same change of variable in the integral as for $g_t$, we get that

Finally, this combined with the previous result on the integral with $g_t$ yields

which leads to the desired equation

E.2.2.   Proof that the dynamics on the angle $\theta$ are closed

The proof follow closely that of Appendix E.1.2 (where we prove closed dynamics), except here we take advantage of the added symmetry of the dynamics. As in Appendix E.1.2, we have

and the only thing to prove is that this quantity can be expressed using only $(\tau_t^+, \tau_t^-)$. As in Appendix E.1.2, we first prove this when $\nu_t^\pm$ has a density, which is the case for $\nu_0^\pm$ and should thus remain so during the dynamics.

Similarly to what occurs in Appendix E.1.2, $\nu_t^\pm$ is invariant by any orthogonal map $T$ which preserves $H$ and $H^\perp$ because $\mu_t$ has those symmetries given by Proposition 2.1, and if $\nu_t^\pm$ has a density $p_{\nu_t^\pm}$ w.r.t. $\tilde{\omega}_d$, then $p^\pm_t$ is also invariant by any such map $T$, and thus depends only on the norms $||u^H||$ and $||u^\perp||$ of its input $u \in \mathbb{S}^{d-1}$. But since its input is on the sphere, those norms are determined by the angle $\theta = \arccos(||u^H||)$ between the input $u$ and $H$. Calling $q_t^\pm$ such that $p_t^\pm(u) = q_t^\pm(\arccos(||u^H||))$, this will lead $\tau_t^\pm$ to have the density $q_t^\pm$ w.r.t. $\tilde{\gamma}$. Then, we show below that similarly to Appendix E.1.2, the integral of any measurable $\varphi:\mathbb{S}^{d-1} \to \mathbb{R}$ against $\nu_t^\pm$ can be expressed as an integral against $\tau_t^\pm$. Indeed, using the disintegration Lemma A.4,

where

which concludes the proof if $\nu_t^\pm$ has a density w.r.t. the uniform measure $\tilde{\omega}_d$ on the sphere $\mathbb{S}^{d-1}$. The general case is obtained by a weak convergence argument (of measures with density) as in the third step of Section E.1.2.

E.2.3.   Proof of the Wasserstein-Fisher-Rao GF

Proof. Recall that $\gamma$ is the measure in Definition 3, and consider the following objective functional over $\mathcal{M}([0, \pi/2])^2$:

where, for any $p \in \mathbb{N}$, $\mathrm{d} \gamma_p(r) = (1-r^2)^{(p-3)/2} \mathrm{d} r$, and $\tilde{\gamma}_p = \gamma_p / |\gamma_p|$ with the normalizing factor $|\gamma_p| = B\left(1/2, (p-1)/2 \right) = \sqrt{\pi} \Gamma((p-1)/2) / \Gamma(p/2) = |\mathbb{S}^{p-1}| / |\mathbb{S}^{p-2}|$. Note that $\tilde{\gamma}_p$ can be simply expressed as the law of $\epsilon \times \sqrt{X}$ where $\epsilon \sim \mathcal{U}(\{-1, +1\})$ and $X \sim \text{Beta}(1/2, (p-1)/2)$.

Computing the first variation or Fréchet derivative of the functional $A$ w.r.t. to its first and second argument yields, for any $\theta \in [0, \pi/2]$,

To conclude one needs only observe that the quantity above is simply equal to $G_t(\theta)$, up to a fixed multiplicative constant. Since we have assumed $\rho$ to be the uniform measure over $\mathbb{S}^{d-1}$ to ensure that the Wasserstein GF (1.3) is well-defined, the constant is one here but in the case of a general $\rho$ with spherical symmetry, the result should also hold (as long as the Wasserstein GF (1.3) is well-defined) but the proof is more technical and different constants might appear.

Simplifying $f(\mu_t; \cdot)$. Using the results from Appendix E.2.2, we have for any $\varphi, z^H, z^\perp \in [0, \pi/2] \times \mathbb{S}^{d_H-1} \times \mathbb{S}^{d_\perp-1}$ (so that $u = \cos(\varphi) z^H + \sin(\varphi) z^\perp \in \mathbb{S}^{d-1}$)

Now, because of the integration against uniform measures on the unit spheres, and the inner products involved, we can use some spherical harmonics theory to simplify those calculations. Using The Funk-Hecke formula (see Atkinson and Han [34,Theorem 2.22], $n = 0$, $d = d_H$ or $d = d_\perp$), we get

Simplifying $f^*(\cos(\varphi) z^H + \sin(\varphi) z^\perp)$. Because $f^*(y) = ||y^H||$, $f^*(\cos(\varphi) z^H + \sin(\varphi) z^\perp)$ is simply $||\cos(\varphi) z^H|| = \cos(\varphi)$ because $z^H \in \mathbb{S}^{d_H-1}$.

With the previous expressions for $f(\mu_t; \cdot)$ and $f^*$ we have that for any function $\Phi: \mathbb{R}^2 \to \mathbb{R}$,

Note that this applies both to $\Phi(y, \hat y) = \ell(y, \hat y)$ and $\Phi(y, \hat y) = -\partial_2 \ell(y, \hat y)$.

Proof that $F(\mu_t) = A(\tau_t^+, \tau_t^-)$.

Using the disintegration Lemma A.4 for the uniform measure on the unit sphere $\mathbb{S}^{d-1}$, we have

where we have used in the last equality the fact that the integrand does not depend on $z^H$ or $z^\perp$ and that $\tilde{\omega}_{d_H}$ and $\tilde{\omega}_{d_\perp}$ are probability measures (and thus their total mass is $1$).

Simplifying $G_t$. Using the disintegration Lemma A.4, we have:

Similarly to what we did for simplifying $f(\mu_t; \cdot)$, we can simplify the integrals against $\tilde{\omega}_{d_H}$ and $\tilde{\omega}_{d_\perp}$ using spherical harmonics theory to get:

This shows that

which proves that Eq (4.5) indeed describes the evolution of the Wasserstein-Fisher-Rao for the objective functional $A$ over $\mathcal{M}([0, \pi/2])^2$, given by the pair $(\tau_t^+, \tau_t^-)$.

F.   Numerical simulations in one dimension

Measure discretization. Discretizing $\mu_t$ via $\mu_{m, t} = \frac{1}{m} \sum_{j = 1}^m \delta_{(a_j(t), b_j(t))}$, we get that $\tau_{m, t} : = \tau_{m, t}^+ - \tau_{m, t}^- = \frac{1}{m} \sum_{j = 1}^m c_j(t) \delta_{\theta(t)}$ where

Initializing through $a_j(0) \sim \mathcal{U}\{-1, +1 \}$ and $b_j(0) \sim \tilde{\omega}_d = \mathcal{U}(\mathbb{S}^{d-1})$, yields $c_j(0) \sim \mathcal{U}\{-1, +1 \}$ and $\theta_j(0) \sim \tilde{\gamma}$, i.i.d. over $j$. The gradient flows of Eq (4.5) translates into the following ODEs on $(c_j)_{j \in [1, m]}$ and $(\theta_j)_{j \in [1, m]}$:

where $\varepsilon_j = a_j(0) \in \{-1, +1\}$ denotes whether the corresponding quantity appears in $\tau_{t, m}^+$ ($\varepsilon = +1$) or $\tau_{t, m}^-$ ($\varepsilon = -1$).

Time discretization. Simulating these ODEs via the discrete Euler scheme with step $\eta > 0$, leads, for any iteration $k \in \mathbb{N}$, to:

Approximating integrals numerically. The only thing that needs to be dealt with numerically is estimating the values of $G_t$ and $V_t$ which are defined by integrals. With the discretization of the measures, we have:

We thus get:

with

Similarly, we have:

with

We use Monte-Carlo estimation through sampling to approximate the integrals against the five variables $(\varphi, r, r^\prime, s, s^\prime)$ by drawing $N$ samples from the corresponding distributions. We get:

and similarly

where we have drawn the samples i.i.d. over $i \in [1, N]$:

Iterations in the numerical simulation. Defining the vectors $c(k) = (c_j(k))_{j \in [1, m]}$, $\theta(k) = (\theta_j)_{j \in [1, m]}$, and $\varepsilon = (\varepsilon_j)_{j \in [1, m]}$, the update Eq (13.3) can then be written in terms of update rules using the matrices $\Psi(k) = (\Psi_{ji}(k))_{j, i \in [1, m] \times [1, N]}$, $\tilde{\Psi}(k) = (\tilde{\Psi}_{ji}(k))_{j, i \in [1, m] \times [1, N]}$, and finally $\chi = (\chi_{ji}(k))_{j, i \in [1, m] \times [1, N]}$, and the vectors $(\Phi, R, R^\prime, S, S^\prime) = (\Phi_i, R_i, R^\prime_i, S_i, S^\prime_i)_{i \in [1, N]}$, which are re-sampled at each iteration $k \in [0, K]$, where $K \in \mathbb{N}$:

where

and $\odot$ denotes the Hadamard (element-wise) product of two vectors. One can compute the loss through sampling in a similar way.

Experimental value for $\alpha$ and parameters of the numerical simulation. For the numerical simulations, we fix the number of atoms of the measure (or equivalently the width of the network) to $m = 1024$, the learning rate to $\eta = 5.10^{-3}$, the number of samples for the Monte-Carlo scheme to $N = 1000$, and the total number of iterations to $K = 20, 000$. The experimental value for $\alpha$ (see Section 4.2) is computed through $\alpha_{\text{exp}} = \tau_{m, K}^+([0, \pi/2])$, that is

As mentioned in the main text, the behaviour of the numerical simulation depends a lot on the step-size $\eta$. Some of the differences between our observations and our intuitive description of the limiting model (infinite-width and continuous time) can come from too big a step-size. We have thus run the numerical simulation with $\eta = 2.10^{-5}$ as well, for $K = 230, 000$ steps but the same differences still appear (e.g., $\tau_{m, k}^+([0, \pi/2])$ still grows larger than the theoretically expected limit $\alpha$ after some time, albeit by a smaller margin) and after the critical $t^*$, some negative particles seem to go slightly beyond $\pi/2$, even with a very small step-size, a fact which cannot happen for the limiting model. Consequently, in Figure 3, the first histogram bin right after $\pi/2$ has been merged with the one before.