SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics

Introduction

Deep learning has emerged as the standard approach to exploiting massive high-dimensional datasets. At the core of its success lies its capability to learn effective features with fairly blackbox architectures without suffering from the curse of dimensionality. To explain this success, two structural properties of data are commonly conjectured: (i) a low-dimensional structure that SGD-trained neural networks are able to adapt to; (ii) a hierarchical structure that neural networks can leverage with SGD training. In particular,

A line of work has investigated the sample complexity of learning with deep neural networks, decoupled from computational considerations. By directly considering global solutions of empirical risk minimization (ERM) problems over arbitrarily large neural networks and sparsity inducing norms, they showed that deep neural networks can overcome the curse of dimensionality on classes of functions with low-dimensional and hierarchical structures. However, this approach does not provide efficient algorithms: instead, a number of works have shown computational hardness of ERM problems and it is unclear how much this line of work can inform practical neural networks, which are trained using SGD and variants.

A line of work in computational learning theory has provided time- and sample-efficient algorithms for learning Boolean functions with low-dimensional structure, based on their sparse Fourier spectrum . However, these algorithms are a priori quite different from SGD-trained neural networks. While unconstrained architectures can emulate any efficient learning algorithms , it is unclear whether more ‘standard’ neural networks can succeed on these same classes of functions or whether they require additional structure that pertains to hierarchical properties.

Thus, an outstanding question emerges from the current state of affairs:

For neural networks satisfying “regularity assumptions” (e.g., fully-connected, isotropically initialized layers), are there structural properties of the data that govern the time complexity of SGD learning? How does SGD exploit these properties in its training dynamics?

Here the key points are: (i) the “regularity” assumption, which prohibits the use of unorthodox neural networks that can emulate generic PAC/SQ learning algorithms as in ; (ii) the requirement on the time complexity, which prohibits direct applications of infinite width, continuous time or infinite time analyses as in . We discuss in Section 1.3 the various works that made progress towards the above, in particular regarding single- and multi-index models. We now specify the setting of this paper.

We focus on the following class of data distributions. First of all, we consider IID inputs, i.e.,

and we focus on the case where $\mu$ is either $\mathcal{N}(0,1)$ or ${\rm Unif}(\{+1,-1\})$, although we expect that other distributions would admit a similar treatment. Incidentally, the latter distribution is of interest in reasoning tasks related to Boolean arithmetic or logic . We now make a key assumption on the target function, that of having a low latent dimension, i.e., $f_{*}({\bm{x}})=h_{*}({\bm{z}})$ where ${\bm{z}}={\bm{M}}{\bm{x}}$ and

with the assumption that $P=O_{d}(1)$. In other words, the target function has a large ambient dimension but depends only on a finite number of latent coordinates. In the Gaussian case the coordinates are not known because of a possible rotation of the input, and in the Boolean case the coordinates are not known because of a possible permutation of the input.

Data with large ambient dimension but low latent dimension have long been a center of focus in machine learning and data science. It is known that kernel methods cannot exploit low latent dimension, i.e., it was proved in that any kernel method needs a number of features $p$ or samples $n$ satisfying

in order to learn a Boolean function as above with degree $D=O_{d}(1)$. In other words, for kernel methods $D$ controls the sample complexity irrespective of any potential additional structural properties of $f_{*}$ (e.g., hierarchical properties). On the other hand, it is known that this is not the limit for deep learning, which can break the $d^{D}$ curse, as discussed next.

Consider the following example: ${\bm{x}}\sim{\rm Unif}(\{+1,-1\}^{d})$ is drawn from the hypercube and the target function is $4$-sparse, either

The first function is called a vanilla staircase of degree 4 . The second is a monomial of degree 4. Each of these functions induces a function class under the permutation of the variables (i.e., one can consider the class of all monomials on any 4 of the $d$ input variables, and similarly for staircases). One can verify that these function classes have similar approximation and statistical complexity because of the low-dimensional structure, but have different computational complexity because of the hierarchical structure. For example, under the Correlational Statistical Query (CSQ) model of computation , the first class has CSQ dimension $\Theta(d)$ versus $\Theta(d^{4})$ for the second classSee Section 2 for more details on CSQ..

1 The leap complexity

We now define the leap complexity. Any function in $L^{2}(\mu^{\otimes P})$ can be expressed in the orthogonal basis of $L^{2}(\mu^{\otimes P})$, i.e., the Hermite or Fourier-Walsh basis for $\mu\sim{\sf N}(0,1)$ and $\mu\sim{\rm Unif}(\{+1,-1\})$ respectively,

In words, a function $h_{*}$ is leap-$k$ if its non-zero monomials can be ordered in a sequence such that each time a monomial is added, the support of $h_{*}$ grows by at most $k$ new coordinates, where each new coordinate is counted with multiplicity in the Gaussian case (and the 1-norm collapses to the cardinality of the difference set in the Boolean case). Note that the definition of leap-$k$ functions on the hypercube generalizes the definition of functions with the merged-staircase property (leap-1 functions) from .

2 Summary of our contributions

This paper puts forward a general conjecture characterizing the time complexity of SGD-learning on regular neural networks with isotropic data of low latent dimension. The key quantity that emerges to govern the complexity is the leap (Definition 1). This gives a formal measure of “hierarchy” in target functions, going beyond spectrum sparsity and emerging from the study of SGD-trained regular networks. The paper then proves a specialization of the conjecture to a representative class of functions on Gaussian inputs, but for 2-layer neural networks and with certain technical assumptions on how SGD is run. The two main innovations of the proof are (i) a full control of the time complexity of SGD learning on a fully-connected network (without infinite width or continuous time approximations); (ii) going beyond one-step gradient analyses and showing that the leap controls the entire learning trajectory due to a sequential learning mechanism (saddle-to-saddle). We also provide experimental evidence towards the more general conjecture with vanilla SGD and derive CSQ lower-bounds for noisy GD that match our achievability bounds.

We believe that the conjecture (in particular the time complexity scaling) holds for more general architectures than those with isotropically-initialized layers, as long as enough ‘regularity’ assumptions are present at initialization (prohibiting the type of ‘emulation networks’ used in ). Note that it is not enough to ask for only the first layer to be initialized with a rotationally-invariant distribution, as this may be handled by using emulation networks on the subsequent layers, but weaker invariances of subsequent layers (e.g., permutation subgroups) may suffice.

The characterization obtained in this paper implies a relatively simple picture for learning low-dimensional functions with SGD on neural networks:

SGD on regular neural networks implicitly implements a form of ‘adaptive curriculum’ learning. SGD first picks up low-level features that are computationally and statistically easier to learn, and by picking up these low level features, it makes the learning of higher-level features in turn easier. As mentioned in the examples of Table 1: learning $z_{1}\cdots z_{2k}$ takes $\widetilde{\Theta}(d^{2k-1})$ sample complexity (leap-$2k$ function). But if we add an intermediary monomial to our target to create $z_{1}\cdots z_{k}+z_{1}\cdots z_{2k}$, then it takes $\widetilde{\Theta}(d^{k-1})$ steps to learn (leap-$k$ function). If we have a full staircase, it only requires $\Theta(d)$ (leap-$1$ function). This thus gives an adaptive learning process that follows a curriculum learning procedure where features of increasing complexity guide the learning.

3 Related works

