The Impact of Neural Network Overparameterization on Gradient Confusion and Stochastic Gradient Descent

Introduction

Stochastic gradient descent (SGD) (Robbins & Monro 1951) and its variants with momentum have become the standard optimization routine for neural networks due to their fast convergence and good generalization properties (Wilson et al. 2017; Sutskever et al. 2013; Smith et al. 2020). Yet the convergence behavior of SGD on neural networks still eludes full theoretical understanding. Furthermore, it is not well understood how design choices on neural network architecture affect training performance. In this paper, we make progress on these open questions.

Classical stochastic optimization theory predicts that the learning rate of SGD needs to decrease over time for convergence to be guaranteed to the minimizer of a convex function (Shamir & Zhang 2013; Bertsekas 2011). For strongly convex functions for example, such results show that a decreasing learning rate schedule of $O(1/k)$ is required to guarantee convergence to within $\epsilon$-accuracy of the minimizer in $O(1/\epsilon)$ iterations, where $k$ denotes the iteration number. Such decay schemes, however, typically lead to poor performance on standard neural network problems.

Neural networks operate in a regime where the number of parameters is much larger than the number of training data. In this “over-parameterized” regime, SGD seems to converge quickly with constant learning rates. Most neural network practitioners use a constant learning rate for the majority of training (with exponential decay only towards the end of training) without seeing the method stall (Krizhevsky et al. 2012; Simonyan & Zisserman 2014; He et al. 2016; Zagoruyko & Komodakis 2016). With constant learning rates, theoretical guarantees show that SGD converges quickly to a neighborhood of the minimizer, but then reaches a noise floor beyond which it stops converging; this noise floor depends on the learning rate and the variance of the gradients (Moulines & Bach 2011; Needell et al. 2014). Recent results show that convergence without a noise floor is possible without decaying the learning rate, provided the model is strongly convex and overfitting occurs (Schmidt & Roux 2013; Ma et al. 2017; Vaswani et al. 2018).

While these results do give important insights, they do not fully explain the dynamics of SGD on neural networks, and how they relate to over-parameterization. Furthermore, training performance is strongly influenced by network architecture. It is common knowledge among practitioners that, under standard Gaussian initialization techniques (Glorot & Bengio 2010; He et al. 2015), deeper networks train slower (Bengio et al. 1994; Saxe et al. 2013). This has led to several innovations over the years to get deeper nets to train more easily, such as careful initialization strategies (Xiao et al. 2018), residual connections (He et al. 2016), and normalization schemes like batch normalization (Ioffe & Szegedy 2015). Furthermore, there is evidence to indicate that wider networks are faster to train (Zagoruyko & Komodakis 2016; Nguyen & Hein 2017), and recent theoretical results suggest that the dynamics of SGD simplify considerably for very wide networks (Jacot et al. 2018; Lee et al. 2019). In this paper, we make progress on theoretically understanding these empirical observations and unifying existing theoretical results. To this end, we identify and analyze a condition that enables us to establish direct relationships between layer width, network depth, problem dimensionality, initialization schemes, and trainability and SGD dynamics for over-parameterized networks.

Our contributions. Typical neural networks are over-parameterized (i.e., the number of parameters exceed the number of training points). In this paper, we ask how this over-parameterization, and more specifically the network architecture, affects the trainability of neural networks and the dynamics of SGD. Through extensive theoretical and experimental studies, we show how layer width, network depth, initialization schemes, and other architecture choices affect the dynamics. The following are our main contributions. To keep the main text of the paper concise, all proofs and several additional experimental results are delegated to the appendix.

We identify a condition, termed gradient confusion, that impacts the convergence properties of SGD on over-parameterized models. We prove that high gradient confusion may lead to slower convergence, while convergence is accelerated (and could be faster than predicted by existing theory) if confusion is low, indicating a regime where constant learning rates work well in practice (sections 2 and 3). We use the gradient confusion condition to study the effect of various architecture choices on trainability and convergence.

We study the effect of neural network architecture on gradient confusion at standard Gaussian initialization schemes (section 4), and prove (a) gradient confusion increases as the network depth increases, and (b) wider networks have lower gradient confusion. These indicate that deeper networks are more difficult to train and wider networks can improve trainability of networks. Directly analyzing the gradient confusion bound enables us to derive results on the effect of depth and width, without requiring restrictive assumptions like large layer widths (Du et al. 2018; Allen-Zhu et al. 2018). Our results hold for a large class of neural networks with different non-linear activations and loss-functions. In section 5, we present a more general result on the effect of depth on the trainability of networks without assuming the network is at initialization.

We prove that for linear neural networks, gradient confusion is independent of depth when using orthogonal initialization schemes (section 6) (Saxe et al. 2013; Schoenholz et al. 2016). This indicates a way forward in developing techniques for training deeper models.

We test our theoretical predictions using extensive experiments on wide residual networks (WRNs) (Zagoruyko & Komodakis 2016), convolutional networks (CNNs) and multi-layer perceptrons (MLPs) for image classification tasks on CIFAR-10, CIFAR-100 and MNIST (section 7 and appendix A). We find that our theoretical results consistently hold across all our experiments. We further show that the combination of batch normalization and skip connections in residual networks help lower gradient confusion, thus indicating why SGD can efficiently train deep neural networks that employ such techniques.

Gradient confusion

Notations. We denote vectors in bold lower-case and matrices in bold upper-case. We use $(\mathbf{W})_{i,j}$ to indicate the $(i,j)$ cell in matrix $\mathbf{W}$ and $(\mathbf{W})_{i}$ for the $i^{\text{th}}$ row of matrix $\mathbf{W}$. $\|\mathbf{W}\|$ denotes the operator norm of $\mathbf{W}$. $[N]$ denotes $\{1,2,\ldots,N\}$ and $[N]_{0}$ denotes $\{0,1,\ldots,N\}$.

Preliminaries. Given $N$ training points (specified by the corresponding loss functions $\{f_{i}\}_{i\in[N]}$), we use SGD to solve empirical risk minimization problems of the form,

using the following iterative update rule for $T$ rounds:

Can we expect a problem to have low gradient confusion in practice? From the logistic regression problem, we have: $|\langle\nabla f_{i}(\mathbf{w}),\nabla f_{j}(\mathbf{w})\rangle|=|\langle\mathbf{x}_{i},\mathbf{x}_{j}\rangle|\cdot|\zeta_{i}\zeta_{j}|.$ This inner product is expected to be small for all $\mathbf{w}$; the logistic loss satisfies $|\zeta_{i}\zeta_{j}|<1$, and for fixed $N$ the quantity $\max_{ij}|\langle\mathbf{x}_{j},\mathbf{x}_{i}\rangle|$ is $O(1/\sqrt{d})$ whenever $\{\mathbf{x}_{i}\}$ are randomly sampled from a sphere (see lemma B.1 for the formal statement). Generally, this is true whenever $\mathbf{x}_{i}=\frac{1}{\sqrt{d}}\mathbf{y}_{i},$ where $\mathbf{y}_{i}$ is an isotropic random vector (Vershynin 2018). Thus, we would expect a random linear model to have nearly orthogonal gradients, when the number of parameters is "large" and the number of training data is "small", i.e., when the model is over-parameterized. This is further evidenced by a toy example in figure 1, where we show a slightly over-parameterized linear regression model can have much faster convergence rates, as well as lower gradient confusion. One can prove a similar result for problems that have random and low-rank Hessians, which suggests that one might expect gradient to be small near the minimizer for many standard neural nets (see appendix C for more discussion).

