The Effect of Neural Net Architecture on Gradient Confusion & Training Performance

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 ith 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 { fi } i∈ [ N ] ) , we use SGD to solve empirical risk minimization problems of the form , minw∈Rd F ( w ) : = minw∈Rd 1 N ∑N i=1 fi ( w ) , ( 1 ) using the following iterative update rule for T rounds : wk+1 = wk − α∇f̃k ( wk ) . ( 2 ) Here α is the learning rate and f̃k is a function chosen uniformly at random from { fi } i∈ [ N ] at iteration k ∈ [ T ] . We use w ? to denote the optimal solution , i.e. , w ? = argminw F ( w ) . 1Due to space constraints , all proofs and several additional experiments are delegated to the appendix . Gradient confusion . SGD works by iteratively selecting a random function f̃k , and modifying the parameters to move in the direction of the negative gradient of the objective term f̃k . It may happen that the selected gradient ∇f̃k is negatively correlated with the gradient of another term ∇fj . 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 { fi } i∈ [ N ] has gradient confusion bound η ≥ 0 if the pair-wise inner products between gradients satisfy , for a fixed w ∈ Rd , 〈∇fi ( w ) , ∇fj ( w ) 〉 ≥ −η , ∀i 6= j ∈ [ N ] . ( 3 ) 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〈∇fi ( w ) , ∇fj ( w ) 〉/N2 ≥ −η . 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/B2 , 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 fi ( w ) = ` ( yix > i w ) , where ` : R→ R is the logistic loss , { xi } i∈ [ N ] is a set of orthogonal training vectors , and yi ∈ { −1 , 1 } is the label for the ith training example . We then have ∇fi ( w ) = ζixi , where ζi = yi ` ′ ( yi · x > i w ) . Note that the gradient confusion is 0 since 〈∇fi ( w ) , ∇fj ( w ) 〉 = ζiζj〈xi , xj〉 = 0 , ∀i , j ∈ [ N ] and i 6= j . Thus , an update in the gradient direction fi has no effect on the loss value of fj for i 6= 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 : |〈∇fi ( w ) , ∇fj ( w ) 〉| = |〈xi , xj〉| · |ζ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 maxij |〈xj , xi〉| is O ( 1/ √ d ) whenever { xi } 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 fi 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 fi that satisfies ‖Hi ( w ) −Hi ( w′ ) ‖ ≤ LH‖w−w′‖ ( note that this is a standard optimization assumption ( Nesterov , 2018 ) , with evidence that it is applicable for neural nets ( Martens , 2016 ) ) . Then ∇fi ( w ) = Hiw + e , where e is an error term bounded as : ‖e‖ ≤ 12LH‖w‖ 2 , and we use the shorthand Hi to denote Hi ( 0 ) . Then we have ( appendix C ) : |〈∇fi ( w ) , ∇fj ( w ) 〉| ≤ ‖w‖2‖Hi‖‖Hj‖+ 12LH‖w‖ 3 ( ‖Hi‖+ ‖Hj‖ ) + 14L 2 H‖w‖4 . 2This is related to both gradient variance and gradient diversity ( Yin et al. , 2017 ) , but with important differences . See section 6 for more discussion . 3Generally , this is true whenever xi = 1√ d yi , where yi is an isotropic random vector ( Vershynin , 2018 ) . If the Hessians are sufficiently random and low-rank ( e.g. , of the form Hi = aia > i where ai ∈ RN×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 .

Summary: This paper introduces gradient confusion, a bound on the negated dot product of gradients at two data points, and studies its effect on the optimization of neural networks with SGD. If gradient confusion is high, reducing the loss on one data point with SGD increases it on another data point. Theoretical results show that (1) with a fixed learning rate, lower gradient confusion results in faster learning, and (2) increasing the width and decreasing the depth of a network reduces gradient confusion. The experiments corroborate these findings and also show that batch normalization and skip connections can reduce gradient confusion and speed up the training.