In parallel, several works have studied the dynamics of SGD in simpler non-convex models in high dimensions . Our analysis relies on a similar drift plus martingale decomposition of online-SGD as in . In particular, the leap complexity is related to the information-exponent introduced in . The latter considers a single-index model trained with online-SGD on a non-convex loss and the information exponent captures the scaling of the correlation between the model at a typical initialization and the global solution. showed that, with information exponent $k$, online-SGD requires $\widetilde{\Theta}(d^{k-1\vee 1})$ steps to converge, similarly to the scaling presented in this paper. However our analysis and the definition of the leap complexity differ from in two major ways. First, our model is not a single parameter model, so a much more involved analysis is required for the dynamics. Second, the information exponent is a coefficient that only applies at initialization, while the leap-complexity is a measure of targets that controls the entire learning trajectory (our neural networks visit several saddles during training).

Lower bounds on learning leap functions

Linear methods such as kernel methods suffer exponentially in the degree of the target function, and cannot use the “hierarchical” structure to learn faster. This was proved in for the Boolean case, and this work extends the result to the Gaussian case:

Let $h_{*}$ be a degree-$D$ polynomial over the Boolean hypercube (resp., Gaussian measure). Then there are $c_{h_{*}},\varepsilon_{h_{*}}>0$, such that any linear method needs $c_{h_{*}}d^{D}$ samples to learn $f_{*}({\bm{x}})=h_{*}({\bm{M}}{\bm{x}})$ to less than $\varepsilon_{h_{*}}>0$ error, where ${\bm{M}}$ is an unknown permutation (resp., rotation) as in (2).

Consider now the Correlational Statistical Query (CSQ) model of computation . A CSQ algorithm accesses the data via expectation queries, plus additive noise. We show that for CSQ methods the query complexity scales exponentially in the leap of the target function, which can be much less than the degree.

We note that the above lower-bounds are for CSQ models or noisy population-GD models, and not for online-SGD since the latter takes a single sample per time step. Our proof does show a correspondence between online-SGD and population-GD, but without the additional noise. It is however intriguing that the regularity in the network model for online-SGD appears to act comparatively in terms of constraints to a noisy population-GD model (on possibly non-regular architectures), and we leave potential investigations of such correspondences to future work (see also discussion in Appendix B.5). Further, we note that the correspondence to CSQ may not hold beyond the finite $P$ regime. First there is the ‘extremal case’ of learning the full parity function, which is efficiently learnable in CSQ (with 0 queries) but not necessarily with online-SGD on regular networks: shows it is efficiently learnable by a i.i.d. Rademacher$(1/2)$ initialization, but not necessarily by a Gaussian isotropic initialization. Further, the positive result of the Rademacher initialization may disappear under proper hyperparameter ‘stability’ assumptions. Beyond this extremal case, a more important nuance arises for large $P$: the fitting of the function on the support may become costly for regular neural networks in certain cases. For example, let $g:[P]\to\{0,1\}$ be a known function and consider learning $f_{*}({\bm{x}})$ which depends on $P$ unknown coordinates as $h_{*}({\bm{z}})=\sum_{i=1}^{P}iz_{i}+\prod_{i=1}^{P}z_{i}^{g(i)}$. This is a leap-1 function where the linear part reveals the support and the permutation, and with a parity term on the indices such that $g(i)=1$. In this case, SGD on a regular network would first pick up the support, and then have to express a potentially large degree monomial on that support, which may be hard if $P$ is large (i.e., $P\gg 1$). The latter part may be non trivial for SGD on a regular network, while, since $g$ is known, it would require 0 queries for a CSQ algorithm once the permutation was determined from learning the linear coefficients.

Learning leap functions with SGD on neural networks

We consider the following assumption on the activation function:This is satisfied, for example, by the shifted sigmoid $\sigma(z)=1/(1+e^{-z+c})$ for almost all shifts $c$.

For the purposes of the analysis, we make two modifications to SGD training. First, we train layerwise: training $\{{\bm{w}}_{j}\}_{j\in[M]}$ and then $\{a_{j}\}_{j\in[M]}$, while keeping the biases $\{b_{j}\}_{j\in[M]}$ frozen during the whole training. Second, during the training of the first layer weights $\{{\bm{w}}_{j}\}_{j\in[M]}$, we project the weights in order to ensure that they remain bounded in magnitude. See Algorithm 1 for pseudocode, and see below for a detailed explanation. These modifications are not needed in practice for SGD to learn, as we demonstrate in our experiments in Figure 1 and Appendix A.

Analyzing layerwise training is a fairly standard tool in the theoretical literature to obtain rigorous analyses; it is used in a number of works, including . In our setting, layerwise training allows us to analyze the complicated dynamics of neural network training, but it also leads to a major issue. During the training of the first layer, the target function $f_{*}$ is not fully fitted because we do not train the second layer concurrently. Therefore the first-layer weights continue to evolve even after they pick up the support of $f_{*}$. This is a challenge since we must train the first-layer weights for a large number of steps, and so they can potentially grow to a very large magnitude, leading to instability.We emphasize that this problem is due to layerwise training, since in practice if we train both layers at the same time the residual quickly goes to zero after the support is picked up, and so the first-layer weights stop evolving and remain bounded in magnitude (see Appendix A).

We correct the issue by projecting each neuron’s first-layer weights ${\bm{w}}_{j}$ to ensure that the coordinates do not blow up. First, we keep the “small” coordinates of ${\bm{w}}_{j}$ on the unit sphere, i.e., for some parameter $r>0$, we define the “small” coordinates for neuron $j$ at time $t$ by $S_{j,0}=[d]$ and

We project these coordinates on the unit sphere using the operator ${\mathsf{P}}^{t}_{j}$ defined by

and use the spherical gradient with respect to the sphere $\|{\mathcal{S}}_{t}({\bm{w}}_{j}^{t})\|_{2}=1$, i.e., for any function $f$,

In the second phase, the training of the second layer weights ${\bm{a}}$ is by standard SGD (without projection) with added ridge-regularization term $\frac{\lambda_{a}}{2}\|{\bm{a}}\|^{2}$ to encourage low-norm solutions.

2 Learning a single monomial

We first consider the case of learning a single monomial with Hermite exponents $k_{1},\ldots,k_{P}\geq 1$:

We assume $D=k_{1}+\ldots+k_{P}\geq 2$ (the case $D=1$ is straightforward). $h_{*}$ is a leap-$D$ function. We start by proving that, during the first phase, the first layer weights grow in the directions of $z_{1},\ldots,z_{P}$ which are the variables in the support of the target function.

Assume $\sigma$ satisfy Assumption 1. Then for $0<r<\Delta$ sufficiently small (depending on $D,K$) and $\rho\leq\Delta$ the following holds. For any constant $C_{*}>0$, there exist $C_{i}$ for $i=0,\ldots,6$, that only depend on $D,K$ and $C_{*}$ such that

and for $d$ large enough that $r\geq C_{0}\log(d)^{C_{0}}/\sqrt{d}$, the following event holds with probability at least $1-Md^{-C_{*}}$. For any neuron $j\in[M]$,

Early stopping: $|w_{j,i}^{t}-w_{j,i}^{0}|\leq C_{3}/\sqrt{d\log(d)}$ for all $i\in[d]$ and $t\leq\overline{T}_{1}/(C_{4}\log(d)^{C_{4}})$.

And for any neuron $j\in[M]$ such that $a_{j}^{0}\mu_{D}(\sigma)(w_{j,1}^{0})^{k_{1}}\cdots(w_{j,P}^{0})^{k_{P}}>0$,

On the support: \big{|}w_{j,i}^{\overline{T}_{1}}-\text{sign}(w_{j,i}^{0})\cdot\Delta\big{|}\leq C_{5}/\sqrt{d\log(d)} for $i=1,\ldots,P$.

Outside the support: $|w_{j,i}^{\overline{T}_{1}}-w_{j,i}^{0}|\leq C_{6}r^{2}/\sqrt{d}$ for $i=P+1,\ldots,d$, and $\sum_{i>P}(w_{j,i}^{\overline{T}_{1}})^{2}=1$.

Theorem 1 shows that after the end of the first phase, the coordinates ${\bm{w}}_{j}^{\overline{T}_{1}}$ aligned with the support ${\bm{z}}$ are all close to $\pm\Delta$ with the same signs as $w_{j,1}^{0},\ldots,w_{j,P}^{0}$ as long as $(w_{j,1}^{0})^{k_{1}}\cdots(w_{j,P}^{0})^{k_{P}}>0$ has the same sign as $a_{j}^{0}\mu_{D}(\sigma)$ at initialization. Furthermore, the correlation with the support only appears at the end of the dynamics, and does not appear if we stop early.

The proof of Theorem 1 follows a similar proof strategy as , namely a decomposition of the dynamics into a drift and martingale terms with information exponent $D$. However, our problem is multi-index, and the analysis will require a tighter control of the different contributions to the dynamics as the dynamics move from saddle to saddle. An heuristic explanation for this result can be found in Appendix B.3. The complete proof of Theorem 1 is deferred to Appendix C.

Let $h_{*}({\bm{z}})=z_{1}\cdots z_{D}$ or $h_{*}({\bm{z}})={\rm He}_{D}(z_{1})$ and assume $\sigma$ satisfies Assumption 1. For any constants $C_{*}>0$ and $\varepsilon>0$, there exist $C_{i}$ for $i=0,\ldots,11$, that only depend on $D,K$ and $C_{*}$ such that taking width $M=C_{0}\varepsilon^{-C_{0}}$, bias initialization scale $\rho=\varepsilon^{C_{1}}/C_{1}$, and $\Delta=\varepsilon^{C_{1}}/C_{1}$ and second-layer initialization scale $\kappa=\frac{1}{C_{2}Md^{C_{2}}}$, and second-layer regularization $\lambda_{a}=M\varepsilon/C_{3}$, and , and $r=\varepsilon^{C_{4}}/C_{4}$, and

for $d\geq C_{9}\varepsilon^{-C_{9}}$ we have with probability at least $1-d^{-C_{*}}-\varepsilon$:

If we train the first layer weights for $\overline{T}_{1}^{\prime}\leq\overline{T}_{1}/(C_{10}\log(d)^{C_{10}})$ steps and for $M\leq C_{10}\log(d)$, then we cannot fit $f_{*}$ using the second-layer weights, i.e.,

This result suggests that the dynamics of SGD with one monomial can be decomposed into a ‘search phase’ (plateau in the learning curve) and a ‘fitting phase’ (rapid decrease of the loss) similarly to . SGD progressively aligns the first layer weights with the support, with little progress, and as soon as SGD picks up the support, the second layer weights can drive the risk quickly to . Because of the layer-wise training, we only show in Corollary 1.(b) that with early stopping on the training of the first layer weights, we cannot approximate the function $f_{*}$ at all using the second layer weights (hence, we cannot learn it even with infinite number of samples). The proof of Corollary 1 is in Appendix E.1.

3 Learning multiple monomials

We now consider $h_{*}$ with several monomials in its decomposition. In order to simplify the statement and the proofs, we will specifically consider the case of nested monomials

where $0=:P_{0}<P_{1}<P_{2}<\ldots<P_{L}=:P$ and $k_{1},\ldots,k_{P}$ are positive integers. For $l\in[L]$, we denote $D_{l}=k_{P_{l-1}+1}+\ldots+k_{P_{l}}$, and $D=\max_{l\in[L]}D_{l}$ the size of the biggest leap (such that $h_{*}$ is a leap-$D$ function), $\overline{D}_{l}=D_{1}+\ldots+D_{L}$ and $\overline{D}:=D_{L}$ the total degree of the polynomial $h_{*}$. We will assume that $\min_{l\in[L]}D_{l}\geq 2$ (i.e., leap of size at least $2$ between monomials). This specific choice for $h_{*}$ allows for a more compact proof, similar to Theorem 1. However, the compositionality of $h_{*}$ is not a required structure for the sequential alignment to hold and we describe in Appendix D.2 how to modify the analysis for more generalHowever, our current proof techniques do not allow for fully general leap functions: e.g., $h_{*}({\bm{z}})={\rm He}_{2}(z_{1}){\rm He}_{3}(z_{2})-{\rm He}_{3}(z_{1}){\rm He}_{2}(z_{2})$ has its two monomials pushing the ${\bm{w}}_{j}$’s in two opposite directions. $h_{*}$.

We first prove that the first-layer weights grow in the relevant directions during training.

On the support: \big{|}w_{j,i}^{\overline{T}_{1}}-\text{sign}(w_{j,i}^{0})\cdot\Delta\big{|}\leq C_{5}/\sqrt{d\log(d)} for $i=1,\ldots,P$.

Outside the support: $|w_{j,i}^{\overline{T}_{1}}-w_{j,i}^{0}|\leq C_{6}r^{2}/\sqrt{d}$ for $i=P+1,\ldots,d$ and $\sum_{i>P}(w_{j,i}^{\overline{T}_{1}})^{2}=1$.

The proof follows by showing the sequential alignment of the weights to the support: with high probability and for each neurons satisfying the sign condition at initialization, it takes between $d^{\frac{D_{l}+D}{2}-1}/(C\log(d)^{C})$ and $d^{\frac{D_{l}+D}{2}-1}C\log(d)^{C}$ steps to align with coordinates $[P_{l}]$, after having picked up coordinates $[P_{l-1}]$. The proof can be found in Appendix D.

While Theorem 2 captures the tight scaling in overall number of steps, it does not capture the number of steps for smaller leaps $D_{l}<D$ shown in Figure 1 in the case of increasing leaps. In Appendix D.2.1, we show that the scaling of $d^{D_{l}-1}$ steps to align to the next monomial can be obtained by varying the step size, in the case of increasing leaps. Note that in practice, neural networks with constant step size seem to achieve this optimal scaling for escaping each saddle (such as in Figure 1). Hence, there might be a mechanism in the SGD training that can implicitly control the martingale part of the dynamics, without rescaling the step sizes. However, understanding such a mechanism would require to study the joint training of both layers, which is currently out of reach of our proof techniques.

As in the single monomial case, we consider fitting the second layer weights only for a specific class of functions (where all monomials are multilinear):

We require extra assumptions on the activation function to prove that the fitting is possible. The following is an informal statement, and we leave the formal statement and proof to Appendix E.2.

Discussion

One direction for future work is to remove the modifications to vanilla SGD used in the analysis (layerwise training and the projection step). Another direction is to prove the conjecture by extending our analysis of the training dynamics to general functions, beyond those of the form (11). Another direction is to study extensions of the leap complexity measure beyond isotropic input distributions.

Acknowledgement

Part of this work was supported by the NSF-Simons Research Collaborations on the Mathematical and Scientific Foundations of Deep Learning (MoDL) Award and the EPFL PhD Exchange Fellowship. EB was also generously supported by Apple with an AI/ML fellowship. TM also acknowledges the NSF grant CCF-2006489 and the ONR grant N00014-18-1-2729.

References

Appendix A Additional numerical simulations

In Figures 2, 3, 4 and 5 we plot the risk versus number of samples for SGD training of 5-layer ResNets with fully-connected layers for various different target functions and for Boolean and Gaussian data. In these plots, the saddle-to-saddle dynamics are visible, which are caused by the neural network sequentially picking up the support using the hierarchical structure of the monomials in the function. In Figures 6 and 7, we study learning a leap-1 function (merged-staircase function), and we experiment with the effect of adding depth to see its effect on fitting. There is also an interesting edge-of-stability behavior during the “second-layer fitting” part, where the loss does not decrease monotonically . We leave understanding this to future work.

Appendix B Additional discussion from the main text

In addition to the references listed in the main text, we further review other relevant papers.

A line of work in computational learning theory studied the complexity of learning Boolean functions under the uniform input distribution. It was realized that functions with concentrated Fourier spectrum can be learned efficiently, both in sample and time complexity using the sparse Fourier algorithm . Namely, under knowledge of a set of basis elements $\mathcal{S}$ such that $\sum_{S\in\mathcal{S}}f^{2}(S)\geq 1-\epsilon/2$ for all $f\in\mathcal{F}$, one can learn $\mathcal{F}$ with error $\epsilon$, sample complexity $O((1/\epsilon)|\mathcal{S}|\log(|\mathcal{S}|/\delta))$ and polynomial time complexity if $|\mathcal{S}|$ is polynomial using the sparse Fourier algorithm that estimates the coefficients in $\mathcal{S}$. Many interesting classes of functions fall under this setting, such as juntas, low degree functions, bounded-size or -depth decision trees . While $\mathcal{S}$ has to be knownThe set knowledge can be relaxed under the query access model using the Kushilevitz-Mansour algorithm (based on the Goldreich-Levin algorithm) that uses a divide-and-conquer procedure to identify the coefficients to be estimated . under the random sample model, no degree constraints are imposed. In particular, the low-degree assumption (degree at most $k$) is just a special case that provides this knowledge (with order $d^{k}$ time complexity), monomials of degree $k$ or $d-k$ are equivalent in the eye of the sparse Fourier algorithm. This is not necessarily the case for SGD-trained neural networks.

A line of work has considered SGD learning on ‘unconstrained’ neural networks (besides polynomial size) and shows that we can emulate any efficient PAC or SQ algorithm . Such networks are far from the practical neural networks used in applications. Against this state of affairs, several works have attempted to derive computational lower bounds on learning with regular neural networks. For example, shows that for fully connected 2-layer networks, if the initial alignment (INAL) of a network with a Boolean target function (measured by the maximal expected correlation between target and neurons) is not significant, then noisy-GD cannot amplify the correlation to any significant level. This is achieved by showing that a low INAL implies a large minimal degree in the target function (thus a large leap) under some additional conditions. Another work uses the permutation, sign-flip, or rotational equivariance of noisy-GD training of fully-connected neural networks to show a lower bound on the number of gradient descent steps required for global convergence, when we have access to population gradients with an additive Gaussian noise. In particular, for leap-$L$ functions on the hypercube and the hypercube, $\Omega(d^{L}\tau^{2})$ steps are to shown to be required, where $\tau^{2}$ is the Gaussian noise variance. This roughly matches the conjecture in this paper in its exponential dependence on the leap – however, the computational model is different (noisy-GD versus online-SGD).

Finally let’s remark that a large body of work in the statistics and machine learning literature has studied the problem of learning multi-index models. These include for example phase retrieval , intersection of halfspaces and subspace juntas . We refer to and references therein for an overview of this line of work. In particular, it is well understood that in order to break the “curse of dimensionality”, the algorithm needs to estimate the low-dimensional support. In contrast with this line of work, we consider learning these multi-index functions with generic SGD on regular neural networks, with no a priori information on the target function. Surprisingly, we show that this generic algorithm can nearly match the computational complexity of the best CSQ algorithm. Note that specialized algorithms can achieve better sample and computational complexity: for example, showed an algorithm that can learn low-rank Gaussian polynomials in $\widetilde{O}_{d}(d)$ samples and $\widetilde{O}_{d}(d^{3})$ runtime, regardless of the leap-complexity, by going beyond CSQ algorithms.

B.2 Discussion on the definition of the leap complexity

It was noted in that some “degenerate” leap-1 functions on the hypercube are not learned in $\Theta(d)$ SGD-steps. Take for example $h_{*}({\bm{z}})=z_{1}+z_{2}+z_{3}+z_{1}z_{2}z_{3}$: by permutation symmetry on the support of $h_{*}$, $O(d)$ steps of SGD will learn first layer weights ${\bm{w}}_{j}$ aligned with $(1,1,1)$ on the support $(z_{1},z_{2},z_{3})$. SGD will require many more steps to break this symmetryWe conjecture $\Theta(d\log(d)^{C})$ steps are required, see following discussion in the Gaussian case. and fit $h_{*}$. circumvents this difficulty under a smoothed complexity analysis, and shows that the set of degenerate leap-$1$ functions has $\{\hat{h}_{*}(S)\}_{S\in{\mathcal{S}}}$ of Lebesgue-measure . Alternatively, a possible approach to learn these degenerate cases (for “axis-aligned” sparse functions) is to use different random learning rates for each coordinates and break the symmetry in learning.

B.3 Intuition for the proof of Theorem 1

In this section, we give some intuition behind the proof of Theorem 1. The complete proof can be found in Appendix C.

We first consider a simple SGD dynamics, with no projection step, and neglect the biases. We later discuss our choice of algorithm and how the analysis needs to be modified to control the projection step. The dynamics on the first layer weights is now simply given by

Recall that we initialize the second layer weights $|a_{j}^{0}|=\kappa$. By Assumption 1, we have

With high probability over a polynomial number of steps, $\|{\bm{x}}^{t}\|_{\infty}\leq C\log(d)$ with $C$ constant chosen sufficiently large. Hence,

and we can chose $\kappa,\eta_{1}$ with $\eta_{1}\kappa^{2}$ sufficiently small, while keeping $\eta_{1}\kappa$ constant, so that we can neglect the interaction term between the different neurons and get:

Let us directly consider the correlation loss and track the dynamics of a unique neuron $(a,{\bm{w}})$. We assume that $w^{0}_{1}=\ldots=w^{0}_{d}=1/\sqrt{d}$, $a^{0}=\kappa$ and $\mu_{D}(\sigma)>0$. We further make the following heuristic simplification: we assume the dynamics is described by only two parameters

with SGD updates $g_{1}^{t}=y^{t}\sigma^{\prime}(\langle{\bm{w}}^{t},{\bm{x}}^{t}\rangle)x_{1}^{t}$ and $g_{2}^{t}=y^{t}\sigma^{\prime}(\langle{\bm{w}}^{t},{\bm{x}}^{t}\rangle)x_{P+1}^{t}$, i.e.,

where ${\bm{u}}=(u_{1},u_{2})$ and ${\bm{g}}=(g_{1},g_{2})\sim{\sf N}(0,{\mathbf{I}}_{2})$. We deduce that to leading term (assuming $\|{\bm{w}}^{t}\|_{2}\approx 1$)

Let us now control the different contributions to the dynamics:

Martingale part: By Doob’s maximal inequality for martingales, we have with high probability

We choose $\eta_{1}\kappa$ so that we can neglect the martingale contribution during the entire dynamics by taking $\eta_{1}\kappa\sqrt{\overline{T}_{1}}\lesssim\alpha_{i}^{0}=d^{-1/2}$.

Drift part for $\alpha_{1}$: We now neglect the martingale term and write for all $0\leq t\leq\overline{T}_{1}-1$

We can study this sequence (see ) and show that

In order for $\alpha^{\overline{T}_{1}}_{1}\approx 1$, we need to take $\eta_{1}\kappa\mu_{D}(\sigma)\overline{T}_{1}\gtrsim(\alpha_{1}^{0})^{-(D-2)}=d^{D/2-1}$.

Drift part for $\alpha_{2}$: Again, by neglecting the martingale contribution and for $0\leq t\leq\overline{T}_{1}-1$,

We can show that this sequence is bounded by

where we used Eq. (14) in the last inequality.

We deduce from Eq. (15) that for $\overline{T}_{1}$ chosen such that $\alpha^{\overline{T}_{1}}_{1}\approx 1$, then $\ln(\alpha_{2}^{t+1}/\alpha_{2}^{0})\lesssim 1$. Hence, during the dynamics, the weights $\alpha_{2}^{t}$ not aligned with the support of $h_{*}$ remain small, of order $1/\sqrt{d}$, while the weights $\alpha_{1}^{t}$ aligned with the support of $h_{*}$ become of order $1$. From the bounds in (i) and (ii), we need to choose $\eta_{1}$ and $\overline{T}_{1}$ such that $\eta_{1}\kappa\sqrt{\overline{T}_{1}}\lesssim d^{-1/2}$ (martingale part) and $\eta_{1}\kappa\overline{T}_{1}\gtrsim d^{D/2-1}$ (drift part), i.e., we can take

While the above heuristic derivation was useful to get intuitions, the assumption that the weights remain equal (or approximately equal) on and outside the support is not valid. Because of the statistical fluctuations over $\widetilde{\Theta}(d^{D-1})$ steps, different coordinates over different neurons will grow to be order 1 on the support at a stochastic time (with high probability between $d^{D-1}/(C\log(d)^{C})$ and $d^{D-1}C\log(d)^{C}$ for some large enough constant $C$). To prevent these coordinates to continue growing (because we neglected the interaction term in the dynamics, which could otherwise prevent this growth), we introduce the projection step

where ${\mathsf{P}}^{t+1}{\mathsf{P}}_{\infty}$ is the projection step defined in Eq. (10), and we use the spherical gradient defined in Eq. (9). Note that because of the choice $\Delta>r$ and the definition of the set $S_{j,t}$ on which we do the projection on the sphere, ${\mathsf{P}}^{t+1}$ and ${\mathsf{P}}_{\infty}$ commute.

Thanks to the spherical gradient, we can show that the projection steps only have a negligible impact on the dynamics (similarly to the analysis in ). By carefully arranging these additional terms, we can essentially recover the drift plus martingale analysis presented heuristically above.

B.4 Going beyond sparsity

In the $P=O_{d}(1)$ regime the complexity scaling in $d$ is dominated by the ‘hard’ part of learning the low-dimensional latent space on which the function depends, and the complexity of fitting the function on the support is secondary and only results in constants. This also makes the conjecture fairly general in terms of architecture choices as long as there is enough expressivity to fit the function on the support. One could also consider functions that depend on a finite number of basis elements, without necessarily involving a finite number of coordinates. For instance the full parity $\prod_{i\in[d]}x_{i}$ function is such an example. For SQ algorithms, the class of monomials of degree 0 (more generally $k$) has equivalent complexity to the class of monomials of degree $d$ (more generally degree $d-k$), and the SQ-dimension is symmetrical for these dual cases. However for SGD learning on regular nets, this is not exactly the case. It is true that the full parity can be learned by regular nets under a specific setting; provides a regular 2-layer neural net that can learn the full parity if the weight measure of the first layer at initialization is i.i.d. Rademacher(1/2) and the activation is a ReLU. A constant number of step can also be sufficient in such cases, as for the 0-degree monomial. It is however conjectured that this is not achievable with a polynomial number of steps for weights that have a Gaussian initialization. Thus, for isotropic layers, it is possible that the full parity is not polytime learnable. This means that the generalized notion of leap to non-coordinate sparse may depend on more specific choices of the parameters. Further, in the non-isotropic case where the full parity is efficiently learnable, one may define the leap with basis sets that can either grow from the 0-monomial or descend from the full-monomial, with the mirror symmetry as for SQ algorithms.

Another notion to factor in when considering non-coordinate sparse function is the fitting of the function once the support is learned. First of all, there may be a non-polynomial number of coefficients to handle, although one can probably cover enough interesting cases with functions that are well-approximated by polynomially many coefficients . Further, there is the fitting of the function by the neural net that may now turn non-trivial. Consider even a function with few basis elements, $h_{*}({\bm{z}})=\sum_{i=1}^{P}ix_{i}+\prod_{i=1}^{P}x_{i}^{g(i)}$, where $g:[P]\to\{0,1\}$ is an arbitrary, but known function, and $P\gg 1$ is large. SGD on a regular neural network would first pick up the $P$ coordinates in the support and then learn the monomial $\prod_{i=1}^{P}z_{i}^{g(i)}$ based on that support. The latter part may not be trivial for SGD on a regular net, while it would require 0 queries for an SQ algorithm (once the linear part is learned, the permutation is identified and the coefficients in front of each variable would allow us to calculate $g(i)$). Thus the complexity of learning the second monomial on the detected support set is likely to factor in for such cases, and this is likely going to depend more on the model hyperparameters and architecture choice. In less contrived cases, the naive generalization of the leap applied verbatim to non-constant $P$ remains likely relevant.

B.5 Lower-bounds: beyond noisy GD

Note that the CSQ and noisy-GD models do not exactly match the SGD learning model; we do prove in this paper that the drift of the population gradient dominates the dynamic on the considered horizon, but the CSQ model also has noise added to the query outputs. It is nonetheless interesting that the regularity of the network model drives us to an achievability result that matches that of CSQ lower-bounds. Since it is known how to go beyond the CSQ/SQ lower-bounds with non-regular networks , e.g., learning dense parities by emulating matrix inversions with irregular networks, our results raise an intriguing question: may the model “regularity” act comparably to a CSQ constraint? We leave this to future work.

Appendix C Proof of Theorem 1: alignment with a single monomial

In this appendix, we prove the alignment of the first layer’s weights with the support of one monomial. The proof will follow from a similar proof strategy as in , namely decomposing the dynamics into drift and martingale terms. However, it will differ in a key aspect: while considers a single-index model, we will need to track for each neuron $P$ parameters (the first $P$ coordinates of ${\bm{w}}_{j}$) and show that the $d-P$ other parameters remain well behaved along their whole trajectories, which requires a tighter control of the different contributions to the dynamics.

Recall that we denote by $K$ a constant that only depends on $\sigma$ (Assumption 1) and the sub-Gaussianity of the label noise $\varepsilon$. Throughout the proofs, we will write $C,c>0$ for generic constants that only depend on $D$ and $K$. The values of these constants are allowed to change from line to line or within the same line.

In the proof, we will consider $0<r\leq\Delta\leq 1$ to be small enough constants that can depend on $D$ and $K$, but are independent of $d$. We will track the dependency in $r,\Delta$ when necessary, and otherwise use that they are bounded by $1$ (in particular, the constants $c,C$ in the proof will be independent of $r,\Delta$). These constants $r,\Delta$ will be fixed in Theorem 1.

We will show that we can take initialization scale $\kappa$ of second layer weights $a^{0}$ and step size $\eta$ such that the dynamics of the first layer training can be approximated by a correlation dynamics, with no interactions between the neurons, so that we can analyze each neuron independently. We consider below an arbitrary neuron $(a_{j},b_{j},{\bm{w}}_{j})$ for $j\in[M]$. In the case that $a_{j}^{0}\mu_{D}(\sigma(\cdot+b_{j}))(w_{j,1})^{k_{1}}\cdots(w_{j,P}^{0})^{k_{P}}>0$ we prove that the event claimed in Theorem 1.(b) and (c) holds with probability at least $1-d^{-C_{*}}$ for neuron $j$. Theorem 1.(a) will follow from a similar analysis. The result for all neurons follows by a union bound.

We further consider $|b_{j}^{0}|\leq\rho\leq\Delta$ small enough such that $1/2\leq|\mu_{k}(\sigma(\cdot+b_{j}))|/|\mu_{k}(\sigma)|\leq 3/2$ for $k=0,\ldots,D+2$ (see comments below Lemma 2). Hence, the biases will not impact the training of the first layer weights and for the simplicity, we will fix $b_{j}=0$ in the proof.

Without loss of generality, we assume that all of the first-layer coordinates of neuron $j$ have positive sign at initialization $w_{j,1}^{0}=\ldots=w_{j,d}^{0}=1/\sqrt{d}$ (and therefore $a^{0}\mu_{D}(\sigma)>0$ by our choice of $(a_{j}^{0},{\bm{w}}_{j}^{0})$). To see why, define $s_{0}=\prod_{i\in[P]}\text{sign}(w_{j,i}^{0})^{k_{i}}$ and consider instead initializing the network at $\breve{{\bm{\Theta}}}^{0}=(\breve{{\bm{a}}}^{0},\breve{{\bm{W}}}^{0})$ where $\breve{{\bm{a}}}^{0}=s_{0}{\bm{a}}^{0}$ and $\breve{{\bm{w}}}_{j^{\prime}}^{0}={\bm{w}}_{j^{\prime}}\odot\text{sign}({\bm{w}}_{j}^{0})$ for all $j^{\prime}$. Then consider training the network with samples $(\breve{y}^{t},\breve{{\bm{x}}}^{t})$ where $\breve{{\bm{x}}}^{t}={\bm{x}}^{t}\odot\text{sign}({\bm{w}}_{j}^{0})$ and $\breve{y}^{t}=f_{*}(\breve{{\bm{x}}}^{t})+\varepsilon_{t}$. The distribution of data $(\breve{y}^{t},\breve{x}^{t})$ is the same as that of $(y^{t},{\bm{x}}^{t})$, and the the training dynamics $\breve{{\bm{\Theta}}}$ match those of ${\bm{\Theta}}$ up to sign flips, and $\breve{{\bm{w}}}_{j}^{0}=[1/\sqrt{d},\ldots,1/\sqrt{d}]$.

with $S_{t+1}=S_{t}\setminus\{i\in[d]:|\widetilde{w}_{i}^{t+1}|\geq r\}$.

Let us introduce the following stopping times on the dynamics:

Note that \{\tau=t\}\in{\mathcal{F}}_{t}:=\sigma\big{(}{\bm{\Theta}}^{0},\{{\bm{x}}^{s},y^{s}\}_{s\leq t}\big{)} for $\tau\in\{\tau^{+},\tau^{-},\tau^{0}\}$ and $\sigma({\bm{w}}^{t+1}),\sigma(S_{t+1})\subseteq{\mathcal{F}}_{t}$. For $t\leq\tau^{+}$ and $r\geq 3/(2\sqrt{d})$, we have $\{P+1,\ldots,d\}\subseteq S_{t}$, and $\|{\bm{w}}^{t}\|_{2}\leq\|{\bm{w}}_{1:P}^{t}\|_{2}+\|{\mathcal{S}}_{t}({\bm{w}}^{t})\|_{2}\leq\sqrt{P}\Delta+1$. We further define for all $i\in[d]$,

where $C_{0}$ is a constant that will be chosen large enough. In particular, at time $\tau_{i}^{r}+1$, the $i$-th coordinate is removed from the set on which we do the projection, i.e., $\{i\}\subseteq S_{\tau_{i}^{r}}\setminus S_{\tau_{i}^{r}+1}$. We will show in the proof that $\tau^{+}\wedge\tau^{-}\wedge\tau^{0}>\overline{T}_{1}$ with high probability.

By concentration of polynomials of Gaussian variables, we have:

Assume that $\Delta\leq 1$. Then for any $C_{*}>0$, there exists $C_{0}$ large enough that only depends on $C_{*},D$ and $K$, such that for $d\geq 2$,

For $t\leq\tau^{+}$, we must have $\|{\bm{w}}^{t}\|_{2}\leq\sqrt{P}+1$. Using the bounds (53) in Lemma 5 and a union bound, there exists a constant $C_{0}$ such that

Note that for $t<\tau_{i}^{\Delta}\wedge\tau_{0}$, we have $|\widetilde{w}_{i}^{t+1}|=|w_{i}^{t}+\eta_{1}\widetilde{v}_{i}^{t}|\leq\Delta$, and therefore $\overline{w}_{i}^{t+1}=\widetilde{w}_{i}^{t+1}$. Let us introduce the truncated spherical gradient ${\bm{g}}^{t}$ defined by

where $\gamma_{i}^{t}\in$ is a multiplicative factor that models the projection step $\overline{{\bm{w}}}^{t+1}={\mathsf{P}}_{\infty}\widetilde{{\bm{w}}}^{t+1}$,

It is easy to check that $\sigma({\bm{g}}^{t})\subseteq{\mathcal{F}}_{t}$ and $\overline{{\bm{w}}}^{t+1}={\mathsf{P}}_{\infty}({\bm{w}}^{t}+\eta_{1}\widetilde{{\bm{v}}}^{t})={\bm{w}}^{t}+\eta_{1}{\bm{g}}^{t}$ for all $t\geq 0$. With these notations, our dynamics are now simply given by

For $i>P$ and $t<\tau^{+}\wedge\tau_{0}$,

We recall the following useful identities (where $G\sim{\sf N}(0,1)$)

In particular, by integration by parts, we have

which gives Eq. (19) by using $\|{\mathcal{S}}_{t}({\bm{w}}^{t})\|_{2}^{2}=1$. Eqs. (20) and (21) are obtained similarly. ∎

From Assumption 1, we can choose $\Delta$ small enough and depending only on $K$ and $D$ such that that for all $u,v\in[-P\Delta,P\Delta]$ and $0\leq k\leq D$

We further assume that $\Delta$ is chosen small enough such that $\Delta^{2}\leq 1/(2D)$ and $\Delta^{2}\leq 1/(4K^{2})$. With this choice of $\Delta$, there exist constants $C,c>0$ that only depend on $D,K$ such that for all $t<\tau_{i}^{\Delta}\wedge\tau^{+}\wedge\tau^{-}$, if $i\in[P]$,

C.2 Bounding the different contributions to the dynamics

The following lemma tracks the contribution of the projection on the sphere ${\bm{w}}^{t+1}={\mathsf{P}}^{t+1}\overline{{\bm{w}}}^{t+1}$:

Assume that $C_{0}\log(d)^{C_{0}}/\sqrt{d}\leq r\leq\Delta/2\leq 1/(\sqrt{8P})$, $|a^{0}|\leq 1$ and $\eta_{1}\leq 1/d$. Then there exist constants $C,C^{\prime}>0$ (that only depend on $D,K$ and $C_{0}$) such that for $d\geq C^{\prime}$ and all $t<\tau^{0}\wedge\tau^{+}$, if $S_{t+1}=S_{t}$,

First consider the case $S_{t+1}=S_{t}$. We have ${\mathcal{S}}_{t+1}(\overline{{\bm{w}}}^{t+1})={\mathcal{S}}_{t}({\bm{w}}^{t})+\eta_{1}{\mathcal{S}}_{t}({\bm{g}}^{t})$. Note that on $i\in S_{t}$, we have $t<\tau_{i}^{\Delta}$ and therefore $\gamma_{i}^{t}({\bm{w}}^{t})=1$ and ${\mathcal{S}}_{t}({\bm{g}}^{t})={\mathcal{S}}_{t}(\widetilde{{\bm{v}}}^{t})$. We therefore have

where we used that $\|{\mathcal{S}}_{t}({\bm{w}}^{t})\|_{2}^{2}=1$ and $\langle{\mathcal{S}}_{t}({\bm{w}}^{t}),{\mathcal{S}}_{t}(\widetilde{{\bm{v}}}^{t})\rangle=0$ by definition of the spherical gradient. Furthermore, $\eta_{1}^{2}\|{\mathcal{S}}_{t}(\widetilde{{\bm{v}}}^{t})\|_{2}^{2}\leq\eta_{1}^{2}\|{\bm{g}}^{t}\|_{2}^{2}\leq d^{-2}\|{\bm{g}}^{t}\|_{2}^{2}\leq d^{-1}\|{\bm{g}}^{t}\|_{\infty}^{2}\leq Cd^{-1}\log(d)^{C}\leq 1/4$ for $t<\tau_{0}$. Therefore, there exists a constant $C>0$ such that bound (25) holds.

In the case $S_{t+1}\neq S_{t}$, we have $|S_{t}\setminus S_{t+1}|\leq P$ for $t<\tau^{+}$ and the coordinates that are removed at time $t+1$ satisfy $w_{i}^{t}+\eta_{1}g_{i}^{t}\leq r+C_{0}d^{-1/2}\log(d)^{C_{0}}\leq 2r$. Hence

We can then use Eq. (27) and that $\eta_{1}^{2}\|{\mathcal{S}}_{t}(\widetilde{{\bm{v}}}^{t})\|_{2}^{2}\leq Pr^{2}$ to derive Eq. (26). ∎

Let us decompose the different contributions to the dynamics. We define $\bm{m}^{t}={\bm{g}}^{t}-\overline{{\bm{g}}}^{t}$ the martingale updates. Let us bound the change of a coordinate after one update. For $t<\tau^{0}\wedge\tau^{+}\wedge\tau^{-}$, if $S_{t+1}=S_{t}$, then by Eq. (25), we have for $i\in S_{t+1}$

(Note that for $t<\tau^{-}$, we have $\text{sign}(w_{i}^{t+1})=\text{sign}(w_{i}^{t})=1$.) For $t<\tau^{0}$, we have $\|{\bm{g}}^{t}\|_{\infty}\leq C|a^{0}|\log(d)$ and $|w_{i}^{t}|/|w_{i}^{t+1}|\leq C$ because $\eta_{1}\leq 1/d$. Hence, we can rearrange Eqs. (28) and obtain

On the other hand, if $S_{t+1}\neq S_{t}$, by Eq. (26), we have for $i\in S_{t+1}$,

Rearranging these equations, we obtain for $t<\tau^{0}\wedge\tau^{-}$ and $S_{t+1}\neq S_{t}$ (using $|Cr^{2}|\leq 1/2$),

On the other hand, if $i\not\in S_{t+1}$, then

Define $X_{t}=\#\{i\in[P]:\tau_{i}^{r}<t\}$, and

Note $\sigma(\underline{p}^{t}),\sigma(\overline{p}^{t})\in{\mathcal{F}}_{t-1}$, so that $\underline{p}^{t}\bm{m}^{t}$ and $\overline{p}^{t}{\bm{w}}^{t}$ are still martingale updates. By induction on Eqs (29) and (30), we deduce that for $t<\tau^{0}\wedge\tau^{+}\wedge\tau^{-}$,

Let us introduce the following quantities:

The term $D^{t,t^{\prime}}$, which is the sum of population gradients, plays the role of a drift term, while the term $M^{t,t^{\prime}}$ is a martingale and corresponds to the comparison between the stochastic and the population gradients.

We can choose $C^{\prime}$ a constant large enough, depending only on $K,D$, such that for

In particular, this implies that for any $t<T\wedge\tau^{+}\wedge\tau^{-}$ and $r$ constant sufficiently small,

Similarly, we can choose $C^{\prime}$ a constant large enough, depending only on $K,D$, such that for

Hence for $\eta_{1}$ and $T$ satisfying Eqs (32) and (34), we get the following bounds on the trajectory for $t<T\wedge\tau^{+}\wedge\tau^{-}\wedge\tau^{0}$: for $i\geq P+1$ or $i\in[P]$, $t\leq\tau_{i}^{r}$,

while for $i\in[P]$ and $t>\tau_{i}^{r}$,

We prove the following bounds on the martingale part:

Assume $\Delta\leq 1$ and $r$ are chosen as in Lemma 3. Fix $T\leq d^{D}$ and $C_{*}>0$. There exists a constant $C$ that only depends on $D$, $K$, and $C_{*}$, such that if we choose

then with probability at least $1-d^{-C_{*}}$, we have

We will show the theorem for $\max_{0<t\leq T\wedge\tau^{+}}\underline{M}_{i}^{0,t}$. The result for $\max_{t<t^{\prime}\leq T\wedge\tau^{+}}\underline{M}_{i}^{t,t^{\prime}}$ will follow from an union bound on all $t\leq T$, which are also martingales (the proofs for $\overline{M}_{i}^{t,t^{\prime}}$ and $M_{i}^{t,t^{\prime}}$ will follow by the same argument).

Denote $\underline{M}^{t}_{i}:=\underline{M}^{0,t}_{i}$. We will use a truncation argument. For some $\widetilde{C}$, define for all $t\geq 1$ and $i\in[d]$,

For $t\leq\tau^{+}$, we have $\|{\bm{w}}\|_{2}\leq\sqrt{P}\Delta+1$ and we can use Lemma 5 to choose $\widetilde{C}$ that only depends on $D,K,C_{*}$ such that

and for all $t\leq T\wedge\tau^{+}$ and $i\in[d]$

Hence with probability at least $1-d^{-C_{*}}/2$,

Let us now apply Doob’s maximal inequality on $U_{i}^{t}$: the increments are bounded by $2|a^{0}|\widetilde{C}\log(d)^{\widetilde{C}}$, hence we have

Choosing $\varepsilon=1/(2\sqrt{d}\log(d))$ and $\eta_{1}$ as in Eq. (37), as well as a union bound, yields the result. ∎

C.3 Proof of Theorem 1