The above arguments are a bit simplistic, considering toy scenarios and ignoring issues like the effect of network structure. In the following sections, we rigorously analyze the effect of gradient confusion on the speed of convergence on non-convex problems, and the effect of width and depth of the neural network architecture on the gradient confusion.

SGD is fast when gradient confusion is low

We first consider functions satisfying the Polyak-Lojasiewicz (PL) inequality (Lojasiewicz 1965), a condition related to, but weaker than, strong convexity, and used in recent work (Karimi et al. 2016; De et al. 2017). We provide bounds on the rate of convergence in terms of the optimality gap. We start with two standard assumptions.

$\{f_{i}\}_{i\in[N]}$ are Lipschitz smooth: $f_{i}(\mathbf{w}^{\prime})\leq f_{i}(\mathbf{w})+\nabla f_{i}(\mathbf{w})^{\top}(\mathbf{w}^{\prime}-\mathbf{w})+\frac{L}{2}\|\mathbf{w}^{\prime}-\mathbf{w}\|^{2}.$

$\{f_{i}\}_{i\in[N]}$ satisfy the PL inequality: $\frac{1}{2}\|\nabla f_{i}(\mathbf{w})\|^{2}\geq\mu(f_{i}(\mathbf{w})-f_{i}^{\star}),$ $f_{i}^{\star}=\min_{\mathbf{w}}f_{i}(\mathbf{w})$.

We now state a convergence result of constant learning rate SGD in terms of the gradient confusion.

If the objective function satisfies (A1) and (A2), and has gradient confusion $\eta$, SGD converges linearly to a neighborhood of the minima of problem (1) as:

where $\alpha<\frac{2}{NL}$, $\rho=1-\frac{2\mu}{N}\big(\alpha-\frac{NL\alpha^{2}}{2}\big)$, $F^{\star}=\min_{{\mathbf{w}}}F({\mathbf{w}})$ and $\mathbf{w}_{0}$ is the initialized weights.

This result shows that SGD converges linearly to a neighborhood of a minimizer, and the size of this neighborhood depends on the level of gradient confusion. When the gradient confusion is small, i.e., $\eta=O(\epsilon)$, SGD has fast convergence to $O(\epsilon)$-accuracy in $T=O(\log(1/\epsilon))$ iterations, without requiring the learning rate to be vanishingly small. We now extend this to general smooth functions.

If the objective satisfies (A1) and has gradient confusion $\eta$, then SGD converges to a neighborhood of a stationary point of problem (1) as:

for $\alpha<\frac{2}{NL}$, $\rho=\frac{2N}{2-NL\alpha}$, and $F^{\star}=\min_{{\mathbf{w}}}F({\mathbf{w}})$.

Thus, as long as $\eta=O(1/T)$, SGD has fast $O(1/T)$ convergence on smooth non-convex functions. Theorems 3.1 and 3.2 predict an initial phase of optimization with fast convergence to the neighborhood of a minimizer or a stationary point. This behavior is often observed when optimizing neural nets (Darken & Moody 1992; Sutskever et al. 2013), where a constant learning rate reaches a high level of accuracy on the model. As we show in subsequent sections, this is expected since for neural networks typically used, the gradient confusion is expected to be low. See section 9 for more discussion on the above results and how they relate to previous work. We stress that our goal is not to study convergence rates per se, nor is it to prove state-of-the-art rate bounds for this class of problems. Rather, we show the direct effect that the gradient confusion bound has on the convergence rate and the noise floor for constant learning rate SGD. As we show in the following sections, this new perspective in terms of the gradient confusion helps us more directly understand how neural network architecture design affects SGD dynamics and why.

Effect of neural network architecture at Gaussian initializations

To draw a connection between neural network architecture and training performance, we analyze gradient confusion for generic (i.e., random) model problems using methods from high-dimensional probability. In this section, we analyze the effect of neural network architecture at the beginning of training, when using standard Gaussian initialization techniques. Analyzing these models at initialization is important to understand which architectures are more easily trainable than others. Our results cover a wide range of scenarios compared to prior work, require minimal additional assumptions, and hold for a large family of neural networks with different non-linear activation functions and loss-functions. In particular, our results hold for fully connected networks (and can be extended to convolutional networks) with the square-loss and logistic-loss functions, and commonly used non-linear activations such as sigmoid, tanh and ReLU. We consider both the case where the input data is arbitrary but bounded (theorem 4.1, part 1), as well as where the input data is randomly drawn from the surface of a unit sphere (theorem 4.1, part 2).

where $\sigma$ denotes the non-linear activation function applied point-wise to its arguments. We assume that the activation is given by a function $\sigma(x)$ with the following properties.

(P1) Boundedness: $|\sigma(x)|\leq 1$ for $x\in$.

(P2) Bounded differentials: Let $\sigma^{\prime}(x)$ and $\sigma^{\prime\prime}(x)$ denote the first and second sub-differentials respectively. Then, $|\sigma^{\prime}(x)|\leq 1$ and $|\sigma^{\prime\prime}(x)|\leq 1$ for all $x\in$.

When $\|\mathbf{x}\|\leq 1$, activation functions such as sigmoid, tanh, softmax and ReLU satisfy these requirements.

Furthermore, in this section, we consider the following Gaussian weight initialization strategy.

This initialization strategy with different settings of $\kappa$ are used almost universally for neural networks (Glorot & Bengio 2010; LeCun et al. 2012; He et al. 2015). For instance, typically $\kappa=\frac{1}{2}$ when ReLU activations are used, and $\kappa=1$ when tanh activations are used.

Let $\mathbf{W}_{0},\mathbf{W}_{1},\ldots,\mathbf{W}_{\beta}$ be weight matrices chosen according to strategy 4.1. There exists fixed constants $c_{1},c_{2}>0$ such that we have the following.

Consider a fixed but arbitrary dataset $\mathbf{x}_{1},\mathbf{x}_{2},\ldots,\mathbf{x}_{N}$ with $\|\mathbf{x}_{i}\|\leq 1$ for every $i\in[N]$. For $\eta>4$, the gradient confusion bound in equation 3 holds with probability at least

If the dataset $\{\mathbf{x}_{i}\}_{i\in[N]}$ is such that each $\mathbf{x}_{i}$ is an i.i.d. sample from the surface of $d$-dimensional unit sphere, then for every $\eta>0$ the gradient confusion bound in equation 3 holds with probability at least

Theorem 4.1 shows that under popular Gaussian initializations used, training becomes harder as networks get deeper. The result however also shows a way forward: layer width improves the trainability of deep networks. Other related work supports this showing that when the layers are infinitely wide, the learning dynamics of gradient descent simplifies considerably (Jacot et al. 2018; Lee et al. 2019). Hanin & Rolnick 2018 also suggest that the width should increase linearly with depth in a neural network to help dynamics at the beginning of training. In section 7 and appendix A, we show substantial empirical evidence that, given a sufficiently deep network, increasing the layer width often helps in lowering gradient confusion and speeding up convergence for a range of models.

A more general result on the effect of depth

While our results in section 4 hold at standard initialization schemes, in this section we derive a more general version of the result. In particular, we assume the setting where the data is drawn uniformly at random from a unit sphere and the weights lie in a ball around a local minimizer. Our results hold for both fully connected networks and convolutional networks with the square-loss and logistic-loss functions, and commonly-used non-linear activations such as sigmoid, tanh, softmax and ReLU.

We consider the same setup as in the previous section, and assume additionally that the data points $\{\mathbf{x}_{i}\}$ are drawn uniformly from the surface of a $d$-dimensional unit sphere. Additionally, instead of studying the network at initialization, we make the following assumption on the weights.

We assume that the operator norm of the weight matrices $\{\mathbf{W}_{i}\}_{i\in[\beta]_{0}}$ are bounded above by $1$, i.e., for every $i\in[\beta]_{0}$ we have $\|\mathbf{W}_{i}\|\leq 1$.

The operator norm of the weight matrices $\|{\mathbf{W}}\|$ being close to 1 is important for the trainability of neural networks, as it ensures that the input signal is passed through the network without exploding or shrinking across layers (Glorot & Bengio 2010). Proving non-vacuous bounds in case of such blow-ups in magnitude of the signal or the gradient is not possible in general, and thus, we consider this restricted class of weights. Most standard neural networks are trained using weight decay regularizers of the form $\sum_{i}\|W_{i}\|_{F}^{2}$. This biases the weights to be small when training neural networks in practice. See appendix F for further discussion on the small weights assumption.

We now present a more general version of theorem 4.1.

Let $\mathbf{W}_{0},\mathbf{W}_{1},\ldots,\mathbf{W}_{\beta}$ satisfy assumption 1. For some fixed constant $c>0$, the gradient confusion bound (equation 3) holds with probability at least

Theorem 5.1 shows that (for fixed dimension $d$ and number of samples $N$) when the depth $\beta$ decreases, the probability that the gradient confusion bound in equation 3 holds increases, and vice versa. Thus, our results indicate that in the general case when the weights are small, increasing the network depth will typically lead to slower model training.

Note that on assuming $\|{\mathbf{W}}\|\leq 1$ for each weight matrix ${\mathbf{W}}$, the dependence of gradient confusion on the width goes away. To see why this, consider an example where each weight matrix in the neural network has exactly one non-zero element, which is set to 1. The operator norm of each such weight matrix is 1, but the forward or backward propagated signals would not depend on the width.

Note that the convergence rate results of SGD in section 3 assume that the gradient confusion bound holds at every point along the path of SGD. On the other hand, theorem 5.1 shows concentration bounds for the gradient confusion at a fixed weight ${\mathbf{W}}$. Thus, to make the above result more relevant for the convergence of SGD on neural networks, we now make the concentration bound in theorem 5.1 uniform over all weights inside a ball $\mathcal{B}_{r}$ of radius $r$.

Select a point $\mathbf{W}=(\mathbf{W}_{0},\mathbf{W}_{1},\ldots,\mathbf{W}_{\beta})$, satisfying assumption 1. Consider a ball $\mathcal{B}_{r}$ centered at $\mathbf{W}$ of radius $r>0$. If the data $\{{\mathbf{x}_{i}}\}_{i\in[N]}$ are sampled uniformly from a unit sphere, then the gradient confusion bound in equation 3 holds uniformly at all points $\mathbf{W}^{\prime}\in\mathcal{B}_{r}$ with probability at least

Corollary 5.1 shows that the probability that the gradient confusion bound holds decreases with increasing depth, for all weights in a ball around the minimizer. The above results automatically hold for convolutional networks, since a convolution operation on ${\mathbf{x}}$ can be represented as a matrix multiplication $\mathbf{U}\mathbf{x}$ for an appropriate Toeplitz matrix $\mathbf{U}$. This explains why, in the general case, training very deep models might always be hard. This raises the question why most deep neural networks used in practice are so efficiently trained using SGD. While careful Gaussian initialization strategies prevent vanishing or exploding gradients, these strategies still suffer from high gradient confusion for very deep networks unless the width is also increased with the depth, as we show in section 4. Practitioners over the years, however, have achieved state-of-the-art results by making networks deeper, without necessarily making networks wider. Thus, in section 7, we empirically study how popular techniques used in these models like skip connections and batch normalization affect gradient confusion. We find that these techniques drastically lower gradient confusion, making deep networks significantly easier to train. Furthermore, in the next section, we show how deep linear nets are trainable when used with orthogonal initialization techniques, indicating a way forward for training deeper models.

Gradient confusion is independent of depth for orthogonal initializations

In this section, we show that for deep linear neural networks, gradient confusion is independent of depth when the weight matrices are initialized as orthogonal matrices. An orthogonal matrix $\mathbf{A}$ satisfies $\mathbf{A}^{T}\cdot\mathbf{A}=\mathbf{A}\cdot\mathbf{A}^{T}=\mathbf{I}$. Consider the following linear neural network:

where the rescaling parameter $\gamma=\frac{1}{\sqrt{2\beta}}$, and assume we use the squared loss function. Then we have the following.

Let $\{\mathbf{W}_{i}\}_{i\in[\beta]}$ be arbitrary orthogonal matrices that satisfy assumption 1. Let the dataset $\{\mathbf{x}_{i}\}_{i\in[N]}$ be such that each $\mathbf{x}_{i}$ is an i.i.d. sample from the surface of $d$-dimensional unit sphere. Consider the linear neural network in equation 4 that minimizes the empirical square loss function. For some fixed constant $c>0$, the gradient confusion bound (equation 3) holds with probability at least

Experimental results

To test our theoretical results and to probe why standard neural networks are efficiently trained with SGD, we now present experimental results showing the effect of the neural network architecture on the convergence of SGD and gradient confusion. It is worth noting that theorems 3.1 and 3.2 indicate that we would expect the effect of gradient confusion to be most prominent closer to the end of training.

Effect of depth. To test our theoretical results, we consider CNNs with a fixed width factor of 2 and varying network depth. From figure 2, we see that our theoretical results are backed by the experiments: increasing depth slows down convergence, and increases gradient confusion. We also notice that with increasing depth, the density of pairwise gradient cosine similarities concentrates less sharply around 0, which makes the network harder to train.

Effect of width. We now consider CNNs with a fixed depth of 16 and varying width factors. From figure 3, we see that increasing width results in faster convergence and lower gradient confusion. We further see that gradient cosine similarities concentrate around 0 with growing width, indicating that SGD decouples across the training samples with growing width. Note that the smallest network considered (CNN-16-2) is still over-parameterized and achieves a high level of performance (see appendix A.3).

Effect of batch normalization and skip connections. Almost all state-of-the-art neural networks currently contain both skip connections and normalization layers. To help understand why such neural networks are so efficiently trained using SGD with constant learning rates, we test the effect of adding skip connections and batch normalization to CNNs of fixed width and varying depth. Figure 4 shows that adding skip connections or batch normalization individually help in training deeper models, but these models still suffer from worsening results and increasing gradient confusion as the network gets deeper. When these techniques are used together, the model has relatively low gradient confusion even for very deep networks, significantly improving trainability of deep models. Note that our observations are consistent with prior work (De & Smith 2020; Yang et al. 2019).

Alternate definitions of gradient confusion

Note that the gradient confusion bound $\eta$ in equation 3 is defined for the worst-case gradient inner product. However, all the results in this paper can be trivially extended to using a bound on the average gradient inner product of the form:

In this case, all theoretical results would remain the same up to constants. We can also define a normalized variant of the gradient confusion condition:

This condition inherently makes an additional assumption that the norm of the stochastic gradients, $\|\nabla f_{i}(\mathbf{w})\|$, is bounded, and thus the gradient variance is also bounded (see discussion in section 9). Thus, while all our theoretical results would qualitatively remain the same under this condition, we can prove tighter versions of our current results.

Finally, note that gradient confusion condition in equation 3 is applicable even when the stochastic gradients are averaged over minibatches of size $B$. The variance of the gradient inner product scales down as $1/B^{2}$ in this case, and thus $\eta$ is expected to decrease as $B$ grows.

Related work

The gradient confusion bound and our theoretical results have interesting connections to prior work. In this section, we briefly discuss some of these connections.

Connections to gradient diversity: Gradient diversity (Yin et al. 2017) also measures the degree to which individual gradients at different data samples are different from each other. However, the gradient diversity measure gets larger as the individual gradients become orthogonal to each other, and further increases as the gradients start pointing in opposite directions. On the other hand, gradient confusion between two individual gradients is zero unless the inner product between them is negative. As we show in this paper, this has important implications when we study the convergence of SGD in the over-parameterized setting: increased width makes gradients more orthogonal to each other improving trainability, while increased depth result in gradients pointing in opposite directions making networks harder to train. Thus, we view our papers to be complementary, providing insights about different issues (large batch distributed training vs. small batch convergence).

Related work on the impact of network architecture: Balduzzi et al. 2017 studied neural networks with ReLU activations at Gaussian initializations, and showed that gradients become increasingly negatively correlated with depth. Hanin 2018 showed that the variance of gradients in fully connected networks with ReLU activations is exponential in the sum of the reciprocals of the hidden layer widths at Gaussian initializations. In a follow-up work, Hanin & Rolnick 2018 showed that this sum of the reciprocals of the hidden layer widths determines the variance of the sizes of the activations at each layer. When this sum of reciprocals is too large, early training dynamics are very slow, suggesting the difficulties of starting training on deeper networks, as well as the benefits of increased width.

Other work on SGD convergence: There has recently been a lot of interest in analyzing conditions under which SGD converges to global minimizers of over-parameterized linear and non-linear neural networks. Arora et al. 2018 shows SGD converges linearly to global minimizers for linear neural networks under certain conditions. Du et al. 2018; Allen-Zhu et al. 2018; Zou et al. 2018; Brutzkus et al. 2017 also show convergence to global minimizers of SGD for non-linear networks. This paper complements these recent results by studying how low gradient confusion contributes to SGD’s success on over-parameterized neural networks used in practice.

Discussion

In this paper, we study how neural network architecture affects the trainability of networks and the dynamics of SGD. To rigorously analyze this, we introduce a concept called gradient confusion, and show that when gradient confusion is low, SGD has fast convergence. We show at standard Gaussian initializations, increasing layer width leads to lower gradient confusion, making the model easier to train. In contrast, increasing depth results in higher gradient confusion, making models harder to train. These results indicate that increasing the layer width with the network depth is important to maintain trainability of the neural network. This is supported by other recent work that suggest that the width should increase linearly with depth in a Gaussian-initialized neural network to help dynamics early in training (Hanin 2018; Hanin & Rolnick 2018).

Many previous results have shown how deeper models are more efficient at modeling higher complexity function classes than wider models, and thus depth is essential for the success of neural networks (Eldan & Shamir 2016; Telgarsky 2016). Indeed, practitioners over the years have achieved state-of-the-art results on various tasks by making networks deeper, without necessarily making networks wider. We thus study techniques that enable us to train deep models without requiring us to increase the width with depth. Most state-of-the-art neural networks currently contain both skip connections and normalization layers. We thus, empirically study the effect of introducing batch normalization and skip connections to a neural network. We show that the combination of batch normalization and skip connections lower gradient confusion and help train very deep models, explaining why many neural networks used in practice are so efficiently trained. Furthermore, we show how orthogonal initialization techniques provide a promising direction for improving the trainability of very deep networks.

Our results provide a number of important insights that can be used for neural network model design. We demonstrate that the gradient confusion condition could be useful as a measure of trainability of networks, and thus could potentially be used to develop algorithms for more efficient training. Additionally, the correlation between gradient confusion and the test set accuracies shown in appendix A suggest that an interesting topic for future work would be to investigate the connection between gradient confusion and generalization (Fort et al. 2019). Our results also suggest the importance of further work on orthogonal initialization schemes for neural networks with non-linear activations that make training very deep models possible.

Acknowledgements

The authors thank Brendan O’Donoghue, Aleksandar Botev, James Martens, Sudha Rao, and Samuel L. Smith for helpful discussions and for reviewing earlier versions of this manuscript. This paper was supported by the ONR MURI program, AFOSR MURI Program, and the National Science Foundation DMS directorate.

References

Appendix

We first briefly outline the different sections in the appendix.

In appendix A, we provide details of our experimental setup, and provide additional empirical results on fully connected networks, convolutional networks and residual networks with the MNIST, CIFAR-10 and CIFAR-100 datasets.

In appendix B, we state and prove a lemma on the near orthogonality of random vectors, which we refer to in the main text. This result is often attributed to Milman & Schechtman 1986.

In appendix C, we provide some intuition on why many standard over-parameterized neural networks with low-rank Hessians might have low gradient confusion for a large set of weights near the minimizer.

In appendix D, we provide the proofs of the theorems presented in the main section. In appendix D.1, we provide proofs of theorems 3.1 and 3.2. In appendix D.2, we provide the proof of lemma D.1, which we refer to in the main text. In appendix D.3, we provide proofs of theorem 5.1 and corollary 5.1. In appendix D.4, we provide the proof of theorem 4.1. In appendix D.5, we provide the proof of theorem 6.1.

In appendix E, we briefly describe a few lemmas that we require in our analysis.

In appendix F, we discuss the small weights assumption (assumption 1), which is required for theorem 5.1, corollary 5.1 and theorem 6.1 in the main text.

Appendix A Additional experimental results

In this section, we present more details about our experimental setup, as well as, additional experimental results on a range of models (MLPs, CNNs and Wide ResNets) and a range of datasets (MNIST, CIFAR-10, CIFAR-100).

To further test the main claims in the paper, we performed additional experiments on an image classification problem on the MNIST dataset using fully connected neural networks. We iterated over neural networks of varying depth and width, and considered both the identity activation function (i.e., linear neural networks) and the tanh activation function. We also considered two different weight initializations that are popularly used and appropriate for these activation functions:

The Glorot normal initializer (Glorot & Bengio 2010) with weights initialized by sampling from the distribution $\mathcal{N}\big(0,2/(\text{fan-in}+\text{fan-out})\big)$, where fan-in denotes the number of input units in the weight matrix, and fan-out denotes the number of output units in the weight matrix.

The LeCun normal initializer (LeCun et al. 2012) with weights initialized by sampling from the distribution $\mathcal{N}\big(0,1/\text{fan-in}\big)$.

This relatively simple model gave us the ability to iterate over a large number of combinations of network architectures of varying width and depth, and different activation functions and weight initializations. Linear neural networks are an efficient way to directly understand the effect of changing depth and width without increasing model complexity over linear regression. Thus, we considered both linear and non-linear neural networks in our experiments.

We used SGD with constant learning rates for training with a mini-batch size of 128 and trained each model for 40000 iterations (more than 100 epochs). The constant learning rate $\alpha$ was tuned over a logarithmically-spaced grid:

We ran each experiment 10 times (making sure at least 8 of them ran till completion), and picked the learning rate that achieved the lowest training loss value on average at the end of training. Our grid search was such that the optimal learning rate never occurred at one of the extreme values tested.

To measure gradient confusion at the end training, we sampled 1000 pairs of mini-batches each of size 128 (the same size as the training batch size). We calculated gradients on each of these pairs of mini-batches, and then calculated the cosine similarity between them. To measure the worse-case gradient confusion, we computed the lowest gradient cosine similarity among all pairs. We explored the effect of changing depth and changing width on the different activation functions and weight initializations. We plot the final training loss achieved for each model and the minimum gradient cosine similarities calculated over the 1000 pairs of gradients at the end of training. For each point, we plot both the mean and the standard deviation over the 10 independent runs.

Figure 5 shows results on neural networks with identity and tanh activation functions for the two weight initializations considered (Glorot normal and LeCun normal). Similar to the experimental results in section 7, and matching our theoretical results in sections 4 and 5, we notice the consistent trend of gradient confusion increasing with increasing depth. This makes the networks harder to train with increasing depth, and this is evidenced by an increase in the final training loss value. By depth $\beta=1000$, the increased gradient confusion effectively makes the network untrainable when using tanh non-linearities.

The effect of width. We explored the effect of width by varying the width of the neural network while keeping the depth fixed at $\beta=300$. We chose a very deep model, which is essentially untrainable for small widths (with standard initialization techniques) and helps better illustrate the effects of increasing width. We varied the width of the network, again on the log scale, as:

A.2 Additional experimental details for CNNs and WRNs

In this section, we review the details of our setup for the image classification experiments on CNNs and WRNs on the CIFAR-10 and CIFAR-100 datasets.

For cleaner figures, we sometimes plot a subset of these results: $\beta\in\{16,28,40,52,76,100\}.$ To study the effect of width, we considered WRNs with depth $\beta=16$ and width factor varying as:

Convolutional neural networks

To study the effect of width, we considered CNNs with depth $\beta=16$ and width factor varying as:

Hyperparameter tuning and other details

We used SGD as the optimizer without any momentum. Following Zagoruyko & Komodakis 2016, we ran all experiments for 200 epochs with minibatches of size 128, and reduced the initial learning rate by a factor of 10 at epochs 80 and 160. We turned off weight decay for all our experiments.

We ran each individual experiment 5 times. We ignored any runs that were unable to decrease the loss from its initial value. We also made sure at least 4 out of the 5 independent runs ran till completion. When the learning rate is close to the threshold at which training is still possible, some runs may converge, while others may fail to converge. Thus, these checks ensure that we pick a learning rate that converges reliably in most cases on each problem. We show the standard deviation across runs in our plots.

We tuned the optimal initial learning rate for each model over a logarithmically-spaced grid:

and selected the run that achieved the lowest final training loss value (averaged over the independent runs). Our grid search was such that the optimal learning rate never occurred at one of the extreme values tested. We used the standard train-valid-test splits of 40000-10000-10000 for CIFAR-10 and CIFAR-100.

To measure gradient confusion, at the end of every training epoch, we sampled 100 pairs of mini-batches each of size 128 (the same size as the training batch size). We calculated gradients on each mini-batch, and then computed pairwise cosine similarities. To measure the worse-case gradient confusion, we computed the lowest gradient cosine similarity among all pairs. We also show the kernel density estimation of the pairwise gradient cosine similarities of the 100 minibatches sampled at the end of training (after 200 epochs), to see the concentration of the distribution. To do this, we combine together the 100 samples for each independent run and then perform kernel density estimation with a gaussian kernel on this data.

A.3 Additional plots for CIFAR-10 on CNNs

In section 7, we showed results for image classification using CNNs on CIFAR-10. In this section, we show some additional plots for this experiment. Figure 7 shows the effect of changing the depth, while figure 8 shows the effect of changing the width factor of the CNN. We see that the final training loss and test set accuracy values show the same trends as in section 7: deeper networks are harder to train, while wider networks are easier to train. As mentioned previously, theorems 3.1 and 3.2 indicate that we would expect the effect of gradient confusion to be more prominent near the end of training. From the plots we see that deeper networks have higher gradient confusion close to minimum, while wider networks have lower gradient confusion close to the minimum.

A.4 CIFAR-100 on CNNs

We now consider image classifications tasks with CNNs on the CIFAR-100 dataset. Figure 9 shows the effect of varying depth, while figure 10 shows the effect of varying width. We notice the same trends as in our results with CNNs on CIFAR-10. Interestingly, from the width results in figure 10, we see that while there is no perceptible change to the minimum pairwise gradient cosine similarity, the distribution still sharply concentrates around 0 with increasing width. Thus more gradients become orthogonal to each other with increasing width, implying that SGD on very wide networks becomes closer to decoupling over the data samples.

A.5 Image classification with WRNs on CIFAR-10 and CIFAR-100

We now show results for image classification problems using wide residual networks (WRNs) on CIFAR-10 and CIFAR-100. The WRNs we consider do not have any batch normalization. Later we show results on the effect of adding batch normalization to these networks.

Figures 11 and 12 show results on the effect of depth using WRNs on CIFAR-10 and CIFAR-100 respectively. We again see the consistent trend of deeper networks having higher gradient confusion, making them harder to train. We further see that increasing depth results in the pairwise gradient cosine similarities concentrating less around 0.

Figures 13 and 14 show results on the effect of width using WRNs on CIFAR-10 and CIFAR-100 respectively. We see that increasing width typically lowers gradient confusion and helps the network achieve lower loss values. The pairwise gradient cosine similarities also typically concentrate around 0 with higher width. We also notice from these figures that in some cases, increasing width might lead to diminishing returns, i.e., the benefits of increased width diminish after a certain point, as one would expect.

A.6 Effect of batch normalization and skip connections

In section 7 we showed results on the effect of adding batch normalization and skip connections to CNNs and WRNs on an image classification task on CIFAR-10. In this section, we present similar results for image classification on CIFAR-100. Similar to section 7, figure 15 shows that adding skip connections or batch normalization individually help in training deeper models, but these models still suffer from worsening results and increasing gradient confusion as the network gets deeper. Both these techniques together keep the gradient confusion relatively low even for very deep networks, significantly improving trainability of deep models.

Appendix B Near orthogonality of random vectors

For completeness, we state and prove below a lemma on the near orthogonality of random vectors. This result is often attributed to Milman & Schechtman 1986.

For vectors $\{\mathbf{x_{i}}\}_{i\in[N]}$ drawn uniformly from a unit sphere in $d$ dimensions, and $\nu>0,$

Given a fixed vector $\mathbf{x},$ a uniform random vector $\mathbf{y}$ satisfies $|\mathbf{x}^{\top}\mathbf{y}|\geq\nu$ only if $\mathbf{y}$ lies in one of two spherical caps: one centered at $\mathbf{x}$ and the other at $-\mathbf{x},$ and both with angular radius $\cos^{-1}(\nu)\leq\frac{\pi}{2}-\nu.$ A simple result often attributed to Milman & Schechtman 1986 bounds the probability of lying in either of these caps as

Because of rotational symmetry, the bound (5) holds if both $\mathbf{x}$ and $\mathbf{y}$ are chosen uniformly at random.

We next apply a union bound to control the probability that $|\mathbf{x}_{i}^{\top}\mathbf{x}_{j}|\geq\nu$ for some pair $(i,j).$ There are fewer than $N^{2}/2$ such pairs, and so the probability of this condition is

Appendix C Low-rank Hessians lead to low gradient confusion

In this section, we show that low-rank random Hessians result in low gradient confusion. For clarity in presentation, suppose each $f_{i}$ has a minimizer at the origin (the same argument can be easily extended to the more general case). Suppose also that there is a Lipschitz constant for the Hessian of each function $f_{i}$ that satisfies $\|\mathbf{H}_{i}(\mathbf{w})-\mathbf{H}_{i}(\mathbf{w^{\prime}})\|\leq L_{H}\|\mathbf{w}-\mathbf{w^{\prime}}\|$ (note that this is a standard optimization assumption (Nesterov 2018), with evidence that it is applicable for neural networks (Martens 2016)). Then $\nabla f_{i}(\mathbf{w})=\mathbf{H}_{i}\mathbf{w}+\mathbf{e}$, where $\mathbf{e}$ is an error term bounded as: $\|\mathbf{e}\|\leq\frac{1}{2}L_{H}\|\mathbf{w}\|^{2},$ and we use the shorthand $\mathbf{H}_{i}$ to denote $\mathbf{H}_{i}(\mathbf{0}).$ Then we have:

There is evidence that the Hessian at the minimizer is very low rank for many standard over-parameterized neural network models (Sagun et al. 2017; Cooper 2018; Chaudhari et al. 2016; Wu et al. 2017; Ghorbani et al. 2019). While a bit non-rigorous, the above result nonetheless suggests that for many standard neural network models, the gradient confusion might be small for a large class of weights near the minimizer.

Appendix D Missing proofs

This section presents proofs for the convergence theorems of SGD presented in section 3, under the assumption of low gradient confusion. For clarity of presentation, we re-state each theorem before its proof.

where the second-last inequality follows from definition 2.1. Let the learning rate $\alpha<2/NL$. Then, using assumption (A2) and subtracting by $F^{\star}=\min_{\mathbf{w}}F(\mathbf{w})$ on both sides, we get

Writing $\rho=1-\frac{2\mu\alpha}{N}+\mu L\alpha^{2}$, and unrolling the iterations, we get

Thus, taking expectation and assuming the step size $\alpha<2/(NL)$, we can rewrite equation 6 as:

Taking an average over $T$ iterations, and using $F^{\star}=\min_{\mathbf{w}}F(\mathbf{w})$, we get:

D.2 Proof of lemma D.1

Consider the set of loss-functions $\{f_{i}(\mathbf{W})\}_{i\in[N]}$ where all $f_{i}$ are either the square-loss function or the logistic-loss function. Recall that $f_{i}(\mathbf{W}):=f(\mathbf{W},\mathbf{x}_{i})$. Consider a feed-forward neural network as defined in equation 4 whose weights $\mathbf{W}$ satisfy assumption 1. Consider the gradient $\nabla_{\mathbf{W}}f_{i}(\mathbf{W})$ of each function $f_{i}$. From definition we have that $\nabla_{\mathbf{W}}f_{i}(\mathbf{W})=\zeta_{\mathbf{x}_{i}}(\mathbf{W})\nabla_{\mathbf{W}}g_{\mathbf{W}}(\mathbf{x}_{i})$, where we define $\zeta_{\mathbf{x}_{i}}(\mathbf{W})=\partial f_{i}(\mathbf{W})/\partial g_{\mathbf{W}}$. Then we have the following properties.

When $\|\mathbf{x}\|\leq 1$ for every $p\in[\beta]$ we have $\|\nabla_{\mathbf{W}_{p}}g_{\mathbf{W}}(\mathbf{x}_{i})\|\leq 1$.

There exists $0<\zeta_{0}\leq 2\sqrt{\beta}$, such that $|\zeta_{\mathbf{x}_{i}}(\mathbf{W})|\leq 2\,$, $\|\nabla_{\mathbf{x}_{i}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})\|_{2}\leq\zeta_{0}\,$, $\|\nabla_{{\mathbf{W}}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})\|_{2}\leq\zeta_{0}$.

The first property is a direct consequence of assumption 1 and property (P2) of the activation function.

Let $\mathbf{W}$ denote the tuple $(\mathbf{W}_{p})_{p\in[\beta]_{0}}$. Consider $|\zeta_{\mathbf{x}_{i}}(\mathbf{W})|=|\partial f_{i}(\mathbf{W})/\partial g_{\mathbf{W}}|$. In the case of square-loss function this evaluates to $|g_{\mathbf{W}}(\mathbf{x})-\mathcal{C}(\mathbf{x})|\leq 2$. In case of logistic regression, this evaluates to $|\frac{-1}{1+\exp(\mathcal{C}(\mathbf{x}_{i})g_{\mathbf{W}}(\mathbf{x}_{i}))}|\leq 1$. Now we consider $\|\nabla_{\mathbf{x}_{i}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})\|$. Consider the squared loss function. We then have the following.

Likewise, consider the logistic-loss function. We then have the following.

Thus, it suffices to bound $\|\nabla_{\mathbf{x}_{i}}g_{\mathbf{W}}(\mathbf{x}_{i})\|$. Using assumption 1 and the properties (P1), (P2) of $\sigma$, this can be upper-bounded by $1$.

Consider $\nabla_{\mathbf{W}_{p}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})$ for some layer index $p\in[\beta]_{0}$. We will show that $\|\nabla_{\mathbf{W}_{p}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})\|_{2}\leq 2$. Then it immediately follows that $\|\nabla_{\mathbf{W}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})\|_{2}\leq 2\sqrt{\beta}$. In the case of a squared loss function. We have the following.

Likewise, consider the logistic-loss function. We then have the following.

Since $\|\nabla_{\mathbf{W}_{p}}g_{\mathbf{W}}(\mathbf{x}_{i})\|\leq 1$, we have that $\|\nabla_{\mathbf{W}_{p}}\zeta_{\mathbf{x}_{i}}(\mathbf{W})\|\leq 2$ in both the cases. Thus, $\zeta_{0}=2\sqrt{\beta}$. ∎

D.3 Proofs of theorem 5.1 and corollary 5.1

In this section, we will present the proofs of theorem 5.1 and corollary 5.1.

We show two key properties, namely bounded gradient and non negative expectation. We will then use both these properties to complete the proof.

Bounded gradient. For every $i\in[n]$ define $\zeta_{\mathbf{x}_{i}}(\mathbf{W}):=f^{\prime}(\mathbf{W})$. For every $p\in[\beta]$ define $\mathbf{H}_{p}$ as follows.

Fix an $i\in[N]$. Then we have the following recurrence

Then the gradients can be written in terms of the above quantities as follows.

We can write, the gradient confusion denote by $h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})$, as follows.

We will now bound $\|\nabla_{(\mathbf{x}_{i},\mathbf{x}_{j})}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|_{2}$. Consider $\nabla_{\mathbf{x}_{i}}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})$. This can be written as follows.

Observe that each of the entries in the diagonal matrix $\operatorname{Diag}(\sigma^{\prime}(\mathbf{H}_{p}(\mathbf{x}_{i})))$ is at most $1$. Thus, we have that $\|\operatorname{Diag}(\sigma^{\prime}(\mathbf{H}_{p}(\mathbf{x}_{i})))\|\leq 1$.

Consider $\|\nabla_{\mathbf{x}_{i}}\left(\mathbf{H}_{p-1}(\mathbf{x}_{i})\cdot\mathbf{g}_{p}(\mathbf{x}_{i})^{\top}\cdot\mathbf{g}_{p}(\mathbf{x}_{j})\cdot\mathbf{H}_{p-1}(\mathbf{x}_{i})\right)\|$ for every $p\in[\beta]_{0}$.

Note that $\nabla_{\mathbf{x}_{i}}\mathbf{H}_{p-1}(\mathbf{x}_{i})=\mathbf{g}_{1}(\mathbf{x}_{i})\cdot\operatorname{Diag}(\sigma^{\prime}(\mathbf{W}_{0}\cdot\mathbf{x}_{i}))\cdot\mathbf{W}_{0}^{\top}\cdot\mathbf{g}_{p}(\mathbf{x}_{i})^{\top}$. Thus, $\|\nabla_{\mathbf{x}_{i}}\mathbf{H}_{p-1}(\mathbf{x}_{i})\|\leq 1$. We will now show that $\|\nabla_{\mathbf{x}_{i}}\mathbf{g}_{p}(\mathbf{x}_{i})\|\leq\beta-p+1$. We prove this inductively. Consider the base case when $p=\beta$.

Thus, using equation 8 and the above arguments, we obtain, $\|\nabla_{\mathbf{x}_{i}}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|_{2}\leq\zeta_{0}^{2}(\beta+1)+\zeta_{0}^{2}(\beta+1)(\beta+2)\leq 2\zeta_{0}^{2}(\beta+2)^{2}$ and thus, $\|\nabla_{(\mathbf{x}_{i},\mathbf{x}_{j})}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|_{2}\leq 4\zeta_{0}^{2}(\beta+2)^{2}$.

We have used the fact that $\nabla f_{i}(\mathbf{W})$ and $\nabla f_{j}(\mathbf{W})$ are identically distributed and independent.

Concentration of Measure. We combine the two properties as follows. From Non-negative Expectation property and equation 26, we have that

To obtain the probability that some value of $h_{\mathbf{w}}(\nabla_{\mathbf{w}}f_{i},\nabla_{\mathbf{w}}f_{j})$ lies below $-\eta,$ we use a union bound. There are $N(N-1)/2<N^{2}/2$ possible pairs of data points to consider, and so this probability is bounded above by $N^{2}\exp\left(\frac{-cd\eta^{2}}{16\zeta_{0}^{4}(\beta+2)^{4}}\right)$. ∎

Before we prove corollary 5.1 we first prove the following helper lemma.

Suppose $\max_{\mathbf{W}}\|\nabla_{\mathbf{W}}f_{i}({\mathbf{W}})\|\leq M,$ and both $\nabla_{{\mathbf{W}}}f_{i}({\mathbf{w}})$ and $\nabla_{{\mathbf{W}}}f_{j}({\mathbf{W}})$ are Lipschitz in ${\mathbf{W}}$ with constant $L$. Then $h_{\mathbf{W}}({\mathbf{x}_{i}},{\mathbf{x}_{j}})$ is Lipschitz in ${\mathbf{W}}$ with constant $2LM.$

We view ${\mathbf{W}}$ as flattened vector. We now prove the above result for these two vectors. For two vectors ${\mathbf{w}},{\mathbf{w}}^{\prime},$

Here the first inequality uses the triangle inequality, the second inequality uses the Cauchy-Schwartz inequality, and the third and fourth inequalities use the assumptions that $\nabla_{\mathbf{w}}f_{i}({\mathbf{w}})$ and $\nabla_{\mathbf{w}}f_{j}({\mathbf{w}})$ are Lipschitz in ${\mathbf{w}}$ and have bounded norm. ∎

We are now ready to prove the corollary, which we restate here. The proof uses a standard "epsilon-net" argument; we identify a fine net of points within the ball $\mathcal{B}_{r}.$ If the gradient confusion is small at every point in this discrete set, and the gradient confusion varies slowly enough with ${\mathbf{W}},$ when we can guarantee small gradient confusion at every point in $\mathcal{B}_{r}.$

Define the function $h^{+}({\mathbf{W}})=\max_{ij}h_{\mathbf{W}}({\mathbf{x}_{i}},{\mathbf{x}_{j}}).$ Our goal is to find conditions under which $h^{+}({\mathbf{W}})>-\eta$ for all ${\mathbf{W}}$ in a large set. To derive such conditions, we will need a Lipschitz constant for $h^{+}({\mathbf{W}}),$ which is no larger than the maximal Lipschitz constant of $h_{\mathbf{W}}({\mathbf{x}_{i}},{\mathbf{x}_{j}})$ for all $i,j.$ We have that $\|\nabla_{{\mathbf{W}}}f_{i}\|=\|\zeta_{\mathbf{x}_{i}}(\mathbf{W})\mathbf{x}_{i}\|\leq\zeta_{0}.$ Now we need to get a ${\mathbf{W}}$-Lipschitz constants for $\nabla_{\mathbf{x}_{i}}f_{i}=\zeta_{\mathbf{x}_{i}}(\mathbf{W}){\mathbf{x}_{i}}.$ By lemma D.1, we have $\|\nabla_{\mathbf{W}}(\zeta_{\mathbf{x}_{i}}(\mathbf{W})\mathbf{x}_{i})\|=\|(\nabla_{\mathbf{W}}\zeta_{\mathbf{x}_{i}}(\mathbf{W}))\mathbf{x}_{i}\|\leq\zeta_{0}.$ Using lemma D.2, we see that $2\zeta_{0}^{2}$ is a Lipschitz constant for $h_{\mathbf{W}}({\mathbf{x}_{i}},{\mathbf{x}_{j}}),$ and thus also $h^{+}({\mathbf{W}}).$

Now, consider a minimizer ${\mathbf{W}}$ of the objective, and a ball $\mathcal{B}_{r}$ around this point of radius $r$. Define the constant $\epsilon=\frac{\eta}{4\zeta_{0}^{2}},$ and create an $\epsilon$-net of points $\mathcal{N}_{\epsilon}=\{{\mathbf{W}}_{i}\}$ inside the ball. This net is sufficiently dense that any ${\mathbf{W}}^{\prime}\in\mathcal{B}_{r}$ is at most $\epsilon$ units away from some ${\mathbf{W}}_{i}\in\mathcal{N}_{\epsilon}.$ Furthermore, because $h^{+}({\mathbf{W}})$ is Lipschitz in ${\mathbf{W}},$ $|h^{+}({\mathbf{W}}^{\prime})-h^{+}({\mathbf{W}_{i}})|\leq 2\zeta_{0}^{2}\epsilon=\eta/2.$

