SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \underset {\xi \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ \overline {{P}} _ {k, d} (\langle u, \xi \rangle) \overline {{P}} _ {k ^ {\prime}, d} (\langle v, \xi \rangle) ] = \frac {1}{\sqrt {N _ {k , d} N _ {k ^ {\prime} , d}}} \sum_ {i = 1} ^ {N _ {k, d}} \sum_ {j = 1} ^ {N _ {k ^ {\prime}, d}} \underset {\xi \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ Y _ {k, i} (u) Y _ {k, i} (\xi) Y _ {k ^ {\prime}, j} (v) Y _ {k ^ {\prime}, j} (\xi) ] (C.18) \\ = \frac {\mathbf {1} [ k = k ^ {\prime} ]}{N _ {k , d}} \sum_ {j = 1} ^ {N _ {k, d}} \underset {\xi \sim \mathfrak {S} ^ {d - 1}} {\mathbb {E}} \left[ Y _ {k, j} (u) Y _ {k, j} (v) \right] (C.19) \\ = \mathbf {1} [ k = k ^ {\prime} ] P _ {k, d} (\langle u, v \rangle). (C.20) \\ \end{array}
$$
Lemma C.6. For any  $k \geq 0, d \geq 3$ , there exists a feature mapping  $\phi: \mathbb{S}^{d-1} \to \mathbb{R}^{N_{k,d}}$  such that for every  $u, v \in \mathbb{S}^{d-1}$ ,
$$
\langle \phi (u), \phi (v) \rangle = N _ {k, d} P _ {k, d} (\langle u, v \rangle), \tag {C.21}
$$
$$
\left\| \phi (u) \right\| _ {2} ^ {2} = N _ {k, d}, \tag {C.22}
$$
$$
\underset {u \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ \phi (u) \phi (u) ^ {\top} ] = I. \tag {C.23}
$$
Proof. Let  $Y_{1}, \dots, Y_{N_{k,d}}: \mathbb{S}^{d-1} \to \mathbb{R}$  be an orthonormal basis for the degree- $k$  spherical harmonics  $\mathbb{Y}_{k,d}$ . We let  $\phi(x) = (Y_l(x))_{l \in [N_{k,d}]}$  be the feature mapping, and in the following we verify Eqs. (C.21)-(C.23). By the addition theorem (Theorem C.4), for every  $u, v \in \mathbb{S}^{d-1}$  we have
$$
\langle \phi (u), \phi (v) \rangle = \sum_ {j = 1} ^ {N _ {k, d}} Y _ {j} (u) Y _ {j} (v) = N _ {k, d} P _ {k, d} (\langle u, v \rangle), \tag {C.24}
$$
which proves Eq. (C.21). Note that for every  $u \in \mathbb{S}^{d-1}$ ,  $P_{k,d}(\langle u, u \rangle) = P_{k,d}(1) = 1$ . Hence Eq. (C.22) follows directly. To prove Eq. (C.23), note that  $Y_1, \dots, Y_{N_{k,d}}$  are orthonormal. Therefore for every  $i, j \in [N_{k,d}]$
$$
\underset {u \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ Y _ {i} (u) Y _ {j} (u) \right] = \mathbf {1} [ i = j ]. \tag {C.25}
$$
Hence, Eq. (C.23) follows directly.
Lemma C.7. Let  $\mu_d$  be the distribution of  $\langle u, v \rangle$  when  $v$  is drawn uniformly at random from the unit sphere  $\mathbb{S}^{d-1}$  and  $u \in \mathbb{S}^{d-1}$  is a fixed vector. For fixed  $n \geq 1$ , let  $x_1, \dots, x_n \in \mathbb{S}^{d-1}$  be (not necessarily independent) random variables drawn from the distribution  $\mu_d$ . Then there exists a universal constant  $c > 0$  such that for any integer  $k \geq 1$ ,
$$
\forall t \geq 4 ^ {1 + k} (k \ln d) ^ {k}, \quad \Pr \left(\max  _ {i \leq n} d ^ {k} P _ {k, d} \left(x _ {i}\right) ^ {2} \geq t\right) \leq 2 n \exp \left(- \frac {t ^ {1 / k}}{3 2}\right). \tag {C.26}
$$
In addition, we also have
$$
\mathbb {E} \left[ \max  _ {i \leq n} d ^ {k} P _ {k, d} \left(x _ {i}\right) ^ {2} \right] \lesssim (3 2 k \ln (d n)) ^ {k}. \tag {C.27}
$$
Proof. To prove the first part of this lemma, we first invoke Proposition C.8, which states that
$$
\forall x \in \left[ \sqrt {\frac {k \ln d}{d}}, 1 \right], \quad | P _ {k, d} (x) | \leq 2 ^ {1 + k} x ^ {k}. \tag {C.28}
$$
When  $x\sim \mu_d$  , Proposition I.5 states that
$$
\forall t > 0, \quad \Pr (| x | \geq t) \leq 2 \exp \left(- \frac {t ^ {2} d}{2}\right). \tag {C.29}
$$
Therefore for every  $t \geq 4^{1 + k}(k \ln d)^k$  we have
$$
\begin{array}{l} \Pr \left(P _ {k, d} (x) ^ {2} d ^ {k} \geq t\right) \leq \Pr \left(| x | \geq t ^ {1 / 2 k} d ^ {- 1 / 2} 2 ^ {- 1 - 1 / k}\right) (C.30) \\ \leq 2 \exp \left(- \frac {1}{2} t ^ {1 / k} 4 ^ {- 1 - 1 / k}\right) (C.31) \\ \leq 2 \exp \left(- t ^ {1 / k} / 3 2\right). (C.32) \\ \end{array}
$$
By union bound we get
$$
\forall t \geq 4 ^ {1 + k} (k \ln d) ^ {k}, \quad \Pr \left(\max  _ {i \leq n} d ^ {k} P _ {k, d} \left(x _ {i}\right) ^ {2} \geq t\right) \leq 2 n \exp \left(- \frac {t ^ {1 / k}}{3 2}\right). \tag {C.33}
$$
which proves the first part of this lemma.
As a corollary, for any fixed  $\epsilon > 4^{1 + k}(k\ln d)^k$  we have
$$
\begin{array}{l} \mathbb {E} \left[ \max  _ {i \leq n} d ^ {k} P _ {k, d} \left(x _ {i}\right) ^ {2} \right] \leq \epsilon + \int_ {\epsilon} ^ {\infty} \Pr \left(\max  _ {i \leq n} d ^ {k} P _ {k, d} \left(x _ {i}\right) ^ {2} \geq t\right) \mathrm {d} t (C.34) \\ \leq \epsilon + n \int_ {\epsilon} ^ {\infty} \Pr \left(d ^ {k} P _ {k, d} (t) ^ {2} \geq t\right) \mathrm {d} t (C.35) \\ \leq \epsilon + 2 n \int_ {\epsilon} ^ {\infty} \exp \left(- \frac {t ^ {1 / k}}{3 2}\right) d t. (C.36) \\ \end{array}
$$
Now we upper bound the integral by changing variables. Letting  $u = t^{1 / k} / 32$ , we get
$$
\begin{array}{l} \int_ {\epsilon} ^ {\infty} \exp \left(- \frac {t ^ {1 / k}}{3 2}\right) \mathrm {d} t \leq \int_ {\frac {\epsilon 1 / k}{3 2}} ^ {\infty} 3 2 ^ {k} k u ^ {k - 1} \exp (- u) \mathrm {d} u (C.37) \\ \leq k ^ {2} \exp \left(- \frac {\epsilon^ {1 / k}}{3 2}\right) 3 2 ^ {k} \left(\frac {\epsilon^ {1 / k}}{3 2}\right) ^ {k - 1} (C.38) \\ \leq 3 2 k ^ {2} \exp \left(- \frac {\epsilon^ {1 / k}}{3 2}\right) \epsilon (C.40) \\ \end{array}
$$
where the last inequality comes from an upper bound for the incomplete gamma functions [31, Theorem 4.4.3]. Therefore we get
$$
\mathbb {E} \left[ \max  _ {i \leq n} d ^ {k} P _ {k, d} \left(x _ {i}\right) ^ {2} \right] \leq \epsilon + 6 4 n \epsilon k ^ {2} \exp \left(- \frac {\epsilon^ {1 / k}}{3 2}\right). \tag {C.41}
$$
Finally taking  $\epsilon = (32k\ln (dn))^k$ , we prove the desired result.
Proposition C.8. For any  $k\geq 0,d\geq 2$  we have
$$
\forall t \in \left[ \sqrt {\frac {k \ln d}{d}}, 1 \right], \quad | P _ {k, d} (t) | \leq 2 ^ {1 + k} t ^ {k}. \tag {C.42}
$$
Proof. Let  $\mu_d(t) \triangleq (1 - t^2)^{\frac{d-3}{2}} \frac{\Gamma(d/2)}{\Gamma((d-1)/2)} \frac{1}{\sqrt{\pi}}$  be the density of  $u_1$  when  $u = (u_1, \dots, u_d)$  is drawn uniformly from  $\mathbb{S}^{d-1}$ . By Atkinson and Han [12, Theorem 2.24] we have
$$
P _ {k, d} (t) = \underset {s \sim \mu_ {d - 1}} {\mathbb {E}} \left[ (t + i (1 - t ^ {2}) ^ {1 / 2} s) ^ {k} \right]. \tag {C.43}
$$
$$
\left| P _ {k, d} (t) \right| \leq \mathbb {E} _ {s \sim \mu_ {d - 1}} \left[ \left| t + i \left(1 - t ^ {2}\right) ^ {1 / 2} s \right| ^ {k} \right] \leq \mathbb {E} _ {s \sim \mu_ {d - 1}} \left[ \left(t ^ {2} + \left(1 - t ^ {2}\right) s ^ {2}\right) ^ {k / 2} \right]. \tag {C.44}
$$
By Proposition I.5 we have,
$$
\Pr (\sqrt {d - 1} | s | \leq 2 \sqrt {k \ln d}) \leq 2 \exp (- 2 k \ln d) \leq d ^ {- k}. \tag {C.45}
$$
$$
\begin{array}{l} \underset {s \sim \mu_ {d - 1}} {\mathbb {E}} \left[ \left(t ^ {2} + (1 - t ^ {2}) s ^ {2}\right) ^ {k / 2} \right] \leq d ^ {- k} + \underset {s \sim \mu_ {d - 1}} {\mathbb {E}} \left[ \mathbf {1} \left[ s \leq 2 \sqrt {k \ln d} / \sqrt {d - 1} \right] \left(t ^ {2} + s ^ {2}\right) ^ {k / 2} \right] (C.46) \\ \leq d ^ {- k} + \left(t ^ {2} + \frac {2 k \ln d}{d - 1}\right) ^ {k / 2}. (C.47) \\ \end{array}
$$
Hence, when  $t \geq \sqrt{\frac{k \ln d}{d}}$  we get
$$
\left| P _ {k, d} (t) \right| \leq d ^ {- k} + \left(t ^ {2} + \frac {2 k \ln d}{d - 1}\right) ^ {k / 2} \leq 2 ^ {1 + k} t ^ {k}. \tag {C.48}
$$
D Proofs for Population Dynamics
D.1 Proof of Lemma 4.2
In the following, we prove Lemma 4.2, which provides a way to simplify  $L(\rho)$  when  $\rho$  is rotationally invariant and symmetric.
Proof of Lemma 4.2. Recall that the population loss is
$$
L (\rho) = \frac {1}{2} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho} (x) - h (\langle x, e _ {1} \rangle)) ^ {2} ], \tag {D.1}
$$
and the activation function  $\sigma : \mathbb{R} \to \mathbb{R}$  and target function  $h: \mathbb{R} \to \mathbb{R}$  has the following decomposition:
$$
\sigma (t) = \sum_ {k = 0} ^ {\infty} \hat {\sigma} _ {k, d} \bar {P} _ {k, d} (t), \tag {D.2}
$$
$$
h (t) = \sum_ {k = 0} ^ {\infty} \hat {h} _ {k, d} \bar {P} _ {k, d} (t). \tag {D.3}
$$
$$
\begin{array}{l} 2 L (\rho) = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (x) - h \left(\langle x, e _ {1} \rangle\right)\right) ^ {2} \right] (D.4) \\ = \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} \left(\mathbb {E} _ {u \sim \rho} [ \sigma (\langle u, x \rangle) ] - h (\langle x, e _ {1} \rangle)\right) ^ {2} (D.5) \\ = \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} \left(\mathbb {E} _ {u \sim \rho} \left[ \sum_ {k = 0} ^ {\infty} \hat {\sigma} _ {k, d} \overline {{P _ {k , d}}} (\langle u, x \rangle) \right] - \sum_ {k = 0} ^ {\infty} \hat {h} _ {k, d} \overline {{P}} _ {k, d} (\langle x, e _ {1} \rangle)\right) ^ {2} (D.6) \\ = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left(\sum_ {k = 0} ^ {\infty} \left(\hat {\sigma} _ {k, d} \underset {u \sim \rho} {\mathbb {E}} \left[ \bar {P} _ {k, d} (\langle u, x \rangle) \right] - \hat {h} _ {k, d} \bar {P} _ {k, d} (\langle x, e _ {1} \rangle)\right)\right) ^ {2}. (D.7) \\ \end{array}
$$
When the neural network  $\rho$  is rotationally invariant, we can simplify the above equation by invoking Lemma C.2, which states that
$$
\underset {u \sim \rho} {\mathbb {E}} [ \bar {P} _ {k, d} (\langle u, x \rangle) ] = \underset {u \sim \rho} {\mathbb {E}} [ P _ {k, d} (\langle e _ {1}, u \rangle) ] \bar {P} _ {k, d} (\langle e _ {1}, x \rangle). \tag {D.8}
$$
Continuing Eq. (D.4) above we have,
$$
\begin{array}{l} 2 L (\rho) = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left(\sum_ {k = 0} ^ {\infty} \hat {\sigma} _ {k, d} \underset {u \sim \rho} {\mathbb {E}} [ P _ {k, d} (\langle e _ {1}, u \rangle) ] \overline {{P}} _ {k, d} (\langle e _ {1}, x \rangle) - \hat {h} _ {k, d} \overline {{P}} _ {k, d} (\langle e _ {1}, x \rangle)\right) ^ {2} (D.9) \\ = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left(\sum_ {k = 0} ^ {\infty} \left(\hat {\sigma} _ {k, d} \underset {u \sim \rho} {\mathbb {E}} \left[ P _ {k, d} \left(\langle e _ {1}, u \rangle\right) \right] - \hat {h} _ {k, d}\right) \bar {P} _ {k, d} \left(\langle e _ {1}, x \rangle\right)\right) ^ {2} (D.10) \\ \end{array}
$$
Now we expand the square by invoking Lemma C.5. In particular, Lemma C.5 states that
$$
\underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \bar {P} _ {k, d} \left(\langle e _ {1}, x \rangle\right) \bar {P} _ {k ^ {\prime}, d} \left(\langle e _ {1}, x \rangle\right) \right] = \mathbf {1} [ k = k ^ {\prime} ] P _ {k, d} \left(\langle e _ {1}, e _ {1} \rangle\right) = \mathbf {1} [ k = k ^ {\prime} ]. \tag {D.11}
$$
$$
2 L (\rho) = \sum_ {k = 0} ^ {\infty} \left(\hat {\sigma} _ {k, d} \underset {u \sim \rho} {\mathbb {E}} \left[ P _ {k, d} \left(\langle e _ {1}, u \rangle\right) \right] - \hat {h} _ {k, d}\right) ^ {2}, \tag {D.12}
$$
To prove the second part of this lemma, in the following, we assume that (1)  $\hat{h}_{k,d} = 0, \forall k \notin \{0,2,4\}$ , (2)  $\hat{h}_{0,d} = \hat{\sigma}_{0,d}$ , and (3) that  $\sigma$  is a degree-4 polynomial. Then we get
$$
L (\rho) = \frac {\hat {\sigma} _ {2 , d} ^ {2}}{2} \left(\underset {u \sim \rho} {\mathbb {E}} [ P _ {2, d} (w) ] - \gamma_ {2}\right) ^ {2} + \frac {\hat {\sigma} _ {4 , d} ^ {2}}{2} \left(\underset {u \sim \rho} {\mathbb {E}} [ P _ {4, d} (w) ] - \gamma_ {4}\right) ^ {2} + \sum_ {k \in \{1, 3 \}} \frac {\hat {\sigma} _ {k , d} ^ {2}}{2} \left(\underset {u \sim \rho} {\mathbb {E}} [ P _ {k, d} (w) ]\right) ^ {2}. \tag {D.13}
$$
Since  $\rho$  is symmetric and  $P_{k,d}(w)$  is an odd function when  $k$  is odd, we get
$$
\forall k \in \{1, 3 \}, \quad \mathbb {E} _ {u \sim \rho} [ P _ {k, d} (w) ] = 0. \tag {D.14}
$$
It follows directly that
$$
L (\rho) = \frac {\hat {\sigma} _ {2 , d} ^ {2}}{2} \left(\underset {u \sim \rho} {\mathbb {E}} [ P _ {2, d} (w) ] - \gamma_ {2}\right) ^ {2} + \frac {\hat {\sigma} _ {4 , d} ^ {2}}{2} \left(\underset {u \sim \rho} {\mathbb {E}} [ P _ {4, d} (w) ] - \gamma_ {4}\right) ^ {2} \tag {D.15}
$$
D.2 Proof of Lemma 4.3 and Lemma 4.4
The goal of this subsection is to show Lemma 4.3, which states that  $\rho_{t}$  remains symmetric and rotationally invariant under the infinite-width population dynamics, and to show Lemma 4.4 which describes the dynamics of the first coordinates of the particles in  $\rho_{t}$ .
The following lemma is a first step. Recall that when  $\rho$  is rotationally invariant, then Lemma 4.2 allows us to rewrite  $L(\rho) = \frac{1}{2}\mathbb{E}_{x\sim \mathbb{S}^{d - 1}}\left(f_{\rho}(x) - y\right)^{2}$  as  $\frac{1}{2}\sum_{k = 0}^{\infty}\left(\hat{\sigma}_{k,d}\mathbb{E}_{u\sim \rho}[P_{k,d}(w)] - \hat{h}_{k,d}\right)^{2}$ . The following lemma states that the gradient from the two formulas for  $L(\rho)$  are the same when  $\rho$  is rotationally invariant.
Lemma D.1. Let  $\rho$  be a distribution on  $\mathbb{S}^{d - 1}$  and let  $\nu$  be a distribution on  $[-1,1]$ . Define
$$
F (\nu) = \frac {1}{2} \sum_ {k = 0} ^ {\infty} \left(\hat {\sigma} _ {k, d} \underset {w \sim \nu} {\mathbb {E}} [ P _ {k, d} (w) ] - \hat {h} _ {k, d}\right) ^ {2}. \tag {D.16}
$$
If  $\rho$  is rotationally invariant as in Definition 4.1 and  $\nu$  is the marginal distribution of  $w$  for  $u = (w,z) \sim \rho$ , then for any particle  $u = (w,z)$ , we have
$$
(1 - w ^ {2}) \nabla_ {w} F (\nu) = \operatorname {g r a d} _ {w} L (\rho) \tag {D.17}
$$
where  $\operatorname{grad}_w L(\rho)$  denotes the first coordinate of  $\operatorname{grad}_u L(\rho)$ .
Proof of Lemma D.1. For any particle  $u = (w,z)$ , we write the first coordinate of  $\nabla_u L(\rho)$  as  $\nabla_w L(\rho)$ , and the vector consisting of the last  $(d - 1)$  coordinates as  $\nabla_z L(\rho)$ . Observe that
$$
\begin{array}{l} \operatorname {g r a d} _ {w} L (\rho) = \left\langle e _ {1}, \operatorname {g r a d} _ {u} L (\rho) \right\rangle (D.18) \\ = \left\langle e _ {1}, (I - u u ^ {\top}) \nabla_ {u} L (\rho) \right\rangle (D.19) \\ = \nabla_ {w} L (\rho) - w \langle u, \nabla_ {u} L (\rho) \rangle (D.20) \\ = \left(1 - w ^ {2}\right) \nabla_ {w} L (\rho) - w \langle z, \nabla_ {z} L (\rho) \rangle (D.21) \\ \end{array}
$$
On the other hand, suppose  $\rho$  satisfies the rotational invariance property, and that  $\nu$  is the marginal distribution of the first coordinate under  $\rho$ . Then, we can write  $\rho$  in terms of  $\nu$  — for any  $w \in [-1, 1]$ , the distribution of  $z$  conditioned on  $w$  is  $\sqrt{1 - w^2}\mathbb{S}^{d-1}$ . Thus, since  $L(\rho) = F(\nu)$  by Lemma 4.2, we have by the chain rule that
$$
\begin{array}{l} \nabla_ {w} F (\nu) = \underset { \begin{array}{c} \xi \sim \mathbb {S} ^ {d - 2} \\ u = w e _ {1} + \sqrt {1 - w ^ {2}} \xi \end{array} } {\mathbb {E}} \left[ \left\langle \nabla_ {u} L (\rho), \frac {\partial u}{\partial w} \right\rangle \right] (D.22) \\ = \nabla_ {w} L (\rho) + \underset {\xi \sim \mathbb {S} ^ {d - 2}} {\mathbb {E}} \left[ \left\langle \nabla_ {z} L (\rho) \right| _ {z = \sqrt {1 - w ^ {2}} \xi}, \frac {\partial z}{\partial w} \right\rangle \Bigg ] (D.23) \\ = \nabla_ {w} L (\rho) + \underset {\xi \sim \mathbb {S} ^ {d - 2}} {\mathbb {E}} \left[ \left\langle \nabla_ {z} L (\rho) \right| _ {z = \sqrt {1 - w ^ {2}} \xi}, - \frac {w}{\sqrt {1 - w ^ {2}}} \xi \right\rangle \Bigg ] (D.24) \\ = \nabla_ {w} L (\rho) + \underset {\xi \sim \mathbb {S} ^ {d - 2}} {\mathbb {E}} \left[ \left\langle \nabla_ {z} L (\rho) \right| _ {z = \sqrt {1 - w ^ {2}} \xi}, - \frac {w}{1 - w ^ {2}} z \right\rangle \bigg ] (D.25) \\ = \nabla_ {w} L (\rho) - \frac {w}{1 - w ^ {2}} \underset {\xi \sim \mathbb {S} ^ {d - 2}} {\mathbb {E}} [ \langle \nabla_ {z} L (\rho), z \rangle ] (D.26) \\ \end{array}
$$
Moreover, for any fixed  $w \in [-1, 1]$ , the quantity  $\langle \nabla_z L(\rho), z \rangle$  does not depend on  $z$ , since by Lemma D.2, for any rotation matrix  $A$  and  $z' = Az$ , we have  $\langle \nabla_{z'} L(\rho), z' \rangle = \langle A \cdot \nabla_z L(\rho), Az \rangle = \langle \nabla_z L(\rho), z \rangle$ . Thus, for any  $w \in [-1, 1]$  and  $z \in \sqrt{1 - w^2} \mathbb{S}^{d-2}$ , we can write
$$
(1 - w ^ {2}) \nabla_ {w} F (\nu) = (1 - w ^ {2}) \nabla_ {w} L (\rho) - w \langle z, \nabla_ {z} L (\rho) \rangle = \operatorname {g r a d} _ {w} L (\rho) \tag {D.27}
$$
In the proof of the above lemma, we used the following fact which we now prove: when  $\rho$  is rotationally invariant,  $\nabla_u L(\rho)$  is rotationally equivariant as a function of  $u$ , i.e. if  $u$  is rotated by a rotation matrix  $A$  which only modifies the last  $(d - 1)$  coordinates, then  $\nabla_u L(\rho)$  is rotated by  $A$  as well.
Lemma D.2 (Gradient is Rotationally Equivariant). Let  $\rho$  be a rotationally invariant distribution as in Definition 4.1. Let  $A\in \mathbb{R}^{d\times d}$  be a rotation matrix with  $Ae_1 = e_1$ , i.e. A only modifies the last  $(d - 1)$  coordinates. Let  $u\in \mathbb{S}^{d - 1}$  and  $u^{\prime} = Au$ . Then,  $\nabla_{u^{\prime}}L(\rho) = A\cdot \nabla_{u}L(\rho)$ , and  $\mathrm{grad}_{u^{\prime}}L(\rho) = A\cdot \mathrm{grad}_{u}L(\rho)$ .
Proof. The proof is by a straightforward calculation:
$$
\begin{array}{l} \nabla_ {u ^ {\prime}} L (\rho) = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho} (x) - y (x)) \sigma^ {\prime} (\langle u ^ {\prime}, x \rangle) x ] (D.28) \\ = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (x) - h \left(\langle e _ {1}, x \rangle\right)\right) \sigma^ {\prime} \left(\langle u ^ {\prime}, x \rangle\right) x \right] (D.29) \\ = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (A x) - h \left(\langle e _ {1}, A x \rangle\right)\right) \sigma^ {\prime} \left(\langle u ^ {\prime}, A x \rangle\right) A x \right] \\ \end{array}
$$
(By rotational invariance of uniform distribution on  $\mathbb{S}^{d - 1}$ )
$$
\begin{array}{l} = A \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho} (A x) - h (\langle e _ {1}, A x \rangle)) \sigma^ {\prime} (\langle u ^ {\prime}, A x \rangle) x ] \tag {D.30} \\ = A \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (A x) - h \left(\langle e _ {1}, A x \rangle\right)\right) \sigma^ {\prime} \left(\langle A u, A x \rangle\right) x \right] \quad (\text {B y}) \\ = A \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho} (A x) - h (\langle e _ {1}, A x \rangle)) \sigma^ {\prime} (\langle u, x \rangle) x ] \\ \end{array}
$$
(B.c.  $A$  is a rotation matrix so  $A^{\top}A = I_{d\times d}$ )
By the rotational invariance of  $\rho$ , we have  $f_{\rho}(Ax) = f_{\rho}(x)$ , meaning
$$
\begin{array}{l} \nabla_ {u ^ {\prime}} L (\rho) = A \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho} (x) - h (\langle e _ {1}, A x \rangle)) \sigma^ {\prime} (\langle u, x \rangle) x ] \\ (\mathrm {B . c .} f _ {\rho} (A x) = f _ {\rho} (x) \text {b y r o t a t i o n a l i n v a r i a n c e o f} \rho) \\ = A \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (x) - h (\langle e _ {1}, x \rangle)\right) \sigma^ {\prime} (\langle u, x \rangle) x \right] \quad \left(\mathrm {B . c .} A e _ {1} = e _ {1} \text {a n d} A ^ {\top} A = I _ {d \times d}\right) \\ = A \cdot \nabla_ {u} L (\rho) \tag {D.31} \\ \end{array}
$$
This proves the first statement of the lemma. The second statement also follows from a similar calculation:
$$
\begin{array}{l} \operatorname {g r a d} _ {u ^ {\prime}} L (\rho) = \left(I - u ^ {\prime} \left(u ^ {\prime}\right) ^ {\top}\right) \cdot \nabla_ {u ^ {\prime}} L (\rho) (D.32) \\ = \left(I - A u u ^ {\top} A ^ {\top}\right) \cdot A \cdot \nabla_ {u} L (\rho) (D.33) \\ = A \cdot \nabla_ {u} L (\rho) - A u u ^ {\top} A ^ {\top} A \cdot \nabla_ {u} L (\rho) (D.34) \\ = A \cdot \nabla_ {u} L (\rho) - A u u ^ {\top} \cdot \nabla_ {u} L (\rho) \quad \left(\mathrm {B . c .} A ^ {\top} A = I _ {d \times d}\right) \\ = A \cdot \left(I - u u ^ {\top}\right) \cdot \nabla_ {u} L (\rho) (D.35) \\ = A \cdot \operatorname {g r a d} _ {u} L (\rho) (D.36) \\ \end{array}
$$
As a corollary, we obtain our first formula for the dynamics of the 1-dimensional particles:
Lemma D.3 (One-Dimensional Velocity). Suppose we are in the setting of Theorem 3.3 and that  $\rho_t$  is rotationally invariant. Then, for any particle  $u_t$  (where we omit the initialization  $\chi$  for convenience), if  $w_t$  is the first coordinate of  $u_t$ , then
$$
\frac {d w _ {t}}{d t} = - \left(1 - w _ {t} ^ {2}\right) \sum_ {k = 0} ^ {4} \left(\hat {\sigma} _ {k, d} \underset {u \sim \rho_ {t}} {\mathbb {E}} [ P _ {k, d} (w) ] - \hat {h} _ {k, d}\right) \cdot \hat {\sigma} _ {k, d} P _ {k, d} ^ {\prime} (w). \tag {D.37}
$$
Proof of Lemma D.3. This follows directly from Lemma D.1, and because the activation  $\sigma$  and the target  $h$  only have nonzero coefficients  $\hat{\sigma}_{k,d}$  and  $\hat{h}_{k,d}$  for  $k = 0,\dots ,4$ .
We now complete the proof of Lemma 4.3, i.e. that  $\rho_{t}$  remains rotationally invariant and symmetric if it is defined according to the population projected gradient flow as in Eq. (3.6) and Eq. (3.7). To prove that  $\rho_{t}$  remains symmetric, we make use of the above formula for the 1-dimensional dynamics.
Proof of Lemma 4.3. First, we show the rotational invariance property. Assume inductively that at a time  $t$ ,  $\rho_t$  satisfies the desired rotational invariance property — it holds at initialization since  $\rho_0$  is the uniform distribution on  $\mathbb{S}^{d-1}$ . Consider a particle  $u_t$  and let  $A \in \mathbb{R}^{d \times d}$  be a rotation matrix which only modifies the last  $(d-1)$  coordinates. Let  $u_t' = Au_t$  be another particle. Then, it follows from Lemma D.2 that
$$
\frac {d u _ {t} ^ {\prime}}{d t} = A \frac {d u _ {t}}{d t} \tag {D.38}
$$
Thus, for a fixed  $w \in [-1, 1]$ , if  $z \sim \sqrt{1 - w^2} \mathbb{S}^{d-2}$ , then the distribution of  $\frac{dz}{dt}$  is uniform on  $c\mathbb{S}^{d-2}$  for a fixed scalar  $c$  that depends on  $w$ . Therefore,  $\rho$  will remain rotationally invariant.
Next, we prove that  $\rho_{t}$  is symmetric. First observe that  $\rho_{t}$  is symmetric at initialization, since it is the uniform distribution on  $\mathbb{S}^{d - 1}$ . Now, suppose that  $\rho_{t}$  is symmetric at time  $t$ . The polynomial  $P_{k,d}$  is odd if  $k$  is odd and even if  $k$  is even — this can be seen from induction on Eq. (C.2) and Eq. (C.3). Thus,  $\mathbb{E}_{u\sim \rho_t}[P_{k,d}(w)] = 0$  for odd  $k$  by our assumption that  $\rho_{t}$  is symmetric. Additionally, recall that we defined  $h$  so that  $\hat{\sigma}_{0,d} = \hat{h}_{0,d}$  and  $\hat{h}_{k,d} = 0$  for odd  $k$ . Thus, by our definition of  $D_{2,t}$  and  $D_{4,t}$  in Section 4,
$$
\frac {d w _ {t}}{d t} = - \left(1 - w _ {t} ^ {2}\right) \left(\hat {\sigma} _ {2, d} ^ {2} D _ {2, t} P _ {2, d} ^ {\prime} \left(w _ {t}\right) + \hat {\sigma} _ {4, d} ^ {2} D _ {4, t} P _ {4, d} ^ {\prime} \left(w _ {t}\right)\right) \tag {D.39}
$$
by Lemma D.3. Since  $P_{2,d}$  and  $P_{4,d}$  are even polynomials, we know  $P_{2,d}'$  and  $P_{4,d}'$  are odd polynomials, and thus  $\frac{dw_t}{dt}$  is an odd polynomial in  $w_t$ . In other words, if  $w_t$  and  $w_t'$  are the first
coordinates of two particles  $u_{t}$  and  $u_{t}^{\prime}$  respectively, such that  $w_{t}^{\prime} = -w_{t}$ , then  $\frac{dw_t'}{dt} = -\frac{dw_t}{dt}$ . Thus, the distribution of  $\frac{dw_t}{dt}$  is symmetric, meaning that  $\rho_{t}$  will remain symmetric.
In the rest of the proofs for the infinite-width population dynamics, we make use of the symmetry property. In particular, many of the lemmas are stated for  $w > 0$  or  $\iota > 0$  — however, the generalizations to  $w < 0$  or  $\iota < 0$  also hold. We use this symmetry implicitly throughout the proofs in this section.
In Lemma 4.4, we are simply rewriting our formula for the 1-dimensional dynamics in a more convenient form. The following proof shows the detailed calculation.
Proof of Lemma 4.4. By Lemma D.3 and the fact that only the second and fourth order terms are nonzero (which follows from Lemma 4.3, the fact that  $P_{k,d}$  is an odd polynomial for odd  $k$  (Eq. (C.2) and Eq. (C.3)) and because  $\hat{h}_{1,d} = \hat{h}_{3,d} = 0$ ), the velocity of  $w$  is equal to the following: