What Can ResNet Learn Efficiently, Going Beyond Kernels?

Introduction

Neural network learning has become a key practical machine learning approach and has achieved remarkable success in a wide range of real-world domains, such as computer vision, speech recognition, and game playing . On the other hand, from a theoretical standpoint, it is less understood that how large-scale, non-convex, non-smooth neural networks can be optimized efficiently over the training data and generalize to the test data with relatively few training examples.

There has been a sequence of research trying to address this question, showing that under certain conditions neural networks can be learned efficiently . These provable guarantees typically come with strong assumptions and the proofs heavily rely on them. One common assumption from them is on the input distribution, usually being random Gaussian or sufficiently close to Gaussian. While providing great insights to the optimization side of neural networks, it is not clear whether these works emphasizing on Gaussian inputs can coincide with the neural network learning process in practice. Indeed, in nearly all real world data where deep learning is applied to, the input distributions are not close to Gaussians; even worse, there may be no simple model to capture such distributions.

The difficulty of modeling real-world distributions brings us back to the traditional PAC-learning language which is distribution-free. In this language, one of the most popular, provable learning methods is the kernel methods, defined with respect to kernel functions $K(x,x^{\prime})$ over pairs of data $(x,x^{\prime})$. The optimization task associated with kernel methods is convex, hence the convergence rate and the generalization error bound are well-established in theory.

Recently, there is a line of work studying the convergence of neural networks in the PAC-learning language, especially for over-parameterized neural networks , putting neural network theory back to the distribution-free setting. Most of these works rely on the so-called Neural Tangent Kernel (NTK) technique , by relating the training process of sufficiently over-parameterized (or even infinite-width) neural networks to the learning process over a kernel whose features are defined by the randomly initialized weights of the neural network. In other words, on the same training data set, these works prove that neural networks can efficiently learn a concept class with as good generalization as kernels, but nothing more is known.Technically speaking, the three-layer learning theorem of is beyond NTK, because the learned weights across different layers interact with each other, while in NTK the learned weights of each layer only interact with random weights of other layers. However, there exist other kernels— such as recursive kernels — that can more or less efficiently learn the same concept class proposed in .

In contrast, in many practical tasks, neural networks give much better generalization error compared to kernels, although both methods can achieve zero training error. For example, ResNet achieves 96% test accuracy on the CIFAR-10 data set, but NTKs achieve 77% and random feature kernels achieve 85% . This gap becomes larger on more complicated data sets.

To separate the generalization power of neural networks from kernel methods, the recent work tries to identify conditions where the solutions found by neural networks provably generalize better than kernels. This approach assumes that the optimization converges to minimal complexity solutions (i.e. the ones minimizing the value of the regularizer, usually the sum of squared Frobenius norms of weight matrices) of the training objective. However, for most practical applications, it is unclear how, when training neural networks, minimal complexity solutions can be found efficiently by local search algorithms such as stochastic gradient descent. In fact, it is not true even for rather simple problems (see Figure 1).Consider the class of degree-6 polynomials over 6 coordinates of the $d$-dimensional input. There exist two-layer networks with F-norm $O(\sqrt{d})$ implementing this function (thus have near-zero training and testing error). By Rademacher complexity, $O(d)$ samples suffice to learn if we are able to find a minimal complexity solution. Unfortunately, due to the non-convexity of the optimization landscape, two-layer networks can not be trained to match this F-norm even with $O(d^{2})$ samples, see Figure 1. Towards this end, the following fundamental question is largely unsolved:

Can neural networks efficiently and distribution-freely learn a concept class,

with better generalization than kernel methods?

In this paper, we give arguably the first positive answer to this question for neural networks with ReLU activations. We show without any distributional assumption, a three-layer residual network (ResNet) can (improperly) learn a concept class that includes three-layer ResNets of smaller size and smooth activations. This learning process can be efficiently done by stochastic gradient descent (SGD), and the generalization error is also small if polynomially many training examples are given.

More importantly, we give a provable separation between the generalization error obtained by neural networks and arbitrary kernel methods. For some $\delta\in(0,1)$, with $N=O(\delta^{-2})$ training samples, we prove that neural networks can efficiently achieve generalization error $\delta$ for this concept class over any distribution; in contrast, there exists rather simple distributions such that any kernel method (including NTK, recursive kernel, etc) cannot have generalization error better than $\sqrt{\delta}$ for this class. To the best of our knowledge, this is the first work that gives provable, efficiently achievable separation between neural networks with ReLU activations and kernels in the distribution-free setting. In the end, we also prove a computation complexity advantage of neural networks with respect to linear regression over arbitrary feature mappings as well.

Roadmap. We present detailed overview of our positive and negative results in Section 2 and 3. Then, we introduce notations in Section 4, formally define our concept class in Section 5, and give proof overviews in Section 6 and 7.

Positive Result: The Learnability of Three-Layer ResNet

We wish to learn a concept class given by target functions that can be written as

To illustrate our result, we first assume for simplicity that $y=\mathcal{H}\left(x\right)$ for some $\mathcal{H}$ of the form (2.2) (so the optimal target has zero regression error). Our main theorem can be sketched as follows.

Let $C_{\mathcal{F}}$ and $C_{\mathcal{G}}$ respectively be the individual “complexity” of $\mathcal{F}$ and $\mathcal{G}$, which at a high level, capture the size and smoothness of $\mathcal{F}$ and $\mathcal{G}$. This complexity notion shall be formally introduced in Section 4, and is used by prior works such as .

For any distribution over $x$, for every \delta\in\big{(}(\alpha C_{\mathcal{G}})^{4},1\big{)}, with probability at least $0.99$, SGD efficiently learns a network $\mathsf{out}(x)$ in the form (2.1) satisfying

The running time of SGD is polynomial in $\poly(C_{\mathcal{G}},C_{\mathcal{F}},\alpha^{-1})$.

In other words, ResNet is capable of achieving population risk $\alpha^{4}$, or equivalently learning the output $\mathcal{H}(x)$ up to $\alpha^{2}$ error. In our full theorem, we also allow label $y$ to be generated from $\mathcal{H}(x)$ with error, thus our result also holds in the agnostic learning framework.

we need $\widetilde{O}(d)$ samples to efficiently learn $\mathcal{H}=\mathcal{F}+\alpha\mathcal{G}(\mathcal{F})$ up to accuracy $\widetilde{O}(\alpha^{2})$.

In contrast, the complexity of $\mathcal{G}(\mathcal{F})$ is $\widetilde{O}((\sqrt{d})^{10})$, so

prior works need $\widetilde{\Omega}(d^{10})$ samples to efficiently learn $\mathcal{H}$ up to any accuracy $o(\alpha)$,

Inductive Bias. Our network is over-parameterized, thus intuitively in the example above, with only $O(d)$ training examples, the learner network could over-fit to the training data since it has to decide from a set of $d^{10}$ many possible coefficients to learn the degree 10 polynomial $\mathcal{G}$. This is indeed the case if we learn the target function using kernels, or possibly even learn it with a two-layer network. However, three-layer ResNet posts a completely different inductive bias, and manages to avoid over-fitting to $\mathcal{G}(\mathcal{F})$ with the help from $\mathcal{F}$.

Implicit Hierarchical Learning using Forward Feature Learning. Since $\mathcal{H}(x)=\mathcal{F}(x)+\alpha\mathcal{G}\left(\mathcal{F}(x)\right)$, if we only learn $\mathcal{F}$ but not $\mathcal{G}\left(\mathcal{F}\right)$, we will have regression error $\approx(\alpha C_{\mathcal{G}})^{2}$. Thus, to get to regression error $(\alpha C_{\mathcal{G}})^{4}$, Theorem 1 shows that ResNet is also capable of learning $\mathcal{G}\left(\mathcal{F}\right)$ up to some good accuracy with relatively few training examples. This is also observed in practice, where with this number of training examples, three-layer fully-connected networks and kernel methods can indeed fail to learn $\mathcal{G}\left(\mathcal{F}\right)$ up to any non-trivial accuracy, see Figure 2.

Intuitively, there is a hierarchy of the learning process: we would like to first learn $\mathcal{F}$, and then we could learn $\mathcal{G}(\mathcal{F})$ much easier with the help of $\mathcal{F}$ using the residual link. In our learner network (2.1), the first hidden layer serves to learn $\mathcal{F}$ and the second hidden layer serves to learn $\mathcal{G}$ with the help of $\mathcal{F}$, which reduces the sample complexity. However, the important message is that $\mathcal{F}$ and $\mathcal{G}$ are not given as separate data to the network, rather the learning algorithm has to disentangle them from the “combined” function $\mathcal{H}=\mathcal{F}+\alpha\mathcal{G}(\mathcal{F})$ automatically during the training process. Moreover, since we train both layers simultaneously, the learning algorithm also has to distribute the learning task of $\mathcal{F}$ and $\mathcal{G}$ onto different layers automatically. We call this process “forward feature learning”:

We point out forward feature learning is different from layer-wise training. For instance, our result cannot be obtained by first training the hidden layer close to the input, and then fixing it and training the hidden layer close to the output. Since it could be the case the first layer incurs some $\alpha$ error (since it cannot learn $\mathcal{G}(\mathcal{F})$ directly), then it could be really hard, or perhaps impossible, for the second layer to fix it only using inputs of the form $\mathcal{F}(x)\pm\alpha$. In other words, it is crucial that the two hidden layers are simultaneously trained. This does not mean that the error of the first layer can be reduced by its own, since it is still possible for the first layer to learn $\mathcal{F}+\alpha\mathcal{R}(x)\pm\alpha^{2}$ and the second layer to learn $\mathcal{G}(\mathcal{F})(x)-\mathcal{R}(x)\pm\alpha$, for an arbitrary (bounded) function $\mathcal{R}$.

A follow-up work. In a follow-up work , this theory of hierarchical learning is strengthened to further incorporate the backward feature correction step when training deep neural networks. In the language of this paper, when the two layers trained together, given enough samples, the accuracy in the first layer can actually be improved from $\mathcal{F}\pm\alpha$ to arbitrarily close to $\mathcal{F}$ during the training process. As a consequence, the final training and generalization error can be arbitrarily small as well, as opposite to $\alpha^{2}$ (or equivalently population risk $\alpha^{4}$) in this work. The new “backward feature correction” is also critical to extend the hierarchical learning process from $3$ layers to arbitrarily number of layers.

Negative Results

For every constant $k\geq 2$, for every sufficiently large $d\geq 2$, there exist concept classes consisting of functions $\mathcal{H}(x)=\mathcal{F}(x)+\alpha\mathcal{G}\left(\mathcal{F}(x)\right)$ with complexities $C_{\mathcal{F}},C_{\mathcal{G}}$ and $\alpha\in(0,\frac{1}{C_{\mathcal{G}}})$ such that, letting

then there exists simple distributions $\mathcal{D}$ over $(x,\mathcal{H}(x))$ such that, for at least $99\%$ of the functions $\mathcal{H}$ in this concept class, even given N=O\big{(}(N_{\mathsf{res}})^{k/2}\big{)} training samples from $\mathcal{D}$, any function $\mathfrak{K}(x)$ of the form (3.1) has to suffer population risk

Contribution and Intuition. Let us compare this to Theorem 1. While both algorithms are efficient, neural networks (trained by SGD) achieve population risk $\alpha^{3.9}$ using $N_{\mathsf{res}}$ samples for any distribution over $x$, while kernel methods cannot achieve any population risk better than $\alpha^{2}$ for some simple distributions even with $N=(N_{\mathsf{res}})^{k/2}\gg N_{\mathsf{res}}$ samples.It is necessary the negative result of kernel methods is distribution dependent, since for trivial distributions where $x$ is non-zero only on the first constantly many coordinates, both neural networks and kernel methods can learn it with constantly many samples. Our two theorems together gives a provable separation between the generalization error of the solutions found by neural networks and kernel methods, in the efficiently computable regime.

More specifically, recall $C_{\mathcal{F}}$ and $C_{\mathcal{G}}$ only depend on individual complexity of $\mathcal{G},\mathcal{F}$, but not on $\mathcal{G}(\mathcal{F})$. In Theorem 2, we will construct $\mathcal{F}$ as linear functions and $\mathcal{G}$ as degree-$k$ polynomials. This ensures $C_{\mathcal{F}}=O(\sqrt{d})$ and $C_{\mathcal{G}}=O(1)$ for $k$ being constant, but the combined complexity of $\mathcal{G}(\mathcal{F})$ is as high as $\Omega(d^{k/2})$. Since ResNet can perform hierarchical learning, it only needs sample complexity $N_{\mathsf{res}}=O(d/\alpha^{8})$ instead of paying (square of) the combined complexity $\Omega(d^{k})$.

In contrast, a kernel method is not hierarchical: rather than discovering $\mathcal{F}$ first and then learning $\mathcal{G}(\mathcal{F})$ with the guidance of it, kernel method tries to learn everything in one shot. This unavoidably requires the sample complexity to be at least $\Omega(d^{k})$. Intuitively, as the kernel method tries to learn $\mathcal{G}(\mathcal{F})$ from scratch, this means that it has to take into account all $\Omega(d^{k})$ many possible choices of $\mathcal{G}(\mathcal{F})$ (recall that $\mathcal{G}$ is a degree $k$ polynomial over dimension $d$). On the other hand, a kernel method with $N$ samples only has $N$-degrees of freedom (for each output dimension). This means, if $N\ll o(d^{k})$, kernel method simply does not have enough degrees of freedom to distinguish between different $\mathcal{G}(\mathcal{F})$, so has to pay $\Omega(\alpha^{2})$ in population risk. Choosing for instance $\alpha=d^{-0.1}$, we have the desired negative result for all N\leq O\big{(}(N_{\mathsf{res}})^{k/2}\big{)}\ll o(d^{k}).

2 Limitation of Linear Regression Over Feature Mappings

for some regularizer $R(w)$. In this paper, we do not make assumptions about how the weighted are found. Instead, we focus on any linear function over such feature mapping in the form (3.3).

For sufficiently large integers $d,k$, there exist concept classes consisting of functions $\mathcal{H}(x)=\mathcal{F}(x)+\alpha\mathcal{G}\left(\mathcal{F}(x)\right)$ with complexities $C_{\mathcal{F}},C_{\mathcal{G}}$ and $\alpha\in(0,\frac{1}{C_{\mathcal{G}}})$ such that, letting

then for at least $99\%$ of the functions $\mathcal{H}$ in this concept class, even with arbitrary $D=(T_{\mathsf{res}})^{2}$ dimensional feature mapping, any function $\mathfrak{F}(x)$ of the form (3.3) has to suffer population risk

Interpretation. Since any algorithm that optimizes linear functions over $D$-dimensional feature mapping has to run in time $\Omega(D)$, this proves a time complexity separation between neural networks (say, for achieving population risk $\alpha^{3.9}$) and linear regression over feature mappings (for achieving even any population risk better than $\alpha^{2}\gg\alpha^{3.9}$). Usually, such an algorithm also has to suffer from $\Omega(D)$ space complexity. If that happens, we also have a space complexity separation. Our hard instance in proving Theorem 3 is the same as Theorem 2, and the proof is analogous.

Notations

For notation simplicity, throughout this paper “with high probability” (or w.h.p.) means with probability $1-e^{-c\log^{2}m}$ for a sufficiently large constant $c$. We use $\widetilde{O}$ to hide $\polylog(m)$ factors.

where $C^{*}$ is a sufficiently large constant (e.g., $10^{4}$).

Concept Class

We denote by $\mathfrak{C}_{\varepsilon}(\mathcal{F})=\max_{r,i}\{\mathfrak{C}_{\varepsilon}(\mathcal{F}_{r,i})\}$ and $\mathfrak{C}_{\mathfrak{s}}(\mathcal{F})=\max_{r,i}\{\mathfrak{C}_{\mathfrak{s}}(\mathcal{F}_{r,i})\}$. Intuitively, $\mathcal{F}$ and $\mathcal{G}$ are both generated by two-layer neural networks with smooth activation functions $\mathcal{F}_{r,i}$ and $\mathcal{G}_{r,i}$.

Borrowing the agnostic PAC-learning language, our concept class consists of all functions $\mathcal{H}(x)$ in the form of Concept 1 with complexity bounded by tuple $(p_{F},C_{F},p_{G},C_{G})$. Let $\mathsf{OPT}$ be the population risk achieved by the best target function in this concept class. Then, our goal is to learn this concept class with population risk $O(\mathsf{OPT})+\varepsilon$ using sample and time complexity polynomial in $p_{F},C_{F},p_{G},C_{G}$ and $1/\varepsilon$. In the remainder of this paper, to simplify notations, we do not explicitly define this concept class parameterized by $(p_{F},C_{F},p_{G},C_{G})$. Instead, we equivalently state our theorem with respect to any (unknown) fixed target function $\mathcal{H}$ with with population risk $\mathsf{OPT}$:

In the analysis we adopt the following notations. For every $(x,y)\sim\mathcal{D}$, it satisfies $\|\mathcal{F}(x)\|_{2}\leq\mathfrak{B}_{\mathcal{F}}$ and $\|\mathcal{G}(\mathcal{F}(x))\|_{2}\leq\mathfrak{B}_{\mathcal{F}\circ\mathcal{G}}$. We assume $\mathcal{G}(\cdot)$ is $\mathfrak{L}_{\mathcal{G}}$-Lipschitz continuous. It is a simple exercise (see Fact A.3) to verify that $\mathfrak{L}_{\mathcal{G}}\leq\sqrt{k}p_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})$, $\mathfrak{B}_{\mathcal{F}}\leq\sqrt{k}p_{\mathcal{F}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{F})$ and $\mathfrak{B}_{\mathcal{F}\circ\mathcal{G}}\leq\mathfrak{L}_{\mathcal{G}}\mathfrak{B}_{\mathcal{F}}+\sqrt{k}p_{\mathcal{G}}\mathfrak{C}(\mathcal{G})\leq kp_{\mathcal{F}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{F})p_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})$.

Overview of Theorem 1

We consider the vanilla SGD algorithm given in Algorithm 1.Performing SGD with respect to ${\mathbf{W}^{(0)}}+\mathbf{W}$ and ${\mathbf{V}^{(0)}}+\mathbf{V}$ is the same as that with respect to $\mathbf{W}$ and $\mathbf{V}$; we introduce ${\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$ notation for analysis purpose. Note also, one can alternatively consider having a training set and then performing SGD on this training set with multiple passes; similar results can be obtained.

Under Concept 1 or Concept 2, for every \alpha\in\big{(}0,\widetilde{\Theta}(\frac{1}{kp_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})})\big{)} and $\delta\geq\mathsf{OPT}+\widetilde{\Theta}\left(\alpha^{4}(kp_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G}))^{4}(1+\mathfrak{B}_{\mathcal{F}})^{2}\right)$. There exist $M=\poly(\mathfrak{C}_{\alpha}(\mathcal{F}),\mathfrak{C}_{\alpha}(\mathcal{G}),p_{\mathcal{F}},\alpha^{-1})$ satisfying that for every $m\geq M$, with high probability over $\mathbf{A},{\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$, for a wide range of random initialization parameters $\sigma_{w},\sigma_{v}$ (see Table 1), choosing

With high probability, the SGD algorithm satisfies

As a corollary, under Concept 1, we can archive population risk

Our Theorem 1 is almost in the PAC-learning language, except that the final error has an additive $\alpha^{4}$ term that can not be arbitrarily small.

In the analysis, let us define diagonal matrices

which satisfy $\mathsf{out}_{1}(x)=\mathbf{A}D_{\mathbf{W}}({\mathbf{W}^{(0)}}+\mathbf{W})(x,1)$ and $\mathsf{out}(x)=\mathbf{A}D_{\mathbf{V},\mathbf{W}}({\mathbf{V}^{(0)}}+\mathbf{V})(\mathsf{out}_{1}(x),1)$.

The proof of Theorem 1 can be divided into three simple steps with parameter choices in Table 1.

In the first step, we prove that for all weight matrices not very far from random initialization (namely, all $\|\mathbf{W}\|_{2}\leq\tau_{w}$ and $\|\mathbf{V}\|_{2}\leq\tau_{v}$), many good “coupling properties” occur. This includes upper bounds on the number of sign changes (i.e., on $\left\|D_{{\mathbf{W}^{(0)}}}-D_{\mathbf{W}}\right\|_{0}$ and $\left\|D_{{\mathbf{V}^{(0)}},\mathbf{W}}-D_{\mathbf{V},\mathbf{W}}\right\|_{0}$) as well as vanishing properties such as $\mathbf{A}D_{\mathbf{W}}{\mathbf{W}^{(0)}},\mathbf{A}D_{\mathbf{V},\mathbf{W}}{\mathbf{V}^{(0)}}$ being negligible. We prove such properties using techniques from prior works . Details are in Section C.1.

In the second step, we prove the existence of $\mathbf{W}^{\divideontimes},\mathbf{V}^{\divideontimes}$ with $\|\mathbf{W}^{\divideontimes}\|_{F}\leq\frac{\tau_{w}}{10}$ and $\|\mathbf{V}^{\divideontimes}\|_{F}\leq\frac{\tau_{v}}{10}$ satisfying $\mathbf{A}D_{{\mathbf{W}^{(0)}}}\mathbf{W}^{\divideontimes}(x,1)\approx\mathcal{F}(x)$ and $\mathbf{A}D_{{\mathbf{V}^{(0)}},\mathbf{W}}\mathbf{V}^{\divideontimes}(\mathsf{out}_{1}(x),1)\approx\alpha\mathcal{G}\left(\mathsf{out}_{1}(x)\right)$. This existential proof relies on an “indicator to function” lemma from ; for the purpose of this paper we have to revise it to include a trainable bias term (or equivalently, to support vectors of the form $(x,1)$). Combining it with the aforementioned vanishing properties, we derive (details are in Section C.2):

In the third step, consider iteration $t$ of SGD with sample $(x_{t},y_{t})\sim\mathcal{D}$. For simplicity we assume $\mathsf{OPT}=0$ so $y_{t}=\mathcal{H}(x_{t})$. One can carefully write down gradient formula, and plug in (6.3) to derive

Overview of Theorem 2 and 3

Consider the class of target functions $\mathcal{H}(x)=\mathcal{F}(x)+\alpha\mathcal{G}(\mathcal{F}(x))$, where

where $\mathbf{W}^{*}=\sqrt{d}(\mathbf{e}_{i_{1}},\mathbf{e}_{i_{2}},\cdots\mathbf{e}_{i_{k}})$ for $i_{1},i_{2},\dots,i_{k}\in[d_{1}]$ are distinct indices chosen from the first $d_{1}$ coordinates. There are clearly $\binom{d_{1}}{k}$ many target functions in this class.

Intuitively, $\mathbf{e}_{1},\cdots,\mathbf{e}_{d_{1}}$ represent the directions where the signal possibly lies, where usually the inputs would have high variance; and $\mathbf{e}_{d_{1}+1},\dots,\mathbf{e}_{d}$ represent the directions that can be view as “background noise”, where the distribution can be arbitrary. For example when $d_{1}\leq d/2$, such distribution $\mathcal{D}$ can be very different from Gaussian distribution or uniform distribution over Boolean cube, yet kernel methods still suffer from high population risk when learning over these distributions comparing to using neural networks.

We first state the population risk for the three-layer ResNet to learn this concept class: Our Theorem 1 implies the following complexity on learning this concept class (after verifying that $\mathfrak{C}_{\mathfrak{s}}(\mathcal{F})=O(\sqrt{d})$, $p_{\mathcal{F}}=1$, $\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})=2^{O(k)}$, $p_{\mathcal{G}}=2^{k}$, see Section D.4).

For every $d\geq d_{1}\geq k\geq 2$, for every \alpha\in\big{(}0,\frac{1}{\widetilde{\Theta}(2^{O(k)})}\big{)}, there exist $M=\poly(d,2^{k},\alpha^{-1})$ satisfying that for every $m\geq M$, for every target functions $\mathcal{H}(x)$ in the class (7.1), with probability at least 0.99 over $\mathbf{A},{\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$ and $\mathcal{X}$, given labels $y^{(n)}=\mathcal{H}(x^{(n)})$ for $n\in[N]$, SGD finds a network $\mathsf{out}(x)$ with population risk

As an example, when $k\geq 2$ is constant, $d=\Theta(d_{1})$ is sufficiently large, and $\alpha=\Theta(d^{-0.1})$,

Corollary 7.1 says that ResNet achieves regression error $\alpha^{3.9}$ on the true distribution, with $N_{\mathsf{res}}=\widetilde{O}(d^{1.8})$ samples to learn any function in (7.1);

Theorem 2 says that kernel methods cannot achieve $\alpha^{2}/16$ error even with $N\leq(N_{\mathsf{res}})^{k/2}\ll o(d^{k})$ samples. Hence, to achieve generalization $\alpha^{2}/16\gg\alpha^{3.9}$, the sample complexity of any kernel method is at least $N\geq(N_{\mathsf{res}})^{k/2}\gg N_{\mathsf{res}}$.

Proof Overview. Our proof of Theorem 2 is relatively simple, and we illustrate the main idea in the case of $d=d_{1}$. At a high level, given $N\ll\binom{n}{d}$ samples, the kernel regression function only has $N$-degrees of freedom (each with respect to a sample point). Now, since there are possibly $\binom{n}{d}$ many target functions, if the kernel regression learns most of these target functions to some sufficient accuracy, then by some rank counting argument, the degree of freedom is not enough.

2 Linear Regression Over Feature Mappings

As an example, there exists sufficiently large constant $c>1$ such that, for every $k\geq 4c$, for every $d_{1}\geq d/2$, for every $d\geq\Omega(2^{k})$, there exists choice $\alpha=2^{-\Theta(k)}\cdot d^{-0.001}$ such that

Corollary 7.1 says that ResNet achieves regression error $\widetilde{O}(\alpha^{4}2^{O(k)})\leq\alpha^{3.9}$ in time $T_{\mathsf{res}}=\poly(d,2^{k},\alpha^{-1})\leq d^{c}$ to learn any function in (7.1);

Theorem 3 says that linear regression over feature mapping cannot achieve regression error $\alpha^{2}/16$ even if D=\Omega\big{(}{d_{1}\choose k}\big{)}\geq d^{2c}.

In particular, this means linear regression over feature mappings cannot achieve regression error $\alpha^{2}/16$ even if $D=(T_{\mathsf{res}})^{2}$. Since a linear regression over $R^{D}$ normally takes at least time/space $D$ to compute/store, this implies that ResNet is also more time/space efficient than linear regression over feature mappings as well.

Theorem 3 can be proved in the same way as Theorem 2, using exactly the same hard instance, since $\mathfrak{F}(x)$ has exactly $D$-degrees of freedom.

Experiments

Neural Networks Algorithms. Recall in our positive result on three-layer ResNet (see Theorem 1 and Footnote 12), to prove the strongest result, we only train hidden weights $\mathbf{W}$ and $\mathbf{V}$ but not the output layer $\mathbf{A}$. One can naturally extend this to show that Theorem 1 also holds when $\mathbf{W},\mathbf{V},\mathbf{A}$ are jointly trained. For such reason, we implement both algorithms: 3resnet(hidden) for training only $\mathbf{W},\mathbf{V}$ and 3resnet(all) for training all $\mathbf{W},\mathbf{V},\mathbf{A}$. This is similar for two-layer and three-layer fully-connected networks, where previously the strongest theoretical work is in terms of training only hidden weights , so we implement both (all) and (hidden) for them.

Kernel Methods. We implement conjugate kernel, which corresponds to training only the last (output) layer ; as well as neural tangent kernel (NTK), in which we train all the layers .

Setup. We choose the network width (i.e., parameter $m$) in the range $m\in\{20,50,100,200,\dots\}$ until the largest possible value $m$ that fits into a 16GB GPU memory. We choose the popular random initialization: entries of $\mathbf{A},\mathbf{V},\mathbf{W}$ (and their corresponding bias terms) are all i.i.d. from $\mathcal{N}(1,\frac{1}{m})$.This corresponds to choosing the standard deviation as $\frac{1}{\sqrt{\texttt{fan\_in}}+\sqrt{\texttt{fan\_out}}}$. Some practitioners also use $\frac{1}{\sqrt{\texttt{fan\_in}}}$ as the standard deviation. We have included an experiment with respect to that choice in our V1/V2 of this paper. We use similar initializations for two and three-layer networks.

Experiment 1: Performance Comparison. Since it is unfair to compare neural network training “with respect to hidden weights only” vs. “with respect to all weights”, we conduct two experiments. The first experiment is on training all layers vs. kernel methods, see Figure 2(a); and the second experiment is on training hidden layers vs. kernel methods, see Figure 2(b). We use $N=500$ training samples for the former case and $N=1000$ samples for the latter case, because training the last layer together gives more power to a neural network.

In both experiments, we choose $\alpha=0.3$ and $k=15$ so that test error $k\alpha^{2}=1.35$ is a threshold for detecting whether the trained model has successfully learned $\alpha\mathcal{G}(\mathcal{F}(x))$ or not. If the model has not learned $\alpha\mathcal{G}(\mathcal{F}(x))$ to any non-trivial accuracy, then the error is $\alpha$ per output coordinate, totaling to $k\alpha^{2}$ in regression error.

From Figure 2, it is clear that for our choice of $N$, training a three-layer ResNet is the only method among the ones we compare that can learn $\alpha\mathcal{G}(\mathcal{F}(x))$ (even only non-trivially). All kernel methods fall far behind even when the network width $m$ is large.

𝛽ℱ𝑥𝛼𝒢ℱ𝑥\mathcal{H}(x)=\beta\mathcal{F}(x)+\alpha\mathcal{G}(\mathcal{F}(x)) with $\alpha=0.3$ and varying $\beta\in$. Experiment 2: Sensitivity on $\alpha$. One key assumption of this paper is to have $\alpha$ to be sufficiently small, so that ResNet can perform hierarchical learning, by first learning the base signal $\mathcal{F}$, which is simpler and contributes more to the target, and then learning the composite signal $\alpha\mathcal{G}\left(\mathcal{F}\right)$, which is more complicated but contributes less.

In Figure 3, we verify that this assumption is indeed necessary. Instead of varying $\alpha$ (which will change the error magnitude), we define $\mathcal{H}(x)=\beta\mathcal{F}(x)+\alpha\mathcal{G}(\mathcal{F}(x))$ and let $\beta$ vary between and $1$. As shown in Figure 3, when $\alpha\lesssim\beta$, the base signal is larger than the composite signal, so indeed ResNet can perform hierachical learning; in contrast, when $\alpha\gtrsim\beta$, learning the composite signal becomes practically impossible.

Other Findings. Although this paper proves theoretical separation between three-layer ResNet and kernel methods (and it is verified by Figure 2), we do not yet have

theoretical separation between two/three-layer fully-connected networks and kernel methods;

theoretical separation between three-layer ResNet and two/three-layer networks.

It seems in practice such separations do exist (as observed in Figure 2). We leave these as future research directions.

2 SGD Does Not Converge To Minimal Norm Solutions

We give a simple experiment to show that optimization methods (such as SGD) do not necessarily converge to minimal complexity solutions.

It is a simple experimental exercise to verify that, for every even $m\geq 200$ and every $d\geq 6$, there exist This can be done by first considering $m=200$ and $d=6$. Experimentally one can easily use SGD to train such two-layer networks to obtain some $W^{*}$ with such test errors. Then, for general $d>6$, one can pad $W^{*}$ with $d-6$ zero columns; and for general $m>200$, one can duplicate the rows of $W^{*}$ and re-scale.

Using simple Rademacher complexity argument, the above existential statement implies if we focus only on matrices $W$ with $\|W\|_{F}\leq 9.7\sqrt{d}$, then given $N$ training samples the Rademacher complexity is at most $\frac{\frac{2}{\sqrt{m}}\sum_{j\in[m]}\|W_{j}\|_{2}}{\sqrt{N}}\leq\frac{2\|W\|_{F}}{\sqrt{N}}$.This can found for instance in . A cleaner one page proof can be found in the lecture notes . This implies, for any $m\geq 200$ and $d\geq 6$, if $N=O(d)$ samples are given and if SGD finds any close-to-minimal complexity solution (i.e. with F-norm within some constant times $\sqrt{d}$) that performs well on the training set, then it also generalizes to give small test error (i.e. test error $<0.3$).

SGD cannot find solution with test error better than 0.69 (see Figure 4(a)), and

SGD cannot find solution with small training error and small Frobenius norm (see Figure 4(a)). Thus, SGD starting from random initialization fails to find the minimal complexity solution.

SGD cannot find solution with test error better than 0.98 (see Figure 4(b)), and

SGD cannot find solution with small training error and small Frobenius norm (see Figure 4(b)).

In Appendix A we give some more information about our concept class and complexity measure.

In Appendix B we review some simple lemmas from probability theory.

In Appendix C we give our full proof to Theorem 1.

In Appendix D we give our full proof to Theorem 2.

In Appendix E we include a variant of the existential lemma from prior work, and include its proof only for completeness’ sake.

Appendix A Complexity and Concept Class

In this section we introduce an alternative (but bigger) concept class.

where where $\mathcal{F}$ and $\mathcal{G}$ respectively have general complexity $(p_{\mathcal{F}},\mathfrak{C}_{\mathfrak{s}}(\mathcal{F}),\mathfrak{C}_{\varepsilon}(\mathcal{G}))$ and $(p_{\mathcal{G}},\mathfrak{C}_{\mathfrak{s}}(\mathcal{G}),\mathfrak{C}_{\varepsilon}(\mathcal{G}))$. We further assume $\|\mathcal{F}(x)\|_{2}\leq\mathfrak{B}_{\mathcal{F}}$ for all $(x,y)\sim\mathcal{D}$.

We have the following lemma which states that Concept 1 is a special case of Concept 2 (with constant factor $2$ blow up).

Under Concept 1, we can construct $\mathcal{F}^{\prime},\mathcal{G}^{\prime}$ satisfying Concept 2 with general complexity $(2p_{\mathcal{F}},\mathfrak{C}_{\mathfrak{s}}(\mathcal{F}),\mathfrak{C}_{\varepsilon}(\mathcal{G}))$ and $(2p_{\mathcal{G}},\mathfrak{C}_{\mathfrak{s}}(\mathcal{G}),\mathfrak{C}_{\varepsilon}(\mathcal{G}))$ and with $\mathfrak{B}_{\mathcal{F}}=1$.

Lemma A.2 is a simple corollary of the following claim.

$\mathcal{F}^{\prime}$ has general complexity $(2p,\mathfrak{C}_{\mathfrak{s}}(\mathcal{F}),\mathfrak{C}_{\varepsilon}(\mathcal{F}))$ where $\mathfrak{C}_{\varepsilon}(\mathcal{F})=\max_{r,i}\{\mathfrak{C}_{\varepsilon}(\mathcal{F}_{r,i})\}$ and $\mathfrak{C}_{\mathfrak{s}}(\mathcal{F})=\max_{r,i}\{\mathfrak{C}_{\mathfrak{s}}(\mathcal{F}_{r,i})\}$.

Below we prove that the above claim holds. For each $\mathcal{F}_{r,i}(\cdot)$ suppose we have $\mathcal{F}_{r,i}(z)=\sum_{i=0}^{\infty}c_{i}z^{i}$ as its Taylor expansion, then we can write

It is a simple exercise to verify that $\mathcal{F}^{\prime}(x)=\mathcal{F}(x)$ for all unit vectors $x$. ∎

We also state some simple properties regarding our complexity measure.

The boundedness of $\|\mathcal{F}(x)\|_{2}$ is trivial so we only focus on $\|\mathcal{F}(x)-\mathcal{F}(y)\|_{2}$. For each component $g(x)=\mathcal{F}_{r,i}\left(\frac{\langle w^{*}_{1,i},(x,1)\rangle}{\|(x,1)\|_{2}}\right)\cdot\langle w^{*}_{2,i},(x,1)\rangle$, denoting by $w^{*}_{1}$ as the first $d$ coordinate of $w^{*}_{1,i}$, and by $w^{*}_{2,i}$ as the first $d$ coordinates of $w^{*}_{2,i}$, we have

As a result, $|\mathcal{F}_{r}(x)-\mathcal{F}_{r}(y)|\leq 3p\mathfrak{C}_{\mathfrak{s}}(\mathcal{F}_{r,i})$. ∎

Appendix B Probability Theory Review

The following concentration of chi-square distribution is standard.

If $g\sim\mathcal{N}(0,\mathbf{I})$ is $m$-dimensional, then for every $t\geq 1$

The following norm bound on random Gaussian matrix is standard.

For any $t\geq 1$, with probability $\geq 1-e^{-\Omega(t^{2})}$ it satisfies $\|\mathbf{M}\|_{2}\leq O(\sqrt{n}+\sqrt{m})+t$.

The first statement can be found for instance in [34, Proposition 2.4]. As for the second statement, it suffices for us to consider all $\binom{m}{s}$ possible $n\times s$ sub-matrices of $\mathbf{M}$, each applying the first statement, and then taking a union bound. ∎

The following concentration is proved for instance in .

Let $(n_{1},\alpha_{1},a_{1,1},a_{2,1}),\cdots,(n_{m},\alpha_{m},a_{1,m},a_{2,m})$ be $m$ i.i.d. samples from some distribution, where within a 4-tuples:

the marginal distribution of $a_{1,i}$ and $a_{2,i}$ is standard Gaussian $\mathcal{N}(0,1)$;

$n_{i}$ and $\alpha_{i}$ are not necessarily independent;

$a_{1,i}$ and $a_{2,i}$ are independent; and

$n_{i}$ and $\alpha_{i}$ are independent of $a_{1,i}$ and $a_{2,i}$.

Let us consider a fixed $n_{1},\alpha_{1},\cdots,n_{m},\alpha_{m}$, then since each $|\mathds{1}[n_{i}\geq 0]h(\alpha_{i})|\leq L$, by Gaussian chaos variables concentration bound (e.g., Example 2.15 in ) we have that

Since this holds for every choice of $\{n_{i},\alpha_{i}\}_{i\in[m]}$ we can complete the proof. The second inequality follows from sub-exponential concentration bounds. ∎

The next proposition at least traces back to and was stated for instance in .

Observe that $(D^{\prime})_{j,j}$ is non-zero for some $j\in[m]$ only if

Therefore, denoting by $x=D^{\prime}g^{(0)}$, for each $j\in[m]$ such that $x_{j}\neq 0$, we must have $|x_{j}|=|(g^{(0)})_{j}|\leq|(g^{\prime})_{j}|$ so we have $\|x\|_{2}\leq\|g^{\prime}\|_{2}\enspace.$

Let $\xi\leq\frac{1}{2\sqrt{m}}$ be a constant parameter to be chosen later.

We denote by $S_{1}\subseteq[m]$ the index sets where $j$ satisfies $|(g^{(0)})_{j}|\leq\xi$. Since we know $(g^{(0)})_{j}\sim\mathcal{N}(0,1/m)$, we have $\operatornamewithlimits{\mathbf{Pr}}[|(g^{(0)})_{j}|\leq\xi]\leq O\left(\xi\sqrt{m}\right)$ for each $j\in[m]$. Using Chernoff bound for all $j\in[m]$, we have with probability at least $1-e^{-\Omega(m^{3/2}\xi)}$,

We denote by $S_{2}\subseteq[m]\setminus S_{1}$ the index set of all $j\in[m]\setminus S_{1}$ where $(D^{\prime})_{j,j}\neq 0$. Using (B.1), we have for each $j\in S_{2}$ it satisfies $|(g^{\prime})_{j}|\geq|(g^{(0)})_{j}|\geq\xi\enspace.$ This means $|S_{2}|\leq\frac{\|g^{\prime}_{1}\|_{2}^{2}}{\xi^{2}}\enspace.$

From above, we have \|D^{\prime}\|_{0}\leq|S_{1}|+|S_{2}|\leq O\big{(}\xi m^{3/2}+\frac{\delta^{2}}{\xi^{2}}\big{)}. Choosing $\xi=\frac{\delta^{2/3}}{2m^{1/2}}$ gives the desired result. ∎

Appendix C Theorem 1 Proof Details

In the analysis, let us define a diagonal matrices

which satisfy $\mathsf{out}_{1}(x)=\mathbf{A}D_{\mathbf{W}}({\mathbf{W}^{(0)}}+\mathbf{W})(x,1)$ and $\mathsf{out}(x)=\mathbf{A}D_{\mathbf{V},\mathbf{W}}({\mathbf{V}^{(0)}}+\mathbf{V})(\mathsf{out}_{1}(x),1)$.

Throughout the proof, we assume $m\geq\poly(\mathfrak{C}_{\alpha}(\mathcal{F}),\mathfrak{C}_{\alpha}(\mathcal{G}),p_{\mathcal{G}},p_{\mathcal{F}},k,\alpha^{-1})$.

In this subsection we present our coupling lemma. It shows that for all weight matrices not very far from random initialization (namely, all $\|\mathbf{W}\|_{2}\leq\tau_{w}$ and $\|\mathbf{V}\|_{2}\leq\tau_{v}$), many good properties occur. This includes upper bounds on the number of sign changes (i.e., on $\left\|D_{{\mathbf{W}^{(0)}}}-D_{\mathbf{W}}\right\|_{0}$ and $\left\|D_{{\mathbf{V}^{(0)}},\mathbf{W}}-D_{\mathbf{V},\mathbf{W}}\right\|_{0}$) as well as vanishing properties such as $\mathbf{A}D_{\mathbf{W}}{\mathbf{W}^{(0)}},\mathbf{A}D_{\mathbf{V},\mathbf{W}}{\mathbf{V}^{(0)}}$ being negligible. We prove such properties using techniques from prior works .

Suppose $\tau_{w}\geq 1$, \tau_{w}\in\big{[}m^{1/8+0.001}\sigma_{w},m^{1/8-0.001}\sigma_{w}^{1/4}\big{]}, and \tau_{v}\in\big{[}\sigma_{v}\cdot(k/m)^{3/8},\sigma_{v}\big{]}. Then, for every fixed $x$, with high probability over $\mathbf{A},{\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$, we have that for all $\mathbf{W},\mathbf{V}$ satisfying $\|\mathbf{W}\|_{2}\leq\tau_{w}$ and $\|\mathbf{V}\|_{2}\leq\tau_{v}$, it holds that

$\left\|D_{{\mathbf{W}^{(0)}}}-D_{\mathbf{W}}\right\|_{0}\leq O((\tau_{w}/\sigma_{w})^{2/3}m^{2/3})$

$\left\|\mathbf{A}D_{\mathbf{W}}\mathbf{W}(x,1)-\mathbf{A}D_{\mathbf{W}}(({\mathbf{W}^{(0)}}+\mathbf{W})(x,1))\right\|_{2}\leq\widetilde{O}\left(\frac{\tau_{w}(\tau_{w}/\sigma_{w})^{1/3}}{m^{1/6}}\right)\leq O(m^{-0.001})$

$\|\mathsf{out}_{1}(x)\|_{2}=\left\|\mathbf{A}D_{\mathbf{W}}({\mathbf{W}^{(0)}}+\mathbf{W})(x,1)\right\|_{2}\leq O\left(\tau_{w}\right)$

$\left\|D_{{\mathbf{V}^{(0)}},\mathbf{W}}-D_{\mathbf{V},\mathbf{W}}\right\|_{0}\leq O((\tau_{v}/\sigma_{v})^{2/3}m)$

$\left\|\mathbf{A}D_{\mathbf{V},\mathbf{W}}\mathbf{V}(\mathsf{out}_{1}(x),1)-\mathbf{A}D_{\mathbf{V},\mathbf{W}}({\mathbf{V}^{(0)}}+\mathbf{V})(\mathsf{out}_{1}(x),1)\right\|_{2}\leq\widetilde{O}\left(\tau_{v}(\tau_{v}/\sigma_{v})^{1/3}\right)\cdot(\|\mathsf{out}_{1}(x)\|_{2}+1)$

$\left\|\mathbf{A}D_{\mathbf{V},\mathbf{W}}{\mathbf{V}^{(0)}}\right\|_{2}\leq\widetilde{O}\left(\tau_{v}(\tau_{v}/\sigma_{v})^{1/3}\right)$

$\left\|\mathbf{A}D_{\mathbf{V},\mathbf{W}}({\mathbf{V}^{(0)}}+\mathbf{V})(\mathsf{out}_{1}(x),1)\right\|_{2}\leq\widetilde{O}\left(\tau_{v}(\|\mathsf{out}_{1}(x)\|_{2}+1)\right)$

Using basic probability argument (appropriately scaling and invoking Proposition B.4) we have

For the first term, we have $\left\|D_{\mathbf{W}^{(0)}}{\mathbf{W}^{(0)}}(x,1)\right\|_{2}\leq\left\|{\mathbf{W}^{(0)}}(x,1)\right\|_{2}\leq O(\sigma_{w}\sqrt{m})$ with high probability due to concentration of chi-square distribution, and then using the randomness of $\mathbf{A}$ and applying concentration of chi-square distribution again, we have $\left\|\mathbf{A}D_{\mathbf{W}^{(0)}}{\mathbf{W}^{(0)}}(x,1)\right\|_{2}\leq\widetilde{O}(\sqrt{k}\sigma_{w})$ with high probability.

For the second term, invoking Proposition B.4 again, we have

Recall for every $s$-sparse vectors $y$, it satisfies $\|\mathbf{A}y\|_{2}\leq\widetilde{O}(\frac{\sqrt{s}}{\sqrt{m}})\cdot\|y\|_{2}$ with high probability (see Proposition B.2). This implies

for s=O\left(\big{(}\frac{\tau_{w}}{\sigma_{w}\sqrt{m}}\big{)}^{2/3}\cdot m\right). Together, we have

We use Lemma lem:couplingb together with $\|\mathbf{A}D_{\mathbf{W}}\mathbf{W}(x,1)\|_{2}\leq\|\mathbf{A}\|_{2}\|\mathbf{W}\|_{2}\leq O(\tau_{w})$, where the property $\|\mathbf{A}\|_{2}\leq O(1)$ holds with high probability using Proposition B.2.

Then, for the first term, we have $\mathbf{A}D_{{\mathbf{V}^{(0)}},\mathbf{W}}{\mathbf{V}^{(0)}}z=\mathbf{A}\sigma({\mathbf{V}^{(0)}}z)$ and by by concentration of chi-square distribution we have $\|{\mathbf{V}^{(0)}}z\|_{2}\leq O(\sigma_{v}\|z\|_{2})$ with probability at least $1-e^{-\Omega(m)}$, and then using the randomness of $\mathbf{A}$ and applying chi-square concentration again (see Proposition B.1), we have with probability at least $1-e^{-\Omega(k\log^{2}m)}$,

For the second term, invoking Proposition B.4, we have

Recall for every $s$-sparse vectors $y$, it satisfies $\|\mathbf{A}y\|_{2}\leq\widetilde{O}(\frac{\sqrt{s}}{\sqrt{m}})\cdot\|y\|_{2}$ with probability at least $1-e^{-\widetilde{\Omega}(s)}$ (see Proposition B.2). This implies

This is a byproduct of the proof of Lemma lem:couplinge.

Combining this with Lemma lem:couplinge gives the proof.

C.2 Existantial

In this subsection, we prove the existence of matrices $\mathbf{W}^{\divideontimes},\mathbf{V}^{\divideontimes}$ with $\|\mathbf{W}^{\divideontimes}\|_{F}\leq\frac{\tau_{w}}{10}$ and $\|\mathbf{V}^{\divideontimes}\|_{F}\leq\frac{\tau_{v}}{10}$ satisfying $\mathbf{A}D_{{\mathbf{W}^{(0)}}}\mathbf{W}^{\divideontimes}(x,1)\approx\mathcal{F}(x)$ and $\mathbf{A}D_{{\mathbf{V}^{(0)}},\mathbf{W}}\mathbf{V}^{\divideontimes}(\mathsf{out}_{1}(x),1)\approx\alpha\mathcal{G}\left(\mathsf{out}_{1}(x)\right)$.

This existential proof relies on an “indicator to function” lemma that was used in prior work ; however, for the purpose of this paper we have to revise it to include a trainable bias term (or equivalently, to support vectors of the form $(x,1)$). We treat that carefully in Appendix E.

for all $x$ and $\mathbf{W}$, $\left\|\mathbf{A}D_{{\mathbf{V}^{(0)}},\mathbf{W}}\mathbf{V}^{\divideontimes}(\mathsf{out}_{1}(x),1)-\alpha\mathcal{G}\left(\mathsf{out}_{1}(x)\right)\right\|_{2}\leq\widetilde{\alpha}^{2}\cdot\|(\mathsf{out}_{1}(x),1)\|_{2}$.

Finally, choosing $z=\mathsf{out}_{1}(x)$ finishes the proof.

Next, we can combine coupling and existential lemmas:

Under the assumptions of Lemma C.1 and Lemma C.2, we have

For every $s$-sparse vectors $y$, it satisfies $\|\mathbf{A}y\|_{2}\leq\widetilde{O}(\frac{\sqrt{s}}{\sqrt{m}})\cdot\|y\|_{2}$ with high probability (see Proposition B.2). We also have $\|\mathbf{W}^{\divideontimes}(x,1)\|_{2}\leq O(\|\mathbf{W}^{\divideontimes}\|_{F})\leq O(\tau_{w})$. Therefore, $\|\mathbf{A}(D_{{\mathbf{W}^{(0)}}}-D_{\mathbf{W}})\mathbf{W}^{\divideontimes}(x,1)\|\leq O(\sqrt{s}\tau_{w}/\sqrt{m})$ where $s$ is the maximum sparsity of $D_{{\mathbf{W}^{(0)}}}-D_{\mathbf{W}}$, which satisfies $s=O((\tau_{w}/\sigma_{w})^{2/3}m^{2/3})$ by Lemma lem:couplinga. This, combining with Lemma lem:exist-priora gives

Again, for every $s$-sparse vectors $y$, it satisfies $\|\mathbf{A}y\|_{2}\leq\widetilde{O}(\frac{\sqrt{s}}{\sqrt{m}})\cdot\|y\|_{2}$ with high probability. We also have $\|\mathbf{V}^{\divideontimes}(\mathsf{out}_{1}(x),1)\|_{2}\leq O(\|\mathbf{V}^{\divideontimes}\|_{F})\cdot\|(\mathsf{out}_{1}(x),1)\|_{2}\leq O(\tau_{v})\cdot\|(\mathsf{out}_{1}(x),1)\|_{2}$. Therefore,

where $s$ is the maximum sparsity of $D_{{\mathbf{V}^{(0)}},\mathbf{W}}-D_{\mathbf{V},\mathbf{W}}$, which satisfies $s=O((\tau_{v}/\sigma_{v})^{2/3}m)$ by Lemma lem:couplingd. This, combining with Lemma lem:exist-priorb gives

This combines Lemma lem:couplingb and Lemma lem:exist-and-couplea, together with our sufficiently large choice of $m$.

C.3 Optimization

In this subsection we give some structural results that shall be later used in the optimization step. The first fact gives an explicit formula of the gradient.

When $\mathsf{Obj}(\mathbf{W},\mathbf{V};(x,y))=\frac{1}{2}\|y-\mathsf{out}(\mathbf{W},\mathbf{V};x)\|_{2}^{2}$, we can write its gradient as follows.

The next claim gives simple upper bound on the norm of the gradient.

For all $(x,y)$ in the support of $\mathcal{D}$, with high probability over $\mathbf{A},{\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$, we have that for all $\mathbf{W},\mathbf{V}$ satisfying $\|\mathbf{W}\|_{F}\leq\tau_{w}$ and $\|\mathbf{V}\|_{F}\leq\tau_{v}$, it holds that

For the gradient in $\mathbf{W}$, we derive using the gradient formula Fact C.4 that

Above, the last inequality uses $\|\mathbf{A}\|_{2}\leq O(1)$ and $\|{\mathbf{V}^{(0)}}\|_{2}\leq O(\sigma_{v})$ with high probability (using random matrix theory, see Proposition B.2), as well as $\tau_{v}\leq\sigma_{v}$. Similarly, using the gradient formula Fact C.4, we derive that

where the last inequality uses Lemma lem:couplingc and $\|\mathbf{A}\|_{2}\leq O(1)$. ∎

The next claim gives a careful approximation to $f(\mathbf{W}^{\divideontimes}-\mathbf{W};x)+g(\mathbf{V}^{\divideontimes}-\mathbf{V};x)$, which according to Fact C.4 is related to the correlation between the gradient direction and $(\mathbf{W}-\mathbf{W}^{\divideontimes},\mathbf{V}-\mathbf{V}^{\divideontimes})$.

In the same setting as Lemma C.1 and Lemma C.2, suppose we set parameters according to Table 1. Then, we can write

and for every $(x,y)\sim\mathcal{D}$, with high probability $\|Err\|_{2}\leq O(\tau_{w})$.

For the $\clubsuit$ term, under expectation over $(x,y)\sim\mathcal{D}$,

where the last inequality uses Lemma lem:couplingf and Lemma lem:exist-and-couplec, together with $\tau_{v}\leq\frac{1}{\polylog(m)}\sigma_{v}$.

For the $\spadesuit$ term, under expectation over $(x,y)\sim\mathcal{D}$,

where the first inequality uses Lemma lem:exist-and-couplea and Lemma lem:exist-and-coupleb, as well as the Lipscthiz continuity of $\mathcal{G}(x)$ (which satisfies $\|\mathcal{G}(x)-\mathcal{G}(y)\|\leq\mathfrak{L}_{\mathcal{G}}\|x-y\|$); and the second inequality uses $\frac{1}{\sigma_{v}}\leq\tau_{v}^{2}$ and the definition of $\widetilde{\alpha}$.

For the $\diamondsuit$ term, under expectation over $(x,y)\sim\mathcal{D}$,

where the inequality uses Lemma lem:couplingb, Lemma lem:couplinge and $\frac{1}{\sigma_{v}}\leq\tau_{v}^{2}$.

Combining this with Claim C.7, and using $\|\mathsf{out}_{1}(x)\|_{2}\leq\|\mathsf{out}_{1}(x)-\mathcal{F}(x)\|_{2}+\mathfrak{B}_{\mathcal{F}}$, we have

As for the absolute value bound, one can naively derive that with high probability $\|f(\mathbf{W}^{\divideontimes}-\mathbf{W};x)\|_{2}\leq O(\tau_{w})$, $\|g(\mathbf{V}^{\divideontimes}-\mathbf{V};x)\|_{2}\leq O(\tau_{w}\tau_{v})$, $\|\mathcal{H}(x)\|_{2}\leq\mathfrak{B}_{\mathcal{F}}+\alpha\mathfrak{B}_{\mathcal{F}\circ\mathcal{G}}$, and $\|\mathsf{out}(\tau_{w})\|_{2}\leq O(\tau_{w})$ (by Lemma lem:couplingc and lem:couplingg). Combining them with $\mathfrak{B}_{\mathcal{F}}\leq\sqrt{k}p_{\mathcal{F}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{F})\leq\tau_{w}$ and \alpha\mathfrak{B}_{\mathcal{F}\circ\mathcal{G}}\leq\alpha\big{(}\mathfrak{B}_{\mathcal{F}}\mathfrak{L}_{\mathcal{G}}+\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})\big{)}\leq\frac{1}{kp_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})}\big{(}\mathfrak{B}_{\mathcal{F}}\mathfrak{L}_{\mathcal{G}}+\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})\big{)}\leq\tau_{w} finishes the proof. ∎

