# How Does Gradient Descent Learn Features – A Local Analysis for Regularized
Two-Layer Neural Networks

Mo Zhou
Duke University
<EMAIL_ADDRESS>Rong Ge
Duke University
<EMAIL_ADDRESS>
###### Abstract

The ability of learning useful features is one of the major advantages of
neural networks. Although recent works show that neural network can operate in
a neural tangent kernel (NTK) regime that does not allow feature learning,
many works also demonstrate the potential for neural networks to go beyond NTK
regime and perform feature learning. Recently, a line of work highlighted the
feature learning capabilities of the early stages of gradient-based training.
In this paper we consider another mechanism for feature learning via gradient
descent through a local convergence analysis. We show that once the loss is
below a certain threshold, gradient descent with a carefully regularized
objective will capture ground-truth directions. Our results demonstrate that
feature learning not only happens at the initial gradient steps, but can also
occur towards the end of training.

##  1 Introduction

Feature learning has long been considered to be a major advantage of neural
networks. However, how gradient-based training algorithms can learn useful
features is not well-understood. In particular, the most widely applied
analysis for overparametrized neural networks is the neural tangent kernel
(NTK) (Jacot et al., 2018; Du et al., 2019; Allen-Zhu et al., 2019b). In this
setting, the neurons don’t move far from their initialization and the features
are determined by the network architecture and random initialization (Chizat
et al., 2019).

While there are empirical and theoretical evidence on the limitation of NTK
regime (Chizat et al., 2019; Arora et al., 2019), extending the analysis
beyond the NTK regime has been challenging. For 2-layer networks, an
alternative framework for analyzing overparametrized neural networks called
mean-field analysis was introduced. Earlier mean-field analysis (e.g., Chizat
and Bach, 2018; Mei et al., 2018) require either infinite or exponentially
many neurons. Later works (e.g., Li et al., 2020; Ge et al., 2021; Bietti et
al., 2022; Mahankali et al., 2024) can analyze the training dynamics of mildly
overparametrized networks with polynomially many neurons with stronger
assumptions on the ground-truth function.

Recently, a growing line of works (Daniely and Malach, 2020; Damian et al.,
2022; Abbe et al., 2021, 2022, 2023; Yehudai and Shamir, 2019; Shi et al.,
2022; Ba et al., 2022; Mousavi-Hosseini et al., 2023; Barak et al., 2022;
Dandi et al., 2023; Wang et al., 2024; Nichani et al., 2024a, b) showed that
early stages of gradient training (either one/a few steps of gradient descent
or a small amount of time of gradient flow) can be useful in feature learning.
These works show that after the early stages of gradient training, the first
layer in a 2-layer neural network already captures useful features (usually in
the form of a low dimensional subspace), and continuing training the second
layer weights will give performance guarantees that are stronger than any
kernel or random feature based models. In this work, we consider the natural
follow-up question:

Does feature learning only happen in the early stages of gradient training?

We show that this is not the case by demonstrating feature learning capability
for the final stage of gradient training – local convergence. In particular,
we prove the following result:

######  Theorem 1.1 (Informal).

If the data is generated by a 2-layer teacher network $f_{*}$, as long as the
width of student network $m$ is at least some quantity $m_{0}$ that only
depends on $f_{*}$, a variant of gradient descent algorithm (Algorithm 1,
roughly gradient descent with decreasing weight decay) can recover the target
network within polynomial time. Moreover, the student neurons align with the
teacher neurons at the end.

Our result highlights the different mechanisms of feature learning: previous
works show that in the early stages of gradient descent, the network learns
the subspace spanned by the neurons in the teacher network. Our local
convergence result shows that at later stages, gradient descent is able to
learn the exact directions of the teacher neurons, which are much more
informative compared to the subspace and lead to stronger guarantees.

Analyzing the entire training dynamics is still challenging so in our
algorithm (see Algorithm 1) we use a convex second stage to “fast-forward” to
the local analysis. Our technique for local convergence is similar to the
earlier work (Zhou et al., 2021), however we consider a more complicated
setting with ReLU activation and allow second-layer weights to be both
positive or negative. This change requires additional regularization in the
form of standard weight decay and new dual certificate analysis.

###  1.1 Related works

##### Neural Tangent Kernel

Early works often studied neural network optimization using NTK theory (Jacot
et al., 2018; Allen-Zhu et al., 2019b; Du et al., 2019). It is shown that
highly-overparametrized neural nets are essentially kernel methods under
certain initialization scale. However, NTK theory cannot explain the
performance of neural nets in practice (Arora et al., 2019) and leads to lazy
training dynamics that neurons stay close to their initialization (Chizat et
al., 2019). Hence, later research efforts (e.g., Allen-Zhu et al., 2019a; Bai
and Lee, 2020; Li et al., 2020), as well as current paper, focus on feature
learning regime where neural nets can learn features and outperform kernel
methods.

##### Early stage feature learning

Researchers have recently tried to understand how neural networks trained with
gradient descent (GD) can learn features, going beyond the kernel/lazy regime
(Jacot et al., 2018; Chizat et al., 2019). A typical setup is to use 2-layer
neural networks to learn certain target function, often equipped with low-
dimensional structure. Examples include learning polynomials (Yehudai and
Shamir, 2019; Damian et al., 2022), single-index models (Ba et al., 2022;
Mousavi-Hosseini et al., 2023), multi-index models (Dandi et al., 2023),
sparse boolean functions (Abbe et al., 2021, 2022, 2023), sparse parity
functions (Daniely and Malach, 2020; Shi et al., 2022; Barak et al., 2022) and
causal graph (Nichani et al., 2024b). Also, few works use 3-layer networks as
learner model (Nichani et al., 2024a; Wang et al., 2024). These works
essentially showed that feature learning happens in the early stage of
training. Specifically, they often use 2-stage layer-wise training procedure:
first-layer weights/features are only trained with one or few steps of
gradient descent/flow and only update the second-layer afterwards. Our results
give a complementary view that feature learning can also happen in the final
stage training that leading student neurons eventually align with ground-truth
directions. This cannot be achieved if first-layer weights are fixed after few
steps.

##### Learning single/multi-index models with neural networks

Single/Multi-index models are the functions that only depend on one or few
directions of the high dimensional input. Many recent works have studied the
problem of using 2-layer networks to learn single-index models
(Soltanolkotabi, 2017; Yehudai and Ohad, 2020; Frei et al., 2020; Wu, 2022;
Bietti et al., 2022; Xu and Du, 2023; Berthier et al., 2023; Mahankali et al.,
2024) and multi-index models (Damian et al., 2022; Bietti et al., 2023; Suzuki
et al., 2024; Glasgow, 2024). These works show the advantages of feature
learning over fixed random features in various settings. In this paper, we
consider target multi-index function that can be represented by a small
2-layer network, and show a variant of GD with weight decay can learn it and,
moreover, recover the ground-truth directions.

##### Local loss landscape

Safran et al. (2021) showed that in the overparametrized case with orthogonal
teacher neurons, even around the local region of global minima, the landscape
neither is convex nor satisfies PL condition. Chizat (2022) considered square
loss with $\ell_{2}$ regularization similar to our setup and showed the local
loss landscape is strongly-convex under certain non-degenerate assumptions.
However, it is not known when such assumptions actually hold and the proof
cannot handle ReLU. Later Akiyama and Suzuki (2021) gives a result for ReLU,
but the non-degeneracy assumption is still required (and also focus on
effective $\ell_{1}$ regularization instead of $\ell_{2}$ regularization).
Zhou et al. (2021) studies a similar local convergence setting but restricts
second-layer weights to be positive and uses absolute activation. In this
paper, we focus on a more natural but technically challenging case that
second-layer can be positive and negative and using ReLU activation. We
develop new techniques to overcome the above challenges (additional
assumption, ReLU, standard second-layer, etc).

##  2 Preliminary

##### Notation

Let $[n]$ be set $\\{1,\ldots,n\\}$. For vector ${\bm{w}}$, we use
$\left\|{\bm{w}}\right\|_{2}$ for its 2-norm and
$\overline{{\bm{w}}}={\bm{w}}/\left\|{\bm{w}}\right\|_{2}$ as its normalized
version. For two vectors ${\bm{w}},{\bm{v}}$ we use
$\angle({\bm{w}},{\bm{v}})=\arccos(|{\bm{w}}^{\top}{\bm{v}}|/(\left\|{\bm{w}}\right\|_{2}\left\|{\bm{v}}\right\|_{2})]\in[0,\pi/2]$
as the angle between them (up to a sign). For matrix ${\bm{A}}$ let
$\left\|{\bm{A}}\right\|_{F}$ be its Frobenius norm. We use standard
$O,\Omega,\Theta$ to hide constants and
$\widetilde{O},\widetilde{\Omega},\widetilde{\Theta}$ to hide polylog factors.
We use $O_{*},\Omega_{*},\Theta_{*}$ to hide problem dependent parameters that
only depend on the target network (see paragraph above (1)).

##### Teacher-student setup

We will consider the teacher-student setup for two-layer neural networks with
Gaussian input ${\bm{x}}\sim N({\bm{0}},{\bm{I}}_{d})$. The goal is to learn
the teacher network of size $m_{*}$

$\displaystyle
f_{*}({\bm{x}})=\sum_{i=1}^{m^{*}}a_{i}^{*}\sigma({\bm{w}}_{i}^{*\top}{\bm{x}})+{\bm{w}}_{0}^{*\top}{\bm{x}}+b_{0}^{*},$

where $\sigma(x):=\max\\{0,x\\}$ is ReLU activation,
$S_{*}:=\operatorname{span}\\{{\bm{w}}_{1}^{*},\ldots,{\bm{w}}_{m^{*}}^{*}\\}$
is the target subspace. Without loss of generality, we will assume
$\left\|{\bm{w}}_{i}^{*}\right\|_{2}=1$ due to the homogeneity of ReLU.

Following the recent line of works in learning single/multi-index models (Ba
et al., 2022; Damian et al., 2022), we assume the target network has low
dimensional structure.

######  Assumption 2.1.

Teacher neurons form a low dimensional subspace in $\mathbb{R}^{d}$, that is

$\dim(S_{*})=\dim(\operatorname{span}\\{{\bm{w}}_{1}^{*},\ldots,{\bm{w}}_{m^{*}}^{*}\\})=r\ll
d.$

We will also assume the teacher neurons are non-degenerate in the following
sense:

######  Assumption 2.2.

Teacher neurons are $\Delta$-separated, that is angle
$\angle({\bm{w}}_{i}^{*},{\bm{w}}_{j}^{*})\geq\Delta$ for all $i\neq j$.

######  Assumption 2.3.

${\bm{H}}:=\sum_{i=1}^{m^{*}}a_{i}^{*}{\bm{w}}_{i}^{*}{\bm{w}}_{i}^{*\top}$ is
non-degenerate in target subspace $S_{*}$, i.e., $\text{rank}(H)=r$. Denote
$\kappa:=|\lambda_{r}({\bm{H}})|$.

Assumption 2.2 simply requires all teacher neurons pointing to different
directions, which is crucial for identifiability (Zhou et al., 2021).
Assumption 2.3 says the target network contains low-order (second-order)
information, which is related with the notion of information exponent (Arous
et al., 2021). Indeed, one can show
$\mathbb{E}_{\bm{x}}[f_{*}({\bm{x}})h_{2}({\bm{v}}^{\top}{\bm{x}})]=\hat{\sigma}_{2}{\bm{v}}^{\top}{\bm{H}}{\bm{v}}$,
where $h_{2}(x)$ is the 2nd-order normalized Hermite polynomial and
$\hat{\sigma}_{2}$ is the 2nd Hermite coefficient of ReLU. See Appendix A for
more details. Many previous works also rely on same or similar assumption to
show neural networks can learn features to perform better than kernels (Damian
et al., 2022; Abbe et al., 2022; Ba et al., 2022).

In this paper, we are interested in the case where the complexity of target
network is small. Therefore, we will use $O_{*},\Omega_{*},\Theta_{*}$ to hide
$\operatorname{poly}(r,m_{*},\Delta,|a_{1}|,\ldots,|a_{m_{*}}|,\kappa)$, which
is the polynomial dependency on relevant parameters of target $f_{*}$ (does
not depend on student network).

We will use the following overparametrized student network:

$\displaystyle
f({\bm{x}};{\bm{\theta}})=\sum_{i=1}^{m}a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})+\alpha+{\bm{\beta}}^{\top}{\bm{x}},$
(1)

where ${\bm{a}}=(a_{1},\ldots,a_{m})^{\top}\in\mathbb{R}^{m}$,
${\bm{W}}=({\bm{w}}_{1}\cdots{\bm{w}}_{m})^{\top}\in\mathbb{R}^{m\times d}$
and ${\bm{\theta}}=({\bm{a}},{\bm{W}},\alpha,{\bm{\beta}})$.

##### Loss and algorithm

Consider the square loss function with $\ell_{2}$ regularization under
Gaussian input

$\displaystyle L_{\lambda}({\bm{\theta}})=\mathbb{E}_{{\bm{x}}\sim
N(0,{\bm{I}}_{d})}[(f({\bm{x}};{\bm{\theta}})-\widetilde{y})^{2}]+\frac{\lambda}{2}\left\|{\bm{a}}\right\|_{2}^{2}+\frac{\lambda}{2}\left\|{\bm{W}}\right\|_{2}^{2}.$
(2)

We will use $L$ to denote the square loss for simplicity. The $\ell_{2}$
regularization is the same as the commonly used weight decay in practice. Our
goal is to find the minima of unregularized problem ($\lambda=0$) to recover
teacher network $f_{*}$.

In above, we use preprocessed data $(x,\widetilde{y})$ in the loss function as
in Damian et al. (2022). Specifically, given any $({\bm{x}},y)$ with
$y=f_{*}({\bm{x}})$, denote $\alpha_{*}=\mathbb{E}_{\bm{x}}[y]$ and
${\bm{\beta}}_{*}=\mathbb{E}_{\bm{x}}[y{\bm{x}}]$, we get

$\displaystyle\widetilde{f}_{*}({\bm{x}})=\widetilde{y}=y-\alpha_{*}-{\bm{\beta}}_{*}^{\top}{\bm{x}}.$
(3)

This preprocessing process essentially removes the 0-th and 1-st order term in
the Hermite expansion of $\sigma$. See Appendix A for a brief introduction of
Hermite polynomials and Claim B.1.

Our algorithm is shown in Algorithm 1. It is roughly the standard GD following
a given schedule of weight decay $\lambda_{t}$ that goes to 0. Due to the
difficulty in analyzing gradient descent training beyond early and final
stage, we choose to only train the norms in Stage 2 as a tractable way to
reach the local convergence regime.

We will use symmetric initialization that $a_{i}=-a_{i+m/2}$,
${\bm{w}}_{i}={\bm{w}}_{i+m/2}$ with
$a_{i}\sim\mathrm{Unif}\\{\pm\sqrt{d}\\}$,
${\bm{w}}_{i}\sim\mathrm{Unif}((1/\sqrt{m}){\mathbb{S}}^{d-1})$, $\alpha=0$,
${\bm{\beta}}={\bm{0}}$. Our analysis is not sensitive to the initialization
scale we choose here. The choice is just for the simplicity of the proof.

Input: initialization ${\bm{\theta}}^{(0)}$, weight decay $\lambda_{t}$ and
stepsize $\eta_{t}$

Data preprocess: get $({\bm{x}},\widetilde{y})$ according to (3)

Stage 1: one step gradient update
${\bm{\theta}}^{(1)}\leftarrow{\bm{\theta}}^{(0)}-\eta_{0}\nabla_{{\bm{\theta}}}L_{\lambda_{0}}({\bm{\theta}}^{(0)})$

Stage 2: norm adjustment by convex program
${\bm{a}}^{(T_{2})},\alpha^{(T_{2})},{\bm{\beta}}^{(T_{2})}\leftarrow\min_{{\bm{a}},\alpha,{\bm{\beta}}}L({\bm{a}},{\bm{W}}^{(1)},\alpha,{\bm{\beta}})+\lambda\sum_{i}\left\|{\bm{w}}_{i}\right\|_{2}|a_{i}|$

Balancing norm between two layers s.t.
$|a_{i}|=\left\|{\bm{w}}_{i}\right\|_{2}$ for all $i$

Stage 3: local convergence

for  _$k\leq K$_ do // for each epoch, run GD until convergence

for  _$T_{3,k-1}\leq t\leq T_{3,k}$_ do

${\bm{\theta}}^{(t+1)}\leftarrow{\bm{\theta}}^{(t)}-\eta\nabla_{\bm{\theta}}L_{\lambda_{3,k}}({\bm{\theta}}^{(t)})$

Output:
${\bm{\theta}}^{(T_{3,K})}=({\bm{a}}^{(T_{3,K})},{\bm{W}}^{(T_{3,K})},\alpha^{(T_{3,K})},{\bm{\beta}}^{(T_{3,K})})$

Algorithm 1 Learning 2-layer neural networks

##  3 Main results

In this section, we give our main result that shows training student network
using Algorithm 1 can recover the target network within polynomial time. We
will focus on the case that $d\geq\Omega_{*}(1)$ when the complexity of target
function is small.

######  Theorem 3.1 (Main result).

Under Assumption 2.1, 2.2, 2.3, consider Algorithm 1 on loss (2). There exists
a schedule of weight decay $\lambda_{t}$ and step size $\eta_{t}$ such that
given $m\geq m_{0}=\widetilde{O}_{*}(1)\cdot(1/{\varepsilon}_{0})^{O(r)}$
neurons with small enough ${\varepsilon}_{0}=\Theta_{*}(1)$, with high
probability we will recover the target network
$L({\bm{\theta}})\leq{\varepsilon}$ within time
$T=O_{*}(\log(1/{\varepsilon})/\eta)$ where
$\eta=\operatorname{poly}({\varepsilon},1/d,1/m)$.

Moreover, when $\varepsilon\to 0$ every student neuron ${\bm{w}}_{i}$ either
aligns with one of teacher neuron ${\bm{w}}_{j}^{*}$ as
$\angle({\bm{w}}_{i},{\bm{w}}_{j}^{*})=0$ or vanishes as
$|a_{i}|=\left\|{\bm{w}}_{i}\right\|=0$.

Note that our results can be extended to only have access to polynomial number
of samples by using standard concentration tools. We omit the sample
complexity for simplicity. We emphasize that the required width $m_{0}$ only
depends on the complexity of target function $f_{*}$ (not on student network
$f$ or error $\varepsilon$), so any mildly overparametrized networks can learn
$f_{*}$ efficiently to arbitrary small error.

Our result improves the previous works that only train the first layer weight
with small number of gradient steps at the beginning (Damian et al., 2022; Ba
et al., 2022; Abbe et al., 2021, 2022, 2023). In these works, neural networks
only learn the target subspace and do random features within it. Intuitively,
these random features need to span the whole space of the target function
class to perform well, which means its number (the width) should be on the
order of the dimension of target function class. For 2-layer networks, random
features in the target subspace need $(1/{\varepsilon})^{O(r)}$ neurons to
achieve desired accuracy ${\varepsilon}$. In contrast, continue training both
layer at the last phase of training allows us to learn not only subspace but
also exactly the ground-truth directions. Moreover, we only use
$(1/{\varepsilon}_{0})^{O(r)}$ neurons that only depends on the complexity of
target network. This highlights the benefit of continue training first layer
weights instead of fixing them after first step.

##  4 Proof overview

In this section, we give the proof overview of these three stages separately.

Denote the optimality gap $\zeta$ at time $t$ as the difference between
current loss and the best loss one could achieve with networks of any size
(including infinite-width networks)

$\displaystyle\zeta_{t}=$ $\displaystyle
L_{\lambda_{t}}({\bm{\theta}}^{(t)})-\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda_{t}}(\mu),$
(4)

where ${\mathcal{M}}({\mathbb{S}}^{d-1})$ is the set of measures on the sphere
${\mathbb{S}}^{d-1}$. As an example, if
$\mu=\sum_{i}a_{i}\left\|{\bm{w}}_{i}\right\|\delta_{\overline{{\bm{w}}}_{i}}$,
then $L_{\lambda}(\mu)$ recovers $L_{\lambda}({\bm{\theta}})$ when linear term
$\alpha,{\bm{\beta}}$ are perfectly fitted and norms are balanced
$|a_{i}|=\left\|{\bm{w}}_{i}\right\|$. We defer the precise definition of
$L_{\lambda}(\mu)$ to (6) in appendix.

###  4.1 Stage 1

For Stage 1, we show in the lemma below that the first step of gradient
descent identifies the target subspace and ensures there always exists student
neuron that is close to every teacher neuron.

######  Lemma 4.1 (Stage 1).

Under Assumption 2.1,2.2,2.3, consider Algorithm 1 with
$\lambda_{0}=\eta_{0}=1$ and $m\geq
m_{0}=\widetilde{O}_{*}(1)\cdot(1/{\varepsilon}_{0})^{O(r)}$ with any
${\varepsilon}_{0}=\Theta_{*}(1)$. After first step, with probability
$1-\delta$ we have

  1. (i)

for every teacher neuron ${\bm{w}}_{i}^{*}$, there exists at least one student
neuron ${\bm{w}}_{j}$ s.t.
$\angle({\bm{w}}_{i}^{*},{\bm{w}}_{j})\leq{\varepsilon}_{0}$.

  2. (ii)

$\left\|{\bm{w}}_{i}^{(1)}\right\|_{2}=\Theta_{*}(1)$, $|a_{i}^{(1)}|\leq
O_{*}(1/\sqrt{m})$ for all $i\in[m_{*}]$, $\alpha_{1}=0$ and
${\bm{\beta}}_{1}={\bm{0}}$.

The key observation here is similar to Damian et al. (2022) that
${\bm{w}}_{i}^{(1)}\approx-2\eta_{0}a_{i}^{(0)}\left(\hat{\sigma}_{2}^{2}{\bm{H}}\overline{{\bm{w}}}_{i}\right)$
so that given ${\bm{H}}$ is non-degenerate in target subspace $S_{*}$ we
essentially sample ${\bm{w}}_{i}^{(1)}$ from the target subspace. It is then
natural to expect that the neurons form an ${\varepsilon}_{0}$-net in the
target subspace given $m_{0}$ neurons.

###  4.2 Stage 2

Given the learned features (first-layer weights) in Stage 1, we now perform
least squares to adjust the norms and reach a low loss solution in Stage 2.

######  Lemma 4.2 (Stage 2).

Under Assumption 2.1,2.2,2.3, consider Algorithm 1 with
$\lambda_{t}=\sqrt{{\varepsilon}_{0}}$. Given Stage 1 in Lemma 4.1, we have
Stage 2 ends within time $T_{2}=\widetilde{O}_{*}(1/\eta{\varepsilon}_{0})$
such that optimality gap $\zeta_{T_{2}}=O_{*}({\varepsilon}_{0})$.

It remains an open problem to prove the convergence when training both layers
simultaneously beyond early and final stage. To overcome this technical
challenge, we choose to use a simple least square for Stage 2. We use the
simple (sub)gradient descent to optimize this loss. There exist many other
algorithms that can solve this Lasso-type problem, but we omit it for
simplicity as this is not the main focus of this paper.

Note that the regularization in Algorithm 1 is the same as standard weight
decay when we train both layers. This regularization leads to several desired
properties at the end of Stage 2: (1) prevent norm cancellation between
neurons: neurons with similar direction but different sign of second layer
weights cancel with each other; (2) neurons mostly concentrate around ground-
truth directions. As we will see later, these nice properties continue to hold
in Stage 3, thanks to the regularization.

###  4.3 Stage 3

After Stage 2 we are in the local convergence regime. The following lemma
shows that we could recover the target network within polynomial time using a
multi-epoch gradient descent with decreasing weight decay $\lambda$ at every
epoch. Note that this result only requires the initial optimality gap is small
and width $m\geq m_{*}$ (target network width, not $m_{0}$).

######  Lemma 4.3 (Stage 3).

Under Assumption 2.1,2.2,2.3, consider Algorithm 1 on loss (2). Given Stage 2
in Lemma 4.2, if the initial optimality gap $\zeta_{3,0}\leq
O_{*}(\lambda_{3,0}^{9/5})$, weight decay $\lambda$ follows the schedule of
initial value $\lambda_{3,0}=O_{*}(1)$, and $k$-th epoch
$\lambda_{3,k}=\lambda_{3,k-1}/2$ and stepsize $\eta_{3k}=\eta\leq
O_{*}(\lambda_{3,k}^{12}d^{-3})$ for all $T_{3,k}\leq t\leq T_{3,k+1}$ in
epoch $k$, then within $K=O_{*}(\log(1/{\varepsilon}))$ epochs and total
$T_{3}-T_{2}=O_{*}(\lambda_{3,0}^{-4}\eta^{-1}\log(1/\varepsilon))$ time we
recover the ground-truth network $L({\bm{\theta}})\leq{\varepsilon}$.

The lemma above relies on the following result that shows the local landscape
is benign in the sense that it satisfies a special case of Łojasiewicz
property (Lojasiewicz, 1963). This means GD can always make progress until the
optimality gap $\zeta$ is small.

######  Lemma 4.4 (Gradient lower bound).

When $\Omega_{*}(\lambda^{2})\leq\zeta\leq O_{*}(\lambda^{9/5})$ and
$\lambda\leq O_{*}(1)$, we have

$\displaystyle\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}\geq\Omega_{*}(\zeta^{4}/\lambda^{2}).$

Note that this generalizes previous result in Zhou et al. (2021) that only
focuses on 2-layer networks with positive second layer weights. This turns out
to be technically challenging as two neurons with different signs can cancel
each other. We discuss how to deal with this challenge in the next section.

##  5 Descent direction in local convergence (Stage 3): the benefit of weight
decay

In this section, we give the high-level proof ideas for the most technical
challenging part of our results — characterize the local landscape in Stage 3
(Lemma 4.4).

The key idea is to construct descent direction — a direction that has positive
correlation with the gradient direction. The gradient lower bound follows from
the existence of such descent direction.

It turns out that the existence of both positive and negative second-layer
weights introduces significant challenge for the analysis: there might exist
neurons with similar directions (e.g., $(a,{\bm{w}})$ and $(-a,{\bm{w}})$)
that can cancel with each other to have no effect on the output of network.
Intuitively, we would hope all of them to move towards 0, but they have no
incentive to do so. Moreover, if they are not exactly symmetric it’s hard to
characterize which directions these neurons will move.

We use standard weight decay to address the above challenge. Specifically,
weight decay helps us to

  * •

Balance norm between neurons. When norm between two layers are balanced, the
$\ell_{2}$ regularization
$\sum_{i}|a_{i}|^{2}+\left\|{\bm{w}}_{i}\right\|^{2}$ would become the
effective $\ell_{1}$ regularization
$2\sum_{i}|a_{i}|\left\|{\bm{w}}_{i}\right\|$ over the distribution of
neurons. Such sparsity penalty ensures most neurons concentrate around the
ground-truth directions, especially preventing norm cancellation between far-
away neurons.

  * •

Reduce cancellation between close-by neurons. For close-by neurons, weight
decay helps to reduce the norm of neurons with the ‘incorrect’ sign (different
sign with the ground-truth neuron). This is because weight decay prefers low
norm solutions, and reducing cancellations between neurons can reduce total
norm (regularization term) while keeping the square loss same.

We will group the neurons (i.e., partitioning ${\mathbb{S}}^{d-1}$) based on
their distance to the closest teacher neurons: denote
${\mathcal{T}}_{i}=\\{{\bm{w}}:\angle({\bm{w}},{\bm{w}}_{i}^{*})\leq\angle({\bm{w}},{\bm{w}}_{j}^{*})\text{
for any $j\neq i$}\\}$ (break the tie arbitrarily) so that
$\cup_{i}{\mathcal{T}}_{i}={\mathbb{S}}^{d-1}$. We will also use $\delta_{j}$
to denote $\angle({\bm{w}}_{j},{\bm{w}}_{i}^{*})$ for $j\in{\mathcal{T}}_{i}$.

As described above, weight decay can always lead to descent direction when
norms are not balanced or norm cancellation happens (see Lemma F.15 and Lemma
F.16). The following lemma shows that in other scenarios we can always improve
features towards the ground-truth directions.

######  Lemma 5.1 (Feature improvement descent direction, informal).

When norms are balanced and no norm cancellation happens, there exists
properly chosen $q_{ij}\geq 0$ and
$\sum_{j\in{\mathcal{T}}_{i}}a_{j}q_{ij}=a_{i}^{*}$ such that

$\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\langle\nabla_{{\bm{w}}_{i}}L_{\lambda},{\bm{w}}_{j}-q_{ij}{\bm{w}}_{i}^{*}\rangle=\Omega(\zeta).$

In words, this descent direction is the following: we move neuron
${\bm{w}}_{j}\in{\mathcal{T}}_{i}$ toward either ground-truth direction
${\bm{w}}_{i}^{*}$ or 0 depending on whether it is in the neighborhood of
teacher neuron ${\bm{w}}_{i}^{*}$. Specifically, we move far-away neurons
towards 0 (and thus setting $q_{ij}=0$) and move close-by neurons towards its
’closest’ minima $q_{ij}{\bm{w}}_{i}^{*}$ (the fraction of ${\bm{w}}_{i}^{*}$
that neuron ${\bm{w}}_{j}$ should target to approximate).

The proof of the above lemma requires a dedicated characterization of the low
loss solution’s structure, which we describe in Section 6.

##  6 Structure of (approximated) minima

In this section, we first highlight the importance of understanding local
geometry by showing the challenges in proving the existence of descent
direction (Lemma 5.1). Then after presenting the main result of this section
to show the structure of (approximated) minima (Lemma 6.1), we discuss several
proof ideas such as dual certificate analysis in the remaining part.

###  6.1 Constructing descent direction requires better understanding of local
geometry

To show the existence of descent direction in Lemma 5.1, we compute the inner
product between gradient and constructed descent direction. We can lower bound
it by (assuming norms are balanced)

$\displaystyle\zeta+2\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\mathbb{E}_{\bm{x}}[R({\bm{x}})a_{j}q_{ij}{\bm{w}}_{i}^{*\top}{\bm{x}}(\sigma^{\prime}({\bm{w}}_{i}^{*\top}{\bm{x}})-\sigma^{\prime}({\bm{w}}_{j}^{\top}{\bm{x}}))],$

where $R({\bm{x}})=f({\bm{x}})-\widetilde{f}_{*}({\bm{x}})$ is the residual.
Thus, in order to get a lower bound, the goal is to show second term above is
small than $\zeta$. As we can see, this term is quite complicated and can be
viewed as the inner product between $R({\bm{x}})$ and
$h({\bm{x}})=\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}a_{j}q_{ij}{\bm{w}}_{i}^{*\top}{\bm{x}}(\sigma^{\prime}({\bm{w}}_{i}^{*\top}{\bm{x}})-\sigma^{\prime}({\bm{w}}_{j}^{\top}{\bm{x}}))$.

##### Average neuron and residual decomposition

To deal with above challenge, we use the idea of average neuron and residual
decomposition. For each teacher neuron ${\bm{w}}_{i}^{*}$, denote
${\bm{v}}_{i}=\sum_{j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}$ as the average
neuron. Intuitively, this average neuron ${\bm{v}}_{i}$ stands for an idealize
case where all neurons belong to ${\mathcal{T}}_{i}$ (closer to
${\bm{w}}_{i}^{*}$ than other ${\bm{w}}_{j}^{*}$) collapse into a single
neuron.

We decompose the residual
$R({\bm{x}})=f({\bm{x}})-\widetilde{f}_{*}({\bm{x}})$ into the 3 terms below:
denote $\hat{{\bm{v}}}_{i}={\bm{v}}_{i}-{\bm{w}}_{i}^{*}$

$\displaystyle R_{1}({\bm{x}})$
$\displaystyle=\frac{1}{2}\sum_{i\in[m_{*}]}\hat{{\bm{v}}}_{i}^{\top}{\bm{x}}\operatorname{sign}({\bm{w}}_{i}^{*\top}{\bm{x}}),R_{2}({\bm{x}})=\frac{1}{2}\sum_{i\in[m_{*}],j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}^{\top}{\bm{x}}(\operatorname{sign}({\bm{w}}_{j}^{\top}{\bm{x}})-\operatorname{sign}({\bm{w}}_{i}^{*\top}{\bm{x}})),$
$\displaystyle R_{3}({\bm{x}})$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\sum_{i\in[m_{*}]}a_{i}^{*}\left\|{\bm{w}}_{i}^{*}\right\|_{2}-\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\right)+\alpha-\hat{\alpha}+({\bm{\beta}}-\hat{{\bm{\beta}}})^{\top}{\bm{x}}.$

$R_{1}$ can be thought as the exact-parametrization setting (use $m_{*}$
neurons to learn $m_{*}$ neurons), where the average neurons
$\\{{\bm{v}}_{i}\\}_{i=1}^{m_{*}}$ are the effective neurons. The difference
between this exact-parametrization and overparametrization setting is then
characterized by the term $R_{2}$, which captures the difference in nonlinear
activation pattern. This term in fact suggests the loss landscape is
degenerate in overparametrized case and slows down the convergence (Zhou et
al., 2021; Xu and Du, 2023). Overall, this residual decomposition is similar
to Zhou et al. (2021), with additional modification of $R_{3}$ to deal with
ReLU activation and linear term $\alpha,{\bm{\beta}}$.

To some extent, our residual decomposition can be viewed as a kind of ‘bias-
variance’ decomposition in the sense that the ‘bias’ term $R_{1}$ captures the
overall average contribution of all neurons, and the ‘variance’ term $R_{2}$
captures the individual contributions of each neuron that are not reflected in
$R_{1}$.

##### High-level proof plan of Lemma 5.1

We now are ready to give a proof plan for Lemma 5.1. The key is to show
properties of minima that can help us to bound $\langle R,h\rangle$.

  1. 1.

Show that neurons mostly concentrate around ground-truth directions.

  2. 2.

Show that average neuron ${\bm{v}}_{i}$ is close to teacher neuron
${\bm{w}}_{i}^{*}$ for all $i\in[m]$.

  3. 3.

Use above structure to bound $\langle R_{i},h\rangle$. Specifically, bounding
$\langle R_{1},h\rangle$ relies on the fact that average neuron is close to
teacher neuron (step 2); a bound on $\langle R_{2},h\rangle$ follows from far-
away neurons are small (step 1); third term $\langle R_{3},h\rangle$ can be
directly bounded using the loss. Detailed calculations are deferred into
Appendix H.3.

We give main result of this section that shows the desired local geometry
properties more precisely ((i)(ii) corresponding to step 1 and (iii)
corresponding to step 2 above).

######  Lemma 6.1 (Informal).

Suppose the optimality gap is $\zeta$, we have

  1. (i)

Total norm of far-away neurons is small:
$\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}\delta_{j}^{2}=O_{*}(\zeta/\lambda),$
where angle $\delta_{j}=\angle({\bm{w}}_{j},{\bm{w}}_{i}^{*})$ for
${\bm{w}}_{j}$ that $j\in{\mathcal{T}}_{i}$.

  2. (ii)

For every ${\bm{w}}_{i}^{*}$, there exists at least one close-by neuron
${\bm{w}}$ s.t.
$\angle({\bm{w}},{\bm{w}}_{i}^{*})\leq\delta_{close}=O_{*}(\zeta^{1/3})$.

  3. (iii)

Average neuron is close to teach neurons: we have
$\left\|{\bm{v}}_{i}-{\bm{w}}_{i}^{*}\right\|_{2}\leq
O_{*}((\zeta/\lambda)^{3/4})$.

These properties give us a sense of what the network should look like when
loss is small: neurons have large norm only if they are around the ground-
truth directions. Moreover, when $\zeta/\lambda\to 0$, student neuron must
align with one of teacher neurons ($\delta_{j}=0$) or norm becomes 0
($|a_{j}|\left\|{\bm{w}}_{j}\right\|=0$). This can be understood from the
$\ell_{1}$ regularized loss (equivalent to $\ell_{2}$ regularization on both
layers) that promotes the sparsity over the distribution of neurons. In the
rest of this section, we discuss new techniques such as dual certificate that
we develop for the proof.

###  6.2 Neurons concentrate around teacher neurons: dual certificate analysis
and test function

We focus on Lemma 6.1(i)(ii) here. We will use a dual certificate technique
similar to Poon et al. (2023) to prove Lemma 6.1(i), and a more general
construction of test function to prove Lemma 6.1(ii). In below, we consider a
relaxed version of original optimization problem (2) by allowing infinite
number of neurons, i.e., distribution of neurons, with $\sigma_{\geq
2}(x)=\text{ReLU}(x)-1/\sqrt{2\pi}-x/2$ instead of ReLU:

$\displaystyle\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu):=L(\mu;\sigma_{\geq_{2}})+\lambda|\mu|_{1},$
(5)

where $\mu_{\lambda}^{*}$ is the minimizer. We use $\sigma_{\geq 2}$
activation because this is the effective activation when linear terms
$\alpha,\beta$ are perfectly fitted (remove 0th and 1st order Hermite
expansion of ReLU, see Claim B.1 and (6) in appendix).

This is the loss function we would have in the idealized setting: (1) linear
term $\alpha,\beta$ reach their global minima (this is easy to achieve as loss
is convex in them); (2) use $\ell_{1}$ regularization instead of $\ell_{2}$
regularization, since this is the case when the first and second layer norm
are balanced (weight decay encourages this to happen). Note that the results
in this part can handle almost all activation as long as its Hermite expansion
is well-defined, generalizing Zhou et al. (2021) that can only handle
absolute/ReLU activation. In below we will focus on the activation
$\sigma_{\geq 2}$ for simplicity.

##### Dual certificate

