Title: THE EFFECT OF NEURAL NET ARCHITECTURE ON GRADIENT CONFUSION & TRAINING PERFORMANCE Anonymous

Abstract: The goal of this paper is to study why typical neural networks train so fast, and how neural network architecture affects the speed of training. We introduce a simple concept called gradient confusion to help formally analyze this. When confusion is high, stochastic gradients produced by different data samples may be negatively correlated, slowing down convergence. But when gradient confusion is low, data samples interact harmoniously, and training proceeds quickly. Through novel theoretical and experimental results, we show how the neural net architecture affects gradient confusion, and thus the efficiency of training. We show that increasing the width of neural networks leads to lower gradient confusion, and thus easier model training. On the other hand, increasing the depth of neural networks has the opposite effect. Finally, we observe empirically that techniques like batch normalization and skip connections reduce gradient confusion, which helps reduce the training burden of very deep networks.

1. INTRODUCTION:

Stochastic gradient descent (SGD) (Robbins & Monro, 1951) and its variants with momentum (Sutskever et al., 2013) have become the standard optimization routine for neural networks due to their fast convergence and good generalization properties (Wilson et al., 2017; Keskar & Socher, 2017; Sutskever et al., 2013 ). Yet the behavior of SGD on high-dimensional neural network models still eludes full theoretical understanding, both in terms of its convergence and generalization properties. In this paper, we study why SGD is so efficient at converging to low loss values on most standard neural networks, and how neural net architecture design affects training performance.

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 -accuracy of the minimizer in O(1/ ) 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 regime, SGD seems to converge quickly with constant learning rates. Most neural net practitioners use a constant learning rate for the majority of training, with exponentially decaying learning rate schedules at the end, 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) . Some more recent results have shown that when models can fit the data completely while being strongly convex, convergence without a noise floor is possible without decaying the learning rate (Schmidt & Roux, 2013; Ma et al., 2017; Bassily et al., 2018; Vaswani et al., 2018) .

While these results do give important insights, they do not fully explain the dynamics of SGD on neural nets, and how they relate to overparameterization. Training performance is also highly affected by the neural network architecture. It is common knowledge among neural network practitioners that deeper networks train slower (Bengio et al., 1994; Glorot & Bengio, 2010) . This has led to several innovations over the years to get deeper networks to train more easily, such as careful initialization strategies (Glorot & Bengio, 2010; He et al., 2015; Zhang et al., 2019) , residual connections (He et al., 2016) , and various normalization schemes like batch normalization (Ioffe & Szegedy, 2015) and weight normalization (Salimans & Kingma, 2016) . Furthermore, there is ample evidence to indicate that wider networks are easier to train (Zagoruyko & Komodakis, 2016; Nguyen & Hein, 2017; Lee et al., 2019) , and recent theoretical results suggest that the dynamics of SGD simplify considerably for very wide networks (Jacot et al., 2018; Lee et al., 2019) . Several prior works have investigated the difficulties of training deep networks (Glorot & Bengio, 2010; Balduzzi et al., 2017) , and the benefits of width (Nguyen & Hein, 2017; Lee et al., 2019; Du et al., 2018; Allen-Zhu et al., 2018) . This work advances the existing literature by identifying and analyzing a condition that enables us to theoretically and empirically establish novel direct relationships between layer width, network depth, problem dimensionality, and SGD dynamics on overparameterized networks.

Our contributions. Typical neural nets are overparameterized (i.e., the number of parameters exceed the number of training points). In this paper, we ask how this overparameterization, and more specifically the architecture of a neural net, affects the dynamics of SGD. We answer this question through extensive theoretical and experimental studies and show how network width, depth, batch normalization and skip connections affect the dynamics. We emphasize that our main contributions are conceptual. In particular, following are our main contributions.1 

• We identify a condition, termed gradient confusion, that impacts the convergence properties of SGD on overparameterized 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 this gradient confusion condition as the main proxy, to study the effect of various architecture choices on convergence.

• We study the effect of neural net architecture on gradient confusion (section 4), and prove (a) gradient confusion increases as the network depth increases, and (b) at initialization, wider networks have lower gradient confusion. This indicates that deeper networks are more difficult to train and wider networks improves trainability of neural nets. Directly analyzing the gradient confusion bound enables us to derive novel and tight results on the direct effect of depth and width, without requiring arguably restrictive assumptions like infinitely wide networks (Schoenholz et al., 2016; Lee et al., 2019) . Our results hold for a large family of neural networks with non-linear activations and a large class of loss-functions.

• 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 5 and appendix A). We find that our theoretical results consistently hold across all our experiments. We further show that innovations like batch normalization and skip connections in residual networks help lower gradient confusion, thus indicating why standard neural networks that employ such techniques are so efficiently trained using SGD.

2. GRADIENT CONFUSION:

Notations. We denote vectors in bold lower-case and matrices in bold upper-case. We use (W) i,j to indicate the (i, j) cell in matrix W and (W) i for the i th row of matrix W. W denotes the operator norm of W.

[N ] denotes {1, 2, . . . , N } and [N ] 0 denotes {0, 1, . . . , N }.

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

EQUATION

using the following iterative update rule for T rounds:

EQUATION

)

Here α is the learning rate and fk is a function chosen uniformly at random from

{f i } i∈[N ] at iteration k ∈ [T ].

We use w to denote the optimal solution, i.e., w = arg min w F (w).

Gradient confusion. SGD works by iteratively selecting a random function fk , and modifying the parameters to move in the direction of the negative gradient of the objective term fk . It may happen that the selected gradient ∇ fk is negatively correlated with the gradient of another term ∇f j . When the gradients of different mini-batches are negatively correlated, the objective terms disagree on which direction the parameters should move, and we say that there is gradient confusion.2 

Definition 2.1. A set of objective functions {f i } i∈[N ] has gradient confusion bound η ≥ 0 if the pair-wise inner products between gradients satisfy, for a fixed w ∈ R d ,

EQUATION

Remarks. Note that while the gradient confusion bound η is defined for the worst-case gradient inner product, all the results in our paper can be trivially extended to using a bound on the average gradient inner product:

N i,j=1 ∇f i (w), ∇f j (w) /N 2 ≥ -η.

All theoretical results would remain the same up to constants. Further, note that definition 2.1 is applicable even when the stochastic gradients are averaged over minibatches of size B. For minibatches of size B, the variance of the gradient inner product scales down as 1/B 2 , and thus η is expected to decrease as B grows. Observations in simplified settings. SGD converges fast when gradient confusion is low along its path. To see why, consider the case of training a logistic regression model on a dataset with orthogonal vectors. We have f i (w) = (y i x i w), where : R → R is the logistic loss, {x i } i∈[N ] is a set of orthogonal training vectors, and y i ∈ {-1, 1} is the label for the i th training example. We then have ∇f i (w) = ζ i x i , where

ζ i = y i (y i • x i w). Note that the gradient confusion is 0 since ∇f i (w), ∇f j (w) = ζ i ζ j x i , x j = 0, ∀i, j ∈ [N ] and i = j.

Thus, an update in the gradient direction f i has no effect on the loss value of f j for i = j. In this case, SGD decouples into (deterministic) gradient descent on each objective term separately, and we can expect to see the fast convergence rates attained by gradient descent.

Can we expect a problem to have low gradient confusion in practice? From the logistic regression problem, we have:

| ∇f i (w), ∇f j (w) | = | x i , x j | • |ζ i ζ j |.

This inner product is expected to be small for all w; the logistic loss satisfies |ζ i ζ j | < 1, and for fixed N the quantity

max ij | x j , x i | is O(1/ √ d)

whenever {x i } are randomly sampled from a sphere (see lemma B.1 for the formal statement).3 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 overparameterized linear regression model can have much faster convergence rates, as well as lower gradient confusion, compared to the underparameterized model. Now consider more general neural net problems. There is evidence that the Hessian at the minimizer is very low rank for many standard overparameterized neural net models (Sagun et al., 2017; Cooper, 2018; Chaudhari et al., 2016; Wu et al., 2017; Ghorbani et al., 2019) . What does this imply for the 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 H i (w) -H i (w ) ≤ L H ww (note that this is a standard optimization assumption (Nesterov, 2018) , with evidence that it is applicable for neural nets (Martens, 2016) ). Then ∇f i (w) = H i w + e, where e is an error term bounded as: e ≤ 1 2 L H w 2 , and we use the shorthand H i to denote H i (0). Then we have (appendix C):

| ∇f i (w), ∇f j (w) | ≤ w 2 H i H j + 1 2 L H w 3 ( H i + H j ) + 1 4 L 2 H w 4 .

If the Hessians are sufficiently random and low-rank (e.g., of the form H i = a i a i where a i ∈ R N ×r are randomly sampled from a unit sphere), then one would expect the terms in this expression to be small for all w within a neighborhood of the minimizer. While a bit non-rigorous, this nonetheless suggests that for many standard neural network models, the gradient confusion might be small for a large class of weights near the minimizer.

The above arguments are rather informal, and ignore issues like the effect of the structure of neural networks. 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 net architecture on the gradient confusion.

3. SGD IS EFFICIENT WHEN GRADIENT CONFUSION IS LOW:

Several prior papers have analyzed the convergence rates of constant learning rate SGD (Nedić & Bertsekas, 2001; Moulines & Bach, 2011; Needell et al., 2014; Dieuleveut et al., 2017) . These results show that for strongly convex and Lipschitz smooth functions, SGD with a constant learning rate α converges linearly to a neighborhood of the minimizer. The noise floor it converges to depends on the learning rate α and the variance of the gradients at the minimizer, i.e., E i ∇f i (w ) 2 . To guarantee convergence to -accuracy in such a setting, the learning rate needs to be small, i.e., α = O( ), and the method requires T = O(1/ ) iterations. Some more recent results show convergence of constant learning rate SGD without a noise floor and without small step sizes for models that can completely fit the data (Schmidt & Roux, 2013; Ma et al., 2017; Bassily et al., 2018; Vaswani et al., 2018) .

The gradient confusion bound is related to these classical results. Cauchy-Swartz inequality implies that if

E i ∇f i (w ) 2 = O( ), then E i,j | ∇f i (w ), ∇f j (w ) | = O( ), ∀i, j.

Thus the gradient confusion at the minimizer is small when the variance of the gradients at the minimizer is small. Further note that when the variance of the gradients at the minimizer is O( ), a direct application of the results in (Moulines & Bach, 2011; Needell et al., 2014) shows that constant learning rate SGD has fast convergence to -accuracy in T = O(log(1/ )) iterations, without the learning rate needing to be small. Generally however, bounded gradient confusion does not provide a bound on the variance of the gradients (see section 6 for more discussion). Thus, it is instructive to derive convergence bounds of SGD explicitly in terms of the gradient confusion to properly understand its effect.

We begin by considering functions satisfying the Polyak-Lojasiewicz (PL) inequality (Lojasiewicz, 1965) , a condition related to, but weaker than, strong convexity, and provide bounds on the rate of convergence in terms of the optimality gap. Then we look at a broader class of smooth non-convex functions, and analyze convergence to a stationary point. We first make two standard assumptions.

(A1) {f i } i∈[N ] are Lipschitz smooth: f i (w ) ≤ f i (w) + ∇f i (w) (w -w) + L 2 w -w 2 . (A2) {f i } i∈[N ] satisfy the PL inequality: 1 2 ∇f i (w) 2 ≥ µ(f i (w) -f i ), f i = min w f i (w). Theorem 3.1.

If the objective function satisfies (A1) and (A2), and has gradient confusion η, SGD with updates of the form (2) converges linearly to a neighborhood of the minima of problem (1) as:

E[F (w T ) -F ] ≤ ρ T (F (w 0 ) -F ) + αη 1-ρ , where α < 2 N L , ρ = 1 -2µ N α -N Lα 2 2

, F = min w F (w) and 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., η = O( ), SGD has fast convergence to O( )-accuracy in T = O(log(1/ )) iterations, without requiring the learning rate to be vanishingly small. We now extend this to general smooth functions.

Theorem 3.2. If the objective satisfies (A1) and has gradient confusion bound η, then SGD converges to a neighborhood of a stationary point as:

min k=1,...,T E ∇F (w k ) 2 ≤ ρ(F (w1)-F ) T + ρη,

for learning rate α < 2 N L , ρ = 2N 2-N Lα , and F = min w F (w). Thus, as long as η = 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. Convergence slows down as the iterates approach the noise floor, and at this point typically practitioners employ exponentially decaying learning rate schedules (Krizhevsky et al., 2012; Simonyan & Zisserman, 2014; He et al., 2016; Zagoruyko & Komodakis, 2016; Ge et al., 2019) . See section 6 for more discussion on theorems 3.1 and 3.2, and how they relate to results in previous works. 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. The main intention is to show the direct effect that the gradient confusion bound has on the convergence rate and the noise floor that constant learning rate SGD converges to. As we show in the following sections, this new perspective in terms of the gradient confusion helps us more directly understand how neural net architecture design affects SGD dynamics and why.

4. EFFECT OF NEURAL NET ARCHITECTURE ON GRADIENT CONFUSION:

To draw a rigorous connection between neural net structure and training performance, we analyze gradient confusion for generic (i.e., random) model problems using methods from high-dimensional probability. In particular, this section considers the following scenarios: (a) Random data drawn from a unit sphere and the weights in a ball around the local minimizer (theorem 4.1 and corollary 4.1). (b) Random weights using standard initialization schemes and both arbitrary bounded data (theorem 4.2, part 1) and random data drawn from a unit sphere (theorem 4.2, part 2). Our results cover a wide range of scenarios compared to prior work (e.g., Chen et al. (2018); Schoenholz et al. (2016) ; Balduzzi et al. (2017) ), require minimal additional assumptions, and hold for a large family of neural nets with non-linear activations and a large class of loss-functions. In particular, our results hold for fully connected networks (and convolutional networks in some cases) with the square-loss and logistic-loss functions, and commonly used non-linear activations such as sigmoid, tanh and ReLU.

(a) Random Data, Bounded Weights Around Minimizer. In this subsection, we consider training data of the form D = {(x i , C(x i ))} i∈[N ] , for some labeling function C : R d → [-1, 1], and with data points {x i } drawn uniformly from the surface of a d-dimensional unit sphere. The labeling function satisfies |C(x)| ≤ 1 and ∇ x C(x) 2 ≤ 1 for x ≤ 1. Note that this automatically holds for every model considered in this paper where the labeling function is realizable (i.e., where the model can express the labeling function using its parameters). More generally, this assumes a Lipschitz condition on the labels (i.e., the labels don't change too quickly with the inputs). In this paper, we consider two loss-functions, namely, square-loss for regression and logistic loss function for classification. The square-loss function is defined as f i (w) = 1 2 (C(x i ) -g w (x i )) 2 and the logistic function is defined as f i (w) = log(1 + exp(-C(x i )g w (x i ))). Here, g w : R d → R denotes the parameterized function we fit to the training data and f i (w) denotes the loss-function of hypothesis g w on data point x i .

Formally, let W 0 ∈ R 1×d and {W p } p∈ [β] such that W p ∈ R p × p-1 be the given weight matrices. Let W denote the tuple (W p ) p∈[β]0 . Define := max p∈[β] p to be the width and β to be the depth of the neural network. Then, the model g W is defined as

EQUATION

where σ denotes the non-linear activation function applied point-wise to its arguments. We assume that the non-linear activation is given by a function σ(x) with the following properties. When x ≤ 1, as in our random data model, activation functions such as sigmoid, tanh, softmax and ReLU satisfy these requirements. Additionally, we make the following assumption on the weights.

EQUATION

Assumption 1 (Small Weights). We assume that the operator norm of the weight matrices {W i } i∈[β]0 are bounded above by 1. In other words, for every i ∈

[β] 0 we have W i ≤ 1.

The operator norm of the weight matrices W being close to 1 is important for the trainability of neural nets, as it ensures that the input signal is passed through the net 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. The small-weights assumption is not just a theoretical concern, but also usually enforced in practice using weight decay regularizers of the form i W i 2 F , which keep the weights small during optimization. See appendix F for further discussion on the small weights assumption.

We now prove concentration bounds for the gradient confusion on neural nets.

Theorem 4.1. Consider the problem of training neural nets (equation 4) using either the square-loss or the logistic-loss function. Let η > 0 be a given constant. Let the weights satisfy assumption 1 and the non-linearities in each layer satisfy properties (P1) and (P2). For some fixed constant c > 0, the gradient confusion bound in equation 3 holds with probability at least

1 -N 2 exp -cdη 2 16ζ 4 0 (β+2) 4 ,

For both the square-loss and the logistic-loss functions, ζ 0 ≤ 2 √ β (from lemma D.1).

Thus, theorem 4.1 shows that, for a given dimension d and number of samples N , when the network depth β decreases, the probability that the gradient confusion bound in equation 3 holds increases, and vice versa. 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 4.1 shows concentration bounds for the gradient confusion at a fixed weight W. Thus, to ensure that the above result is relevant for the convergence of SGD on overparameterized models, we now make the concentration bound in theorem 4.1 uniform over all weights inside a ball B r of radius r.

Corollary 4.1 (Uniform concentration for all weights around the minimizer).

Select a point W = (W 0 , W 1 , . . . , W β ), satisfying assumption 1. Consider a ball B r centered at W of radius r > 0. If the data {x i } i∈[N ]

are sampled uniformly from a unit sphere, then the gradient confusion bound in equation 3 holds uniformly at all points W ∈ B r with probability at least

1 -N 2 exp -cdη 2 64ζ 4 0 (β+2) 4 , if r ≤ η/4ζ 2 0 , 1 -N 2 exp -cdη 2 64ζ 4 0 (β+2) 4 + 8dζ 2 0 r η , otherwise.

Thus, corollary 4.1 shows that the probability that the gradient confusion bound holds decreases with increasing depth, for all weights in a ball around the minimizer. This explains why training very deep models is hard and typically slow with SGD (Bengio et al., 1994; Glorot & Bengio, 2010) . Note that this is also related to the shattered gradients phenomenon (Balduzzi et al., 2017) that arises with depth (see appendix 6 for more discussion). This naturally raises the question why modern deep neural networks are so efficiently trained using SGD. While careful initialization strategies prevent vanishing or exploding gradients making deeper networks trainable, these strategies still suffer from high gradient confusion for very deep networks (as we show below in theorem 4.2). Thus, in section 5, we empirically study how popular techniques like skip connections (He et al., 2016) and batch normalization (Ioffe & Szegedy, 2015) affect gradient confusion. We find that these techniques drastically lower gradient confusion, making very deep networks significantly easier to train. Note that the above results automatically hold for convolutional nets, since a convolution operation on x can be represented as a matrix multiplication Ux for an appropriate Toeplitz matrix U.

(b) Standard Weight Initializations and the Effect of Layer Width. Note that on assuming W ≤ 1 for each weight matrix W in our results in part 1 of section 4, the dependence of gradient confusion on the layer width goes away in general. A simple example that illustrates this is to consider the case 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. Thus, to better understand the effect of the layer width, in this subsection we focus on the behavior of neural nets at initialization by considering standard weight initialization strategies used when training neural nets. For completeness, we consider both the case where the data is arbitrary but bounded, as well as where the data is randomly drawn from a unit sphere. A key point used in the following results is that typical weight initialization techniques ensure that the operator norm is bounded by 1 with high probability, thus enabling us to derive tight bounds on the gradient confusion. We consider the following weight initialization strategy. Strategy 4.1.

W 0 ∈ R ×d has independent N (0, 1 d ) entries. For every p ∈ [β], the weights W p ∈ R p × p-1 have independent N 0, 1 κ p-1 entries for some constant κ > 0.

This initialization strategy with different settings of κ are used almost universally for neural networks (Glorot & Bengio, 2010; He et al., 2015) . The following theorem shows how the width := max p∈[β] p and the depth β affect the gradient confusion condition. In particular, as width increases or depth decreases the probability that the gradient confusion bound (equation 3) holds increases.

Theorem 4.2 (Neural nets with randomly chosen weights). Let W 0 , W 1 , . . . , W β be weight matrices chosen according to strategy 4.1. There exists fixed constants c 1 , c 2 > 0 such that we have:

1. Consider a fixed but arbitrary dataset

x 1 , x 2 , . . . , x N with x i ≤ 1 for every i ∈ [N ].

For η > 4, the gradient confusion bound in equation 3 holds with probability at least

1 -β exp -c 1 κ 2 2 -N 2 exp -c 2 β(η-4) 2 64ζ 4 0 (β+2) 4 . 2. If the dataset {x i } i∈[N ] is such that each x i is an i.i.d.

sample from the surface of ddimensional unit sphere, then for every η > 0 the gradient confusion bound in equation 3 holds with probability at least

1 -β exp -c 1 κ 2 2 -N 2 exp -c2( d+ 2 β)η 2 16ζ 4 0 (β+2) 4

.

Thus, theorem 4.2 shows that layer width improves the trainability of deep networks under most standard initialization techniques. Note that for reasons described above, almost all recent theoretical results on neural nets have analyzed the effect of width at initialization. Further, one can prove that when the layers are very wide, the weights move very little, and most of the conditions at initialization persist during training, considerably simplifying the learning dynamics of gradient descent (Jacot et al., 2018; Lee et al., 2019) . Under this simplifying assumption of very wide networks, several authors have analyzed the effect of width during optimization on very wide networks (Allen-Zhu et al., 2018; Du et al., 2018) . Under this assumption and using a standard application of the techniques used in prior work, we can likely show that theorem 4.2 also holds during optimization. In the next section (and in 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, and that these effects persist throughout optimization for most models.

5. EXPERIMENTAL RESULTS:

To test our theoretical results and to probe why standard neural nets 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.

We performed 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. We present results for CNNs on CIFAR-10 in this section, and present all other results in appendix A. We use CNN-β-to denote WRNs that have no skip connections or batch normalization, with a depth β and width factor . 4 We turned off dropout and weight decay for all our experiments. 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 used the MSRA initializer (He et al., 2015) for the weights as is standard for this model, and used the same preprocessing steps for the CIFAR-10 images as described in Zagoruyko & Komodakis (2016) . We ran each experiment 5 times, and 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 training loss value. 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. See appendix A.2 for more details on the experimental setup and architectures used.

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 (indicating higher variance), 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 is still overparameterized and achieves a high level of performance (see appendix A.3).

Effect of batch normalization and skip connections. To help understand why many standard deep nets are so efficiently trained using SGD, 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. Both these techniques together keep the gradient confusion relatively low even for very deep networks, significantly improving trainability of deep models. Note that all these observations are consistent with previous work (Balduzzi et al., 2017; Santurkar et al., 2018; Yang et al., 2019) .

6. CONNECTIONS TO RELATED WORK:

The interpolation condition and connections to the gradient variance: The gradient confusion condition has interesting connections to the variance of the gradients. If we assume bounded gradient variance σ 2 , we can bound the gradient confusion parameter η in terms of other quantities. For example, suppose the true gradient g is defined as g = g 1 /2 + g 2 /2. Then we have:

| g 1 , g 2 | ≤ g 1 g 2 ≤ 1/2•( g 1 2 + g 2 2 ) = 1/2•( g 1 -g 2 + g 2 -g 2 )+ g 2 = σ 2 + g 2 ,

where all gradients are defined at the same w.

More interestingly however, it is not possible in general to bound the gradient variance in terms of the gradient confusion parameter. As a counter-example, consider a problem with the following distribution on the stochastic gradients: 1 1-p samples with gradient 1 and 1 p samples with gradient , where p = → 0. In this case, the gradients are positive, so gradient confusion η = 0. The mean of the gradients is given by 1 + (1 -), which remains bounded as → 0. On the other hand, the variance (and thus the squared norm of the stochastic gradients) is actually unbounded (O(1/ ) as → 0). While this is a contrived example, it does show why such a bound does not exist.

A consequence of this is that in theorems 3.1 and 3.2 in section 3, the "noise term" (i.e., the second term in the RHS of the convergence bounds) does not depend on the learning rate in the general case. If gradients have unbounded variance, lowering the learning rate does not reduce the variance of the SGD updates, and thus does not reduce the noise term. This further indicates a regime of high gradient variance but low gradient confusion where SGD would converge faster than predicted by classical convergence rate results (Moulines & Bach, 2011; Needell et al., 2014) .

In practice, typically the gradient variance is bounded. In this case, tighter convergence rate bounds in terms of η are possible, which follow as simple corollaries from classical convergence bounds in terms of the gradient variance such as those in Moulines & Bach (2011) ; Needell et al. (2014) . The interpolation and strong growth conditions, as analyzed in Ma et al. (2017) ; Bassily et al. (2018) ; Vaswani et al. (2018) ; Schmidt & Roux (2013) , are related to this. These conditions imply that the variance of the gradients at the minimizer is small. As shown above, this would imply that the gradient confusion is also small, leading to fast convergence of SGD. This is also similar to the condition of the minimal risk being small as in Srebro et al. (2010) ; Zhang & Zhou (2019) .

That being said, it is worth noting that the main contribution of this paper is to better understand how the network architecture impacts SGD. In general, it is not tractable to prove the concentration bounds in section 4 using the covariance matrix of the gradients alone without further unrealistic assumptions, such as infinitely wide networks (Schoenholz et al., 2016) . A key contribution of this paper is to identify a suitable surrogate (i.e., the gradient confusion bound) to help study this relationship using new tight and clean bounds.

Connections to gradient diversity: In Yin et al. (2017) , the authors define a property called gradient diversity. This quantity is related to gradient confusion, but with important differences. Gradient diversity 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. In a large batch, higher gradient diversity is desirable, and this leads to improved convergence rates in distributed settings, as shown in Yin et al. (2017) . On the other hand, gradient confusion between two individual gradients is zero unless the inner product between them is negative. This makes gradient confusion useful for studying convergence of small minibatch SGD. This is because different possible SGD updates do not conflict with each other unless they are negatively correlated with each other.

The choice of the definition of gradient diversity in Yin et al. (2017) has important implications when its behavior is studied in overparameterized settings. Chen et al. (2018) extends the work of Yin et al. (2017) , where the authors prove on 2-layer neural nets (and multi-layer linear neural nets) that gradient diversity increases with increased width and decreased depth. This metric does not however distinguish between the cases where gradients become more orthogonal vs. more negatively correlated. As we show in this paper, this can have very different effects on the convergence of SGD in overparameterized settings. Specifically, we show that increased width and decreased depth decreases gradient confusion, and makes these networks easier to train. In fact, we see that the gradients become more orthogonal to each other in this case as shown in section 5. Thus, we view our papers to be complementary to each other, providing insights about different issues (large batch distributed training vs. small minibatch convergence).

Other work on the impact of neural net structure: Other works have also studied the impact of structured gradients on SGD. Balduzzi et al. (2017) study the effects of shattered gradients at initialization for ReLU networks, which is when stochastic gradients at initialization become negatively correlated. The authors show how gradients get increasingly shattered with depth in ReLU networks. Hanin (2018) shows 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 initialization. Further, Hanin & Rolnick (2018) show that this sum of the reciprocals of the hidden layer widths determines the variance of the sizes of the activations at each layer during initialization. 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 on overparameterized neural nets: The convergence of SGD on over-parameterized models has received a lot of attention recently. Arora et al. (2018b) study the behavior of SGD on over-parameterized problems, and show that SGD on over-parameterized linear neural nets is similar to applying a certain preconditioner while optimizing. This can sometimes lead to acceleration when overparameterizing by increasing the depth of linear neural networks. In this paper, we show that this property does not hold in general (as mentioned briefly in Arora et al. 2017) also show convergence to global minimizers of SGD for non-linear nets. While all these results require the network to be sufficiently wide, they represent an important step in the direction of better understanding optimization on neural nets. This paper complements these recent results by studying how low gradient confusion contributes to SGD's success on modern overparameterized neural nets.

7. CONCLUSIONS:

In this paper, we investigate how overparameterization and model architecture affect the dynamics of SGD on neural networks. To help formally analyze this, we introduce a concept called gradient confusion, and show that when gradient confusion is low, SGD experiences fast convergence. We then show that increasing layer width leads to lower gradient confusion, making the model easier to train. In contrast, increasing network depth results in higher gradient confusion, making deeper models harder to train. We further show how techniques like batch normalization and skip connections help in tackling this problem.

Our results provide a number of important insights that can be used for better neural net model design, as well as for better algorithms for better training and generalization. Our results on the test set accuracies 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) . The difference in the gradient confusion within the same class and across classes could also be an interesting tool to study adversarial training. Note that 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; Raghu et al., 2017) . Our results indicate that, given a sufficiently deep network, increasing the network width is important for the trainability of the model, and will lead to faster convergence rates. This is further supported by other recent research (Hanin, 2018; Hanin & Rolnick, 2018) that suggest that the width should increase linearly with depth in a neural network to help dynamics at the beginning of training. Our results also suggest the importance of further investigation into good initialization schemes for neural networks that make training very deep models possible (Zhang et al., 2019) .

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).

A.1 MLPS ON MNIST:

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 N 0, 2/(fan-in + fan-out) , 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 N 0, 1/fan-in .

We considered the simplified case where all hidden layers have the same width . Thus, the first weight matrix W 0 ∈ R ×d , where d = 784 for the 28 × 28-sized images of MNIST; all intermediate weight matrices {W p } p∈[β-1] ∈ R × ; and the final layer W β ∈ R 10× for the 10 image classes in MNIST. We added biases to each layer, which we initialized to 0. We used softmax cross entropy as the loss function. We use MLP-β-to denote this fully connected network of depth β and width . We used the standard train-valid-test splits of 40000-10000-10000 for MNIST.

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 nets 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 α was tuned over a logarithmically-spaced grid:

α ∈ {10 0 , 10 -1 , 10 -2 , 10 -3 , 10 -4 , 10 -5 , 10 -6 }.

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.

The effect of depth. We first present results showing the effect of network depth. We considered a fixed width of = 100, and varied the depth of the neural network, on the log scale, as:

β ∈ {3, 10, 30, 100, 300, 1000}.

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 5, and matching our theoretical results in section 4, 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 β = 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 β = 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:

∈ {10, 30, 100, 300, 1000}.

Crucially, note that the smallest network considered here, MLP-300-10, still has more than 50000 parameters (i.e., more than the number of training samples), and the network with width = 30 has almost three times the number of parameters as the high-performing MLP-3-100 network considered in the previous section. Figure 6 show results on linear neural nets and neural nets with tanh activations for both the Glorot normal and LeCun normal initializations. As in the experimental results of section 5, we see the consistent trend of gradient confusion decreasing with increasing width. Thus, wider networks become easier to train and improve the final training loss value. We further see that when the width is too small ( = 30), the gradient confusion becomes drastically high and the network becomes completely untrainable. 

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.

WIDE RESIDUAL NETWORKS:

The Wide ResNet (WRN) architecture (Zagoruyko & Komodakis, 2016) for CIFAR datasets is a stack of three groups of residual blocks. There is a downsampling layer between two blocks, and the number of channels (width of a convolutional layer) is doubled after downsampling. In the three groups, the width of convolutional layers is {16 , 32 , 64 }, respectively. Each group contains β r residual blocks, and each residual block contains two 3 × 3 convolutional layers equipped with ReLU activation, batch normalization and dropout. There is a 3 × 3 convolutional layer with 16 channels before the three groups of residual blocks. And there is a global average pooling, a fully-connected layer and a softmax layer after the three groups. The depth of WRN is β = 6β r + 4.

For our experiments, we turned off dropout. Unless otherwise specified, we also turned off batch normalization. added biases to the convolutional layers when not using batch normalization to maintain model expressivity. We used the MSRA initializer (He et al., 2015) for the weights as is standard for this model, and used the same preprocessing steps for the CIFAR images as described in Zagoruyko & Komodakis (2016) . This preprocessing step involves normalizing the images and doing data augmentation (Zagoruyko & Komodakis, 2016) . We denote this network as WRN-β-, where β represents the depth and represents the width factor of the network.

To study the effect of depth, we considered WRNs with width factor = 2 and depth varying as: 22, 28, 34, 40, 52, 76, 100} .

β ∈ {16,

For cleaner figures, we sometimes plot a subset of these results: β ∈ {16, 28, 40, 52, 76, 100}. To study the effect of width, we considered WRNs with depth β = 16 and width factor varying as: ∈ {2, 3, 4, 5, 6}.

CONVOLUTIONAL NEURAL NETS:

The WRN architecture contains skip connections that, as we show, help in training deep networks.

To consider VGG-like convolutional networks, we consider a family of networks where we remove the skip connections from WRNs. Following the WRN convention, we denote these networks as CNN-β-, where β denotes the depth and denotes the width factor.

To study the effect of depth, we considered CNNs with width factor = 2 and depth varying as:

β ∈ {16, 22, 28, 34, 40}.

To study the effect of width, we considered CNNs with depth β = 16 and width factor varying as: ∈ {2, 3, 4, 5, 6}.

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.

α ∈ {10 1 , 3×10 0 , 10 0 , 3×10 -1 , 10 -1 , 3×10 -2 , 10 -2 , 3×10 -3 , 10 -3 , 3×10 -4 , 10 -4 , 3×10 -5 },

To measure gradient confusion, at the end of every training epoch, we sampled 100 pairs of minibatches 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 5, 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 5: 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:

In section 5 we showed results on the effect of adding batch normalization 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 5, 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.

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) . Lemma B.1 (Near orthogonality of random vectors). For vectors {x i } i∈[N ] drawn uniformly from a unit sphere in d dimensions, and ν > 0, Pr ∃i,

j |x i x j | > ν ≤ N 2 π 8 exp -d-1 2 ν 2 .

Proof. Given a fixed vector x, a uniform random vector y satisfies |x y| ≥ ν only if y lies in one of two spherical caps: one centered at x and the other at -x, and both with angular radius cos -1 (ν) ≤ π 2 -ν. A simple result often attributed to Milman & Schechtman (1986) bounds the probability of lying in either of these caps as

EQUATION

)

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

We next apply a union bound to control the probability that |x i x j | ≥ ν for some pair (i, j). There are fewer than N 2 /2 such pairs, and so the probability of this condition is

Pr[|x i x j | ≥ ν, for some i, j] ≤ N 2 2 π 2 exp - d -1 2 ν 2 .

C LOW-RANK HESSIANS LEAD TO LOW GRADIENT CONFUSION:

In this section, we slightly elaborate on the argument 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 H i (w) -H i (w ) ≤ L H ww .

Then ∇f i (w) = H i w + e, where e is an error term bounded as: e ≤ 1 2 L H w 2 , and we use the shorthand H i to denote H i (0). Then we have: If the Hessians are sufficiently random and low-rank (e.g., of the form H i = a i a i where a i ∈ R N ×r are randomly sampled from a unit sphere), then one would expect the terms in this expression to be small for all w within a neighborhood of the minimizer. Given that it has been observed that many standard neural nets have low rank Hessians at the minimizer, this indicates that the gradient confusion might be low for a large class of weights near the minimizer for these models.

| ∇f i (w), ∇f j (w) | = | H i w, H j w | + e, H i w + H j w + e 2 ≤ w 2 H i H j + e w ( H i + H j ) + e 2 ≤ w 2 H i H j + 1 2 L H w 3 ( H i + H j ) + 1 4 L 2 H w 4 .

D MISSING PROOFS:

D.1 PROOFS OF THEOREMS 3.1 AND 3.2

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.

Theorem 3.1. If the objective function satisfies (A1) and (A2), and has gradient confusion η, SGD with updates of the form (2) converges linearly to a neighborhood of the minima of problem (1) as:

E[F (w T ) -F ] ≤ ρ T (F (w 0 ) -F ) + αη 1-ρ , where α < 2 N L , ρ = 1 -2µ N α -N Lα 2 2

, F = min w F (w) and w 0 is the initialized weights.

Proof. Let ĩ ∈ [N ] denote the index of the realized function fk in the uniform sampling from {f i } i∈[N ] at step k. From assumption (A1), we have ≤

F (w k+1 ) ≤ F (w k ) + ∇F (w k ), w k+1 -w k + L 2 w k+1 -w k 2 = F (w k ) -α ∇F (w k ), ∇ fk (w k ) + Lα 2 2 ∇ fk (w k ) 2 = F (w k ) - α N - Lα 2 2 ∇ fk (w k ) 2 - α N ∀i:i = ĩ ∇f i (w k ), ∇ fk (w k ) ≤ F (w k ) - α N - Lα 2 2 ∇ fk (w k ) 2 + α(N -1)η N ,

F (w k ) - α N - Lα 2 2 ∇ fk (w k ) 2 + αη,

where the second-last inequality follows from definition 2.1. Let the learning rate α < 2/N L. Then, using assumption (A2) and subtracting by F = min w F (w) on both sides, we get

F (w k+1 ) -F ≤ F (w k ) -F -2µ α N - Lα 2 2 ( fk (w k ) -f k ) + αη,

where f k = min w fk (w). It is easy to see that by definition we have,

E i [f i ] ≤ F . Moreover, from assumption that α < 2 N L , it implies that α N -Lα 2 2 > 0.

Therefore, taking expectation on both sides we get,

E[F (w k+1 ) -F ] ≤ 1 - 2µα N + µLα 2 E[F (w k ) -F ] + αη.

Writing ρ = 1 -2µα N + µLα 2 , and unrolling the iterations, we get

E[F (w k+1 ) -F ] ≤ ρ k+1 (F (w 0 ) -F ) + k i=0 ρ i αη ≤ ρ k+1 (F (w 0 ) -F ) + ∞ i=0 ρ i αη = ρ k+1 (F (w 0 ) -F ) + αη 1 -ρ .

Theorem 3.2. If the objective satisfies (A1) and has gradient confusion bound η, then SGD converges to a neighborhood of a stationary point as:

min k=1,...,T E ∇F (w k ) 2 ≤ ρ(F (w1)-F ) T + ρη,

for learning rate α < 2 N L , ρ = 2N 2-N Lα , and F = min w F (w).

Proof. From theorem 3.1, we have:

EQUATION

Now we know that:

E ∇ fk (w k ) 2 = E ∇ fk (w k ) -∇F (w k ) 2 + E ∇F (w k ) 2 ≥ E ∇F (w k ) 2 .

Thus, taking expectation and assuming the step size α < 2/(N L), we can rewrite equation 6 as:

E ∇F (w k ) 2 ≤ 2N 2α -N Lα 2 E[F (w k ) -F (w k+1 )] + 2N η 2 -N Lα .

Taking an average over T iterations, and using F = min w F (w), we get:

min k=1,...,T E ∇F (w k ) 2 ≤ 1 T k=1 E ∇F (w k ) 2 ≤ 2N 2α -N Lα 2 F (w 1 ) -F T + 2N η 2 -N Lα .

D.2 PROOFS OF HELPER LEMMAS:

Lemma D.1. Consider the set of loss-functions {f i (W)} i∈ [N ] where all f i are either the square-loss function or the logistic-loss function. Consider a feed-forward neural network as defined in equation 4 whose weights W satisfy assumption 1,. Consider the gradient ∇ W f i (W) of each function f i . Note that we can write

∇ W f i (W) = ζ xi (W)∇ W g W (x i ), where we define ζ xi (W) = ∂f i (W)/∂g W .

Then we have the following properties.

1. When x ≤ 1 we have ∇ W g W (x i ) ≤ 1.

2. There exists a constant

ζ 0 > 0 such that |ζ xi (W)| ≤ 2 , ∇ xi ζ xi (W) 2 ≤ ζ 0 , ∇ W ζ xi (W) 2 ≤ ζ 0 .

We show that ζ 0 ≤ 2 √ β, where β is the depth of the neural network.

Proof. Let W denote the tuple

(W p ) p∈[β]0 . Consider |ζ xi (W)| = |∂f i (W)/∂g W |.

In the case of square-loss function this evaluates to |g W (x) -C(x)| ≤ 2. In case of logistic regression, this evaluates to

| -1 1+exp(C(xi)g W (xi)) | ≤ 1. Now we consider ∇ xi ζ xi (W)

. Consider the squared loss function. We then have the following.

∇ xi ζ xi (W) = ∇ xi f (W) = ∇ xi g W (x i ) -C(x i ) ≤ ∇ xi g W (x i ) + 1.

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

∇ xi ζ xi (W) ≤ C(x i ) 2 (1 + exp(C(x i )g W (x i ))) 2 i )g W (x i )) ∇ xi g (x i ) ≤ ∇ xi g W (x i ) .

Thus, it suffices to bound ∇ xi g W (x i ) . Using assumption 1 and the properties (P1), (P2) of σ, this can be upper-bounded by 1.

Consider ∇ Wp ζ xi (W) for some layer index p ∈ [β] 0 . We will show that ∇ Wp ζ xi (W) 2 ≤ 2. Then it immediately follows that ∇ W ζ xi (W) 2 ≤ 2 √ β.

In the case of a squared loss function. We have the following.

∇ Wp ζ xi (W) = ∇ Wp f (W) = ∇ Wp g W (x i ) -C(x i ) ≤ ∇ Wp g W (x i ) + 1.

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

∇ Wp ζ xi (W) ≤ C(x i ) 2 (1 + exp(C(x i )g W (x i ))) 2 exp(C(x i )g W (x i )) ∇ Wp g W (x i ) ≤ ∇ Wp g W (x i ) . Since ∇ Wp g W (x i ) ≤ 1, we have that ∇ Wp ζ xi (W) ≤ 2 in both the cases. Thus, ζ 0 = 2 √ β.

D.3 PROOFS OF THEOREM 4.1 AND COROLLARY 4.1

In this section, we will analyze the case when we have non-linear activations at each layer.

Theorem 4.1. Consider the problem of training neural nets (equation 4) using either the square-loss or the logistic-loss function. Let η > 0 be a given constant. Let the weights satisfy assumption 1 and the non-linearities in each layer satisfy properties (P1) and (P2). For some fixed constant c > 0, the gradient confusion bound in equation 3 holds with probability at least 1 -N 2 exp -cdη 2 16ζ 4 0 (β+2) 4 , For both the square-loss and the logistic-loss functions, ζ 0 ≤ 2 √ β (from lemma D.1).

Proof. 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 ∈ [n] define ζ xi (W) := f (W). For every p ∈ [β] define H p as follows. H p (x) := σ(W p • σ(W p-1 • σ(. . . • σ(W 0 • x) . . .).

Fix an i ∈ [N ]. Then we have the following recurrence

g β (x i ) := σ (H β (x i )) g p (x i ) := (W p+1 • g p+1 (x i )) • Diag(σ (H p (x i ))) ∀p ∈ {0, 1, . . . , β -1}.

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

∇ Wp f i (W) = g p (x i ) • H p-1 i ) ∀p ∈ [β] 0 .

We can write h W (x i , x j ) as follows.

EQUATION

We will now bound

∇ (xi,xj ) h W (x i , x j ) 2 . Consider ∇ xi h W (x i , x j ).

This can be written as follows.

EQUATION

Observe that each of the entries in the diagonal matrix Diag(σ (H p (x i ))) is at most 1. Thus, we have that Diag(σ (H p (x i ))) ≤ 1.

We have the following relationship.

g β (x i ) ≤ 1 g p (x i ) ≤ W p+1 g p+1 (x i )) Diag(σ (H p (x i ))) ≤ 1 ∀p ∈ {0, 1, . . . , β -1}.

Moreover we have,

Tr[H p-1 (x i )•g p (x i ) •g p (x j )•H p-1 (x i ) ] ≤ H p-1 (x i ) g p (x i ) g p (x j ) H p-1 (x i ) ≤ 1. Consider ∇ xi H p-1 (x i ) • g p (x i ) • g p (x j ) • H p-1 (x i ) for every p ∈ [β] 0 .

This can be upper-bounded by,

∇ xi H p-1 (x i ) g p (x i ) g p (x j ) H p-1 (x i ) + H p-1 (x i ) ∇ xi g p (x i ) g p (x j ) H p-1 (x i ) . Note that ∇ xi H p-1 (x i ) = g 1 (x i ) • Diag(σ (W 0 • x i )) • W 0 • g p (x i ) . Thus, ∇ xi H p-1 (x i ) ≤ 1.

We will now show that ∇ xi g p (x i ) ≤ β -p + 1. We prove this inductively. Consider the base case when p = β.

∇ xi g β (x i ) = ∇ xi σ (H β (x i )) ≤ 1 = β -β + 1.

Now, the inductive step.

∇ xi g p (x i ) ≤ ∇ xi g p+1 (x i ) + ∇ xi Diag(σ (H p (x i ))) ≤ β -p ≤ β -p + 1.

Thus, using equation 8 and the above arguments, we obtain,

∇ xi h W (x i , x j ) 2 ≤ ζ 2 0 (β + 1) + ζ 2 0 (β + 1)(β + 2) ≤ 2ζ 2 0 (β + 2) 2

and thus, ∇ (xi,xj ) h W (x i , x j ) 2 ≤ 4ζ 2 0 (β + 2) 2 .

Proof. Define the function h + (W) = max ij h W (x i , x j ). Our goal is to find conditions under which h + (W) > -η for all W in a large set. To derive such conditions, we will need a Lipschitz constant for h + (W), which is no larger than the maximal Lipschitz constant of h W (x i , x j ) for all i, j. We have that ∇ Using lemma D.2, we see that 2ζ 2 0 is a Lipschitz constant for h W (x i , x j ), and thus also h + (W). Now, consider a minimizer W of the objective, and a ball B r around this point of radius r. Define the constant = η 4ζ 2 0 , and create an -net of points N = {W i } inside the ball. This net is sufficiently dense that any W ∈ B r is at most units away from some W i ∈ N . Furthermore, because h + (W) is Lipschitz in W, |h + (W ) -h + (W i )| ≤ 2ζ 2 0 = η/2. We now know the following: if we can guarantee that

EQUATION

then we also know that h + (W ) ≥ -η for all W ∈ B r . For this reason, we prove the result by bounding the probability that (10) holds. It is known that N can be constructed so that |N | ≤ (2r/ + 1) d = (8ζ 2 0 r/η + 1) d (see Vershynin (2018) , corollary 4.1.13). Theorem 4.1 provides a bound on the probability that each individual point in the net satisfies condition (10). Using a union bound, we see that all points in the net satisfy this condition with probability at least 1 -N 2 8ζ 2 0 r η + 1 .

Both parts in theorem 4.2 depend on the following argument. From theorem 2.3.8 and Proposition 2.3.10 in Tao (2012) with appropriate scaling 5 , we have for every p = 1, . . . , β we have that the matrix norm W p ≤ 1 with probability at least 1 -β exp -c 1 κ 2 2 and W 0 ≤ 1 with probability at least 1 -exp -c 1 κ 2 d 2 when the weight matrices are initialized according to strategy 4.1. Thus, conditioning on this event it implies that these matrices satisfy assumption 1. The proof strategy is similar to that of theorem 4.1. We will first show that the gradient of the function h(., .) as defined in equation ( 7) with respect to the weights is bounded. Note that in part (1) the random variable is the set of weight matrices {W p } p∈ [β] . Thus, the dimension used to invoke theorem E.1 is at most 2 β. In part (2) along with the weights, the data x ∈ R d is also random. Thus, the dimension used to invoke theorem E.1 is at most d + 2 β. Combining this with theorem E.1, the bound on the gradient of h(., .) and taking a union bound, we get the respective parts of the theorem. Thus, all it remains to prove is the bound on the gradient of the function h(., .) as defined in equation ( 7) with respect to the weights conditioning on the event that W p ≤ 1 for every p ∈ {0, 1, . . . , β}.

We obtain the following analogue of equation ( 8).

EQUATION

As in the case of the proof for theorem 4.1, we will upper-bound the 2 -norm of the above expression.

In particular, we show the following. (18)

EQUATION

Equations ( 16) and 17 follow from the the fact that (∇ W ζ xi (W)) 2 ≤ ζ 0 and the arguments in the proof for theorem 4.1. We will now show the proof sketch for equation (18). For every p ∈ [β] 0 , consider ∇ W H p-1 (x i ) • g p (x i ) • g p (x j ) • H p-1 (x i ) . Using the symmetry between x i and x j , the expression can be upper-bounded by,

2 ∇ W H p-1 (x i ) g p (x i ) g p (x j ) H p-1 (x i ) +2 H p-1 (x i ) ∇ W g p (x i ) g p (x j ) H p-1 (x i ) .

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

∇ W H p-1 (x i ) • g p (x i ) • g p (x j ) • H p-1 (x i ) ≤ 4(β + 2) 2 .

Next, we show that the expected value can be lower-bounded by -4 as in the case of theorem 4.2 above. Combining these two gives us the desired result. Consider E W [h W (x i , x j )]. We compute this expectation iteratively as follows.

E W [h W (x i , x j )] = E W0 [E W1 [. . . E W β [h W (x i , x j )] ≥ -4E W0   E W1   . . . E W β   p∈[β]0

Tr(H p-1 (x i ) • g p (x i ) • g p (x j ) • H p-1 (x i ) )

      .

In the proof of theorem 4.2 we showed that assumption 1 holds at initialization with high probability. 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 nets (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.

:

Due to space constraints, all proofs and several additional experiments are delegated to the appendix.

This is related to both gradient variance and gradient diversity(Yin et al., 2017), but with important differences. See section 6 for more discussion.

Generally, this is true whenever xi = 1 √ d yi, where yi is an isotropic random vector(Vershynin, 2018).

The width factor denotes the number of filters relative to the original ResNet model.

In particular, each entry has to be scaled by 1 for matrices {Wp} p∈[β] and 1 d for the matrix W0.

F , which keep the weights small during optimization. The operator norm of convolutional layers have also recently been used as an effective regularizer as well for image classification tasks bySedghi et al. (2018).

Figures and Tables:

Figure1: Linear regression on an over-parameterized (d = 120) and under-parameterized (d = 80) model with N = 100 samples generated randomly from a Gaussian, trained using SGD with minibatch size 1. Plots are averaged over 3 independent runs. Gradient cosine similarities were calculated over all pairs of gradients.

P2) Bounded differentials: Let σ (x) and σ (x) denote the first and second subdifferentials respectively. Then, |σ (x)| ≤ 1 and |σ (x)| ≤ 1 for all x ∈ [-1, 1].

Figure 2: The effect of network depth with CNN-β-2 on CIFAR-10. Left plot: convergence curves of SGD, Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training (over all independent runs).

Figure 3: The effect of width with CNN-16-on CIFAR-10. Left plot: convergence curves of SGD (for cleaner figures, we plot results for width factors 2, 4 and 6 here), Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training (over all independent runs).

Figure 4: The effect of adding skip connections and batch normalization to CNN-β-2 on CIFAR-10. Plots show the optimal training loss (left plot), minimum pairwise gradient cosine similarities (middle plot), and test set accuracies (right plot) at the end of training.

(2018b)), and that convergence typically slows down because of gradient confusion when training very deep networks. There has recently also been interest in analyzing conditions under which SGD converges to global minimizers of overparameterized linear and non-linear neural networks. Arora et al. (2018a) shows SGD converges linearly to global minimizers for linear neural nets under certain conditions. Du et al. (2018); Allen-Zhu et al. (2018); Zou et al. (2018); Brutzkus et al. (

Figure 5: Effect of varying depth on MLP-β-100.

Figure 6: Effect of varying width on MLP-300-.

Figure 7: The effect of network depth with CNN-β-2 on CIFAR-10. Left plot: final training loss values at the end of training, Middle plot: final test set accuracy values at the end of training. Right plot: curves showing the minimum of pairwise gradient cosine similarities during training.

Figure 8: The effect of width with CNN-16-on CIFAR-10. Left plot: final training loss values at the end of training, Middle plot: final test set accuracy values at the end of training. Right plot: curves showing the minimum of pairwise gradient cosine similarities during training.

Figure 9: The effect of network depth with CNN-β-2 on CIFAR-100. Left plot: training loss values at the end of training. Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training.

Figure 10: The effect of width with CNN-16-on CIFAR-100. Left plot: training loss values at the end of training. Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training.

Figure 11: The effect of depth with WRN-β-2 (no batch normalization) on CIFAR-10. Left plot: training loss values at the end of training. Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training.

Figure 13: The effect of width with WRN-16-(no batch normalization) on CIFAR-10. Left plot: training loss values at the end of training. Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training.

Figure 14: The effect of width with WRN-16-(no batch normalization) on CIFAR-100. Left plot: training loss values at the end of training. Middle plot: minimum of pairwise gradient cosine similarities at the end of training, Right plot: kernel density estimate of the pairwise gradient cosine similarities at the end of training.

Figure 15: The effect of adding skip connections and batch normalization to CNN-β-2 on CIFAR-100. Plots show the training loss, minimum pairwise gradient cosine similarities, and test accuracies at the end of training.

W f i = ζ xi (W)x i ≤ ζ 0 . Now we need to get a W-Lipschitz constants for ∇ xi f i = ζ xi (W)x i . By lemma D.1, we have ∇ W (ζ xi (W)x i ) = (∇ W ζ xi (W))x i ≤ ζ 0 .

that, if r < , then we can form a net with |N | = 1. In this case, the probability of satisfying (10) is at least 1 -N 2 exp -Theorem 4.2 (Neural nets with randomly chosen weights). Let W 0 , W 1 , . . . , W β be weight matrices chosen according to strategy 4.1. There exists fixed constants c 1 , c 2 > 0 such that we have:1. Consider a fixed but arbitrary dataset x 1 , x 2 , . . . , x N with x i ≤ 1 for every i ∈ [N ]. For η > 4, the gradient confusion bound in equation 3 holds with probability at least1 -β exp -c 1 κ 2 2 -N 2 exp -c 2 β(η-4) the dataset {x i } i∈[N ] is such that each x i is an i.i.d.sample from the surface of ddimensional unit sphere, then for every η > 0 the gradient confusion bound in equation 3 holds with probability at least 1 -β exp -c 1 κ 2 2 -N 2 exp -c2( d+ 2 β)η 2

)(∇ W ζ xj (W))ζ xi (W) p-1 (x i ) • g p (x i ) • g p (x j ) • H p-1 (x i ) ] W H p-1 (x i ) • g p (x i ) • g p (x j ) • H p-1 (x i )

Peer review (iclr2020):
This paper introduces the concept of *gradient confusion* to explain why neural networks train fast with SGD. They also study the effects of width, depth on gradient confusion. - The theoretical results assume that the data is sampled from a sphere and do not really give much insight into the effect of width and depth. - There are some confounding factors in the experiments and there needs to be a better comparison to some related work. Detailed review below: Section 1: - Please clarify how *gradient confusion* relate to the interpolation conditon of Ma et al, 2017 and the strong growth condition of Vaswani et al? - If we run SGD with a constant step-size, it will bounce around the optimal point in a ball with radius that depends on the step-size. If I keep decreasing the step-size, this radius shrinks. How does gradient confusion relate to the step-size? Is it upper-bounded by a quantity that depends on the step-size and the batch-size? Section 2: - Definition 2.1: Why should this condition hold for *all* points w? Isn*t it necessary only at w^* or in a small neighborhood around it? - The gradient confusion parameter eta should depend on the batch-size. Please clarify this. - Figure 1: Previous work (Ma et al, Vaswani et al, Gunasekhar, 2017) all have shown that fast convergence can be obtained using SGD with a constant step-size with over-parametrized models and explained it using interpolation. What is the additional insight from gradient confusion? - *Suppose that there is a Lipschitz constant for the Hessian* - This is a strong assumption and a vague argument, that is confusing rather than insightful. Please justify why this is a valid assumption for neural network models. Section 3: - If E_i || 
abla f_i(w) ||^2 = O(epsilon) => gradient confusion = O(epsilon). Isn*t the E_i || 
abla f_i(w) ||^2 exactly the strong growth condition in Vaswani, et al. Can the gradient confusion results be directly derived from the results in that paper? Please compare. Also compare and cite *Stochastic Approximation of Smooth and Strongly Convex Functions: Beyond the O(1/T ) Convergence Rate*, COLT 2019. Section 4: - Please compare against the previous results that assumed the data to be sampled from a sphere. - Thm 4.1: The theorem bounds the probability that gradient confusion holds for a given eta. But the bounds of section 3 are vacuous even if the theorem holds with probability one, but for a large value of eta. There needs to be an upper bound on eta. Please clarify this. - Please compare against the results of this paper by Arora et al: *On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization* - For the effect of layer width, the analysis is only for the initializated weights and does not consider the optimization, which is what the paper claimed in the introduction. Am I missing something? Please justify this. The gradient confusion can decrease as the optimization progresses? Section 5: - *We reduce the learning rate by a factor of 10* But all the theory is for a constant step-size. Please explain this discrepancy. - in Figure 2, in the second figure, why is there a sharp full in the pairwise cosine similarities. - In all these experiments, explain why the batch-size and the step-size is not a confounding factor?