Finally, we state a simple claim that bounds the norm of $\|\mathsf{out}_{1}(x)-\mathcal{F}(x)\|_{2}$ given the norm of $\|\mathsf{out}(x)-\mathcal{H}(x)\|_{2}$.

In the same setting as Lemma C.1, if we additionally have $\tau_{v}\leq\frac{1}{\polylog(m)}$, for every fixed $x$, with high probability over $\mathbf{A},{\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$,

Using Lemma lem:couplingg we have $\|\mathbf{A}D_{\mathbf{V},\mathbf{W}}({\mathbf{V}^{(0)}}+\mathbf{V})(\mathsf{out}_{1}(x),1)\|\leq\widetilde{O}(\tau_{v}(\|\mathsf{out}_{1}(x)\|_{2}+1))$, and using the boundedness we have $\|\mathcal{G}(\mathcal{F}(x))\|_{2}\leq\mathfrak{B}_{\mathcal{F}\circ\mathcal{G}}$. We also have $\|\mathsf{out}_{1}(x)\|_{2}\leq\|\mathsf{out}_{1}(x)-\mathcal{F}(x)\|_{2}+\mathfrak{B}_{\mathcal{F}}$. Together, we have

Using $\tau_{v}\leq\frac{1}{\polylog(m)}$ we finish the proof. ∎

C.4 Proof of Theorem 1

Under Concept 1 or Concept 2, for every \alpha\in\big{(}0,\widetilde{\Theta}(\frac{1}{kp_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G})})\big{)} and $\delta\geq\mathsf{OPT}+\widetilde{\Theta}\left(\alpha^{4}(kp_{\mathcal{G}}\mathfrak{C}_{\mathfrak{s}}(\mathcal{G}))^{4}(1+\mathfrak{B}_{\mathcal{F}})^{2}\right)$. There exist $M=\poly(\mathfrak{C}_{\alpha}(\mathcal{F}),\mathfrak{C}_{\alpha}(\mathcal{G}),p_{\mathcal{F}},\alpha^{-1})$ satisfying that for every $m\geq M$, with high probability over $\mathbf{A},{\mathbf{W}^{(0)}},{\mathbf{V}^{(0)}}$, for a wide range of random initialization parameters $\sigma_{w},\sigma_{v}$ (see Table 1), choosing

With high probability, the SGD algorithm satisfies

We first assume that throughout the SGD algorithm, it satisfies

We shall prove in the end that (C.3) holds throughout the SGD algorithm.

On one hand, using Claim C.6, at any point $\mathbf{W}_{t},\mathbf{V}_{t}$, we have