This optimization problem (5) can be viewed as a natural extension of the
classical compressed sensing problem (Donoho, 2006; Candès et al., 2006) and
Lasso-type problem (Tibshirani, 1996) in the infinite dimensional space, which
has been studied in recent years (Bach, 2017; Poon et al., 2023). One common
way is to study its dual problem. The dual solution $p_{0}({\bm{x}})$ (maps
$\mathbb{R}^{d}$ to $\mathbb{R}$) of (5) when $\lambda=0$ satisfies
$\mathbb{E}_{\bm{x}}[p_{0}({\bm{x}})\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}})]\in\partial|\mu_{*}|({\mathbb{S}}^{d-1})$ (more
detailed discussions on this dual problem can be found in e.g., Poon et al.
(2023)). Here $\eta({\bm{w}})=\mathbb{E}_{\bm{x}}[p({\bm{x}})\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}})]$ is often called dual certificate, as it serves
as a certificate of whether a solution $\mu$ is optimal. Its meaning will be
clear in the discussions below.

We now introduce the notion of non-degenerate dual certificate, motivated by
Poon et al. (2023). Note that the condition
$\eta({\bm{w}})\in\partial|\mu_{*}|({\mathbb{S}}^{d-1})$ implies that
$\eta({\bm{w}}_{i}^{*})=\operatorname{sign}(a_{i}^{*})$ and
$\left\|\eta\right\|_{\infty}\leq 1$. The following definition is a slightly
stronger version of the above implications as it requires $\eta$ to decay at
least quadratic when moves away from ${\bm{w}}_{i}^{*}$.

+1-10$\eta({\bm{w}})$${\bm{w}}_{1}^{*}$${\bm{w}}_{2}^{*}$${\bm{w}}_{3}^{*}$

Figure 1: Dual certificate $\eta$.

######  Definition 1 (Non-degenerate dual certificate).

$\eta({\bm{w}})$ is called a non-degenerate dual certificate if there exists
$p({\bm{x}})$ such that
$\eta({\bm{w}})=\mathbb{E}_{\bm{x}}[p({\bm{x}})\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}})]$ for ${\bm{w}}\in{\mathbb{S}}^{d-1}$ and

  1. (i)

$\eta({\bm{w}}_{i}^{*})=\operatorname{sign}(a_{i}^{*})$ for
$i=1,\ldots,m_{*}$.

  2. (ii)

$|\eta({\bm{w}})|\leq 1-\rho_{\eta}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}$ if
${\bm{w}}\in{\mathcal{T}}_{i}$, where
$\delta({\bm{w}},{\bm{w}}_{i}^{*})=\angle({\bm{w}},{\bm{w}}_{i}^{*})$.

The existence and construction of the non-degenerate dual certificate is
deferred to Appendix G. We focus on the implications of such non-degenerate
dual certificate below.

Roughly speaking, the dual certificate only focuses on the position of ground-
truth directions ${\bm{w}}_{i}^{*}$ as it decays fast when moving away from
these directions (Figure 1). Thus, if $\mu$ exactly recovers ground-truth
$\mu_{*}$, then we have $\langle\eta,\mu_{*}\rangle=|\mu_{*}|_{1}$. The gap
between $\langle\eta,\mu\rangle$ and $|\mu|_{1}$ is large when $\mu$ is away
from $\mu_{*}$. Therefore, $\eta$ can be viewed as a certificate to test the
optimality of $\mu$. The lemma below makes it more precise.

######  Lemma 6.2.

Given a non-degenerate dual certificate $\eta$, then

  1. (i)

$\langle\eta,\mu^{*}\rangle=|\mu^{*}|_{1}$ and
$|\langle\eta,\mu-\mu^{*}\rangle|\leq\left\|p\right\|_{2}\sqrt{L(\mu)}$.

  2. (ii)

For any measure $\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})$,
$|\langle\eta,\mu\rangle|\leq|\mu|_{1}-\rho_{\eta}\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu|({\bm{w}})$.

In the finite width case, we have
$\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu|({\bm{w}})=\sum_{i}|a_{i}|\left\|{\bm{w}}_{i}\right\|\delta_{i}^{2}$.
This is exactly the quantity that we are interested in Lemma 6.1.

To see the usefulness of Lemma 6.2, we show a proof for total norm bound of
the optimal solution $\mu_{\lambda}^{*}$. The proof for general $\mu$ with
optimality gap $\zeta$ is similar (Lemma F.5).

######  Claim 6.1 (Lemma 6.1(i) for $\mu_{\lambda}^{*}$).

$\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu_{\lambda}^{*}|({\bm{w}})\leq
O_{*}(\lambda)$

###### Proof.

It is not hard to show $|\mu_{\lambda}^{*}|_{1}\leq|\mu^{*}|_{1}$ (Lemma F.3)
so we have

$\displaystyle|\mu_{\lambda}^{*}|_{1}-|\mu^{*}|_{1}-\langle\eta,\mu_{\lambda}^{*}-\mu^{*}\rangle\leq-\langle\eta,\mu_{\lambda}^{*}-\mu^{*}\rangle.$

Using Lemma 6.2 and the fact $L(\mu_{\lambda}^{*})=O_{*}(\lambda^{2})$ from
Lemma F.3,

$\displaystyle\operatorname{LHS}=$
$\displaystyle|\mu_{\lambda}^{*}|_{1}-\langle\eta,\mu_{\lambda}^{*}\rangle\geq\rho_{\eta}\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu_{\lambda}^{*}|({\bm{w}}),\quad\operatorname{RHS}\leq\left\|p\right\|_{2}\sqrt{L(\mu_{\lambda}^{*})}=O_{*}(\lambda).$

We have
$\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu_{\lambda}^{*}|({\bm{w}})=O_{*}(\lambda)$.
∎

##### Test function

The idea of using test function is to identify certain properties of the
target function/distribution that we are interested in. Specifically, we
construct test function so that it only correlates well with the target
function that has the desired property. Generally speaking, the dual
certificate above can be consider as a specific case of a test function: the
correlation between dual certificate $\eta$ and distribution of neurons $\mu$
is large (reach $|\mu|_{1}$) only when $\mu\approx\mu_{*}$.

In below, we use this test function idea to show that every ground-truth
direction has close-by neuron (Lemma 6.1(ii)). Denote
${\mathcal{T}}_{i}(\delta):=\\{j:\angle({\bm{w}}_{j},{\bm{w}}_{i})\leq\delta\\}\cap{\mathcal{T}}_{i}$
as the neurons that are $\delta$-close to ${\bm{w}}_{i}^{*}$.

######  Lemma 6.3 (Lemma 6.1(ii), informal).

Given the optimality gap $\zeta$, we have the total mass near each target
direction is large, i.e.,
$\mu({\mathcal{T}}_{i}(\delta))\operatorname{sign}(a_{i}^{*})\geq|a_{i}^{*}|/2$
for all $i\in[m_{*}]$ and any $\delta\geq\Theta_{*}\left(\zeta^{1/3}\right)$.

0$g({\bm{x}})$${\bm{w}}_{1}^{*}$${\bm{w}}_{2}^{*}$${\bm{w}}_{3}^{*}$

Figure 2: Test function $g$.

Note that although the results in the dual certificate part (Lemma 6.2(ii))
can imply that there are neurons close to teacher neurons, the bound we get
here using carefully designed test function are sharper ($\zeta^{1/3}$ vs.
$\zeta^{1/4}$). This is in fact important to the descent direction
construction (Lemma 5.1).

In the proof, we view the residual
$R({\bm{x}})=f_{\mu}({\bm{x}})-f_{*}({\bm{x}})$ as the target function and
construct test function that will only have large correlation if there is a
teacher neuron that have no close student neurons. Specifically, the test
function $g$ only consists of high-order Hermite polynomial such that it is
large around the ground-truth direction and decays fast when moving away
(Figure 2). It looks like a single spike in dual certificate $\eta$, but in
fact decays much faster than $\eta$ when moving away. It is more flexible to
choose test function than dual certificate, so test function $g$ can focus
only on a local region of one ground-truth direction and give a better
guarantee than dual certificate analysis.

###  6.3 Average neuron is close to teacher neuron: residual decomposition and
average neuron

We give the proof idea for Lemma 6.1(iii) that shows average neuron
${\bm{v}}_{i}$ is close to teacher neuron ${\bm{w}}_{i}^{*}$ using the
residual decomposition $R=R_{1}+R_{2}+R_{3}$.

The key is to observe that $R_{1}$ is an analogue to exact-parametrization
case where loss is often strongly-convex, so we have
$\left\|R_{1}\right\|_{2}^{2}=\Omega_{*}(1)\sum_{i}\left\|{\bm{v}}_{i}-{\bm{w}}_{i}^{*}\right\|_{2}^{2}$.
Then the goal is to upper bound $\left\|R_{1}\right\|$. Given the
decomposition $R=R_{1}+R_{2}+R_{3}$, it is easy to bound
$\left\|R_{1}\right\|\leq\left\|R\right\|+\left\|R_{2}\right\|+\left\|R_{3}\right\|$.
We focus on $\left\|R_{2}\right\|$ as the other two are not hard to bound
(loss is small in local regime). $R_{2}$ is in fact closely related with the
total weighted norm bound in Lemma 6.1: we show
$\left\|R_{2}\right\|=O_{*}(1)\left(\sum_{j\in{\mathcal{T}}_{i}}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}\delta_{j}^{2}\right)^{3/2}=O_{*}((\zeta/\lambda)^{3/2})$.
Thus, we get a bound for $\left\|{\bm{v}}_{i}-{\bm{w}}_{i}^{*}\right\|$. See
Appendix F.1.4 for details.

##  7 Conclusion

In this paper we showed that gradient descent converges in a large local
region depending on the complexity of the teacher network, and the local
convergence allows 2-layer networks to perform a strong notion of feature
learning (matching the directions of ground-truth teacher networks). We hope
our result gives a better understanding of why gradient-based training is
important for feature learning in neural networks. Our results rely on adding
standard weight decay and new constructions of dual certificate and test
functions, which can be helpful in understanding local optimization landscape
in other problems. A natural but challenging next step is to understand
whether the intermediate steps are also important for feature learning.

## Acknowledgement

Rong Ge and Mo Zhou are supported by NSF Award DMS-2031849 and CCF-1845171
(CAREER).

## References

  * Abbe et al. (2021) Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj.  The staircase property: How hierarchical structure can guide deep learning.  _Advances in Neural Information Processing Systems_ , 34:26989–27002, 2021.
  * Abbe et al. (2022) Emmanuel Abbe, Enric Boix Adsera, and Theodor 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.
  * Abbe et al. (2023) Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz.  Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics.  In _The Thirty Sixth Annual Conference on Learning Theory_ , pages 2552–2623. PMLR, 2023.
  * Absil et al. (2008) P-A Absil, Robert Mahony, and Rodolphe Sepulchre.  _Optimization algorithms on matrix manifolds_.  Princeton University Press, 2008.
  * Absil et al. (2013) P-A Absil, Robert Mahony, and Jochen Trumpf.  An extrinsic look at the riemannian hessian.  In _International conference on geometric science of information_ , pages 361–368. Springer, 2013.
  * Akiyama and Suzuki (2021) Shunta Akiyama and Taiji Suzuki.  On learnability via gradient method for two-layer relu neural networks in teacher-student setting.  In _International Conference on Machine Learning_ , pages 152–162. PMLR, 2021.
  * Allen-Zhu et al. (2019a) Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang.  Learning and generalization in overparameterized neural networks, going beyond two layers.  _Advances in neural information processing systems_ , 32, 2019a.
  * Allen-Zhu et al. (2019b) Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International conference on machine learning_ , pages 242–252, 2019b.
  * Arora et al. (2019) Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  _Advances in neural information processing systems_ , 32, 2019.
  * Arous et al. (2021) Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath.  Online stochastic gradient descent on non-convex losses from high-dimensional inference.  _The Journal of Machine Learning Research_ , 22(1):4788–4838, 2021.
  * Ba et al. (2022) Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang.  High-dimensional asymptotics of feature learning: How one gradient step improves the representation.  _Advances in Neural Information Processing Systems_ , 35:37932–37946, 2022.
  * Bach (2017) Francis Bach.  Breaking the curse of dimensionality with convex neural networks.  _Journal of Machine Learning Research_ , 18(19):1–53, 2017.
  * Bai and Lee (2020) Yu Bai and Jason D. Lee.  Beyond linearization: On quadratic and higher-order approximation of wide neural networks.  In _International Conference on Learning Representations_ , 2020.
  * Barak et al. (2022) Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang.  Hidden progress in deep learning: Sgd learns parities near the computational limit.  _Advances in Neural Information Processing Systems_ , 35:21750–21764, 2022.
  * Berthier et al. (2023) Raphaël Berthier, Andrea Montanari, and Kangjie Zhou.  Learning time-scales in two-layers neural networks.  _arXiv preprint arXiv:2303.00055_ , 2023.
  * Bietti et al. (2022) Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song.  Learning single-index models with shallow neural networks.  _Advances in Neural Information Processing Systems_ , 35:9768–9783, 2022.
  * Bietti et al. (2023) Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien.  On learning gaussian multi-index models with gradient flow.  _arXiv preprint arXiv:2310.19793_ , 2023.
  * Candès et al. (2006) Emmanuel J Candès, Justin Romberg, and Terence Tao.  Robust uncertainty principles: Exact signal reconstruction from highly incomplete frequency information.  _IEEE Transactions on information theory_ , 52(2):489–509, 2006.
  * Chizat (2022) Lenaic Chizat.  Sparse optimization on measures with over-parameterized gradient descent.  _Mathematical Programming_ , 194(1):487–532, 2022.
  * Chizat and Bach (2018) Lenaic Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  In _Advances in neural information processing systems_ , pages 3036–3046, 2018.
  * Chizat et al. (2019) Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  In _Advances in Neural Information Processing Systems_ , pages 2933–2943, 2019.
  * Damian et al. (2022) Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi.  Neural networks can learn representations with gradient descent.  In _Conference on Learning Theory_ , pages 5413–5452. PMLR, 2022.
  * Dandi et al. (2023) Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan.  How two-layer neural networks learn, one (giant) step at a time.  In _NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning_ , 2023.
  * Daniely and Malach (2020) Amit Daniely and Eran Malach.  Learning parities with neural networks.  _Advances in Neural Information Processing Systems_ , 33:20356–20365, 2020.
  * Donoho (2006) David L Donoho.  Compressed sensing.  _IEEE Transactions on information theory_ , 52(4):1289–1306, 2006.
  * Du et al. (2019) Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In _International Conference on Learning Representations_ , 2019.
  * Frei et al. (2020) Spencer Frei, Yuan Cao, and Quanquan Gu.  Agnostic learning of a single neuron with gradient descent.  _Advances in Neural Information Processing Systems_ , 33:5417–5428, 2020.
  * Ge et al. (2018) Rong Ge, Jason D Lee, and Tengyu Ma.  Learning one-hidden-layer neural networks with landscape design.  In _International Conference on Learning Representations_ , 2018.
  * Ge et al. (2021) Rong Ge, Yunwei Ren, Xiang Wang, and Mo Zhou.  Understanding deflation process in over-parametrized tensor decomposition.  _Advances in Neural Information Processing Systems_ , 34:1299–1311, 2021.
  * Glasgow (2024) Margalit Glasgow.  SGD finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the XOR problem.  In _The Twelfth International Conference on Learning Representations_ , 2024.
  * Goel et al. (2020) Surbhi Goel, Aravind Gollakota, Zhihan Jin, Sushrut Karmalkar, and Adam Klivans.  Superpolynomial lower bounds for learning one-layer neural networks using gradient descent.  In _International Conference on Machine Learning_ , pages 3587–3596. PMLR, 2020.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Advances in neural information processing systems_ , pages 8571–8580, 2018.
  * Li et al. (2020) Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang.  Learning over-parametrized two-layer neural networks beyond ntk.  In _Conference on learning theory_ , pages 2613–2682. PMLR, 2020.
  * Lojasiewicz (1963) Stanislaw Lojasiewicz.  Une propriété topologique des sous-ensembles analytiques réels.  _Les équations aux dérivées partielles_ , 117:87–89, 1963.
  * Mahankali et al. (2024) Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma.  Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Mei et al. (2018) Song Mei, Andrea Montanari, and Phan-Minh Nguyen.  A mean field view of the landscape of two-layer neural networks.  _Proceedings of the National Academy of Sciences_ , 115(33):E7665–E7671, 2018.
  * Mousavi-Hosseini et al. (2023) Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu.  Neural networks efficiently learn low-dimensional representations with SGD.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Nichani et al. (2024a) Eshaan Nichani, Alex Damian, and Jason D Lee.  Provable guarantees for nonlinear feature learning in three-layer neural networks.  _Advances in Neural Information Processing Systems_ , 36, 2024a.
  * Nichani et al. (2024b) Eshaan Nichani, Alex Damian, and Jason D Lee.  How transformers learn causal structure with gradient descent.  _arXiv preprint arXiv:2402.14735_ , 2024b.
  * O’Donnell (2021) Ryan O’Donnell.  Analysis of boolean functions.  _arXiv preprint arXiv:2105.10386_ , 2021.
  * Poon et al. (2023) Clarice Poon, Nicolas Keriven, and Gabriel Peyré.  The geometry of off-the-grid compressed sensing.  _Foundations of Computational Mathematics_ , 23(1):241–327, 2023.
  * Safran et al. (2021) Itay M Safran, Gilad Yehudai, and Ohad Shamir.  The effects of mild over-parameterization on the optimization landscape of shallow relu neural networks.  In _Conference on Learning Theory_ , pages 3889–3934. PMLR, 2021.
  * Shi et al. (2022) Zhenmei Shi, Junyi Wei, and Yingyu Liang.  A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features.  In _International Conference on Learning Representations_ , 2022.
  * Soltanolkotabi (2017) Mahdi Soltanolkotabi.  Learning relus via gradient descent.  _Advances in neural information processing systems_ , 30, 2017.
  * Suzuki et al. (2024) Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda.  Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Tibshirani (1996) Robert Tibshirani.  Regression shrinkage and selection via the lasso.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ , 58(1):267–288, 1996.
  * Wang et al. (2024) Zihao Wang, Eshaan Nichani, and Jason D. Lee.  Learning hierarchical polynomials with three-layer neural networks.  In _The Twelfth International Conference on Learning Representations_ , 2024.
  * Wu (2022) Lei Wu.  Learning a single neuron for non-monotonic activation functions.  In _International Conference on Artificial Intelligence and Statistics_ , pages 4178–4197. PMLR, 2022.
  * Xu and Du (2023) Weihang Xu and Simon Du.  Over-parameterization exponentially slows down gradient descent for learning a single neuron.  In _The Thirty Sixth Annual Conference on Learning Theory_ , pages 1155–1198. PMLR, 2023.
  * Yehudai and Ohad (2020) Gilad Yehudai and Shamir Ohad.  Learning a single neuron with gradient methods.  In _Conference on Learning Theory_ , pages 3756–3786. PMLR, 2020.
  * Yehudai and Shamir (2019) Gilad Yehudai and Ohad Shamir.  On the power and limitations of random features for understanding neural networks.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Zhou et al. (2021) Mo Zhou, Rong Ge, and Chi Jin.  A local convergence theory for mildly over-parameterized two-layer neural network.  In _Conference on Learning Theory_ , pages 4577–4632. PMLR, 2021.

##  Appendix A Some properties of Hermite polynomials

In this section, we give several properties of Hermite polynomials that are
useful in our analysis. See O’Donnell (2021) for a more complete discussion on
Hermite polynomials. Let $H_{k}$ be the probabilists’ Hermite polynomial where

$\displaystyle
H_{k}(x)=(-1)^{k}e^{x^{2}/2}\frac{\,\mathrm{d}^{k}}{\,\mathrm{d}x^{k}}(e^{-x^{2}/2})$

and $h_{k}=\frac{1}{\sqrt{k!}}H_{k}$ be the normalized Hermite polynomials.

Hermite polynomials are classical orthogonal polynomials, which means
$\mathbb{E}_{x\sim N(0,1)}[h_{m}(x)h_{n}(x)]=1$ if $m=n$ and otherwise 0.
Given a function $\sigma$, we call
$\sigma(x)=\sum_{k=0}^{\infty}\hat{\sigma}_{k}h_{k}(x)$ as the Hermit
expansion of $\sigma$ and $\hat{\sigma}_{k}=\mathbb{E}_{x\sim
N(0,1)}[\sigma(x)h_{k}(x)]$ as the $k$-th Hermite coefficient of $\sigma$.

The following is a useful property of Hermite polynomial.

######  Claim A.1 ((O’Donnell, 2021), Section 11.2).

Let $(x,y)$ be $\rho$-correlated standard normal variables (that is, both
$x,y$ have marginal distribution $N(0,1)$ and $\mathbb{E}[xy]=\rho$). Then,
$\mathbb{E}[h_{m}(x)h_{n}(y)]=\rho^{n}\delta_{mn}$, where $\delta_{mn}=1$ if
$m=n$ and otherwise 0.

The following lemma gives the Hermite coefficients for absolute value function
and ReLU.

######  Lemma A.1.

Let $\hat{\sigma}_{k}=\mathbb{E}_{x\sim N(0,1)}[\sigma(x)h_{k}(x)]$ be the
Hermite coefficient of $\sigma$. For $\sigma$ is ReLU or absolute function, we
have $|\hat{\sigma}_{k}|=\Theta(k^{-5/4})$.

###### Proof.

From Goel et al. (2020); Zhou et al. (2021) we have

$\displaystyle\hat{\sigma}_{abs,k}=\left\\{\begin{array}[]{ll}0&\text{, $k$ is
odd}\\\ \sqrt{2/\pi}&\text{, $k=0$}\\\
(-1)^{\frac{k}{2}-1}\sqrt{\frac{2}{\pi}}\frac{(k-2)!}{\sqrt{k!}2^{k/2-1}(k/2-1)!}&\text{,
$k$ is even and $k\geq 2$}\end{array}\right.$

$\displaystyle\hat{\sigma}_{relu,k}=\left\\{\begin{array}[]{ll}0&\text{, $k$
is odd and $k\geq 3$}\\\ \sqrt{1/2\pi}&\text{, $k=0$}\\\ 1/2&\text{, $k=1$}\\\
(-1)^{\frac{k}{2}-1}\sqrt{\frac{1}{2\pi}}\frac{(k-2)!}{\sqrt{k!}2^{k/2-1}(k/2-1)!}&\text{,
$k$ is even and $k\geq 2$}\end{array}\right.$

Using Stirling’s formula, we get
$|\hat{\sigma}_{abs,k}|,|\hat{\sigma}_{relu,k}|=\Theta(k^{-5/4})$. ∎

##  Appendix B Useful facts and proof of Theorem 3.1

In this section we provide several useful facts and present the proof of
Theorem 3.1.

The following claim shows that the square loss can be decomposed into 3 terms,
where $\alpha,{\bm{\beta}}$ are corresponding to 0th and 1st order of Hermite
expansion. The effective activation is in fact $\sigma_{\geq 2}$ as defined
below.

######  Claim B.1.

Denote
$\hat{\alpha}=-(1/\sqrt{2\pi})\sum_{i=1}^{m}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}$,
$\hat{{\bm{\beta}}}=-(1/2)\sum_{i=1}^{m}a_{i}{\bm{w}}_{i}$.We have square loss

$\displaystyle
L({\bm{\theta}})=|\alpha-\hat{\alpha}|^{2}+\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}+\mathbb{E}_{\bm{x}}[(f_{\geq
2}({\bm{x}})-\widetilde{f}_{*}({\bm{x}}))^{2}]$

where $f_{\geq 2}({\bm{x}};{\bm{\theta}})=\sum_{i\in[m]}a_{i}\sigma_{\geq
2}({\bm{w}}_{i}^{\top}{\bm{x}})$ and $\sigma_{\geq
2}(x)=\sigma(x)-1/\sqrt{2\pi}-x/2$ is the activation that after removing 0th
and 1st order term in Hermite expansion.

As a result, when $\alpha,{\bm{\beta}}$ are perfectly fitted and norms are
balanced we have

$L_{\lambda}({\bm{\theta}})=\mathbb{E}_{\bm{x}}[(f_{\geq
2}({\bm{x}})-\widetilde{f}_{*}({\bm{x}}))^{2}]+\lambda\sum_{i\in[m]}.|a_{i}|\left\|{\bm{w}}_{i}\right\|_{2}$

###### Proof.

Following Ge et al. (2018), we can write the loss $L({\bm{\theta}})$ as a sum
of tensor decomposition problem using Hermite expansion as in Section A
(recall $\left\|{\bm{w}}_{i}^{*}\right\|_{2}=1$ and preprocessing procedure
removes the 0-th and 1-st order term in the Hermite expansion of $\sigma$):

$\displaystyle L({\bm{\theta}})=$
$\displaystyle\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\sum_{k\geq
0}\hat{\sigma}_{k}h_{k}(\overline{{\bm{w}}}_{i}^{\top}{\bm{x}})+\alpha+h_{1}({\bm{\beta}}^{\top}{\bm{x}})-\sum_{i\in[m_{*}]}a_{i}^{*}\left\|{\bm{w}}_{i}^{*}\right\|_{2}\sum_{k\geq
2}\hat{\sigma}_{k}h_{k}({\bm{w}}_{i}^{*\top}{\bm{x}})\right)^{2}\right]$
$\displaystyle=$
$\displaystyle\left|\alpha+\hat{\sigma}_{0}\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\right|^{2}+\left\|{\bm{\beta}}+\hat{\sigma}_{1}\sum_{i\in[m]}a_{i}{\bm{w}}_{i}\right\|_{2}^{2}$
$\displaystyle+\sum_{k\geq
2}\hat{\sigma}_{k}^{2}\left\|\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\overline{{\bm{w}}}_{i}^{\otimes
k}-\sum_{i\in[m_{*}]}a_{i}^{*}\left\|{\bm{w}}_{i}^{*}\right\|_{2}{\bm{w}}_{i}^{*\otimes
k}\right\|_{F}^{2}.$

Note that $\hat{\sigma}_{0}=1/\sqrt{2\pi}$, $\hat{\sigma}_{1}=1/2$ as in Lemma
A.1, we get the result. ∎

The proof of main result Theorem 3.1 is simply a combination of few lemmas
appear in other sections. We refer the detailed proof and discussion to their
corresponding sections. See 3.1

###### Proof.

Combine Lemma 4.1 (Stage 1), Lemma 4.2 (Stage 2) and Lemma 4.3 (Stage 3)
together and follow the choice of $\lambda_{t}$ and $\eta_{t}$ we get the
result.

For the student neurons’ alignment, it is a direct corollary from Lemma F.6
and Lemma F.5. ∎

##  Appendix C Stage 1: first gradient step

In this section, we show that after the first gradient update the first layer
weights ${\bm{w}}_{1},\ldots,{\bm{w}}_{m}$ form a ${\varepsilon}_{0}$-net of
the target subspace $S_{*}$, given $m=(1/{\varepsilon}_{0})^{O(r)}$ neurons.
The proof is deferred to Section C.1.

See 4.1

The proof relies on the following lemma from Damian et al. (2022) that shows
after the first step update ${\bm{w}}_{i}$’s are located at positions as if
they are sampled within the target subspace $S_{*}$.

######  Lemma C.1 (Lemma 4, Damian et al. (2022)).

Under Assumption 2.3, we have with high probability in the $\ell_{2}$ norm
sense

$\displaystyle{\bm{w}}_{i}^{(1)}$
$\displaystyle=-\eta_{0}\nabla_{{\bm{w}}_{i}}L({\bm{a}}^{(0)},{\bm{W}}^{(0)})=-2\eta_{0}a_{i}^{(0)}\left(\hat{\sigma}_{2}^{2}{\bm{H}}\overline{{\bm{w}}}_{i}\pm\widetilde{O}(\frac{\sqrt{r}}{d})\right),$

where $\hat{\sigma}_{k}:=\mathbb{E}_{\bm{x}}[\sigma({\bm{x}})h_{k}({\bm{x}})]$
is the $k$-th Hermite polynomial coefficient.

###  C.1 Proofs in Section C

We now are ready to give the proof of Lemma 4.1. See 4.1

###### Proof.

We show them one by one.

##### Part (i)

From Lemma C.1 and the fact that $\overline{{\bm{w}}}_{i}^{(0)}$ samples
uniformly from unit sphere, we know the probability of
$\angle(\overline{{\bm{w}}}_{i}^{(1)},{\bm{w}})$ for any given ${\bm{w}}$ is
at least $\Omega_{*}({\varepsilon}_{0}^{r})$. Applying union bound we get the
desired result.

##### Part (ii)

We have

$\displaystyle{\bm{w}}_{i}^{(1)}$
$\displaystyle=-\eta_{0}\nabla_{{\bm{w}}_{i}}L({\bm{a}}^{(0)},{\bm{W}}^{(0)})=a_{i}^{(0)}\mathbb{E}_{\bm{x}}[\widetilde{f}_{*}({\bm{x}})\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}]$

For the norm bound, using Lemma C.1 we know

$\displaystyle\sqrt{d}\left(\left\|{\bm{H}}\overline{{\bm{w}}}_{i}^{(0)}\right\|_{2}-\widetilde{O}(\frac{\sqrt{r}}{d})\right)\leq\left\|{\bm{w}}_{i}^{(1)}\right\|_{2}\leq\sqrt{d}\left(\left\|{\bm{H}}\overline{{\bm{w}}}_{i}^{(0)}\right\|_{2}+\widetilde{O}(\frac{\sqrt{r}}{d})\right).$

Since ${\bm{w}}_{i}^{(0)}$ initializes from Gaussian distribution, we know the
desired bound hold. Similarly, one can bound $|a_{i}^{(1)}|$.

Since we use a symmetric initialization and have preprocessed the data, it is
easy to see $\alpha,{\bm{\beta}}$ remains at 0. ∎

##  Appendix D Stage 2: reaching low loss

In Stage 2, we show that given the features learned in Stage 1 one can adjust
the norms on top of it to reach low loss that enters the local convergence
regime in Stage 3.

##### Procedure

We first specify the procedure to solve
$\min_{\bm{a}}\min_{\alpha,{\bm{\beta}}}L({\bm{\theta}})+\lambda\sum_{i}\left\|{\bm{w}}_{i}\right\|_{2}|a_{i}|$.
For ${\bm{a}}$ at current point, we first solve the inner optimization
problem, which is a linear regression on $\alpha,{\bm{\beta}}$. From Claim B.1
we know the global minima is $(\hat{\alpha},\hat{{\bm{\beta}}})$. For
simplicity of the proof, we just directly set
$(\alpha,{\bm{\beta}})=(\hat{\alpha},\hat{{\bm{\beta}}})$. Then given the
$\alpha,{\bm{\beta}}$, the outer optimization is a convex optimization for
${\bm{a}}$, which can also be solved efficiently. Specifically, we perform 1
step of (sub)gradient on the loss function. We repeat the above 2 steps until
convergence.

From Claim B.1 we know the actual objective that we optimize is

$\widetilde{L}_{1,\lambda}({\bm{a}})=\mathbb{E}_{\bm{x}}[({\bm{a}}^{\top}\sigma_{\geq
2}({\bm{W}}{\bm{x}})-\widetilde{y})^{2}]+\lambda\sum_{i}\left\|{\bm{w}}_{i}\right\|_{2}|a_{i}|.$

The following lemma shows that after Stage 2 we reach a low loss solution
given the first layer features learned after first gradient step. The proof
requires $\eta$ to be small enough that depends on $1/m$, mostly due to the
large gradient norm. We believe using more advance algorithm for this type of
problem can alleviate this issue. However, as this is not the focus of this
paper, we omit it for simplicity.

See 4.2

###### Proof.

Denote $\widetilde{{\bm{a}}}_{*}$ as the minima of
$\widetilde{L}_{1,\lambda}$. Then, we have

$\displaystyle\left\|{\bm{a}}^{(t+1)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}=$
$\displaystyle\left\|{\bm{a}}^{(t)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}-2\eta\langle\nabla_{\bm{a}}\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)}),{\bm{a}}^{(t)}-\widetilde{{\bm{a}}}_{*}\rangle+\eta^{2}\left\|\nabla_{\bm{a}}\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)})\right\|_{2}^{2}$
$\displaystyle\stackrel{{\scriptstyle(\text{a})}}{{\leq}}$
$\displaystyle\left\|{\bm{a}}^{(t)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}-2\eta(\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*}))+\eta^{2}O_{*}(m)$
$\displaystyle=$
$\displaystyle\left\|{\bm{a}}^{(t)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}-2\eta(\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*}))+\eta{\varepsilon}_{0}/2,$

where (a) we use idea loss $\widetilde{L}_{1,\lambda}$ is convex in
${\bm{a}}$.

Iterating the above inequality over all $t$ we have

$\displaystyle\left\|{\bm{a}}^{(T)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}\leq$
$\displaystyle\left\|{\bm{a}}^{(1)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}-2\eta\sum_{t\leq
T}(\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*}))+\eta
T{\varepsilon}_{0}/2,$

which means

$\displaystyle\min_{t\leq
T}\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})\leq\frac{1}{T}\sum_{t\leq
T}(\widetilde{L}_{1,\lambda}({\bm{a}}^{(t)})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*}))\leq$
$\displaystyle\frac{\left\|{\bm{a}}^{(1)}-\widetilde{{\bm{a}}}_{*}\right\|_{2}^{2}}{\eta
T}+{\varepsilon}_{0}/2.$

It is easy to see
$\left\|{\bm{a}}^{(1)}\right\|_{2},\left\|\widetilde{{\bm{a}}}_{*}\right\|_{1}=O_{*}(1)$.
Thus, when $T\geq O_{*}(1/\eta{\varepsilon}_{0})$ we know
$\widetilde{L}_{1,\lambda}({\bm{a}}^{(T_{2})})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})\leq
3{\varepsilon}_{0}/4$.

This suggests the optimality gap after balancing the norm (so that
$L_{\lambda}({\bm{\theta}}^{(T_{2})})=\widetilde{L}_{1,\lambda}({\bm{a}}^{(T_{2})})$)

$\displaystyle\zeta_{T_{2}}=$ $\displaystyle
L_{\lambda}({\bm{\theta}}^{(T_{2})})-\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu)$
$\displaystyle=$
$\displaystyle\widetilde{L}_{1,\lambda}({\bm{a}}^{(T_{2})})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})+\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})-\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu).$

For
$\widetilde{L}_{1,\lambda}({\bm{a}}^{(T_{2})})-\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})$,
we just show above that it is less than $3{\varepsilon}_{0}/4$.

For
$\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})-\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu)$,
we have

$\displaystyle\widetilde{L}_{1,\lambda}(\widetilde{{\bm{a}}}_{*})-\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu)\leq$
$\displaystyle\widetilde{L}_{1,\lambda}(\hat{{\bm{a}}}_{*})-\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu)$
$\displaystyle\leq$ $\displaystyle
O_{*}(\varepsilon_{0}^{2})+\lambda\left\|{\bm{a}}_{*}\right\|_{1}-\lambda|\mu_{\lambda}^{*}|_{1}\leq
O_{*}(\lambda^{2}),$

where in the last inequality we use Lemma F.3 and
$\mu_{\lambda}^{*}=\arg\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu)$.
Here $\hat{{\bm{a}}}_{*}$ is a rescaled version of ${\bm{a}}_{*}$ and is
constructed as: for every teacher neuron ${\bm{w}}_{i}^{*}$ choose the closest
neuron ${\bm{w}}_{j}$ s.t.
$\angle({\bm{w}}_{j},{\bm{w}}_{i}^{*})\leq\varepsilon_{0}$ and set
$\hat{{\bm{a}}}_{*,j}=a_{i}^{*}/\left\|{\bm{w}}_{j}\right\|_{2}$. Set all
other $\hat{{\bm{a}}}_{*,k}=0$.

Together with above calculations, we have $\zeta_{T_{2}}\leq
O_{*}({\varepsilon}_{0})$. ∎

##  Appendix E Stage 3: local convergence for regularized 2-layer neural
networks

In this section we show the local convergence that loss eventually goes to 0
within polynomial time and recovers teacher neurons’ direction.

The results in this section only need the width $m\geq m_{*}$ as long as its
initial loss is small. See 4.3

The goal of each epoch is to minimize the loss $L_{\lambda}$ with a fix
$\lambda$. The lemma below shows that as long as the initial optimality gap is
$O_{*}(\lambda^{9/5})$, then at the end of each epoch, $L_{\lambda}$ could
decrease to $O_{*}(\lambda^{2})$. Therefore, using a slow decay of weight
decay parameter $\lambda$ for each epoch we could stay in the local
convergence regime for each epoch and eventually recovers the target network.

######  Lemma E.1 (Loss improve within one epoch).

Suppose $|a_{i}^{(0)}|\leq\left\|{\bm{w}}_{i}^{(0)}\right\|_{2}$ for all
$i\in[m]$. If $\zeta_{0}\leq O_{*}(\lambda^{9/5})$ and $\lambda\leq O_{*}(1)$
and $\eta\leq O_{*}(\lambda^{12}d^{-3})$, then within
$O_{*}(\lambda^{-4}\eta^{-1})$ time the optimality gap becomes
$L_{\lambda}-L_{\lambda}(\mu_{\lambda}^{*})=O_{*}(\lambda^{2})$.

The above result relies on the following characterization of local landscape
of regularized loss. We show the gradient is large whenever the optimality gap
is large. This is the main contribution of this paper, see Section F for
detailed proofs. See 4.4

In order to use the above landscape result with standard descent lemma, we
also need certain smoothness condition on the loss function. We show below
that this regularized loss indeed satisfies certain smoothness condition
(though weaker than standard smoothness condition) to allow the convergence
analysis.

######  Lemma E.2 (Smoothness).

Suppose $|a_{i}|\leq\left\|{\bm{w}}_{i}\right\|_{2}$ and
$\left\|\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}=O_{*}(d)$
for all $i\in[m]$. If $\eta=O_{*}(1/d)$, then

$\displaystyle
L_{\lambda}({\bm{\theta}}-\eta\nabla_{\bm{\theta}}L_{\lambda})\leq
L_{\lambda}({\bm{\theta}})-\eta\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}+O_{*}(\eta^{3/2}d^{3/2})$

###  E.1 Proofs in Section E

We now are ready to show the convergence of Stage 3 by using Lemma E.1 to show
the loss makes progress every epoch. See 4.3

###### Proof.

Since $|a_{i}^{(0)}|\leq\left\|{\bm{w}}_{i}^{(0)}\right\|_{2}$ for all
$i\in[m]$ at the beginning of Stage 3, from Lemma E.3 we know they will remain
hold for all epoch and all time $t$.

From Lemma E.1 we know for epoch $k$ it finishes within
$O_{*}(\lambda_{k}^{-4}\eta^{-1})$ time and achieves
$L_{\lambda_{k}}-L_{\lambda_{k}}(\mu_{\lambda_{k}}^{*})=O_{*}(\lambda_{k}^{2})$.
To proceed to next epoch $k+1$, we only need to show the solution at the end
of epoch $k$ ${\bm{\theta}}^{(k)}$ gives the optimality gap
$\zeta=O_{*}(\lambda_{k+1}^{9/5})$ for the next $\lambda_{k+1}$. We have

$\displaystyle
L_{\lambda_{k+1}}({\bm{\theta}}^{(k)})-L_{\lambda_{k+1}}(\mu_{\lambda_{k+1}}^{*})=$
$\displaystyle
L({\bm{\theta}}^{(k)})-L(\mu_{\lambda_{k+1}}^{*})+\frac{\lambda_{k+1}}{2}\left\|{\bm{a}}^{(k)}\right\|_{2}^{2}+\frac{\lambda_{k+1}}{2}\left\|{\bm{W}}^{(k)}\right\|_{F}^{2}-\lambda_{k+1}|\mu_{\lambda_{k+1}}^{*}|_{1}$
$\displaystyle\stackrel{{\scriptstyle(\text{a})}}{{\leq}}$ $\displaystyle
O_{*}(\lambda_{k}^{2})+\frac{\lambda_{k+1}}{\lambda_{k}}\left(\frac{\lambda_{k}}{2}\left\|{\bm{a}}^{(k)}\right\|_{2}^{2}+\frac{\lambda_{k}}{2}\left\|{\bm{W}}^{(k)}\right\|_{F}^{2}-\lambda_{k}|\mu_{\lambda_{k+1}}^{*}|_{1}\right)$
$\displaystyle\stackrel{{\scriptstyle(\text{b})}}{{\leq}}$ $\displaystyle
O_{*}(\lambda_{k}^{2})+\frac{\lambda_{k+1}}{\lambda_{k}}\left(O_{*}(\lambda_{k}^{2})+L(\mu_{\lambda_{k}}^{*})-L({\bm{\theta}}^{(k)})\right)$
$\displaystyle\stackrel{{\scriptstyle(\text{c})}}{{\leq}}$ $\displaystyle
O_{*}(\lambda_{k}^{2})\leq O_{*}(\lambda_{k+1}^{9/5})$

where (a) due to Lemma F.4 that $L({\bm{\theta}}^{(k)})$ is small; (b) the
optimality gap at the end of epoch $k$ is $O_{*}(\lambda_{k}^{2})$ and
$|\mu_{\lambda_{k}}^{*}|_{1}-|\mu_{\lambda_{k+1}}^{*}|_{1}=O_{*}(\lambda_{k})$
from Lemma F.3; (c) due to Lemma F.3 that $L(\mu_{\lambda_{k}}^{*})$ is small.
In this way, we can apply Lemma E.1 again for epoch $k+1$.

From Lemma F.4 we know at the end of epoch $k$ the square loss
$L({\bm{\theta}}^{(k)})=O_{*}(\lambda_{k}^{2})$. Thus, to reach
${\varepsilon}$ square loss, we need $\lambda_{k}=O_{*}({\varepsilon}^{1/2})$,
which means we need to take $O_{*}(\log(1/{\varepsilon}))$ epoch. Since epoch
$k$ it finishes within $O_{*}(\lambda_{k}^{-4}\eta^{-1})$ time, we know the
total time is at most $O_{*}(\lambda_{0}^{-4}\eta^{-1}\log(1/\varepsilon))$
time. ∎

To show the lemma below that loss makes progress within every epoch, we rely
on the gradient lower bound (Lemma 4.4) and smoothness condition of loss
function (Lemma E.2). See E.1

###### Proof.

Since $|a_{i}^{(0)}|\leq\left\|{\bm{w}}_{i}^{(0)}\right\|_{2}$ for all
$i\in[m]$ at the beginning of current epoch, from Lemma E.3 we know they will
remain hold for all time $t$. Then combine Lemma E.4 and Lemma E.2 we know

$\displaystyle
L_{\lambda}({\bm{\theta}}-\eta\nabla_{\bm{\theta}}L_{\lambda})\leq
L_{\lambda}({\bm{\theta}})-\eta\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}+O_{*}(\eta^{3/2}d^{3/2}).$

Recall
$\zeta_{t}=L_{\lambda}({\bm{\theta}}^{(t)})-L_{\lambda}(\mu_{\lambda}^{*})$.
Using gradient lower bound Lemma 4.4 and consider the time before $\zeta_{t}$
reach $O_{*}(\lambda^{2})$ we have

$\displaystyle\zeta_{t+1}\leq\zeta_{t}-\eta\Omega_{*}(\zeta_{t}^{4}/\lambda^{2})+O_{*}(\eta^{3/2}d^{3/2})\leq\zeta_{t}-\Omega_{*}(\eta\zeta_{t}^{4}/\lambda^{2}),$

where we use $\eta=O_{*}(\lambda^{12}d^{-3})$ to be small enough.

The above recursion implies that

$\displaystyle\zeta_{t}=O_{*}((t/\lambda^{2}+\zeta_{0}^{-3})^{-1/3}).$

Thus, within $O_{*}(1/\lambda^{4})$ the optimality gap $\zeta_{t}$ reaches
$O_{*}(\lambda^{2})$. ∎

The lemma below shows a regularity condition on the norm between two layers.

######  Lemma E.3.

If we start at $|a_{i}^{(0)}|\leq\left\|{\bm{w}}_{i}^{(0)}\right\|_{2}$ and
$\eta=O_{*}(1)$, then we have
$|a_{i}^{(t)}|^{2}\leq\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}$ for all
$i\in[m_{*}]$ and all time $t$.

###### Proof.

Denote $R({\bm{x}})=f({\bm{x}})-f_{*}({\bm{x}})$. Assume
$|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}\leq 0$ we show it
remains at $t+1$. We have

$\displaystyle|a_{i}^{(t+1)}|^{2}-\left\|{\bm{w}}_{i}^{(t+1)}\right\|_{2}^{2}$
$\displaystyle=$
$\displaystyle|a_{i}^{(t)}-\eta\nabla_{a_{i}}L_{\lambda}({\bm{\theta}}^{(t)})|^{2}-\left\|{\bm{w}}_{i}^{(t)}-\eta\nabla_{{\bm{w}}_{i}}L_{\lambda}({\bm{\theta}}^{(t)})\right\|_{2}^{2}$
$\displaystyle=$
$\displaystyle|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}+\eta^{2}|\nabla_{a_{i}}L_{\lambda}({\bm{\theta}}^{(t)})|^{2}-\eta^{2}\left\|\nabla_{{\bm{w}}_{i}}L_{\lambda}({\bm{\theta}}^{(t)})\right\|_{2}^{2}$
$\displaystyle=$
$\displaystyle|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}+\eta^{2}|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma({\bm{w}}_{i}^{{(t)}\top}{\bm{x}})]+\lambda
a_{i}^{(t)}|^{2}-\eta^{2}\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})a_{i}^{(t)}\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]+\lambda{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}$

We first focus on the last 2 terms. We have

$\displaystyle|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma({\bm{w}}_{i}^{{(t)}\top}{\bm{x}})]+\lambda
a_{i}^{(t)}|^{2}-\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})a_{i}^{(t)}\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]+\lambda{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}$
$\displaystyle=$
$\displaystyle\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}})]|^{2}+\lambda^{2}|a_{i}^{(t)}|^{2}-|a_{i}^{(t)}|^{2}\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}-\lambda^{2}\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}$
$\displaystyle\stackrel{{\scriptstyle(\text{a})}}{{\leq}}$
$\displaystyle\left(|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}\right)\left(\lambda^{2}-\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}\right),$

where (a) due to
$|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}})]|^{2}\leq\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}$.

Therefore, plug it back to the above equation, we have

$\displaystyle|a_{i}^{(t+1)}|^{2}-\left\|{\bm{w}}_{i}^{(t+1)}\right\|_{2}^{2}\leq$
$\displaystyle\left(|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}\right)\left(1+\eta^{2}\lambda^{2}-\eta^{2}\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}\right)$
$\displaystyle\stackrel{{\scriptstyle(\text{a})}}{{\leq}}$ $\displaystyle 0,$

where (a) due to
$|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}\leq 0$ and we use
$\left\|2\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}})]{\bm{x}}\right\|_{2}^{2}=O_{*}(d)$
from Lemma E.4 and $\eta$ is small enough.

Therefore, we can see that
$|a_{i}^{(t)}|^{2}-\left\|{\bm{w}}_{i}^{(t)}\right\|_{2}^{2}\leq 0$ remains
for all $t$. ∎

This lemma shows the smoothness of loss function. The proof requires a careful
calculations to bound the error terms. See E.2

###### Proof.

Denote $R_{\bm{\theta}}({\bm{x}})=f_{\bm{\theta}}({\bm{x}})-f_{*}({\bm{x}})$
to denote the dependency on ${\bm{\theta}}$. For simplicity, we will use
$\widetilde{\nabla}_{\bm{\theta}}=-\eta\nabla_{\bm{\theta}}L_{\lambda}$ and
same for others. Since
$\left\|\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}=O_{*}(d)$,
we know
$|\widetilde{\nabla}_{a_{i}}|=O_{*}(\eta\left\|{\bm{w}}_{i}\right\|_{2}d)$ and
$\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}=O_{*}(\eta|a_{i}|d)$

We have

$\displaystyle
L_{\lambda}({\bm{\theta}}-\eta\nabla_{\bm{\theta}})-L_{\lambda}({\bm{\theta}})+\eta\left\|\nabla_{\bm{\theta}}\right\|_{F}^{2}$
$\displaystyle=$ $\displaystyle
L_{\lambda}({\bm{\theta}}-\eta\nabla_{\bm{\theta}})-L_{\lambda}({\bm{\theta}})-\langle\nabla_{\bm{\theta}},-\eta\nabla_{\bm{\theta}}\rangle$
$\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{x}}[R_{{\bm{\theta}}+\widetilde{\nabla}_{\bm{\theta}}}({\bm{x}})^{2}]+\frac{\lambda}{2}\left\|{\bm{a}}+\widetilde{\nabla}_{\bm{a}}\right\|_{2}^{2}+\frac{\lambda}{2}\left\|{\bm{W}}+\widetilde{\nabla}_{\bm{W}}\right\|_{F}^{2}-\mathbb{E}_{\bm{x}}[R_{\bm{\theta}}({\bm{x}})^{2}]-\frac{\lambda}{2}\left\|{\bm{a}}\right\|_{2}^{2}-\frac{\lambda}{2}\left\|{\bm{W}}\right\|_{F}^{2}$
$\displaystyle-\sum_{i\in[m]}\mathbb{E}_{\bm{x}}[R_{\bm{\theta}}({\bm{x}})\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\widetilde{\nabla}_{a_{i}}]-\sum_{i\in[m]}\mathbb{E}_{\bm{x}}[R_{\bm{\theta}}({\bm{x}})a_{i}\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{w}}_{i}}]-\mathbb{E}_{\bm{x}}[R_{\bm{\theta}}({\bm{x}})\widetilde{\nabla}_{\alpha}]-\mathbb{E}_{\bm{x}}[R_{\bm{\theta}}({\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{\beta}}}]$
$\displaystyle-\lambda\langle{\bm{a}},\widetilde{\nabla}_{\bm{a}}\rangle-\lambda\langle{\bm{W}},\widetilde{\nabla}_{\bm{W}}\rangle$
$\displaystyle=$
$\displaystyle\underbrace{\mathbb{E}_{\bm{x}}[(R_{{\bm{\theta}}+\widetilde{\nabla}_{\bm{\theta}}}({\bm{x}})-R_{{\bm{\theta}}}({\bm{x}}))^{2}]}_{(I)}$
$\displaystyle+2\underbrace{\mathbb{E}_{\bm{x}}\left[R_{{\bm{\theta}}}({\bm{x}})\left(R_{{\bm{\theta}}+\widetilde{\nabla}_{\bm{\theta}}}({\bm{x}})-R_{{\bm{\theta}}}({\bm{x}})-\sum_{i\in[m]}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\widetilde{\nabla}_{a_{i}}-\sum_{i\in[m]}a_{i}\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{w}}_{i}}-\widetilde{\nabla}_{\alpha}-{\bm{x}}^{\top}\widetilde{\nabla}_{\bm{\beta}}\right)\right]}_{(II)}$
$\displaystyle+\frac{\lambda}{2}\left\|\widetilde{\nabla}_{\bm{a}}\right\|_{2}^{2}+\frac{\lambda}{2}\left\|\widetilde{\nabla}_{\bm{W}}\right\|_{F}^{2}.$

The last line is easy to see on $O_{*}(\eta^{2}d^{2})$ using norm bound in
Lemma F.12, so in below we are going to bound (I) and (II) one by one. The
goal is to show they are small in the sense of on order $o(\eta)$.

##### Bound (I)

For (I), we can write out the expression as

$\displaystyle\mathbb{E}_{\bm{x}}[(R_{{\bm{\theta}}+\widetilde{\nabla}_{\bm{\theta}}}({\bm{x}})-R_{{\bm{\theta}}}({\bm{x}}))^{2}]=$
$\displaystyle\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}(a_{i}+\widetilde{\nabla}_{a_{i}})\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})+\widetilde{\nabla}_{\alpha}+{\bm{x}}^{\top}\widetilde{\nabla}_{\bm{\beta}}\right)^{2}\right]$
$\displaystyle\leq$ $\displaystyle
2\underbrace{\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}(a_{i}+\widetilde{\nabla}_{a_{i}})\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\right)^{2}\right]}_{(I.i)}$
$\displaystyle+2\underbrace{\mathbb{E}_{\bm{x}}\left[\left(\widetilde{\nabla}_{\alpha}+{\bm{x}}^{\top}\widetilde{\nabla}_{\bm{\beta}}\right)^{2}\right]}_{(I.ii)}$

For (I.i), we can split into 2 terms as

$\displaystyle\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}(a_{i}+\widetilde{\nabla}_{a_{i}})\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\right)^{2}\right]$
$\displaystyle\leq$ $\displaystyle
2\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}\widetilde{\nabla}_{a_{i}}\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})\right)^{2}\right]+2\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}a_{i}\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\right)^{2}\right]$
$\displaystyle\leq$ $\displaystyle
2\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}|\widetilde{\nabla}_{a_{i}}||({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}}|\right)^{2}\right]+2\mathbb{E}_{\bm{x}}\left[\left(\sum_{i\in[m]}|a_{i}||\widetilde{\nabla}_{{\bm{w}}_{i}}^{\top}{\bm{x}}|\right)^{2}\right].$

We then can bound them separately as

$\displaystyle(I.i)\stackrel{{\scriptstyle(\text{a})}}{{\leq}}$ $\displaystyle
O(1)\left(\sum_{i\in[m]}|\widetilde{\nabla}_{a_{i}}|\left\|{\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}\right)^{2}+O(1)\left(\sum_{i\in[m]}|a_{i}|\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}\right)^{2}$
$\displaystyle\stackrel{{\scriptstyle(\text{b})}}{{\leq}}$ $\displaystyle
O_{*}(d^{2})\left(\sum_{i\in[m]}\eta\left\|{\bm{w}}_{i}\right\|_{2}^{2}+\eta^{2}|a_{i}|\left\|{\bm{w}}_{i}\right\|_{2}d\right)^{2}+O_{*}(d^{2})\left(\sum_{i\in[m]}\eta
a_{i}^{2}\right)^{2}$
$\displaystyle\stackrel{{\scriptstyle(\text{c})}}{{\leq}}$ $\displaystyle
O_{*}(\eta^{2}d^{2}),$

where (a) we use Lemma E.5; (b) recall
$|\widetilde{\nabla}_{a_{i}}|=O_{*}(\eta\left\|{\bm{w}}_{i}\right\|_{2}d)$ and
$\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}=O_{*}(\eta|a_{i}|d)$;
(c)
$\left\|{\bm{a}}\right\|,\left\|{\bm{W}}\right\|_{F},\sum_{i\in[m]}|a_{i}|\left\|{\bm{w}}_{i}\right\|_{2}=O_{*}(1)$
from Lemma F.12 and Lemma F.4, as well as $\eta$ is small enough.

For (I.ii), we have

$\displaystyle\mathbb{E}_{\bm{x}}\left[\left(\widetilde{\nabla}_{\alpha}+{\bm{x}}^{\top}\widetilde{\nabla}_{\bm{\beta}}\right)^{2}\right]\leq
O(|\widetilde{\nabla}_{\alpha}|^{2}+\left\|\widetilde{\nabla}_{\bm{\beta}}\right\|_{2}^{2})=O_{*}(\eta^{2}d^{2}),$

where we use Lemma F.4.

Combine (I.i) and (I.ii) we know (I)=$O_{*}(\eta^{2}d^{2})$.

##### Bound (II)

For (II), we have

$\displaystyle\mathbb{E}_{\bm{x}}\left[R_{{\bm{\theta}}}({\bm{x}})\left(R_{{\bm{\theta}}+\widetilde{\nabla}_{\bm{\theta}}}({\bm{x}})-R_{{\bm{\theta}}}({\bm{x}})-\sum_{i\in[m]}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\widetilde{\nabla}_{a_{i}}-\sum_{i\in[m]}a_{i}\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{w}}_{i}}-\widetilde{\nabla}_{\alpha}-{\bm{x}}^{\top}\widetilde{\nabla}_{\bm{\beta}}\right)\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{x}}\left[R_{{\bm{\theta}}}({\bm{x}})\left(\sum_{i\in[m]}\underbrace{(a_{i}+\widetilde{\nabla}_{a_{i}})\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})-\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\widetilde{\nabla}_{a_{i}}-a_{i}\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{w}}_{i}}}_{I_{i}({\bm{x}})}\right)\right]$
$\displaystyle\leq$
$\displaystyle\sum_{i\in[m]}\left\|R_{\bm{\theta}}\right\|\left\|I_{i}\right\|$

We focus on bound $\left\|I_{i}\right\|$ below. The goal is to show it is
$o(\eta)$. For $I_{i}({\bm{x}})$, we have

$\displaystyle\left\|I_{i}\right\|_{2}^{2}=$
$\displaystyle\mathbb{E}_{\bm{x}}\left[\left((a_{i}+\widetilde{\nabla}_{a_{i}})\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-a_{i}\sigma({\bm{w}}_{i}^{\top}{\bm{x}})-\sigma({\bm{w}}_{i}^{\top}{\bm{x}})\widetilde{\nabla}_{a_{i}}-a_{i}\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{w}}_{i}}\right)^{2}\right]$
$\displaystyle\leq$
$\displaystyle\mathbb{E}_{\bm{x}}\left[2\left(\widetilde{\nabla}_{a_{i}}(\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-\sigma({\bm{w}}_{i}^{\top}{\bm{x}}))\right)^{2}+2\left(a_{i}(\sigma(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-\sigma({\bm{w}}_{i}^{\top}{\bm{x}})-\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}){\bm{x}}^{\top}\widetilde{\nabla}_{{\bm{w}}_{i}})\right)^{2}\right]$
$\displaystyle\leq$ $\displaystyle
2\underbrace{\mathbb{E}_{\bm{x}}\left[|\widetilde{\nabla}_{a_{i}}|^{2}|\widetilde{\nabla}_{{\bm{w}}_{i}}^{\top}{\bm{x}}|^{2}\right]}_{(II.i)}+2a_{i}^{2}\underbrace{\mathbb{E}_{\bm{x}}\left[|({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}}|^{2}(\sigma^{\prime}(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}))^{2}\right]}_{(II.ii)}$