We consider the dynamics (18) up to time $T\wedge\overline{\tau}$ where $\overline{\tau}:=\tau^{0}\wedge\tau^{+}\wedge\tau^{-}$. We assume $T$ and $\eta_{1}$ satisfy conditions (32), (34) and (37). In particular, with probability at least $1-d^{-C_{*}}$, the dynamics of $w_{i}^{t}$ satisfies the bounds in Eqs (35) and (36), with $M_{i}^{t,t^{\prime}},\overline{M}_{i}^{t,t^{\prime}},\underline{M}_{i}^{t,t^{\prime}}$ satisfying the bounds (38). In the rest of the proof, we show that on this high probability event, we can choose $\eta_{1}$ and $T$ such that $T<\overline{\tau}$ and Theorem 1.(a), (b) and (c) are satisfied.

Step 1: Controlling the coordinates $i\in[P]$ at the end of the dynamics.

Let us first show that as soon as $t>\tau^{\Delta}_{i}$, then $w_{i}^{t}$ stays close to $\Delta$. Note that we can choose $C>0$ constant large enough independent of $d$, such that if $w_{i}^{t}\leq\Delta-|a^{0}|C\eta_{1}\log(d)^{C}$, then by (24)

Hence, for any $\overline{\tau}>t>\tau_{i}^{\Delta}$, consider $t^{\prime}=\sup\{t^{\prime}\leq t:w_{i}^{t^{\prime}}\geq\Delta-|a^{0}|C\eta_{1}\log(d)^{C}\}$ (in particular, $t^{\prime}\geq\tau^{\Delta}_{i}+1$). From Eq. (36) and by Lemma 4, we have

where we used that $w_{i}^{s}<\Delta-|a^{0}|C\eta_{1}\log(d)^{C}$ for $t^{\prime}+1\leq s<t$ and therefore by Eq. (39), we have $D_{i}^{t^{\prime}+1,t}\geq 0$. We deduce that

Similarly, we show that for any $t\geq\tau^{r}_{i}+1$, we have $w_{i}^{t}\geq r/2$. Indeed, for any $\tau^{r}_{i}+1\leq t\leq\tau_{i}^{\Delta}\wedge\overline{\tau}\wedge T$, we have

where we used that $w_{i}^{\tau_{i}^{r}+1}\geq r$ by definition of $\tau_{i}^{r}$, and $\overline{g}_{i}^{t}\geq 0$ for all $s\leq\tau_{i}^{\Delta}\wedge\overline{\tau}\wedge T$. We deduce that

Step 2: Bounding the growth of $w_{i}^{t}$ for $i\in[P]$.

Define $\alpha_{t}=\min\{w_{i}^{t}:i\in S_{t}\cap[P]\}$ (i.e., the minimum of $w_{i}^{t}$ that have $\tau_{i}^{r}\geq t$). Note that $w_{i}^{t}\geq r/2$ for $i\not\in S_{t}$ by Eq. (41) and therefore $w_{i}^{t}\geq\alpha_{t}/2$. By Eq. (33) and Lemma 2, we have for $s\leq\tau_{i}^{r}\wedge\overline{\tau}\wedge T$,

Combining this lower bound with Eq. (35) and the bound on the martingale in Lemma 4, we get that for all $t\leq\tau_{i}^{r}\wedge\overline{\tau}\wedge T$

Similarly, consider $\alpha_{t}=\min\{w_{i}^{t}:i\in[P],\tau^{\Delta}_{i}\geq t\}$ for all $\tau^{r}+1\leq t\leq\tau^{\Delta}\wedge\overline{\tau}$. By Eq. (40), we have $w_{i}^{t}\geq\Delta/2$ for $t>\tau_{i}^{\Delta}$. Hence by Eq. (36), we get

and we deduce that if $\tau^{\Delta}<\overline{\tau}$ then

Combining the above bounds, we deduce that

On the other hand, consider $\beta_{t}=\max\{w_{i}^{t}:i\in[P]\}$ and let us lower bound the time $\overline{t}=\inf\{t:\beta_{t}\geq 2/\sqrt{d}\}$. For $t\leq\overline{\tau}\wedge\overline{t}$, we have $S_{t}=[d]$ and $\frac{1}{\underline{p}^{t}}\leq 1+\frac{C}{\log(d)}$. Hence by Eq. (35),

and therefore by Lemma 6, we get for $D>2$,

Step 3: Bounding the coordinates $P+1\leq i\leq d$.

From Eq. (35) and Lemma 2, we have for all $i\geq P+1$ and $t<\overline{\tau}\wedge T$,

Consider $j\in[P]$ such that $\tau^{r}_{j}=\tau^{r}$. Then by Eq. (35), we have, for any $t<(\tau^{r}+1)\wedge\overline{\tau}\wedge T$, that

Using that $w_{j}^{t}\leq Cr$ for $t\leq\tau_{j}^{r}$ and $w_{i}^{s}\leq 3/(2\sqrt{d})$ for $s<\overline{\tau}$ in Eq. (48), we get that, for any $t<\overline{\tau}\wedge T$,

We deduce that for all $P+1\leq i\leq d$ and $t<\overline{\tau}$ and $t+1<\tau^{0}$

and therefore taking $r$ sufficiently small, $\tau^{+}\wedge\tau^{-}\geq\tau^{0}$.

Choose $\overline{T}_{1}:=T$ and $\eta_{1}$ that satisfy Eqs. (32), (34), (37) and (43). We have with probability at least $1-Cd^{-C_{*}}$ by Lemma 1 and Lemma 4 that $\tau^{0}>\tau^{+}\wedge\tau^{-}\wedge T$. And Eqs. (50) and (42) imply that $\tau^{+}\wedge\tau^{-}\wedge\tau^{0}>\overline{T}_{1}$. Furthermore, by Eq. (43), we have $\tau^{\Delta}<\overline{T}_{1}$ which implies Theorem 1.(b) by Eq. (40). Theorem 1.(c) follows from Eq. (50).

Step 5: Upper bound for all neurons with early stopping.

Theorem 1.(a) follows from Eqs. (45) and (47) for neurons with initialization satisfying

For neurons that do not satisfy this condition, the analysis in Section C.2 still holds and we get bounds on the dynamics similar to the ones in Eq. (35), with the difference that $a^{0}\mu_{D}(\sigma)<0$, and therefore the drift has a negative contribution to the dynamics, and $\tau^{+},\tau^{-}$ are now defined on all coordinates instead of only $i=P+1,\ldots,d$.

We can upper bound the drift contribution using that $\beta_{t}=\max_{i\in[P]}w_{i}^{t}$ satisfy for $t<\overline{\tau}$

and therefore, taking the same bounds (44) and (46), we get for $t\leq\overline{\tau}\wedge\overline{T}_{1}/(C\log(d)^{C})$ for $C$ a constant sufficiently large that

Furthermore, denoting $\alpha_{t}=\min_{i\in[P]}w_{i}^{t}$, we have

Using the analysis of Lemma 6, we get that the drift has the same upper bound as $\beta_{t}$ in Eq. (52),

for $t\leq\overline{\tau}\wedge\overline{T}_{1}/(C\log(d)^{C})$ and therefore

The bound on coordinates $i>P$ follows similarly to step 3. In particular, we deduce that for $t<\overline{\tau}\wedge\overline{T}_{1}/(C\log(d)^{C})$ and $t+1<\tau^{0}$, we have for all $i\in[d]$