SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- For some  $\bar{R}$ , there exists  $\mu^{*} \in \mathcal{P}$  such that  $\mathrm{KL}(\nu, \mu^{*}) \leq R$  and  $L(\mu^{*}) \leq \ell(0) - c_{0}$ .
- For any  $\lambda < c_0 / R$ , the regularized expected risk minimizer  $\mu_{[\lambda]} := \operatorname{argmin} L(\mu) + \lambda \mathrm{KL}(\nu, \mu)$  satisfies  $Yf_{\mu_{[\lambda]}}(X) \geq c_0$  almost surely.
Importantly, we can apply the same general analysis for different classification problems, as long as Assumption 2 is verified. The advantage of this generality is that we do not need to tailor our convergence proof for individual learning problems. Note that the convergence rate of MFLD is strongly affected by the values of  $\bar{R}$  and  $R$ ; therefore, it is crucial to establish this condition using the smallest possible values of  $\bar{R}$  and  $R$ , in order to obtain a tight bound of the classification error. As an illustrative example, we now show that the  $k$ -sparse parity estimation satisfies the above assumption.
Example:  $k$ -sparse parity estimation. In the  $k$ -sparse parity setting (Example 1), Assumption 2 is satisfied with constants specified in the following propositions. The proofs are given in Appendix A.
Proposition 3 ( $k$ -sparse parity). Under Assumption 1 and for  $\bar{R} = k$ , there exists  $\mu^{*} \in \mathcal{P}$  such that
$$
\mathrm {K L} (\nu , \mu^ {*}) \leq c _ {1} k \log (k) ^ {2} d (= R),
$$
and  $L(\mu^{*})\leq \ell (0) - c_{2}$ , where  $c_{1},c_{2} > 0$  are absolute constants.
Proposition 4. Under Assumption 1 and the settings of  $R$  and  $\bar{R}$  given in Proposition 3, if  $\lambda < c_2 / (2R)$ , then  $\mu_{[\lambda]}$  satisfies
$$
\max  \{\bar {L} (\mu_ {[ \lambda ]}), L (\mu_ {[ \lambda ]}) \} \leq \ell (0) - c _ {2} + \lambda R <   \ell (0) - \frac {c _ {2}}{2},
$$
and  $f_{\mu [\lambda ]}$  is a perfect classifier with margin  $c_{2}$ , i.e.,  $Yf_{\mu [\lambda ]}(X)\geq \frac{c_2}{2}$ .
In other words, Assumption 2 is achieved with  $R = O(k\log (k)^{2}d)$ ,  $\bar{R} = k$  and  $c_{0} = c_{2} / 2$ . By substituting these values of  $R$  and  $\bar{R}$  to our general results presented below, we can easily derive a bound for the classification error of the MFLD estimator.
4 Main Result: Annealing Procedure and Classification Error Bound
4.1 Annealing Procedure
The convergence rate of MFLD is heavily dependent on the LSI constant which may impose a large computational cost. We alleviate this dependency by employing a novel annealing scheme where we gradually decrease the regularization parameter  $\lambda$ . In particular, at the  $\kappa$ -th round, we run the MFLD until (near) convergence with a regularization parameter  $\lambda^{(\kappa)} = 2^{-\kappa}\lambda^{(0)}$ :  $\mu_0 = \mu^{(\kappa -1)}$ ,
$$
\mathrm {d} X _ {t} = - \left[ \nabla \frac {\delta L (\mu_ {t})}{\delta \mu} (X _ {t}) \mathrm {d} t + 2 \lambda^ {(\kappa)} \lambda_ {1} X _ {t} \right] + \sqrt {2 \lambda^ {(\kappa)}} \mathrm {d} W _ {t}, \tag {6}
$$
which corresponds to minimizing  $\mathcal{L}^{(\kappa)}(\mu) \coloneqq L(\mu) + \lambda^{(\kappa)}\mathrm{KL}(\nu, \mu)$ . Then, we obtain a near optimal solution  $\mu^{(\kappa)}$  as  $\mathcal{L}^{(\kappa)}(\mu^{(\kappa)}) \leq \min_{\mu \in \mathcal{P}}\mathcal{L}^{(\kappa)}(\mu) + \epsilon^*$  for a given  $\epsilon^* > 0$ . We terminate the procedure after  $K$  rounds and obtain  $\mu^{(K)}$  as the output.
Suppose that there exists  $\mu^{*}$  such that  $\mathrm{KL}(\nu, \mu^{*}) \leq R$  and  $L(\mu^{*}) \leq \delta^{*}$ . Then, as long as  $\lambda^{(\kappa)} \geq \delta^{*}$  and  $\epsilon^{*} < \delta^{*}$ , we have that
$$
L (\mu^ {(\kappa)}) + \lambda^ {(\kappa)} \mathrm {K L} (\nu , \mu^ {(\kappa)}) \leq L (\mu^ {*}) + \lambda^ {(\kappa)} \mathrm {K L} (\nu , \mu^ {*}) + \epsilon^ {*} \leq 2 \delta^ {*} + \lambda^ {(\kappa)} R \leq (R + 2) \lambda^ {(\kappa)}.
$$
Since  $\mathcal{L}^{(\kappa)}(\mu_t)$  is monotonically decreasing during the optimization, we always have  $L(\mu_t) \leq \mathcal{L}^{(\kappa)}(\mu^{(\kappa - 1)}) \leq \mathcal{L}^{(\kappa - 1)}(\mu^{(\kappa - 1)}) \leq (R + 2)\lambda^{(\kappa - 1)}$ . Now we utilize the following structure on common classification losses to ensure that  $\| \delta L(\mu) / \delta \mu \|_{\infty}$  is small when the loss  $L(\mu)$  is small.
Assumption 3. There exists  $c_{\mathrm{L}} > 0$  such that, for any  $\mu \in \mathcal{P}$ , it holds that  $\| \frac{\delta L(\mu)}{\delta \mu} \|_{\infty} \leq c_{\mathrm{L}} \bar{R} L(\mu)$ .
For example, the logistic loss satisfies this assumption:
$$
\left| \partial_ {u} \log (1 + \exp (- u)) \right| = \frac {\exp (- u)}{1 + \exp (- u)} \leq \log (1 + \exp (- u)).
$$
Hence, the condition holds for  $c_{\mathrm{L}} = 1$  because  $\| \frac{\delta L(\mu)}{\delta \mu} \|_{\infty} \leq \frac{1}{n} \sum_{i=1}^{n} |\ell'(y_i, f_\mu(z_i))| \| h_{z_i}(\cdot) \|_{\infty} \leq \frac{1}{n} \sum_{i=1}^{n} \bar{R}\ell(y_i, f_\mu(z_i)) = \bar{R}L(\mu)$ . Under this assumption, the Holley-Strook argument yields that the log-Sobolev constant during the optimization can be bounded as
$$
\alpha \geq \lambda_ {1} \exp \left(- \frac {4 c _ {\mathrm {L}} \bar {R} (R + 2) \lambda^ {(\kappa - 1)}}{\lambda^ {(\kappa)}}\right) \geq \lambda_ {1} \exp \left(- 8 c _ {\mathrm {L}} \bar {R} (R + 2)\right),
$$
which is independent of  $\lambda^{(\kappa)}$  and the final accuracy  $\delta^*$ . Hence, after  $K = \log_2(R / (\lambda^{(0)}\epsilon^*))$  round, we achieve that  $L(\mu^{(K)}) \leq \delta^* + 2\epsilon^*$  where each round takes  $T_{\kappa} = \log (1 / \epsilon^{*}) / (2\alpha \lambda) = O(\log (1 / \epsilon^{*})\exp (8c_{\mathrm{L}}\bar{R} (R + 2)) / (\lambda_1\lambda^{(k)}))$  for the continuous time setting.
For the discrete setting,  $T_{\kappa} = O\left(\left(\frac{C_1}{\alpha \epsilon^*} + \frac{\sqrt{C_1}}{\sqrt{\lambda^{(\kappa)} \alpha \epsilon^*}}\right) \log(1 / \epsilon^*) \exp(8c_{\mathrm{L}} \bar{R}(R + 2)) / (\lambda_1 \lambda^{(\kappa)})\right)$  iterations in each round is sufficient (where  $C_1$  is given for  $\lambda^{(K)}$ ), where the step size is  $\eta = O\left(C_1^{-1} \alpha \epsilon^* \wedge C_1^{-1/2} \sqrt{\lambda^{(\kappa)} \alpha \epsilon^*}\right)$ . As long as  $\epsilon^* = O(\lambda^{(\kappa)} / \alpha)$  for all  $\kappa$ , the total iteration number can be simplified as  $O\left(\log(1 / \epsilon^*) \exp(16c_{\mathrm{L}} \bar{R}(R + 2)) C_1 / (\lambda^{(K)} \epsilon^*)\right)$ , and the width  $N$  (number of particles) can be taken as  $N = O\left[\left(\lambda^{(K)} \alpha\right)^{-2} + (\lambda^{(K)})^{-2} + d\right] / (\epsilon^*)^2 = O\left(\left[\exp(16c_{\mathrm{L}} \bar{R}(R + 2)) / (\lambda^{(K)})^2 + d\right] / (\epsilon^*)^2\right)$ .
Remark 1. We make the following remarks on the annealing procedure.
- The main advantage of this annealing approach is that the exponential factor induced by the LSI constant  $\alpha$  is not dependent on the choice of the regularization parameter  $\lambda^{(k)}$  (note that the LSI is solely determined by the intermediate solution  $\mu_t$ ). In contrast, the naive Holley-Stroock argument of Eq. (5) imposes exponential dependency on the regularization parameter.
- Our annealed algorithm differs from the procedure considered in Chizat (2022, Section 4); importantly, we make use of the structure of classification loss functions to obtain a refined computational complexity analysis.
4.2 Generalization Error Analysis
We utilize the local Rademacher complexity (Mendelson, 2002; Bartlett et al., 2005; Koltchinskii, 2006; Giné and Koltchinskii, 2006) to obtain a faster generalization error rate. For the function class of mean-field neural networks, we introduce  $\mathcal{F} \coloneqq \{f_{\mu} \mid \mu \in \mathcal{P}\}$ , and the KL-constrained model class  $\mathcal{F}_M(\mu^\circ) \coloneqq \{f_\mu \mid \mu \in \mathcal{P}, \mathrm{KL}(\mu^\circ, \mu) \leq M\}$  for  $\mu^\circ \in \mathcal{P}$  and  $M > 0$ . The Rademacher complexity of a function class  $\tilde{\mathcal{F}}$  is defined as
$$
\mathrm {R a d} (\tilde {\mathcal {F}}) := \mathbb {E} _ {\varepsilon_ {i}, z _ {i}} \left[ \sup _ {f \in \tilde {\mathcal {F}}} \frac {1}{n} \sum_ {i = 1} ^ {n} \epsilon_ {i} f (z _ {i}) \right],
$$
where  $(z_{i})_{i = 1}^{n}$  are i.i.d. observations from  $P_Z$  and  $(\varepsilon_i)_i^n$  is an i.i.d. Rademacher sequence  $(P(\varepsilon_{i} = 1) = P(\varepsilon_{i} = -1) = 1 / 2))$ . We have the following bound on the Rademacher complexity of the function class  $\mathcal{F}_M(\mu^\circ)$ .
Lemma 1 (Local Rademacher complexity of  $\mathcal{F}_M(\mu^\circ)$ , Chen et al. (2020) adapted). For any fixed  $\mu^\circ \in \mathcal{P}$  and  $M > 0$ , it holds that  $\mathrm{Rad}(\mathcal{F}_M(\mu^\circ)) \leq 2\bar{R}\sqrt{\frac{M}{n}}$ .
The proof is given in Appendix B.2 in the supplementary material. Combining this local Rademacher complexity bound with the peeling device argument (van de Geer, 2000), we can roughly obtain the following estimate (note that this is an informal derivation):
$$
\underbrace {\bar {L} (\hat {\mu}) - \bar {L} \left(\mu^ {*}\right) - \left(\hat {\mu} - \mu^ {*}\right) \frac {\delta \bar {L} \left(\mu^ {*}\right)}{\delta \mu}} _ {\text {(I I)}} + \underbrace {\lambda \mathrm {K L} \left(\mu^ {*} , \hat {\mu}\right)} _ {\text {(I)}} \lesssim \sqrt {\frac {\mathrm {K L} \left(\mu^ {*} , \hat {\mu}\right)}{n}} \lesssim \frac {1}{n \lambda} + \lambda \mathrm {K L} \left(\mu^ {*}, \hat {\mu}\right), \tag {7}
$$
with high probability, where  $\hat{\mu} = \mathrm{argmin}_{\mu \in \mathcal{P}}\mathcal{L}(\mu)$ ,  $\mu^{*} = \mathrm{argmin}_{\mu \in \mathcal{P}}\bar{L} (\mu) + \lambda \mathrm{KL}(\nu ,\mu)$ ,  $R$  and  $\bar{R}$  are regarded as constants, and the last inequality is by the AM-GM relation. Observe that on the left hand side, we have two non-negative terms (I) and (II). Corresponding to each term, we obtain different types of classification error bounds (Type I and Type II in the following subsection, respectively). Note that there appears a  $O(1 / n)$  factor in the right hand side, which cannot be
obtained by a vanilla Rademacher complexity evaluation because it only yields an  $O(1 / \sqrt{n})$  bound. In other words, localization is essential to obtain our fast convergence rate.
We remark that a local Rademacher complexity technique is also utilized by Telgarsky (2023) to derive a  $O(1/n)$  rate. They adapted a technique developed for a smooth loss function by Srebro et al. (2010), which requires the training loss  $L(\hat{\mu})$  to be sufficiently small, that is, of order  $L(\hat{\mu}) = O(1/n)$ . In our setting, to achieve such a small training loss, we need to take large  $\bar{R}$  such as  $\bar{R} = \Omega(\log(n))$ . Unfortunately, such a large  $\bar{R}$  induces exponentially small log-Sobolev constant like  $\alpha \lesssim \exp(-cd\log(n)) = n^{-cd}$ . In contrast, our analysis focuses on the local Rademacher complexity around  $\mu^{*}$ , and hence we do not require the training loss to be close to 0; instead, it suffices to have a training loss that is close to or smaller than that of  $\mu^{*}$ .
4.2.1 Type I: Perfect Classification with Exponentially Decaying Error
In the regime of  $n = \Omega (1 / \lambda^{(K)2})$ , we can prove that the MFLD estimator attains a perfect classification with an exponentially converging probability by evaluating the term (I). From Eq. (7) we can establish  $\mathrm{KL}(\mu^*,\hat{\mu})\leq O_p(1 / (n\lambda^{(K)2}))$ ; this KL divergence bound can be used to control the  $L^{\infty}$ -norm between  $f_{\hat{\mu}}$  and  $f_{\mu^{*}}$ . Indeed, we can show that  $\| f_{\hat{\mu}} - f_{\mu^{*}}\|_{\infty}^{2}\leq 2\bar{R}^{2}\mathrm{KL}(\mu^{*},\hat{\mu})$  (see the proof of Theorem 1). Then, under the margin assumption of  $f_{\mu^{*}}$  (Assumption 2), we have that  $f_{\hat{\mu}}$  also yields a Bayes optimal classifier. More precisely, we have the following theorem.
Theorem 1. Suppose Assumptions 1 and 2 hold. Let  $M_0 = (\epsilon^* + 2(\bar{R} + 1)) / \lambda^{(K)}$ . Moreover, suppose that  $\lambda^{(K)} < c_0 / R$  and
$$
Q := c _ {0} ^ {2} - \frac {4 \bar {R} ^ {2}}{n \lambda^ {(K) 2}} \left[ \lambda^ {(K)} \left(4 \bar {R} + \frac {\lambda^ {(K)}}{3 2 \bar {R} ^ {2} n}\right) + 8 \bar {R} ^ {2} \left(4 + \log \log_ {2} \left(8 n ^ {2} M _ {0} \bar {R}\right)\right) + n \lambda^ {(K)} \epsilon^ {*} \right] > 0,
$$
then  $f_{\hat{\mu}}$  yields perfect classification, i.e.,  $P(Yf_{\hat{\mu}}(Z) > 0) = 1$ , with probability  $1 - \exp \left(-\frac{n\lambda^{(K)2}}{32R^4} Q\right)$ .
The proof is given in Appendix C.1 in the supplementary material. From Proposition 3 we see that the requirement  $n \asymp 1 / \lambda^{(K)2}$  implies that in the  $n \asymp d^2$  regime, Theorem 1 gives a perfect classification guarantee with a failure probability decaying exponentially fast.
4.2.2 Type II: Polynomial Order Classification Error
Next we evaluate the classification error bound from term (II) in Eq. (7). In this case, we do not require an  $L^{\infty}$ -norm bound as in the Type I analysis above; this results in a milder dependency on  $\lambda^{(K)}$  and hence a better sample complexity.
Theorem 2. Suppose Assumptions 1 and 2 hold. Let  $\lambda^{(K)} < c_0 / R$  and  $M_0 = (\epsilon^* + 2(\bar{R} + 1)) / \lambda^{(K)}$ . Then, with probability  $1 - \exp(-t)$ , the classification error of  $f_{\mu^{(K)}}$  is bounded as
$$
P (Y f _ {\mu^ {(K)}} (Z) \leq 0) \leq 2 \psi (c _ {0}) \left[ \frac {8 \bar {R} ^ {2}}{n \lambda^ {(K)}} \left(4 + t + \log \log_ {2} (8 n ^ {2} M _ {0} \bar {R})\right) + \frac {1}{n} \left(4 \bar {R} + \frac {\lambda^ {(K)}}{3 2 \bar {R} ^ {2} n}\right) + \epsilon^ {*} \right].
$$
The proof is given in Appendix C.2 in the supplementary material. We notice that the right hand side scales with  $O(1 / (n\lambda^{(K)}))$ , which is better than  $O(1 / (n\lambda^{(K)2}))$  in Theorem 1; this implies that a sample size linear in the dimensionality is sufficient to achieve small classification error. The reason for such improvement in the  $\lambda^K$ -dependence is that the stronger  $L^\infty$ -norm convergence is not used in the proof; instead, only the convergence of the loss is utilized. On the other hand, this analysis does not guarantee a perfect classification.
4.2.3 Computational Complexity of MFLD
From the general result in Section 4.1, we can evaluate the computational complexity to achieve the statistical bounds derived above. In both cases (Theorems 1 and 2), we may set the optimization error  $\epsilon^{*} = O(1 / (n\lambda^{(K)}))$ . Then, the total number of iteration can be
$$
\sum_ {\kappa = 1} ^ {K} T _ {\kappa} \leq O \left((d + \lambda^ {(K) - 1}) n \exp (1 6 c _ {\mathrm {L}} \bar {R} (R + 2)) \log (n \lambda^ {(K)})\right).
$$
The width  $N$  (the number of particles) can be taken as  $N = O((\epsilon^{*}\lambda^{(K)}\alpha)^{-2}) = O(n^{2}\alpha^{-2}) = O\left(n^{2}\exp (16c_{\mathrm{L}}\bar{R} (R + 2))\right)$ .
Corollary 1 ( $k$ -sparse parity setting). In the  $k$ -sparse parity setting, we may take  $R = O(k \log (k)^2 d)$ ,  $\bar{R} = k$  and  $\lambda^{(K)} = O(1 / R) = O(1 / (k \log (k)^2 d))$ . Therefore, the classification error is bounded by
$$
P \left(Y f _ {\mu^ {(K)}} <   0\right) \leq O \left(\frac {k ^ {2} \log (k) ^ {2} d}{n} (\log (1 / \delta) + \log \log (n))\right),
$$
with probability  $1 - \delta$ . Moreover, if  $n = \Omega(k^6 \log(k)^4 d^2)$ , then  $P(Y f_{\mu^{(K)}} > 0) = 1$  with probability
$$
1 - \exp (- \Omega (n k ^ {6} \log (k) ^ {4} / d ^ {2})).
$$
As for the computational complexity, we require  $O(k \log (k)^2 dn \log (nd) \exp [O(k^2 \log (k)^2 d)])$  iterations, and the number of particles is  $O(\exp (O(k^2 \log (k)^2 d)))$ .
Comparison with prior results. In the 2-sparse parity setting, neural network in the kernel (NTK) regime achieves only  $O(d^{2} / n)$  convergence in the classification error, as shown in Ji and Telgarsky (2019) and Telgarsky (2023, Theorem 2.3), whereas we demonstrate that mean-field neural network can improve the rate to  $O(d / n)$  via feature learning. Telgarsky (2023) also analyzed the learning of 2-sparse parity problem beyond the kernel regime, and showed that 2-layer ReLU neural network can achieve the best known classification error  $O(d / n)$ . However, their analysis considered a low-rotation dynamics and assumed convergence at  $t\to \infty$ , whereas our framework also provides a concrete estimate of the computational complexity. Indeed, the number of iterations can be bounded as  $O(dn\log (nd)\exp [O(d)])$ . In addition, while we still require exponential width  $N = O(n^{2}\exp (O(d)))$ , such a condition is an improvement over  $N = O(d^{d})$  in Telgarsky (2023).
Barak et al. (2022) considered a learning method in which one-step gradient descent is performed for the purpose of feature learning, and then a network with randomly re-initialized bias units is used to fit the data. For the  $k$ -sparse parity problem, they derived a classification error bound of  $O(d^{(k + 1) / 2} / \sqrt{n})$ . In contrast, our analysis yields a much better statistical complexity of  $O(k^2\log (k)^2 d / n\wedge \exp (-\Omega (nk^6\log (k)^4 /d^2)))$ , which "decouples" the degree  $k$  in the exponent of the dimension dependence.
We validate our theoretical results by numerical experiment on synthetic data. Specifically, we consider the classification of 2-sparse parity with varying dimensionality  $d$  and sample size  $n$ .
Recall that the samples  $\{(z_i, y_i)\}_{i=1}^n$  are independently generated so that  $z_i$  follows the uniform distribution on  $\{\pm 1/\sqrt{d}\}^d$  and  $y_i = d\zeta_{i,1}\zeta_{i,2} \in \{\pm 1\}$  ( $z_i = (\zeta_{i,1}, \dots, \zeta_{i,d})$ ). A finite-width approximation of the mean-field neural network  $\frac{1}{N}\sum_{j=1}^{N}h_{x_j}(z)$  is employed with the width  $N = 2,000$ . For each neuron of the network, all parameters are initialized to independently follow the standard normal distribution (meaning that the network is rotation invariant at initialization) and the scaling parameter  $\bar{R}$  is set to 15. We set  $d$  to take values 5, 10, ..., 150, and  $n = 50, 100, \dots, 2000$ . We trained the network using noisy gradient descent with  $\eta = 0.2$ ,  $\lambda_1 = 0.1$ , and  $\lambda = 0.1/d$  (fixed during the whole training) until  $T = 10,000$ . The logistic loss is used for the training objective.
Figure 1 shows the average test accuracy over five trials. We make the following observations.
- The red line corresponds to the sample size  $n = \Theta(d^2)$ , above which we observe that almost-perfect classification is achieved. According to Section 4.2 (Type I), the classification error gets exponentially small with respect to  $n / d^2$ , which predicts very small classification error above the line of  $n = \Theta(d^2)$ , which matches our experimental result.
- The boundary of test accuracy above  $50\%$  is almost linear, as indicated by the blue line. This matches the theoretical conditions (Type II) to obtain the polynomial order classification error in Section 4.2.
Figure 1: Test accuracy of two-layer neural network optimized by the MFLD to learn a  $d$ -dimensional 2-sparse parity (XOR) problem.
We provided a general framework to evaluate the classification error of a two-layer neural network trained by the mean-field Langevin dynamics. Thanks to the generality of our framework, an error bound for specific settings can be derived by directly specifying the parameters in Assumption 2 such as  $R$ ,  $\bar{R}$ , and  $c_{0}$ . We also proposed an annealing procedure to alleviate the exponential dependencies in the LSI constant. As a special (but important) example, we investigated the  $k$ -sparse parity problem, for which we obtained more general and better sample complexity than existing works.
A limitation of our approach is that the required width and number of iterations are exponential with respect to the dimensionality  $d$ ; in contrast, certain tailored algorithms for learning low-dimensional target functions such as Abbe et al. (2023); Chen and Meka (2020) do not require such exponential computation. An important open problem is whether we can reduce the width and number of iterations to  $\mathrm{poly}(d)$  for the vanilla noisy gradient descent algorithm. Another interesting direction is to investigate the interplay between structured data and the efficiency of feature learning, as done in Ghorbani et al. (2020); Refinetti et al. (2021); Ba et al. (2023); Mousavi-Hosseini et al. (2023).
The authors thank Matus Telgarsky for discussions and feedback. TS was partially supported by JSPS KAKENHI (20H00576) and JST CREST (JPMJCR2115, JPMJCR2015). KO was partially supported by JST, ACT-X Grant Number JPMJAX23C4, JAPAN. AN was partially supported by JSPS KAKENHI (22H03650).
E. Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302.11055, 2023.
J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=akddwRG6EGi.
J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, and D. Wu. Learning in the presence of low-dimensional structure: a spiked random matrix perspective. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
D. Bakry and M. Émery. Diffusions hypercontractives. In J. Azéma and M. Yor, editors, Séminaire de Probabilités XIX 1983/84, pages 177-206, Berlin, Heidelberg, 1985. Springer Berlin Heidelberg. ISBN 978-3-540-39397-9.
D. Bakry, I. Gentil, M. Ledoux, et al. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
B. Barak, B. L. Edelman, S. Goel, S. M. Kakade, eran malach, and C. Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=8XWP2ewX-im.
P. Bartlett, O. Bousquet, and S. Mendelson. Local Rademacher complexities. The Annals of Statistics, 33:1487-1537, 2005.
P. L. Bartlett, M. I. Jordan, and J. D. McAuliffe. Convexity, classification, and risk bounds. Journal of the American Statistical Association, 101(473):138-156, 2006.
A Proofs of Propositions 3 and 4
Proof of Proposition 3. Remember that
$$
h _ {x} (z) = \bar {R} \left[ \tanh  \left(z ^ {\top} x _ {1} + x _ {2}\right) + 2 \tanh  \left(x _ {3}\right) \right] / 3.
$$
Let  $b_{i} = 2i - k$  for  $i = 0, \dots, k$ , let  $\zeta > 0$  be the positive real such that  $\mathbb{E}_{u \sim N(0,1)}[2 \tanh(\zeta + u)] = 1$  (note that, this also yields  $\mathbb{E}_{u \sim N(0,1)}[2 \tanh(-\zeta + u)] = -1$  by the symmetric property of tanh and the Gaussian distribution), and let
$$
\xi = [ \underbrace {\sqrt {d} , \sqrt {d} , \ldots , \sqrt {d}} _ {k \text {- d i m e n s i o n}}, \underbrace {0 , \ldots , 0} _ {d - k \text {- d i m e n s i o n}} ] ^ {\top} \in \mathbb {R} ^ {d}.
$$
$$
\Sigma := \left( \begin{array}{c c c} I / (2 \lambda_ {1}) & 0 & 0 \\ 0 & 1 / (2 \lambda_ {1}) & 0 \\ 0 & 0 & 1 \end{array} \right) \in \mathbb {R} ^ {(d + 1 + 1) \times (d + 1 + 1)},
$$
and  $\rho > 1$  be a constant which will be adjusted later on. Then, for  $\xi_{2i} := [\log (\rho k)\xi^{\top}, -\log (\rho k)(b_i - 1), \zeta]^{\top} \in \mathbb{R}^{\bar{d}}$  and  $\xi_{2i+1} := -[\log (\rho k)\xi^{\top}, -\log (\rho k)(b_i + 1), \zeta]^{\top} \in \mathbb{R}^{\bar{d}}$  for  $i = 0, \dots, k$ , we define
$$
\hat {\mu} _ {2 i} := N \left(\xi_ {2 i}, \Sigma\right), \quad \hat {\mu} _ {2 i + 1} := N \left(\xi_ {2 i + 1}, \Sigma\right).
$$
Then, we can see that, for  $z \in \{\pm 1 / \sqrt{d}\}^d$ , it holds that