For (II.i), recall
$|\widetilde{\nabla}_{a_{i}}|=O_{*}(\eta\left\|{\bm{w}}_{i}\right\|_{2}d)$ and
$\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}=O_{*}(\eta|a_{i}|d)$ we
have

$\displaystyle\mathbb{E}_{\bm{x}}\left[|\widetilde{\nabla}_{a_{i}}|^{2}|\widetilde{\nabla}_{{\bm{w}}_{i}}^{\top}{\bm{x}}|^{2}\right]\leq|\widetilde{\nabla}_{a_{i}}|^{2}\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|^{2}=O_{*}(\eta^{4}|a_{i}|^{2}\left\|{\bm{w}}_{i}\right\|_{2}^{2}d^{4}).$

For (II.ii), we have

$\displaystyle\mathbb{E}_{\bm{x}}\left[|{\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}}^{\top}{\bm{x}}|^{2}(\sigma^{\prime}(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})-\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}}))^{2}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{x}}\left[|({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}}|^{2}\mathbbm{1}_{\operatorname{sign}(({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}})^{\top}{\bm{x}})\neq\operatorname{sign}({\bm{w}}_{i}^{\top}{\bm{x}})}\right]$
$\displaystyle\leq$ $\displaystyle
O(\left\|{\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}^{2}\delta^{3}),$

where
$\delta=\angle({\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}},{\bm{w}}_{i})$
is the angle between ${\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}}$ and
${\bm{w}}_{i}$. Since
$\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}=O_{*}(\eta|a_{i}|d)=O_{*}(\eta\left\|{\bm{w}}_{i}\right\|_{2}d)$,
we know $\delta=O(\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|)$ given
$\eta=O_{*}(1/d)$ to be small enough.

Combine (II.i) and (II.ii) we have

$\displaystyle\left\|I_{i}\right\|_{2}^{2}\leq
O_{*}(\eta^{4}a_{i}^{2}\left\|{\bm{w}}_{i}\right\|_{2}^{2}d^{4})+O(a_{i}^{2}\left\|{\bm{w}}_{i}+\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}^{2}\left\|\widetilde{\nabla}_{{\bm{w}}_{i}}\right\|_{2}^{3})\leq
O_{*}(\eta^{3}a_{i}^{2}\left\|{\bm{w}}_{i}\right\|_{2}^{2}d^{3}).$

Since $\left\|R_{\bm{\theta}}\right\|=O_{*}(1)$, this implies

$\displaystyle(II)\leq\sum_{i\in[m]}O_{*}(\eta^{3/2}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}d^{3/2})=O_{*}(\eta^{3/2}d^{3/2}).$

##### Combine (I)(II)

Finally, combing (I) and (II) we have

$\displaystyle
L_{\lambda}({\bm{\theta}}-\eta\nabla_{\bm{\theta}})-L_{\lambda}({\bm{\theta}})+\eta\left\|\nabla_{\bm{\theta}}\right\|_{F}^{2}=O_{*}(\eta^{3/2}d^{3/2}).$

Going back to the beginning of this proof, we get the desired result. ∎

###  E.2 Technical Lemma

We present technical lemmas that are used in the proof of this section. They
mostly follow from direct calculations.

######  Lemma E.4.

We have
$\left\|\mathbb{E}_{\bm{x}}[R({\bm{x}})\sigma^{\prime}(\overline{{\bm{w}}}_{i}^{{(t)}\top}{\bm{x}}){\bm{x}}]\right\|_{2}^{2}=O_{*}(d)$

###### Proof.

It is easy to see given $\left\|R\right\|=O_{*}(1)$. ∎

######  Lemma E.5 (Lemma D.4 in Zhou et al. (2021)).

Consider $\alpha_{i}\in\mathbb{R}^{d}$ for $i\in[n]$. We have

$\displaystyle\mathbb{E}_{x\sim
N(0,I)}\left[\left(\sum_{i=1}^{n}|\alpha_{i}^{\top}x|\right)^{2}\right]\leq
c_{0}\left(\sum_{i=1}^{n}\left\|\alpha_{i}\right\|\right)^{2},$

where $c_{0}$ is a constant.

##  Appendix F Local landscape of population loss

In this section, we are going to show Lemma 4.4 that characterizing the
population local landscape with a fixed $\lambda$ by giving the lower bound of
gradient.

##### Outline

We generally follow the high-level proof plan that outlines in Section 6. In
Section F.1 and Section F.2, we characterize the local geometry as in Lemma
6.1. Then, we use it to construct descent direction in Section F.3. Finally we
give the proof of Lemma 4.4 in Section F.4.

We start by identifying the structure of (approximated) solution of a closely-
related problem in Section F.1 (rewrite (5)):

$\displaystyle\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu):=$
$\displaystyle
L(\mu)+\lambda|\mu|_{1}:=\mathbb{E}_{{\bm{x}},\widetilde{y}}[(f_{\mu}({\bm{x}})-\widetilde{y})^{2}]+\lambda|\mu|_{1}$
(6) $\displaystyle=$
$\displaystyle\mathbb{E}_{{\bm{x}}}\left[\left(\int_{\bm{w}}\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}}){\textnormal{d}}\
\mu-\mu_{*}\right)^{2}\right]+\lambda|\mu|_{1},$ (7)

where ${\mathcal{M}}({\mathbb{S}}^{d-1})$ is the measure space over unit
sphere ${\mathbb{S}}^{d-1}$,
$\mu_{*}=\sum_{i\in[m_{*}]}a_{i}^{*}\delta_{{\bm{w}}_{i}^{*}}$ and
$\sigma_{\geq 2}(x)=\sigma(x)-1/\sqrt{2\pi}-x/2$ is the activation that after
removing 0th and 1st order term in Hermite expansion. Note that when $\mu$
represents a finite-wdith network, we have
$\mu=\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\delta_{\overline{{\bm{w}}}_{i}}$
is a empirical measure over the neurons. In particular, when $\mu=\mu_{*}$,
model $f_{\mu}$ recovers the target $\widetilde{f}_{*}$.

We call (5) as the ideal loss because the original problem (2) would become
the above (5) when we balance the norms
($\left\|{\bm{w}}_{i}\right\|_{2}=|a_{i}|$), perfectly fit $\alpha,\beta$ and
relax the finite-width constraints to allow infinite-width (see Claim B.1).
This is why we slightly abused the notation to use $L_{\lambda}$ in both (2)
and (5).

In Section F.3 we will use the solution structure to construct descent
direction that are positively correlated with gradient and also handle the
case when norms are not balanced or $\alpha,\beta$ are not fitted well.

##### Notation

Denote the optimality gap between the loss at $\mu$ and the optimal
distribution $\mu_{\lambda}^{*}$ as

$\displaystyle\zeta(\mu):=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*}),$

where $\mu_{\lambda}^{*}$ is the optimal measure that minimize (5). For
simplicity denote $\widetilde{a}_{i}=a_{i}\left\|{\bm{w}}_{i}\right\|_{2}$ so
that $|\mu|_{1}=\left\|\widetilde{{\bm{a}}}\right\|_{1}$ when
$\mu=\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\delta_{\overline{{\bm{w}}}_{i}}$.
Often we use $\zeta_{t}=\zeta(\mu_{t})$ to denote the optimality gap at time
$t$ and just $\zeta$ for simplicity. We slightly abuse the notation to also
use $\zeta=L_{\lambda}(\theta)-L_{\lambda}(\mu_{\lambda}^{*})$. Finally denote
$\mu^{*}=\sum_{i\in[m_{*}]}a_{i}^{*}\delta_{{\bm{w}}_{i}^{*}}$ (assuming
$\left\|{\bm{w}}_{i}^{*}\right\|_{2}=1$) so that
$f_{\mu^{*}}({\bm{x}})=\mathbb{E}_{{\bm{w}}\sim\mu^{*}}[\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}})]$.

###  F.1 Structure of the ideal loss solution

In this section, we will focus on the structure of approximated solution for
the $\ell_{1}$ regularized regression problem (5).

In the rest of this section, we will first introduce the idea of non-
degenerate dual certificate and then use it as a tool to characterize the
structure of the solutions. The proofs are deferred to Section H.

####  F.1.1 Non-degenerate dual certificate

We first recall the definition of non-degenerate dual certificate, which is
similar as in (Poon et al., 2023) but slightly adapted for fit our need. See 1

We first show that there exist such non-degenerate dual certificate. More
discussion and a detailed proof are deferred to Section G.

######  Lemma F.1.

There exists a non-degenerate dual certificate
$\eta=\mathbb{E}_{\bm{x}}[p({\bm{x}})\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}})]$ with $\rho_{\eta}=\Theta(1)$ and
$\left\|p\right\|_{2}\leq\operatorname{poly}(m_{*},\Delta)$

The following lemma (restate of Lemma 6.2) gives the properties that will be
used in the later proofs: the non-degenerate dual certificate $\eta$ allows us
to capture the gap between the current position $\mu$ and the target
$\mu^{*}$.

######  Lemma F.2.

Given a non-degenerate dual certificate $\eta$, then

  1. (i)

$\langle\eta,\mu^{*}\rangle=|\mu^{*}|_{1}$

  2. (ii)

For any measure $\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})$,
$|\langle\eta,\mu\rangle|\leq|\mu|_{1}-\rho_{\eta}\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu|({\bm{w}})$.

  3. (iii)

$\langle\eta,\mu-\mu^{*}\rangle=\langle p,f_{\mu}-f_{\mu^{*}}\rangle$, where
$f_{\mu}({\bm{x}})=\mathbb{E}_{{\bm{w}}\sim\mu}[\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}})]$. Then
$|\langle\eta,\mu-\mu^{*}\rangle|\leq\left\|p\right\|_{2}\sqrt{L(\mu)}$.

####  F.1.2 Properties of $\mu_{\lambda}^{*}$

Given the non-degenerate dual certificate $\eta$, we now are ready to identify
several useful properties of $\mu_{\lambda}^{*}$. The lemma below essentially
says that $\mu_{\lambda}^{*}$ is similar to $\mu^{*}$ in the sense that most
of the norm are concentrated in the ground-truth direction and the square loss
is small. The proof relies on comparing $\mu_{\lambda}^{*}$ with $\mu^{*}$
using the optimality conditions.

######  Lemma F.3.

We have the following hold

  1. (i)

$|\mu_{*}|_{1}-\lambda\left\|p\right\|_{2}^{2}\leq|\mu_{\lambda}^{*}|_{1}\leq|\mu^{*}|_{1}=\left\|{\bm{a}}^{*}\right\|_{1}$

  2. (ii)

$L(\mu_{\lambda}^{*})\leq\lambda^{2}\left\|p\right\|_{2}^{2}=O_{*}(\lambda^{2})$

  3. (iii)

$\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu_{\lambda}^{*}|({\bm{w}})\leq\lambda\left\|p\right\|_{2}^{2}/\rho_{\eta}=O_{*}(\lambda)$

####  F.1.3 Properties of $\mu$ with optimality gap $\zeta$

We now characterize the structure of $\mu$ when the optimality gap is $\zeta$.
The proof mostly relies on comparing $\mu$ with $\mu_{\lambda}^{*}$ and the
structure of $\mu_{\lambda}^{*}$ in previous section.

The following lemma shows the square loss is bounded by the optimality gap and
norms are always bounded. Note that the conditions are true under Lemma 4.4.

######  Lemma F.4.

Recall the optimality gap
$\zeta=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*})$. Then, the following
holds:

  1. (i)

$L(\mu)\leq 5\lambda^{2}\left\|p\right\|^{2}+4\zeta=O_{*}(\lambda^{2}+\zeta)$.

  2. (ii)

if $\zeta\leq\lambda|\mu^{*}|_{1}$ and
$\lambda\leq|\mu^{*}|_{1}/\left\|p\right\|_{2}^{2}$, then $|\mu|_{1}\leq
3|\mu^{*}|_{1}=3\left\|{\bm{a}}^{*}\right\|_{1}$.

The following two lemma characterize the structure of $\mu$ using the fact
that the square loss is small in previous lemma. The lemma below says that the
total norm of far away neuron is small.

######  Lemma F.5.

Recall the optimality gap
$\zeta=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*})$. Then, we have

$\displaystyle\sum_{i\in[m_{*}]}\int_{{\mathcal{T}}_{i}}\delta({\bm{w}},{\bm{w}}_{i}^{*})^{2}\,\mathrm{d}|\mu|({\bm{w}})\leq(\zeta/\lambda+2\lambda\left\|p\right\|_{2}^{2})/\rho_{\eta}=O_{*}(\zeta/\lambda+\lambda).$

In particular, when
$\mu=\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\delta_{\overline{{\bm{w}}}_{i}}$
represents finite number of neurons, we have

$\displaystyle\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}\delta_{j}^{2}\leq(\zeta/\lambda+2\lambda\left\|p\right\|_{2}^{2})/\rho_{\eta}=O_{*}(\zeta/\lambda+\lambda),$

where $\delta_{j}=\angle({\bm{w}}_{j},{\bm{w}}_{i}^{*})$ for
$j\in{\mathcal{T}}_{i}$.

The lemma below shows there are neurons close to the teacher neurons once the
gap is small. The proof idea is similar to Section 5.3 in Zhou et al. (2021)
that use test function to lower bound the loss, but now we can handle almost
all activation.

######  Lemma F.6.

Under Lemma 4.4, if the Hermite coefficient of $\sigma$ decays as
$|\hat{\sigma}_{k}|=\Theta(k^{-c_{\sigma}})$ with some constant
$c_{\sigma}>0$, then the total mass near each target direction is large, i.e.,
$\mu({\mathcal{T}}_{i}(\delta))\operatorname{sign}(a_{i}^{*})\geq|a_{i}^{*}|/2$
for all $i\in[m_{*}]$ and any
$\delta_{close}\geq\widetilde{\Omega}\left((\frac{L(\mu)}{a_{\min}^{2}})^{1/(4c_{\sigma}-2)}\right)$
with large enough hidden constant. In particular, for $\sigma$ is ReLU or
absolute function,
$\delta_{close}\geq\widetilde{\Omega}\left((\frac{L(\mu)}{a_{\min}^{2}})^{1/3}\right)$.
Here $a_{\min}=\min|a_{i}|$ is the smallest entry of ${\bm{a}}_{*}$ in
absolute value.

As a corollary, if the optimality gap
$\zeta=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*})$, then
$\delta_{close}\geq\widetilde{\Omega}_{*}\left((\zeta+\lambda^{2})^{1/(4c_{\sigma}-2)}\right)$
and for ReLU or absolute
$\delta_{close}\geq\widetilde{\Omega}_{*}\left((\zeta+\lambda^{2})^{1/3}\right)$.

####  F.1.4 Residual decomposition and average neuron

In this section, we introduce the residual decomposition and average neuron as
in (Zhou et al., 2021) that will be used when proving the existence of descent
direction.

Denote the decomposition
$R({\bm{x}})=f_{\mu}({\bm{x}})-f_{\mu^{*}}({\bm{x}})=R_{1}({\bm{x}})+R_{2}({\bm{x}})+R_{3}({\bm{x}})$
(this can be directly verified noticing that $\sigma_{\geq
2}(x)=|x|/2-1/\sqrt{2\pi}$),

$\displaystyle R_{1}({\bm{x}})$
$\displaystyle=\frac{1}{2}\sum_{i\in[m_{*}]}\left(\sum_{j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}-{\bm{w}}_{i}^{*}\right)^{\top}{\bm{x}}\operatorname{sign}({\bm{w}}_{i}^{*\top}{\bm{x}}),$
(8) $\displaystyle R_{2}({\bm{x}})$
$\displaystyle=\frac{1}{2}\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}^{\top}{\bm{x}}(\operatorname{sign}({\bm{w}}_{j}^{\top}{\bm{x}})-\operatorname{sign}({\bm{w}}_{i}^{*\top}{\bm{x}})),$
$\displaystyle R_{3}({\bm{x}})$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\sum_{i\in[m_{*}]}a_{i}^{*}\left\|{\bm{w}}_{i}^{*}\right\|_{2}-\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\right).$

In the following we characterize $R_{1},R_{2},R_{3}$ separately. In Lemma F.7
we relate $R_{1}$ with the average neuron. In Lemma F.8 and Lemma F.9 we bound
$R_{2}$ and $R_{3}$ respectively.

######  Lemma F.7 (Zhou et al. (2021), Lemma 11).

$\left\|R_{1}\right\|_{2}^{2}=\Omega(\Delta^{3}/m_{*}^{3})\sum_{i\in[m_{*}]}\left\|\sum_{j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}-{\bm{w}}_{i}^{*}\right\|_{2}^{2}$.

######  Lemma F.8.

Under Lemma 4.4, recall the optimality gap
$\zeta=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*})$. Then

$\displaystyle\left\|R_{2}\right\|_{2}^{2}=O_{*}((\zeta/\lambda+\lambda)^{3/2}).$

######  Lemma F.9.

Under Lemma 4.4 and recall the optimality gap
$\zeta=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*})$. If
$\hat{\sigma}_{0}=0$ and $\hat{\sigma}_{k}>0$ with some
$k=\Theta((1/\Delta^{2})\log(\zeta/\left\|{\bm{a}}_{*}\right\|_{1}))$, then

$\displaystyle\left\|R_{3}\right\|_{2}=$
$\displaystyle\widetilde{O}_{*}((\zeta+\lambda^{2})^{1/2}/\hat{\sigma}_{k}+(\zeta/\lambda+\lambda)+\zeta).$

Now we are ready to bound the difference between average neuron with its
corresponding ground-truth neuron.

######  Lemma F.10.

Under Lemma 4.4, recall the optimality gap
$\zeta=L_{\lambda}(\mu)-L_{\lambda}(\mu_{\lambda}^{*})$. Then for any
$i\in[m_{*}]$, $\zeta=\Omega(\lambda^{2})$ and $\zeta,\lambda\leq
1/\operatorname{poly}(m_{*},\Delta,\left\|{\bm{a}}_{*}\right\|_{1})$

$\displaystyle\left\|\sum_{j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}-{\bm{w}}_{i}^{*}\right\|_{2}\leq\left(\sum_{i\in[m_{*}]}\left\|\sum_{j\in{\mathcal{T}}_{i}}a_{j}{\bm{w}}_{j}-{\bm{w}}_{i}^{*}\right\|_{2}^{2}\right)^{1/2}=O_{*}((\zeta/\lambda)^{3/4}).$

###  F.2 From ideal loss solution to real loss solution

In previous section, we consider the ideal loss solution that assumes the
norms are perfectly balanced ($|a_{i}|=\left\|{\bm{w}}_{i}\right\|_{2}$) and
$\alpha,{\bm{\beta}}$ are perfectly fitted. However, during the training we
are not able to guarantee achieve these exactly but only approximately. This
section is devoted to show that the results in previous section still hold
though the conditions are only approximately satisfied. Recall that the
original loss

$\displaystyle
L_{\lambda}({\bm{\theta}})=L({\bm{\theta}})+\frac{\lambda}{2}\left\|{\bm{a}}\right\|_{2}^{2}+\frac{\lambda}{2}\left\|{\bm{W}}\right\|_{F}^{2}$

so that when norm are balanced and $\alpha,{\bm{\beta}}$ are perfectly fitted,
$L_{\lambda}({\bm{\theta}})=L({\bm{\theta}})+\lambda\sum_{i}|a_{i}|\left\|{\bm{w}}_{i}\right\|_{2}=L_{\lambda}(\mu)$.

The lemma below shows that the properties of ideal loss solution in previous
section still hold for the solution of original loss, when
$\alpha,{\bm{\beta}}$ are approximately fitted.

######  Lemma F.11.

Given any ${\bm{\theta}}=({\bm{a}},{\bm{W}},\alpha,{\bm{\beta}})$ satisfying
$|\alpha-\hat{\alpha}|^{2}=O(\zeta)$,
$\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}=O(\zeta)$, where
$\hat{\alpha}=-(1/\sqrt{2\pi})\sum_{i=1}^{m}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}$
and $\hat{{\bm{\beta}}}=-(1/2)\sum_{i=1}^{m}a_{i}{\bm{w}}_{i}$. Let its
corresponding balanced version
${\bm{\theta}}_{bal}=({\bm{a}}_{bal},{\bm{W}}_{bal},\alpha_{bal},{\bm{\beta}}_{bal})$
as
$a_{bal,i}=\operatorname{sign}(a_{i})\sqrt{|a_{i}|\left\|{\bm{w}}_{i}\right\|_{2}}$,
${\bm{w}}_{bal,i}=\overline{{\bm{w}}}_{i}\sqrt{|a_{i}|\left\|{\bm{w}}_{i}\right\|_{2}}$,
$\alpha_{bal}=\hat{\alpha}$ and ${\bm{\beta}}_{bal}=\hat{{\bm{\beta}}}$. Then,
we have

$\displaystyle
L_{\lambda}({\bm{\theta}})-L_{\lambda}({\bm{\theta}}_{bal})=|\alpha-\hat{\alpha}|^{2}+\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}+\frac{\lambda}{2}\sum_{i\in[m]}(|a_{i}|-\left\|{\bm{w}}_{i}\right\|_{2})^{2}\geq
0.$

Moreover, let the optimality gap
$\zeta=L_{\lambda}({\bm{\theta}})-L_{\lambda}(\mu_{\lambda}^{*})$, we have
results in Lemma F.4, Lemma F.5, Lemma F.6, Lemma F.7, Lemma F.8, Lemma F.9
and Lemma F.10 still hold for $L_{\lambda}({\bm{\theta}})$, with the change of
$R_{3}$ in (8) as

$\displaystyle R_{3}({\bm{x}})$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\left(\sum_{i\in[m_{*}]}a_{i}^{*}\left\|{\bm{w}}_{i}^{*}\right\|_{2}-\sum_{i\in[m]}a_{i}\left\|{\bm{w}}_{i}\right\|_{2}\right)+\alpha-\hat{\alpha}+({\bm{\beta}}-\hat{{\bm{\beta}}})^{\top}{\bm{x}}.$

The following lemma shows the norm remains bounded.

######  Lemma F.12.

Under Lemma 4.4, suppose optimality gap
$\zeta=L_{\lambda}({\bm{\theta}})-L_{\lambda}(\mu_{\lambda}^{*})$. Then
$\left\|{\bm{a}}\right\|_{2}^{2}+\left\|{\bm{W}}\right\|_{F}^{2}\leq
3\left\|{\bm{a}}_{*}\right\|_{1}$.

###  F.3 Descent direction

In this section, we show that there is a descent direction as long as the
optimality gap is small until it reaches $O(\lambda^{2})$. We will assume
$\zeta=\Omega(\lambda^{2})$ in this section for simplicity.

We first show gradient is always large whenever $\alpha,{\bm{\beta}}$ are not
fitted well. This is a direct corollary of Claim B.1.

######  Lemma F.13 (Descent direction, $\alpha$ and ${\bm{\beta}}$).

We have

$\displaystyle|\nabla_{\alpha}L_{\lambda}|^{2}=4(\alpha-\hat{\alpha})^{2},\quad\left\|\nabla_{\bm{\beta}}L_{\lambda}\right\|_{2}^{2}=4\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}.$

Before proceeding to the following descent direction, we first make a
simplification assumption that

######  Assumption F.1.

For every ${\mathcal{T}}_{i}$, for all neuron
${\bm{w}}_{j}\in{\mathcal{T}}_{i}$, assume
${\bm{w}}_{j}^{\top}{\bm{w}}_{i}^{*}\geq 0$.

This is because due to the linear term ${\bm{\beta}}$, the effective
activation is symmetry $\sigma_{\geq 2}(x)=\sigma_{\geq 2}(-x)$. This
introduce the ambiguity of the sign of neurons. Such assumption clarifies the
ambiguity of neurons’ direction.

As the lemma below shows, there always exists a set of parameter (by flipping
the sign of neurons) that satisfy the assumption and gives almost same
gradient norm. Thus, making such assumption will not cause any issue when
$\alpha,\beta$ are perfectly fitted.

######  Lemma F.14.

Suppose
$(\alpha-\hat{\alpha})^{2},\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}\leq\tau$
to be small enough and
$\left\|{\bm{a}}\right\|_{2},\left\|{\bm{W}}\right\|_{F}=O_{*}(1)$. Then,
given any parameter ${\bm{\theta}}$, there exists another set of parameter
$\widetilde{{\bm{\theta}}}$ that satisfies Assumption F.1 such that
$f_{\bm{\theta}}=f_{\widetilde{{\bm{\theta}}}}$ and
$|\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|-\left\|\nabla_{\widetilde{{\bm{\theta}}}}L_{\lambda}\right\|_{F}|\leq
O_{*}(\sqrt{\tau})$.

###### Proof.

Denote
${\bm{\theta}}=({\bm{a}},{\bm{w}}_{1},\ldots,{\bm{w}}_{m},\alpha,{\bm{\beta}})$.
We first construct
$\widetilde{{\bm{\theta}}}=(\widetilde{{\bm{a}}},\widetilde{{\bm{w}}}_{1},\ldots,\widetilde{{\bm{w}}}_{m},\widetilde{\alpha},\widetilde{{\bm{\beta}}})$.

Let $\widetilde{{\bm{a}}}={\bm{a}}$. For $\widetilde{{\bm{w}}}_{i}$, there
exists such sign vector ${\bm{s}}=(s_{1},\ldots,s_{m})\in\\{\pm 1\\}^{m}$ so
that by flipping the sign of neurons we have
$\widetilde{{\bm{w}}}_{i}=s_{i}{\bm{w}}_{i}$ satisfies Assumption F.1. Let
$\widetilde{\alpha}=\alpha$ and
$\widetilde{{\bm{\beta}}}={\bm{\beta}}++\sum_{i:s_{i}=-1}a_{i}{\bm{w}}_{i}$.

One can verify that $f_{\bm{\theta}}=f_{\widetilde{{\bm{\theta}}}}$. Moreover,
for the gradient of $\alpha,{\bm{\beta}}$ we have

$\displaystyle\nabla_{\alpha}L_{\lambda}=\nabla_{\widetilde{\alpha}}L_{\lambda},\nabla_{\bm{\beta}}L_{\lambda}=\nabla_{\widetilde{{\bm{\beta}}}}L_{\lambda},$

For gradient of ${\bm{a}},{\bm{w}}_{i}$, when $s_{i}=1$ we know they are the
same. When $s_{i}=-1$, note that

$\displaystyle\nabla_{a_{i}}L_{\lambda}-\nabla_{\widetilde{a}_{i}}L_{\lambda}=2\mathbb{E}_{\bm{x}}[R({\bm{x}})(\sigma({\bm{w}}_{i}^{\top}{\bm{x}})-\sigma(\widetilde{{\bm{w}}}_{i}^{\top}{\bm{x}}))]=2({\bm{\beta}}-\hat{{\bm{\beta}}})^{\top}{\bm{w}}_{i}$
$\displaystyle\nabla_{{\bm{w}}_{i}}L_{\lambda}+\nabla_{\widetilde{{\bm{w}}}_{i}}L_{\lambda}=2a_{i}\mathbb{E}_{\bm{x}}[R({\bm{x}})(\sigma^{\prime}({\bm{w}}_{i}^{\top}{\bm{x}})+\sigma^{\prime}(\widetilde{{\bm{w}}}_{i}^{\top}{\bm{x}})){\bm{x}}]=2a_{i}({\bm{\beta}}-\hat{{\bm{\beta}}}).$

Therefore, we get the desired result by noting the norm bound.

∎

We then show that if norms are not balanced or norm cancellation happens for
neurons with similar direction, then one can always adjust the norm to
decrease the loss due to the regularization term.

######  Lemma F.15 (Descent direction, norm balance).

We have

$\displaystyle\sum_{i}\sum_{j\in
T_{i}}\left|\langle\nabla_{a_{j}}L_{\lambda},-a_{j}\rangle+\langle\nabla_{{\bm{w}}_{j}}L_{\lambda},{\bm{w}}_{j}\rangle\right|=$
$\displaystyle\lambda\sum_{i\in[m_{*}]}\left|a_{i}^{2}-\left\|{\bm{w}}_{i}\right\|_{2}^{2}\right|$
$\displaystyle\geq$
$\displaystyle\max\left\\{\lambda|\left\|{\bm{a}}\right\|_{2}^{2}-\left\|{\bm{W}}\right\|_{F}^{2}|,\lambda\sum_{i\in[m_{*}]}(|a_{i}|-\left\|{\bm{w}}_{i}\right\|_{2})^{2}\right\\}$

######  Lemma F.16 (Descent direction, norm cancellation).

Under Lemma 4.4 and Assumption F.1, suppose the optimality gap
$\zeta=L_{\lambda}({\bm{\theta}})-L_{\lambda}(\mu_{\lambda}^{*})$. For any
${\bm{w}}_{i}^{*}$, consider $\delta_{\operatorname{sign}}$ such that
$\delta_{close}<\delta_{\operatorname{sign}}=O(\lambda/\zeta^{1/2})$ with
small enough hidden constant ($\delta_{close}$ defined in Lemma F.6), then

$\displaystyle\sum_{s\in\\{+,-\\}}\sum_{j\in
T_{i,s}(\delta_{\operatorname{sign}})}\left\langle\nabla_{a_{j}}L_{\lambda},\frac{a_{j}}{\sum_{j\in
T_{i,s}(\delta_{\operatorname{sign}})}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}}\right\rangle+\left\langle\nabla_{{\bm{w}}_{j}}L_{\lambda},\frac{{\bm{w}}_{j}}{\sum_{j\in
T_{i,s}(\delta_{\operatorname{sign}})}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}}\right\rangle=\Omega(\lambda).$

where $T_{i,+}(\delta_{\operatorname{sign}})=\\{j\in
T_{i}:\delta({\bm{w}}_{j},{\bm{w}}_{i}^{*})\leq\delta_{\operatorname{sign}},\operatorname{sign}(a_{j})=\operatorname{sign}(a_{i}^{*})\\}$,
$T_{i,-}(\delta_{\operatorname{sign}})=\\{j\in
T_{i}:\delta({\bm{w}}_{j},{\bm{w}}_{i}^{*})\leq\delta_{\operatorname{sign}},\operatorname{sign}(a_{j})\neq\operatorname{sign}(a_{i}^{*})\\}$
are the set of neurons that close to ${\bm{w}}_{i}^{*}$ with/without same sign
of $a_{i}^{*}$.

As a result,

$\displaystyle\left\|\nabla_{{\bm{a}}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{{\bm{W}}}L_{\lambda}\right\|_{F}^{2}\geq$
$\displaystyle\lambda^{2}\sum_{j\in
T_{i,-}(\delta_{\operatorname{sign}})}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}$

Now given the above lemmas, it suffices to consider the remaining case that
$\alpha,{\bm{\beta}}$ are well fitted, norms are balanced and no cancellation.
In this case, the loss landscape is roughly the same as the ideal loss (5)
from Lemma F.11. Thus, we could leverage these detailed characterization of
the solution (far-away neurons are small and average neuron is close to
corresponding ground-truth neuron) to construct descent direction.

######  Lemma F.17 (Descent direction).

Under Lemma 4.4 and Assumption F.1, suppose the optimality gap
$\zeta=L_{\lambda}({\bm{\theta}})-L_{\lambda}(\mu_{\lambda}^{*})$. Suppose

  1. (i)

norms are (almost) balanced:
$|\left\|{\bm{W}}\right\|_{F}^{2}-\left\|{\bm{a}}\right\|_{2}^{2}|\leq\zeta/\lambda$,
$\sum_{i\in[m]}(|a_{j}|-\left\|{\bm{w}}_{j}\right\|_{2})^{2}=O_{*}(\zeta^{2}/\lambda^{2})$

  2. (ii)

(almost) no norm cancellation: consider all neurons ${\bm{w}}_{j}$ that are
$\delta_{\operatorname{sign}}$-close w.r.t. teacher neuron ${\bm{w}}_{i}^{*}$
but has a different sign, i.e.,
$\operatorname{sign}(a_{j})\neq\operatorname{sign}(a_{i}^{*})$ with
$\delta_{\operatorname{sign}}=\Theta_{*}(\lambda/\zeta^{1/2})$, we have
$\sum_{j\in
T_{i,-}(\delta_{\operatorname{sign}})}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}\leq\tau=O_{*}(\zeta^{5/6}/\lambda)$
with small enough hidden constant, where $T_{i,-}(\delta)$ defined in Lemma
F.16.

  3. (iii)

$\alpha,{\bm{\beta}}$ are well fitted:
$|\alpha-\hat{\alpha}|^{2}=O_{*}(\zeta)$,
$\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}=O_{*}(\zeta)$ with
small enough hidden factor.

Then, we can construct the following descent direction

$\displaystyle(\alpha+\alpha_{*})\nabla_{\alpha}L_{\lambda}+\langle\nabla_{\bm{\beta}}L_{\lambda},{\bm{\beta}}+{\bm{\beta}}_{*}\rangle+\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\langle\nabla_{{\bm{w}}_{i}}L_{\lambda},{\bm{w}}_{j}-q_{ij}{\bm{w}}_{i}^{*}\rangle=\Omega(\zeta),$

where $q_{ij}$ satisfy the following conditions with
$\delta_{close}<\delta_{\operatorname{sign}}$ and
$\delta_{close}=O_{*}(\zeta^{1/3})$: (1)
$\sum_{j\in{\mathcal{T}}_{i}}a_{j}q_{ij}=a_{i}^{*}$; (2) $q_{ij}\geq 0$; (3)
$q_{ij}=0$ when $\operatorname{sign}(a_{j})\neq\operatorname{sign}(a_{i}^{*})$
or $\delta_{j}>\delta_{close}$. (4)
$\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}q_{ij}^{2}=O_{*}(1)$.

###  F.4 Proof of Lemma 4.4

Now we are ready to prove the gradient lower bound (Lemma 4.4) by combining
all descent direction lemma in the previous section together.

See 4.4

###### Proof.

We check the assumption of Lemma F.17 one by one. We first assume Assumption
F.1 holds to get a gradient lower bound.

For assumption (i) (norm balance) in Lemma F.17, whenever
$\sum_{i\in[m_{*}]}\left|a_{i}^{2}-\left\|{\bm{w}}_{i}\right\|_{2}^{2}\right|=\Omega_{*}(\zeta^{2}/\lambda^{2})$,
by Lemma F.15 we know

$\displaystyle\sum_{i}\sum_{j\in
T_{i}}\left|\langle\nabla_{a_{j}}L_{\lambda},-a_{j}\rangle+\langle\nabla_{{\bm{w}}_{j}}L_{\lambda},{\bm{w}}_{j}\rangle\right|\geq$
$\displaystyle\Omega_{*}(\zeta^{2}/\lambda).$

With Lemma F.12, this implies

$\displaystyle\sqrt{\left\|\nabla_{\bm{a}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{W}}L_{\lambda}\right\|_{F}^{2}}\cdot
O(\left\|{\bm{a}}_{*}\right\|_{1})\geq\sqrt{\left\|\nabla_{\bm{a}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{W}}L_{\lambda}\right\|_{F}^{2}}\sqrt{\left\|{\bm{a}}\right\|_{2}^{2}+\left\|{\bm{W}}\right\|_{F}^{2}}=\Omega_{*}(\zeta^{2}/\lambda),$

which means

$\displaystyle\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}\geq\left\|\nabla_{\bm{a}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{W}}L_{\lambda}\right\|_{F}^{2}\geq\Omega_{*}(\zeta^{4}/\lambda^{2})$

For assumption (ii) (norm cancellation) in Lemma F.17, whenever it does not
hold, by Lemma F.16 we know

$\displaystyle\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}\geq\left\|\nabla_{{\bm{a}}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{{\bm{W}}}L_{\lambda}\right\|_{F}^{2}\geq$
$\displaystyle\lambda^{2}\sum_{j\in
T_{i,-}(\delta_{\operatorname{sign}})}|a_{j}|\left\|{\bm{w}}_{j}\right\|_{2}\geq\Omega_{*}(\zeta^{5/6}\lambda).$

For assumption (iii) ($\alpha,{\bm{\beta}}$) in Lemma F.17, whenever it does
not hold, by Lemma F.13 we know

$\displaystyle|\nabla_{\alpha}L_{\lambda}|^{2}=(\alpha-\hat{\alpha})^{2}=\Omega_{*}(\zeta^{2}),\quad\left\|\nabla_{\bm{\beta}}L_{\lambda}\right\|_{2}^{2}=4\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}=\Omega_{*}(\zeta^{2}),$

which implies

$\displaystyle\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}\geq|\nabla_{\alpha}L_{\lambda}|^{2}+\left\|\nabla_{\bm{\beta}}L_{\lambda}\right\|_{2}^{2}=\Omega_{*}(\zeta^{2}).$

Thus, the remaining case is the one that all assumption (i)-(iii) in Lemma
F.17 hold and also
$\sum_{i\in[m_{*}]}\left|a_{i}^{2}-\left\|{\bm{w}}_{i}\right\|_{2}^{2}\right|=O_{*}(\zeta^{2}/\lambda^{2})$,
we choose

$\displaystyle
q_{ij}=\left\\{\begin{array}[]{ll}\frac{a_{j}a_{i}^{*}}{\sum_{j\in
T_{i,+}(\delta_{close})}a_{j}^{2}}&\text{, if $j\in
T_{i,+}(\delta_{close})$}\\\ 0&\text{, otherwise}\end{array}\right.$

so that condition (1)-(4) on $q_{ij}$ all hold: condition (1)-(3) are easy to
check, Lemma H.4 shows condition (4) holds. Now we know from Lemma F.17 that

$\displaystyle(\alpha+\alpha_{*})\nabla_{\alpha}L_{\lambda}+\langle\nabla_{\bm{\beta}}L_{\lambda},{\bm{\beta}}+{\bm{\beta}}_{*}\rangle+\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\langle\nabla_{{\bm{w}}_{i}}L_{\lambda},{\bm{w}}_{j}-q_{ij}{\bm{w}}_{i}^{*}\rangle=\Omega(\zeta).$

Note that

$\displaystyle(\alpha+\alpha_{*})\nabla_{\alpha}L_{\lambda}+\langle\nabla_{\bm{\beta}}L_{\lambda},{\bm{\beta}}+{\bm{\beta}}_{*}\rangle+\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\langle\nabla_{{\bm{w}}_{i}}L_{\lambda},{\bm{w}}_{j}-q_{ij}{\bm{w}}_{i}^{*}\rangle$
$\displaystyle\leq$
$\displaystyle\sqrt{|\nabla_{\alpha}L_{\lambda}|^{2}+\left\|\nabla_{\bm{\beta}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{a}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{W}}L_{\lambda}\right\|_{F}^{2}}\sqrt{(\alpha+\alpha_{*})^{2}+\left\|{\bm{\beta}}+{\bm{\beta}}_{*}\right\|_{2}^{2}+\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\left\|{\bm{w}}_{j}-q_{ij}{\bm{w}}_{i}^{*}\right\|_{2}^{2}}$

and

$\displaystyle|\alpha+\alpha_{*}|\leq|\hat{\alpha}|+|\alpha_{*}|+O_{*}(\zeta)\stackrel{{\scriptstyle(\text{a})}}{{\leq}}O_{*}(1)$
$\displaystyle\left\|{\bm{\beta}}+{\bm{\beta}}_{*}\right\|_{2}\leq\left\|\hat{{\bm{\beta}}}\right\|_{2}+\left\|{\bm{\beta}}_{*}\right\|_{2}+O_{*}(\zeta)\stackrel{{\scriptstyle(\text{b})}}{{\leq}}O_{*}(1)$
$\displaystyle\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\left\|{\bm{w}}_{j}-q_{ij}{\bm{w}}_{i}^{*}\right\|_{2}^{2}\leq
2\sum_{i\in[m_{*}]}\sum_{j\in{\mathcal{T}}_{i}}\left\|{\bm{w}}_{j}\right\|_{2}^{2}+q_{ij}^{2}\left\|{\bm{w}}_{i}^{*}\right\|_{2}^{2}\stackrel{{\scriptstyle(\text{c})}}{{\leq}}O_{*}(1),$

where (a)(b) by Lemma F.4; (c) we use Lemma F.12 and condition (4) on
$q_{ij}$.

Therefore, we get

$\displaystyle\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}=|\nabla_{\alpha}L_{\lambda}|^{2}+\left\|\nabla_{\bm{\beta}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{a}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{W}}L_{\lambda}\right\|_{F}^{2}=\Omega_{*}(\zeta^{2}).$

Combine all cases above, we know

$\displaystyle\left\|\nabla_{\bm{a}}L_{\lambda}\right\|_{2}^{2}+\left\|\nabla_{\bm{W}}L_{\lambda}\right\|_{F}^{2}=\Omega_{*}(\min\\{\zeta^{4}/\lambda^{2},\zeta^{5/6}\lambda,\zeta^{2}\\})=\Omega_{*}(\zeta^{4}/\lambda^{2}),$

as long as
$\zeta=O(\lambda^{9/5}/\operatorname{poly}(r,m_{*},\Delta,\left\|{\bm{a}}_{*}\right\|_{1},a_{\min}))$.

We now use Lemma F.14 to show when Assumption F.1 is not true, we can get
similar gradient lower bound. Denote the above gradient lower bound as
$\tau_{0}=\Omega_{*}(\zeta^{4}/\lambda^{2})$. Let $\tau=\tau_{0}/2$.

When $(\alpha-\hat{\alpha})^{2}\geq\tau$ or
$\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}\geq\tau$, from Lemma
F.13 we know $\left\|\nabla_{\bm{\theta}}L_{\lambda}\right\|_{F}^{2}\geq\tau$.

When
$(\alpha-\hat{\alpha})^{2},\left\|{\bm{\beta}}-\hat{{\bm{\beta}}}\right\|_{2}^{2}\leq\tau$,
using Lemma F.14 we know there exists $\widetilde{{\bm{\theta}}}$ such that
$\left\|\nabla_{\widetilde{{\bm{\theta}}}}L_{\lambda}\right\|_{F}^{2}\geq\tau_{0}$
and
$|\left\|\nabla_{\widetilde{{\bm{\theta}}}}L_{\lambda}\right\|_{F}-\left\|\nabla_{{\bm{\theta}}}L_{\lambda}\right\|_{F}|\leq\sqrt{\tau}$.
Thus, we know $\left\|\nabla_{{\bm{\theta}}}L_{\lambda}\right\|_{F}^{2}\geq
0.1\tau$.

Therefore, combine above we can show
$\left\|\nabla_{\theta}L_{\lambda}\right\|_{F}^{2}=\Omega_{*}(\zeta^{4}/\lambda^{2})$.
∎

##  Appendix G Non-degenerate dual certificate

In this section, we show that there indeed exists a non-degenerate dual
certificate that satisfies Definition 1 and therefore proving Lemma F.1.

See F.1

Recall that we want to use the dual certificate $\eta$ to characterize the
(approximate) solution for the following regression problem:

$\displaystyle\min_{\mu\in{\mathcal{M}}({\mathbb{S}}^{d-1})}L_{\lambda}(\mu)=\mathbb{E}_{{\bm{x}},\widetilde{y}}[(f_{\mu}({\bm{x}})-\widetilde{y})^{2}]+\lambda|\mu|_{1}=\mathbb{E}_{{\bm{x}}}\left[\left(\int_{\bm{w}}\sigma_{\geq
2}({\bm{w}}^{\top}{\bm{x}}){\textnormal{d}}\
\mu-\mu_{*}\right)^{2}\right]+\lambda|\mu|_{1},$

where $\sigma_{\geq 2}$ is the ReLU activation after removing 0th and 1st
order (corresponding to $\alpha$ and $\beta$ terms) and
$\mu_{*}=\sum_{i\in[m_{*}]}a_{i}^{*}\delta_{{\bm{w}}_{i}^{*}}$ is the ground-
truth.

##### Notation

We need to first introduce few notations before proceeding to the proof.
Denote the kernel $K_{\geq\ell}({\bm{w}},{\bm{u}})=\mathbb{E}_{{\bm{x}}\sim
N(0,{\bm{I}})}[\overline{\sigma_{\geq\ell}}(\overline{{\bm{w}}}^{\top}{\bm{x}})\overline{\sigma_{\geq\ell}}(\overline{{\bm{u}}}^{\top}{\bm{x}})]$
as the kernel induced by activation $\sigma_{\geq\ell}(x)$, where
$\overline{\sigma_{\geq\ell}}(x)=\sum_{k\geq\ell}\hat{\sigma}_{k}h_{k}(x)/Z_{\sigma}$,
$Z_{\sigma}=\left\|\sigma_{\geq\ell}\right\|_{2}=\sqrt{\sum_{k\geq\ell}\hat{\sigma}_{k}^{2}}=\Theta(\ell^{-3/4})$
is the normalizing factor, $h_{k}(x)$ is the normalized $k$-th (probabilistic)
Hermite polynomial and $\hat{\sigma}_{k}$ is the corresponding Hermite
coefficient. We will specify the value of $\ell$ later and use $K$ instead of
$K_{\geq\ell}$ for simplicity.

We will construct the dual certificate $\eta$ following the proof strategy in
Poon et al. (2023) with the form below (the difference is that we now only
keep high order terms that are at least $\ell$):

$\displaystyle\eta({\bm{w}})=\sum_{j\in[m_{*}]}\alpha_{1,j}K({\bm{w}}_{j}^{*},{\bm{w}})+\sum_{j\in[m_{*}]}{\bm{\alpha}}_{2,j}^{\top}\nabla_{1}K({\bm{w}}_{j}^{*},{\bm{w}})$

such that it satisfies

$\displaystyle\eta({\bm{w}}_{i}^{*})=\operatorname{sign}(a_{i}^{*})\text{ and
}\nabla\eta({\bm{w}}_{i}^{*})=0\text{ for all $i\in[m_{*}]$.}$ (9)

Here
${\bm{\alpha}}_{1}=(\alpha_{1},\ldots,\alpha_{m_{*}})^{\top}\in\mathbb{R}^{m_{*}},{\bm{\alpha}}_{2}=({\bm{\alpha}}_{2,1}^{\top},\ldots,{\bm{\alpha}}_{2,m_{*}}^{\top})^{\top}\in\mathbb{R}^{m_{*}d}$
are the parameters that we are going to solve and $\nabla_{i}$ means the
gradient w.r.t. $i$-th variable (for example, $\nabla_{1}K({\bm{x}},{\bm{y}})$
means gradient with respect to ${\bm{x}}$).

One can rewrite the above constraints (9) into the matrix form:

$\displaystyle{\bm{\Upsilon}}\begin{pmatrix}{\bm{\alpha}}_{1}\\\
{\bm{\alpha}}_{2}\end{pmatrix}={\bm{b}},$ (10)

where
${\bm{b}}=(\operatorname{sign}(a_{1}^{*}),\ldots,\operatorname{sign}(a_{m_{*}}^{*}),{\bm{0}}_{m^{*}d}^{\top})^{\top}\in\mathbb{R}^{m_{*}(d+1)}$,
${\bm{\Upsilon}}=\mathbb{E}_{\bm{x}}[{\bm{\gamma}}({\bm{x}}){\bm{\gamma}}({\bm{x}})^{\top}]\in\mathbb{R}^{m_{*}(d+1)\times
m_{*}(d+1)}$,

$\displaystyle{\bm{\gamma}}({\bm{x}})=(\overline{\sigma_{\geq\ell}}({\bm{w}}_{1}^{*\top}{\bm{x}}),\ldots,\overline{\sigma_{\geq\ell}}({\bm{w}}_{m_{*}}^{*\top}{\bm{x}}),\nabla_{\bm{w}}\overline{\sigma_{\geq\ell}}(\overline{{\bm{w}}}_{1}^{*\top}{\bm{x}})^{\top},\ldots,\nabla_{\bm{w}}\overline{\sigma_{\geq\ell}}(\overline{{\bm{w}}}_{m_{*}}^{*\top}{\bm{x}})^{\top})^{\top}\in\mathbb{R}^{m_{*}(d+1)}.$

Here
$\nabla_{\bm{w}}\overline{\sigma_{\geq\ell}}(\overline{{\bm{w}}}_{i}^{*\top}{\bm{x}})={\bm{P}}_{{\bm{w}}_{i}^{*}}\overline{\sigma_{\geq\ell}}^{\prime}({\bm{w}}_{i}^{*\top}{\bm{x}}){\bm{x}}\in\mathbb{R}^{d}$,
where ${\bm{P}}_{{\bm{w}}_{i}^{*}}$ is the projection matrix defined below.

##### Notions on the unit sphere

As we could see, the kernel $K$ is invariant under the change of norms, so it
suffices to focus on the input on the unit sphere ${\mathbb{S}}^{d-1}$. On the
unite sphere, we could compute the gradient and hessian of a function
$f({\bm{w}})$ on the sphere (e.g., Absil et al. (2013))

$\displaystyle\operatorname{grad}f({\bm{w}})$
$\displaystyle={\bm{P}}_{\bm{w}}\nabla f({\bm{w}}),$
$\displaystyle\operatorname{H}f({\bm{w}})[{\bm{z}}]$
$\displaystyle={\bm{P}}_{\bm{w}}(\nabla^{2}f({\bm{w}})-\overline{{\bm{w}}}^{\top}\nabla
f({\bm{w}}){\bm{I}}){\bm{z}}\quad\text{for all tangent vector ${\bm{z}}$ that
${\bm{z}}^{\top}{\bm{w}}=0$},$

where ${\bm{P}}_{\bm{w}}={\bm{I}}-{\bm{w}}{\bm{w}}^{\top}$ is the projection
matrix.

Then, we could define the derivative as in Poon et al. (2023); Absil et al.
(2008): for tangent vectors ${\bm{z}},{{\bm{z}}^{\prime}}$

$\displaystyle\operatorname{D}_{0}f({\bm{w}})$ $\displaystyle:=f({\bm{w}})$
$\displaystyle\operatorname{D}_{1}f({\bm{w}})[{\bm{z}}]$
$\displaystyle:=\langle{\bm{z}},\operatorname{grad}f({\bm{w}})\rangle={\bm{z}}^{\top}{\bm{P}}_{\bm{w}}\nabla
f({\bm{w}})$
$\displaystyle\operatorname{D}_{2}f({\bm{w}})[{\bm{z}},{{\bm{z}}^{\prime}}]$
$\displaystyle:=\langle\operatorname{H}f({\bm{w}})[{\bm{z}}],{{\bm{z}}^{\prime}}\rangle={\bm{z}}^{\top}{\bm{P}}_{\bm{w}}(\nabla^{2}f({\bm{w}})-\overline{{\bm{w}}}^{\top}\nabla
f({\bm{w}}){\bm{I}}){\bm{P}}_{\bm{w}}{{\bm{z}}^{\prime}},$

and their associated norms

$\displaystyle\left\|\operatorname{D}_{1}f({\bm{w}})\right\|_{\bm{w}}$
$\displaystyle:=\sup_{\left\|{\bm{z}}\right\|_{\bm{w}}=1}\operatorname{D}_{1}f({\bm{w}})[{\bm{z}}]=\left\|{\bm{P}}_{\bm{w}}\nabla
f({\bm{w}})\right\|_{2},$
$\displaystyle\left\|\operatorname{D}_{2}f({\bm{w}})\right\|_{\bm{w}}$
$\displaystyle:=\sup_{\left\|{\bm{z}}\right\|_{\bm{w}},\left\|{{\bm{z}}^{\prime}}\right\|_{\bm{w}}=1}\operatorname{D}_{2}f({\bm{w}})[{\bm{z}},{{\bm{z}}^{\prime}}]=\left\|{\bm{P}}_{\bm{w}}\operatorname{H}f({\bm{w}}){\bm{P}}_{\bm{w}}\right\|_{2},$

where
$\left\|{\bm{z}}\right\|_{\bm{w}}=\left\|{\bm{P}}_{\bm{w}}{\bm{z}}\right\|_{2}$.

For simplicity, we will use $K^{(ij)}({\bm{w}},{\bm{u}})$ to denote
$\nabla_{1}^{i}\nabla_{2}^{j}K({\bm{w}},{\bm{u}})$. One can check that this is
in fact the same as the one defined Poon et al. (2023) under our specific
kernel $K$, $i+j\leq 3$ and $i,j\leq 2$. Let

$\displaystyle\left\|K^{(ij)}({\bm{w}},{\bm{u}})\right\|_{{\bm{w}},{\bm{u}}}:=\sup_{\begin{subarray}{c}\left\|{\bm{z}}_{\bm{w}}^{(p)}\right\|_{\bm{w}}=\left\|{\bm{z}}_{\bm{u}}^{(q)}\right\|_{\bm{u}}=1,\\\
{\bm{w}}^{\top}{\bm{z}}_{\bm{w}}^{(p)}={\bm{u}}^{\top}{\bm{z}}_{\bm{u}}^{(q)}=0\
\forall
p\in[i],q\in[j]\end{subarray}}K^{(ij)}({\bm{w}},{\bm{u}})[{\bm{z}}_{\bm{w}}^{(1)},\ldots,{\bm{z}}_{\bm{u}}^{(j)}],$

where ${\bm{z}}_{\bm{w}}^{(p)}$ applies to the dimension corresponding to
${\bm{w}}$ and similarly ${\bm{z}}_{\bm{u}}^{(q)}$ for ${\bm{u}}$.

Before solving (10), we first present some useful proprieties of kernel $K$
that will be used later (see Section I for the proofs). The lemma below shows
that kernel $K({\bm{w}},{\bm{u}})$ is non-degenerate in the sense that it
decays at least quadratic at each ground-truth direction
(${\bm{w}}\approx{\bm{u}}\approx{\bm{w}}_{i}^{*}$) and contributes almost
nothing when ${\bm{w}},{\bm{u}}$ are away.

######  Lemma G.1 (Non-degeneracy of kernel $K$).

For any $h>0$, let $\ell\geq\Theta(\Delta^{-2}\log(m_{*}\ell/h\Delta))$,
kernel $K_{\geq\ell}$ is non-degenerate in the sense that there exists
$r=\Theta(\ell^{-1/2}),\rho_{1}=\Theta(1),\rho_{2}=\Theta(\ell)$ such that
following hold:

  1. (i)

$K({\bm{w}},{\bm{u}})\leq 1-\rho_{1}$ for all
$\delta({\bm{w}},{\bm{u}}):=\angle({\bm{w}},{\bm{u}})\geq r$.

  2. (ii)

$K^{(20)}({\bm{w}},{\bm{u}})[{\bm{z}},{\bm{z}}]\leq-\rho_{2}\left\|{\bm{z}}\right\|^{2}$
for tangent vector ${\bm{z}}$ that ${\bm{z}}^{\top}{\bm{w}}=0$ and
$\delta({\bm{w}},{\bm{u}})\leq r$.

  3. (iii)

$\left\|K^{(ij)}({\bm{w}}_{1}^{*},{\bm{w}}_{k}^{*})\right\|_{{\bm{w}}_{i}^{*},{\bm{w}}_{k}^{*}}\leq
h/m_{*}^{2}$ for $(i,j)\in\\{0,1\\}\times\\{0,1,2\\}$

The following lemma shows that $K$ and its derivatives are bounded.

######  Lemma G.2 (Regularity conditions on kernel $K$).

Let
$B_{ij}:=\sup_{{\bm{w}},{\bm{u}}}\left\|K^{(ij)}({\bm{w}},{\bm{u}})\right\|_{{\bm{w}},{\bm{u}}}$
and $B_{0}=B_{00}+B_{10}+1$, $B_{2}=B_{20}+B_{21}+1$. We have $B_{00}=O(1)$,
$B_{10}=O(\ell^{1/2})$, $B_{11}=O(\ell)$, $B_{20}=O(\ell)$,
$B_{21}=O(\ell^{3/2})$, and therefore $B_{0}=O(\ell^{1/2})$,
$B_{2}=O(\ell^{3/2})$.

The following lemma from Poon et al. (2023) connects the non-degeneracy of
kernel $K$ to the dual certificate $\eta$ that we are interested in.

######  Lemma G.3 (Lemma 2, Poon et al. (2023), adapted in our setting).

Let $a\in\\{\pm 1\\}$. Suppose that for some $\rho>0$, $B>0$ and $0<r\leq
B^{-1/2}$ we have: for all $\delta({\bm{w}},{\bm{w}}_{0})$ and
${\bm{z}}\in\mathbb{R}^{d}$ with ${\bm{z}}^{\top}{\bm{w}}=0$, it holds that