We now know the following: if we can guarantee that

then we also know that $h^{+}({\mathbf{W}}^{\prime})\geq-\eta$ for all ${\mathbf{W}}^{\prime}\in\mathcal{B}_{r}$. For this reason, we prove the result by bounding the probability that (11) holds. It is known that $\mathcal{N}_{\epsilon}$ can be constructed so that $|\mathcal{N}_{\epsilon}|\leq(2r/\epsilon+1)^{d}=(8\zeta_{0}^{2}r/\eta+1)^{d}$ (see Vershynin 2018, corollary 4.1.13). Theorem 5.1 provides a bound on the probability that each individual point in the net satisfies condition (11). Using a union bound, we see that all points in the net satisfy this condition with probability at least

Finally, note that, if $r<\epsilon,$ then we can form a net with $|\mathcal{N}_{\epsilon}|=1$. In this case, the probability of satisfying (11) is at least

D.4 Proof of theorem 4.1

We obtain the following analogue of equation (8).

Equations (17) and 18 follow from the the fact that $\|(\nabla_{\mathbf{W}}\zeta_{\mathbf{x}_{i}}(\mathbf{W}))\|_{2}\leq\zeta_{0}$ and the arguments in the proof for theorem 5.1. We will now show the proof sketch for equation (19). For every $p\in[\beta]_{0}$, consider $\|\nabla_{\mathbf{W}}\left(\mathbf{H}_{p-1}(\mathbf{x}_{i})\cdot\mathbf{g}_{p}(\mathbf{x}_{i})^{\top}\cdot\mathbf{g}_{p}(\mathbf{x}_{j})\cdot\mathbf{H}_{p-1}(\mathbf{x}_{i})\right)\|$. Using the symmetry between $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$, the expression can be upper-bounded by,

As before we can use an inductive argument to find the upper-bound and thus, we obtain the following which implies equation (19).

The inequality combines equation 7 with Lemma D.1. We now prove the following inequality.

Consider the inner-most expectation. Note that the only random variable is $\mathbf{W}_{\beta}$. Moreover, the term inside the trace is scalar. Note that the activation function $\sigma$ satisfies $|\sigma^{\prime}(x)|\leq 1$. Using the linearity of expectation, the LHS in equation (20) can be upper-bounded by the following.

The first sum in the above expression can be upper-bounded by $1$, since $|\sigma(x)|\leq 1$. We will now show that the second sum is $0$. Consider the inner-most expectation. The weights $\mathbf{W}_{\beta}$ appears only in the expression $\mathbf{g}_{p}(\mathbf{x}_{i})^{\top}\cdot\mathbf{g}_{p}(\mathbf{x}_{j})$. Moreover, note that every entry in $\mathbf{W}_{\beta}$ is an i.i.d. normal random variable with mean $0$. Thus, the second summand simplifies to,

Applying the above argument repeatedly we obtain that the second summand (equation (22)) is $0$.

D.5 Proof of Theorem 6.1

In this section, we prove Theorem 6.1. The proof follows similar to those in previous sub-sections; we prove a bound on the gradient of the gradient inner-product and show that the expectation is non-negative. Combining these two with an argument similar to equation 10 we get the theorem.

Note that the dataset is obtained by considering i.i.d. samples from a $d$-dimensional unit sphere. Thus, the lower-bound on the expectation (i.e., non-negative expectation of the gradient inner-product) follows from equation 9. Thus, it remains to prove an upper-bound on the norm of the gradient of the gradient inner-product term.

Throughout this proof, we will use $g(\mathbf{x})$ as a short-hand to denote $g_{\mathbf{W}}(\mathbf{x})$. Consider the gradient $\nabla_{\mathbf{W}}g(\mathbf{x})$. The the $i^{th}$ component of this can be written as follows.

Now consider, the gradient inner-product $h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})$. We want to upper-bound the quantity $\|\nabla_{(\mathbf{x}_{i},\mathbf{x}_{j})}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|$. From symmetry, this can be upper-bounded by $2\|\nabla_{\mathbf{x}_{i}}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|$. Consider the $k^{th}$ coordinate of $\nabla_{\mathbf{x}_{i}}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})$. Using equation 23, the assumption that $\{\mathbf{W}_{i}\}_{i\in[\beta]}$ are orthogonal matrices and taking the gradient, this can be written as,

Combining assumption 1 with equation 24 we have that $\|\nabla_{\mathbf{x}_{i}}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|$ is at most $2\gamma^{2}\beta\|\mathbf{x}_{j}\|\leq 2\gamma^{2}\beta$. For the definition of the scaling factor $\gamma=\frac{1}{\sqrt{2\beta}}$, we have that $2\gamma^{2}\beta=1$. Thus, $\|\nabla_{(\mathbf{x}_{i},\mathbf{x}_{j})}h_{\mathbf{W}}(\mathbf{x}_{i},\mathbf{x}_{j})\|\leq 2$.

Appendix E Technical lemmas

We will briefly describe some technical lemmas we require in our analysis. The following Chernoff-style concentration bound is proved in Chapter 5 of Vershynin 2018.

We will rely on the following generalization of lemma E.1. We would like to point out that the underlying metric is the Euclidean metric and thus we use the $\|.\|_{2}$-norm.

Additionally, we will use the following facts about a normalized Gaussian random variable.

For a normalized Gaussian $\mathbf{x}$ (i.e., an $\mathbf{x}$ sampled uniformly from the surface of a unit $d$-dimensional sphere) the following statements are true.

We use the following well-known Gaussian concentration inequality in our proofs (e.g., Chapter 5 in Boucheron et al. 2013).

Appendix F Additional discussion of the small weights assumption (assumption 1)

Without the small-weights assumption, the signal propagated forward or the gradients $\nabla_{\mathbf{W}}f_{i}$ could potentially blow up in magnitude, making the network untrainable. Proving non-vacuous bounds in case of such blow-ups in magnitude of the signal or the gradient is not possible in general, and thus, we assume this restricted class of weights.

Note that the small-weights assumption is not just a theoretical concern, but also usually holds in practice. Neural networks are often trained with weight decay regularizers of the form $\sum_{i}\|W_{i}\|_{F}^{2}$, which keep the weights small during optimization. The operator norm of convolutional layers have also recently been used as an effective regularizer for image classification tasks by Sedghi et al. 2018.

In the proof of theorem 4.1 we showed that assumption 1 holds with high probability at standard Gaussian initializations used in practice. While, in general, there is no reason to believe that such a small-weights assumption would continue to hold during optimization without explicit regularizers like weight decay, some recent work has shown evidence that the weights do not move too far away during training from the random initialization point for overparameterized neural networks (Neyshabur et al. 2018; Dziugaite & Roy 2017; Nagarajan & Kolter 2019; Zou et al. 2018; Allen-Zhu et al. 2018; Du et al. 2018; Oymak & Soltanolkotabi 2018). It is worth noting though that all these results have been shown under some restrictive assumptions, such as the width requiring to be much larger than generally used by practitioners.

27, 37, 31, 42, 9, 8, 56, 17, 38, 14, 34, 61, 1, 13, 40, 54, 51, 6, 46, 32