SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Taiji Suzuki $^{1,2}$ , Denny Wu $^{3,4}$ , Kazusato Oko $^{1,2}$ , Atsushi Nitanda $^{2,5}$
4 Flatiron Institute, 5 Kyushu Institute of Technology
taiji@mist.i.u-tokyo.ac.jp, dennywu@nyu.edu,
oko-kazusato@g.ecc.u-tokyo.ac.jp, nitanda@ai.kyutech.ac.jp
Neural network in the mean-field regime is known to be capable of feature learning, unlike the kernel (NTK) counterpart. Recent works have shown that mean-field neural networks can be globally optimized by a noisy gradient descent update termed the mean-field Langevin dynamics (MFLD). However, all existing guarantees for MFLD only considered the optimization efficiency, and it is unclear if this algorithm leads to improved generalization performance and sample complexity due to the presence of feature learning. To fill this important gap, in this work we study the sample complexity of MFLD in learning a class of binary classification problems. Unlike existing margin bounds for neural networks, we avoid the typical norm control by utilizing the perspective that MFLD optimizes the distribution of parameters rather than the parameter itself; this leads to an improved analysis of the sample complexity and convergence rate. We apply our general framework to the learning of  $k$ -sparse parity functions, where we prove that unlike kernel methods, two-layer neural networks optimized by MFLD achieves a sample complexity where the degree  $k$  is "decoupled" from the exponent in the dimension dependence.
Mean-field Langevin dynamics. The optimization dynamics of two-layer neural networks in the mean-field regime can be described by a nonlinear partial differential equation of the distribution of parameters (Nitanda and Suzuki, 2017; Chizat and Bach, 2018; Mei et al., 2018; Rotskoff and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020). Such a description has multiple advantages:  $(i)$  global convergence guarantees can be obtained by exploiting convexity of the loss function, and  $(ii)$  the parameters are allowed to evolve away from initialization and learn informative features, in contrast to the neural tangent kernel ("lazy") regime (Jacot et al., 2018).
Among the gradient-based optimization algorithms for mean-field neural networks, the mean-field Langevin dynamics (MFLD) (Mei et al., 2018; Hu et al., 2019) is particularly attractive due to the recently established quantitative optimization guarantees. MFLD arises from a noisy gradient descent update on the parameters, where Gaussian noise is injected to the gradient to encourage "exploration". It has been shown that MFLD globally optimizes an entropy-regularized convex functional in the space of measures, and for the infinite-width and continuous-time dynamics, the convergence rate is exponential under suitable isoperimetric conditions (Nitanda et al., 2022; Chizat, 2022). Furthermore, uniform-in-time estimates of the particle discretization error have also been established (Suzuki et al., 2023a; Chen et al., 2022), meaning that optimization guarantees for the infinite-dimensional problem can be effectively translated to a finite-width neural network.
However, existing analyses of MFLD only considered the optimization of neural networks; this alone does not demonstrate the benefit of mean-field regime nor the presence of feature learning. Therefore,
an important problem is to characterize the generalization of the learned models, and prove efficient sample complexity guarantees. The goal of this work is to address the following question.
Can we show that neural network + simple noisy gradient descent (MFLD) efficiently learns an interesting class of functions, with a better rate of convergence compared to the "lazy" regime?
Learning sparse parity functions. One particularly relevant learning task is the  $k$ -sparse parity problem, where the response  $y$  is given by the sign of the product of  $k$  coordinates of the input (on hypercube); as a special case, setting  $k = 2$  recovers the classical XOR problem. When  $k \ll d$ , this target function is low-dimensional, and hence we expect feature learning to be beneficial in that it can "zoom in" to relevant subspace. In contrast, for kernel methods (including neural networks in the lazy regime) which cannot adapt to such structure, it has been shown that a sample complexity of  $n = \Omega(d^k)$  is unavoidable (Ghorbani et al., 2019; Hsu; Abbe et al., 2022).
For the XOR case ( $k = 2$ ), recent works have shown that neural networks in the mean-field regime can achieve a sample complexity of  $n = \mathcal{O}(d / \epsilon)$  (Wei et al., 2019; Chizat and Bach, 2020; Telgarsky, 2023), which indeed improves upon the NTK complexity (Ji and Telgarsky, 2019). However, all these results directly assumed convergence of the dynamics ( $t \to \infty$ ) with no iteration complexity. Moreover, Wei et al. (2019); Chizat and Bach (2020) directly analyzed the infinite-width limit, and while Telgarsky (2023) provided a finite-width characterization, the dynamics is restricted to the low-rotation regime, and a very large number of particles  $N = \mathcal{O}(d^d)$  is required. Lastly, these analyses are specialized to XOR, and do not directly generalize to the  $k$ -parity setting.
In this work, we bridge the aforementioned gap by presenting a simple and general framework to establish sample complexity of MFLD in learning binary classification problems. We then apply this framework to the sparse  $k$ -parity problem, and obtain improved rate of convergence for the fully time-and space-discretized algorithm. More specifically, our contributions can be summarized as follows.
- We present a general framework to analyze MFLD in the learning of binary classification tasks. Our framework has two main ingredients:  $(i)$  an annealing procedure that applies to common classification losses that removes the exponential dependence on regularization parameters in the logarithmic Sobolev inequality, and  $(ii)$  a novel local Rademacher complexity analysis for the distribution of parameters optimized by MFLD. As a result, we can obtain generalization guarantee for the learned neural network in discrete-time and finite-width settings.
- We apply our general framework to the  $k$ -sparse parity problem, and derived learning guarantees with improved rate of convergence and dimension dependence, as shown in Table 1. Specially, in the  $n \asymp d^2$  regime we obtain exponentially converging classification error, whereas in the  $n \asymp d$  regime we achieve linear dimension dependence. Note that this improves upon the NTK analysis (which gives a sample complexity of  $n = \Omega(d^k)$ ) in that it "decouple" the degree  $k$  from the exponent in the dimension dependence. Our theoretical results are supported by empirical findings.
Authors	regime/method	k-parity	class error	width	# iterations
Ji and Telgarsky (2019)	NTK/SGD	No	d2/n	d8	d2/ε
Telgarsky (2023)	NTK/SGD	No	d2/n	d2	d2/ε
Barak et al. (2022)	Two phase SGD	Yes	d(k+1)/2/√n	O(1)	d/ε2
Wei et al. (2019)	mean-field/GF	No	d/n	∞	∞
Telgarsky (2023)	mean-field/GF	No	d/n	d4	∞
Ours	mean-field/MFLD	Yes	exp(-O(√n/d))	eO(d)	eO(d)
Ours	mean-field/MFLD	Yes	d/n	eO(d)	eO(d)
Table 1: Statistical and computational complexity (omitting poly-log terms) for the  $k$ -sparse (or 2-sparse) parity problem. Column "k-parity" indicates applicability to the general  $k$ -sparse parity setting, and for references that does not handle  $k$ -parity we state the complexity for 2-parity (XOR).  $d$  is the input dimensionality, and  $n$  is the sample size. For SGD,  $n$  is the total sample size given by product of mini-batch size and iterations.
1.2 Additional Related Works
In addition to the mean-field analysis, parity-like functions can be learned via other feature learning procedures such as the one gradient step analysis in Daniely and Malach (2020); Ba et al. (2022);
Damian et al. (2022); Barak et al. (2022). Such analysis requires nontrivial gradient concentration at initialization, which translates to a sample complexity that scales as  $n = \Theta(d^k)$  (Barak et al., 2022). For "narrow" neural networks, Abbe et al. (2023) showed that a modified projected (online) gradient descent algorithm can learn the  $k$ -parity problem with a sample complexity of  $n = O(d^{k-1} + \epsilon^{-c})$  (for some unknown  $c$ ). Also, Refinetti et al. (2021); Ben Arous et al. (2022) showed that a two-layer ReLU network with more than 4 neurons can learn the Gaussian XOR problem with gradient descent.
Throughout this paper, we consider a classification problem given by the following model:
$$
Y = \mathbf {1} _ {A} (Z) - \mathbf {1} _ {A ^ {c}} (Z) \in \{\pm 1 \}
$$
where  $Z = (Z_{1},\ldots ,Z_{d})$  is the input random variable on  $\mathbb{R}^d$  and  $\mathbf{1}_A$  is the indicator function corresponding to a measurable set  $A\in \mathbb{B}(\mathbb{R}^d)$ , i.e.,  $\mathbf{1}_A(Z) = 1$  if  $Z\in A$  and  $\mathbf{1}_A(Z) = 0$  if  $Z\notin A$ . Let  $P_Z$  be the distribution of  $Z$ . We are given input-output pairs  $D_{n} = (z_{i},y_{i})_{i = 1}^{n}$  independently identically distributed from this model as training data. Then, we construct a binary classifier that predicts the label for the test input data as accurate as possible. To achieve this, we learn a two-layer neural network model in the mean-field regime via the mean-field Langevin dynamics.
One important problem setting for our analysis is the  $k$ -sparse parity problem defined as follows.
Example 1 ( $k$ -sparse parity problem).  $P_Z$  is the uniform distribution on the grid  $\{\pm 1/\sqrt{d}\}^d$  and  $A = \{\zeta = (\zeta_1, \dots, \zeta_d) \in \{\pm 1/\sqrt{d}\}^d \mid \zeta_1 \dots \zeta_k > 0\}^1$ .
As a special case,  $k = 2$  (XOR) has been extensively studied (Wei et al., 2019; Telgarsky, 2023).
Mean-field two-layer network. Given input  $z$ , let  $h_x(z)$  be one neuron in a two-layer neural network with parameter  $x = (x_1, x_2, x_3) \in \mathbb{R}^{d + 1 + 1}$  defined as
$$
h _ {x} (z) = \bar {R} \left[ \tanh  \left(z ^ {\top} x _ {1} + x _ {2}\right) + 2 \tanh  \left(x _ {3}\right) \right] / 3,
$$
where  $\bar{R} \in \mathbb{R}$  is a hyper-parameter determining the scale of the network $^2$ . We place an extra tanh activation for the bias term  $x_{3} \in \mathbb{R}$  because the boundedness of  $h_x$  is required in the convergence analysis. Let  $\mathcal{P}$  be the set of probability measures on  $(\mathbb{R}^{\bar{d}}, \mathcal{B}(\mathbb{R}^{\bar{d}}))$  where  $\bar{d} = d + 2$  and  $\mathcal{B}(\mathbb{R}^{\bar{d}})$  is the Borel  $\sigma$ -algebra on  $\mathbb{R}^{\bar{d}}$  and  $\mathcal{P}_p$  be the subset of  $\mathcal{P}$  such that its  $p$ -th moment is bounded:  $\mathbb{E}_{\mu}[\|X\|^p] < \infty$  ( $\mu \in \mathcal{P}$ ). The mean-field neural network is defined as an integral over neurons  $h_x$ ,
$$
f _ {\mu} (\cdot) = \int h _ {x} (\cdot) \mu (\mathrm {d} x),
$$
for  $\mu \in \mathcal{P}$ . To evaluate the performance of  $f_{\mu}$ , we define the empirical risk and the population risk as
$$
L (\mu) := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(y _ {i} f _ {\mu} \left(z _ {i}\right)\right), \quad \bar {L} (\mu) := \mathbb {E} \left[ \ell \left(Y f _ {\mu} (Z)\right) \right],
$$
respectively, where  $\ell : \mathbb{R} \to \mathbb{R}_{\geq 0}$  is a convex loss function. In particular, we consider the logistic loss  $\ell(f, y) = \log(1 + \exp(-y f))$  for  $y \in \{\pm 1\}$  and  $f \in \mathbb{R}$ . To avoid overfitting, we consider a regularized empirical risk  $F(\mu) \coloneqq L(\mu) + \lambda \mathbb{E}_{X \sim \mu}[\lambda_1 \| X\|^2]$ , where  $\lambda, \lambda_1 \geq 0$  are regularization parameters. One advantage of this mean-field definition is that  $f_{\mu}$  is a linear with respect to  $\mu$ , and hence the functional  $L(\mu)$  becomes a convex functional.
Mean-field Langevin dynamics. We optimize the training objective via MFLD, which is given by the following stochastic differential equation:
$$
\mathrm {d} X _ {t} = - \nabla \frac {\delta F (\mu_ {t})}{\delta \mu} (X _ {t}) \mathrm {d} t + \sqrt {2 \lambda} \mathrm {d} W _ {t}, \quad \mu_ {t} = \operatorname {L a w} (X _ {t}), \tag {1}
$$
where  $X_0 \sim \mu_0$ ,  $\operatorname{Law}(X)$  denotes the distribution of the random variable  $X$  and  $(W_t)_{t \geq 0}$  is the  $d$ -dimensional standard Brownian motion. Readers may refer to Theorem 3.3 of Huang et al. (2021) for the existence and uniqueness of the solution. Here,  $\frac{\delta F(\mu_t)}{\delta \mu}$  is the first variation of  $F$ .
Definition 1. For a functional  $G: \mathcal{P} \to \mathbb{R}$ , the first-variation  $\frac{\delta G}{\delta \mu}(\mu)$  at  $\mu \in \mathcal{P}$  is a continuous functional  $\mathcal{P} \times \mathbb{R}^d \to \mathbb{R}$  satisfying  $\lim_{\epsilon \to 0} \frac{G(\epsilon \nu + (1 - \epsilon) \mu)}{\epsilon} = \int \frac{\delta G}{\delta \mu}(\mu)(x) \, \mathrm{d}(\nu - \mu)$  for any  $\nu \in \mathcal{P}$ .
<sup>1</sup>We present the axis-aligned setting for conciseness, but the same result holds under orthogonal transforms.
2We may also deal with more general activation function that is smooth and bounded.
In our setting, we have  $\frac{\delta F(\mu)}{\delta\mu}(x) = \frac{1}{n}\sum_{i=1}^{n}\ell'(y_if_\mu(z_i))y_ih_x(z_i) + \lambda(\lambda_1\|x\|^2)$ . It is known that the Fokker-Planck equation of the SDE (1) is given by
$$
\partial_ {t} \mu_ {t} = \lambda \Delta \mu_ {t} + \nabla \cdot \left[ \mu_ {t} \nabla \frac {\delta F (\mu_ {t})}{\delta \nu} \right] = \nabla \cdot \left[ \mu_ {t} \nabla \left(\lambda \log (\mu_ {t}) + \frac {\delta F (\mu_ {t})}{\delta \nu}\right) \right]. \tag {2}
$$
Then, we can verify that this is equivalent to the Wasserstein gradient flow to optimize the following entropy regularized risk (Mei et al., 2018; Hu et al., 2019):
$$
\mathcal {L} (\mu) = F (\mu) + \lambda \operatorname {E n t} (\mu) = L (\mu) + \lambda \mathrm {K L} (\nu , \mu) + (\text {c o n s t .}) \tag {3}
$$
where  $\mathrm{KL}(\nu, \mu) = \int \log (\mu / \nu) \mathrm{d}\mu$  is the KL divergence between  $\nu$  and  $\mu$ , and  $\nu$  is the Gaussian distribution with mean 0 and variance  $I / (2\lambda_1)$ , i.e.,  $\nu = \mathcal{N}(0, I / (2\lambda_1))$ .
For a practical algorithm, we need to consider a space- and time-discretized version of the MFLD, that is, we approximate the solution  $\mu_t$  by an empirical measure  $\mu_{\mathcal{X}} = \frac{1}{N}\sum_{i=1}^{N}\delta_{X_i}$  corresponding to a set of finite particles  $\mathcal{X} = (X^i)_{i=1}^N \subset \mathbb{R}^{\bar{d}}$ . Let  $\mathcal{X}_{\tau} = (X_{\tau}^i)_{i=1}^N \subset \mathbb{R}^{\bar{d}}$  be  $N$  particles at the  $\tau$ -th update ( $\tau \in \{0,1,2,\ldots\}$ ), and define  $\mu_{\tau} = \mu_{\mathcal{X}_{\tau}}$  as a finite particle approximation of the population counterpart. Then, the discretized MFLD is defined as follows:  $X_0^i \sim \mu_0$ , and  $\mathcal{X}_{\tau}$  is updated as
$$
X _ {\tau + 1} ^ {i} = X _ {\tau} ^ {i} - \eta \nabla \frac {\delta F (\mu_ {\tau})}{\delta \mu} \left(X _ {\tau} ^ {i}\right) + \sqrt {2 \lambda \eta} \xi_ {\tau} ^ {i}, \tag {4}
$$
where  $\eta > 0$  is the step size, and  $\xi_{\tau}^{i} \sim_{i.i.d.} N(0, I)$ . This is the Euler-Maruyama approximation of the MFLD with a discretized measure; we present the discretization error bounds in the next section.
In this section, we introduce the basic assumptions and technical tools for our analysis.
Condition on the loss function To derive the convergence of the classification error, we assume that the loss function satisfies the following condition.
Assumption 1. The convex loss function  $\ell : \mathbb{R} \to \mathbb{R}_{\geq 0}$  satisfies the following conditions:
-  $\ell$  is first order differentiable, its derivative is Lipschitz continuous and its derivative is bounded by  $1$ :  $|\ell^{\prime}(x) - \ell^{\prime}(x^{\prime})| \leq C|x - x^{\prime}|$  and  $\sup_x |\ell'(x)| \leq 1$ .
-  $\ell$  is monotonically decreasing, and is classification calibrated:  $\ell'(0) < 0$  (Bartlett et al., 2006).
-  $\psi(u)^{-1} \coloneqq \ell(0) - (\ell(u) - u\ell'(u)) > 0$  for any  $u > 0$ .
This standard assumption is satisfied by several loss functions such as the logistic loss. We remark that the first assumption is used to show the well-definedness of the mean-field Langevin dynamics and derive its discretization error, and also to obtain a uniform generalization error bound through the classical contraction argument (Boucheron et al., 2013; Ledoux and Talagrand, 1991). The second and third assumptions are used to show the convergence of classification error of our estimator.
Logarithmic Sobolev inequality. Nitanda et al. (2022); Chizat (2022) showed that the convergence of MFLD crucially relies on properties of the proximal Gibbs distribution whose density is given by
$$
p _ {\mu} (X) \propto \exp \left(- \frac {1}{\lambda} \frac {\delta F (\mu)}{\delta \mu} (X)\right),
$$
where  $\mu \in \mathcal{P}$ . By the smoothness of the loss function (Assumption 1) and the tanh activation, we can show that the objective  $\mathcal{L}$  has a unique solution  $\mu^{*}$  which is also a proximal Gibbs measure of itself.
Proposition 1 (Proposition 2.5 of Hu et al. (2019)). The functional  $\mathcal{L}$  has a unique minimizer in  $\mathcal{P}_2$  that is absolutely continuous with respect to the Lebesgue measure. Moreover,  $\mu^{*} \in \mathcal{P}_{2}$  is the optimal solution if and only if  $\mu^{*}$  is absolutely continuous and its density function is given by  $p_{\mu^*}$ .
The next question is how fast the solution  $\mu_t$  converges to the optimal solution  $\mu^*$ . As we will see, the convergence of MFLD heavily depends on a logarithmic Sobolev inequality (LSI) on  $p_\mu$ .
This should be interpreted in a weak sense, that is, for any continuously differentiable function  $\phi$  with a compact support,  $\int \phi \mathrm{d}\mu_t - \int \phi \mathrm{d}\mu_s = -\int_s^t\int \nabla \phi \cdot (\nabla \log (\mu_t) - \nabla \frac{\delta F(\mu_t)}{\delta\nu})\mathrm{d}\mu_\tau \mathrm{d}\tau$ .
Definition 2 (Logarithmic Sobolev inequality). Let  $\mu$  be a probability measure on  $(\mathbb{R}^d,\mathcal{B}(\mathbb{R}^d))$ .  $\mu$  satisfies the LSI with constant  $\alpha >0$  if for any smooth function  $\phi :\mathbb{R}^d\to \mathbb{R}$  with  $\mathbb{E}_{\mu}[\phi^{2}] < \infty$ ,
$$
\mathbb {E} _ {\mu} [ \phi^ {2} \log (\phi^ {2}) ] - \mathbb {E} _ {\mu} [ \phi^ {2} ] \log (\mathbb {E} _ {\mu} [ \phi^ {2} ]) \leq \frac {2}{\alpha} \mathbb {E} _ {\mu} [ \| \nabla \phi \| _ {2} ^ {2} ].
$$
This is equivalent to the condition that the KL divergence from  $\mu$  is bounded by the Fisher divergence:  $\int \log (\mathrm{d}\nu /\mathrm{d}\mu)\mathrm{d}\nu \leq \frac{2}{\alpha}\int \| \nabla \log (\mathrm{d}\nu /\mathrm{d}\mu)\| ^2\mathrm{d}\mu$ , for any  $\nu \in \mathcal{P}$  which is absolutely continuous with respect to  $\mu$ . The LSI of proximal Gibbs measure can be established via standard perturbation criteria. For  $L(\mu)$  with bounded first-variation, we may apply the classical Bakry-Emery and Holley-Stroock arguments (Bakry and Emery, 1985; Holley and Stroock, 1987) (Corollary 5.7.2 and 5.1.7 of Bakry et al. (2014)): If  $\| \frac{\delta L(\mu)}{\delta\mu}\|_{\infty}\leq B$  is satisfied for any  $\mu \in \mathcal{P}_2$ , then  $\mu^{*}$  and  $p_{\mathcal{X}}$  satisfy the LSI with
$$
\alpha \geq \lambda_ {1} \exp (- 4 B / \lambda). \tag {5}
$$
With the LSI condition on the proximal Gibbs distribution, it is known that the MFLD converges to the optimal solution in an exponential order by using a so-called Entropy sandwich technique.
Proposition 2 (Entropy sandwich (Nitanda et al., 2022; Chizat, 2022)). Suppose that  $\mu_0$  satisfies  $\mathcal{L}(\mu_0) < \infty$  and the proximal Gibbs measure  $p_{\mu_t}$  corresponding to the solution  $\mu_t$  has the LSI constant  $\alpha$  for all  $t \geq 0$ , then the solution  $\mu_t$  of MFLD satisfies
$$
\lambda \mathrm {K L} (\mu^ {*}, \mu_ {t}) \leq \mathcal {L} (\mu_ {t}) - \mathcal {L} (\mu^ {*}) \leq \exp (- 2 \alpha \lambda t) (\mathcal {L} (\mu_ {0}) - \mathcal {L} (\mu^ {*})),
$$
where  $\mu^{*} = \operatorname{argmin}_{\mu \in \mathcal{P}} \mathcal{L}(\mu)$  (the existence and uniqueness of  $\mu^{*}$  is guaranteed by Proposition 1).
Hence we know that time horizon  $T = O\left(\frac{1}{\lambda \alpha} \log \left(1 / \tilde{\epsilon}\right)\right)$  is sufficient to achieve  $\tilde{\epsilon} > 0$  accuracy.
Convergence of the discretized algorithm. While Proposition 2 only established the convergence rate of the continuous dynamics, similar guarantee can be shown for the discretized setting. Let
$$
\mathcal {L} ^ {N} (\mu^ {(N)}) = N \mathbb {E} _ {\mathcal {X} \sim \mu^ {(N)}} [ F (\mu_ {\mathcal {X}}) ] + \lambda \mathrm {E n t} (\mu^ {(N)}),
$$
where  $\mu^{(N)}$  is a distribution of  $N$  particles  $\mathcal{X} = (X^i)_{i=1}^N \subset \mathbb{R}^{\bar{d}}$ . Let  $\mu_{\tau}^{(N)}$  be the distribution of the particles  $\mathcal{X}_{\tau} = (X_{\tau}^i)_{i=1}^N$  at the  $\tau$ -th iteration. Suzuki et al. (2023b) showed that, if  $\lambda \alpha \eta \leq 1/4$  and  $\eta \leq 1/4$ , then for  $\bar{B}^2 \coloneqq \mathbb{E}[\|X_0^i\|^2] + \frac{1}{\lambda \lambda_1} \left[ \left( \frac{1}{4} + \frac{1}{\lambda \lambda_1} \right) \bar{R}^2 + \lambda d \right] = O(d + \lambda^{-2})$  and  $\delta_{\eta} \coloneqq C_1 \bar{L}^2 (\eta^2 + \lambda \eta)$ , where  $\bar{L} = 2\bar{R} + \lambda \lambda_1 = O(1)$  and  $C_1 = 8(\bar{R}^2 + \lambda \lambda_1 \bar{B}^2 + d) = O(d + \lambda^{-1})$ ,
$$
\frac {1}{N} \mathbb {E} \left[ \mathcal {L} ^ {N} \left(\mu_ {\tau} ^ {(N)}\right) \right] - \mathcal {L} \left(\mu^ {*}\right) \leq \exp \left(- \frac {\lambda \alpha \eta \tau}{2}\right) \left(\frac {\mathbb {E} \left[ \mathcal {L} ^ {N} \left(\mu_ {0} ^ {(N)}\right) \right]}{N} - \mathcal {L} \left(\mu^ {*}\right)\right) + \frac {4}{\lambda \alpha} \bar {L} ^ {2} C _ {1} \left(\lambda \eta + \eta^ {2}\right) + \frac {4 C _ {\lambda}}{\lambda \alpha N},
$$
where  $C_{\lambda}$  is a constant depending on  $\lambda$ . In particular, for a given  $\tilde{\epsilon} > 0$ , the right hand side can be bounded by  $\tilde{\epsilon} + \frac{4C_{\lambda}}{\lambda\alpha N}$  after  $T = O\left(\frac{\bar{L}^2C_1}{\alpha\tilde{\epsilon}} + \frac{\bar{L}\sqrt{C_1}}{\sqrt{\lambda\alpha\tilde{\epsilon}}}\right)\frac{1}{\lambda\alpha}\log (1 / \tilde{\epsilon})$  iterations with the step size  $\eta = O\bigl (\bigl (\frac{\bar{L}^2C_1}{\alpha\tilde{\epsilon}} +\frac{\bar{L}\sqrt{C_1}}{\sqrt{\lambda\alpha\tilde{\epsilon}}}\bigr)^{-1}\bigr)$ . Furthermore, the convergence of the loss function can be connected to the convergence of the function value of the neural network as follows,
$$
\begin{array}{l} \mathbb {E} _ {\mathcal {X} _ {\tau} \sim \mu_ {k} ^ {(N)}} \left[ \sup  _ {z \in \operatorname {s u p p} (P _ {Z})} \left(f _ {\mu_ {\mathcal {X} _ {\tau}}} (z) - f _ {\mu^ {*}} (z)\right) ^ {2} \right] \\ \leq \frac {4 \bar {L} ^ {2}}{\lambda \alpha} \left(\frac {\mathcal {L} ^ {N} (\mu_ {\tau} ^ {(N)})}{N} - \mathcal {L} (\mu^ {*})\right) + 2 \mathbb {E} _ {\mathcal {X} _ {*} \sim (\mu^ {*}) ^ {\otimes N}} \bigg [ \sup _ {z \in \operatorname {s u p p} (P _ {Z})} \left(\frac {1}{N} \sum_ {i = 1} ^ {N} h _ {X _ {*} ^ {i}} (z) - \int h _ {x} (z) \mathrm {d} \mu^ {*} (z)\right) ^ {2} \bigg ]. \\ \end{array}
$$
Here the second term in the right hand side can be bounded by  $\frac{32\bar{R}^2}{N}\left[1 + 2\left(\frac{2\bar{R}^2}{(\lambda\lambda_1)^2} +\frac{\bar{d}}{\lambda_1}\right)\right]$  via Lemma 2, if  $\| z\| \leq 1$  for any  $z\in \mathrm{supp}(P_Z)$  as in the  $k$  -sparse parity problem. Hence, by taking the number of particles as  $N = \epsilon^{-2}[(\lambda \alpha)^{-2} + (\lambda \lambda_1)^{-2} + d / \lambda_1]$  and letting  $\tilde{\epsilon} = \lambda \alpha \epsilon^{2}$  with the choice of  $T$  and  $\eta$  as described above, we have  $\sup_{z\in \mathrm{supp}(P_Z)}|f_{\mu_{\mathcal{X}_T}}(z) - f_{\mu^*}(z)| = O_p(\epsilon).$
Assumptions on the model specification. We restrict ourselves to the situation where a perfect classifier with margin  $c_{0}$  is included in our model class, stated in the following assumption.
Assumption 2. There exists  $c_0 > 0$  and  $R > 0$  such that the following conditions are satisfied: