Understanding the Generalization of Adam in Learning Neural Networks with Proper Regularization

1 INTRODUCTION . Adaptive gradient methods ( Duchi et al. , 2011 ; Hinton et al. , 2012 ; Kingma & Ba , 2015 ; Reddi et al. , 2018 ) such as Adam are very popular optimizers for training deep neural networks . By adjusting the learning ratethis coordinate-wisely based on historical gradient information , they are known to be able to automatically choose appropriate learning rates to achieve fast convergence in training . Because of this advantage , Adam and its variants are widely used in deep learning . Despite their fast convergence , adaptive gradient methods have been observed to achieve worse generalization performance compared with gradient descent and stochastic gradient descent ( SGD ) ( Wilson et al. , 2017 ; Luo et al. , 2019 ; Chen et al. , 2020 ; Zhou et al. , 2020 ) in many deep learning tasks such as image classification ( we have done some simple deep learning experiments to justify this , the results are reported in Table 1 ) . Even with proper regularization , achieving good test error with adaptive gradient methods seems to be challenging . Several recent works provided theoretical explanations of this generalization gap between Adam and GD . Wilson et al . ( 2017 ) ; Agarwal et al . ( 2019 ) considered a setting of linear regression , and showed that Adam can fail when learning an overparameterized linear model on certain specifically designed data , while SGD can learn the linear model to achieve zero test error . This ex- ample in linear regression offers valuable insights into the difference between SGD and Adam . However , it is under a convex optimization setting , and as we will show in this paper ( Theorem 4.2 ) , the performance difference between Adam and GD can be easily avoided by adding an arbitrarily small regularization term , because the regularized training loss function is strongly convex and all algorithms will converge to the same unique global optimum . For this reason , we argue that the example in the convex setting can not capture the fundamental differences between GD and Adam . More recently , Zhou et al . ( 2020 ) studied the expected escaping time of Adam and SGD from a local basin , and utilized this to explain the difference between SGD and Adam . However , their results do not take NN architecture into consideration , and do not provide an analysis of test errors either . In this paper , we aim at answering the following question Why is there a generalization gap between Adam and gradient descent in learning neural networks , even with proper regularization ? Specifically , we study Adam and GD for training neural networks with weight decay regularization on an image-like data model , and demonstrate the difference between Adam and GD from a feature learning perspective . We consider a model where the data are generated as a combination of feature and noise patches under certain sparsity conditions , and analyze the convergence and generalization of Adam and GD for training a two-layer convolutional neural network ( CNN ) . The contributions of this paper are summarized as follows . • We establish global convergence guarantees for Adam and GD with proper weight decay regularization . We show that , starting at the same random initialization , Adam and GD can both train a two-layer convolutional neural network to achieve zero training error after polynomially many iterations , despite the nonconvex optimization landscape . • We further show that GD and Adam in fact converge to different global solutions with different generalization performance : GD can achieve nearly zero test error , while the generalization performance of the model found by Adam is no better than a random guess . In particular , we show that the reason for this gap is due to the different training behaviors of Adam and GD : Adam is more likely to fit noises in the data and output a model that is largely contributed by the noise patches of the training data ; GD prefers to fit training data based on their feature patch and finds a solution that is mainly composed by the true features . We also illustrate such different training processes in Figure 1 , where it can be seen that the model trained by Adam is clearly more “ noisy ” than that trained by SGD . • We also show that for convex settings with weight decay regularization , both Adam and gradient descent converge to the exact same solution and therefore have no test error difference . This suggests that the difference between Adam and GD can not be fully explained by linear models or neural networks trained in the “ almost convex ” neural tangent kernel ( NTK ) regime Jacot et al . ( 2018 ) ; Allen-Zhu et al . ( 2019b ) ; Du et al . ( 2019a ) ; Zou et al . ( 2019 ) ; Allen-Zhu et al . ( 2019a ) ; Arora et al . ( 2019a ; b ) ; Cao & Gu ( 2019 ) ; Ji & Telgarsky ( 2020 ) ; Chen et al . ( 2021 ) . It also demonstrates that the inferior generalization performance of Adam is fundamentally tied to the nonconvex landscape of deep learning optimization , and can not be solved by adding regularization . Notation . For a scalar x , we use [ x ] + to denote max { x , 0 } . For a vector v = ( v1 , . . . , vd ) > , we denote by ‖v‖2 : = ( ∑d j=1 v 2 j ) 1/2 its ` 2-norm , and use supp ( v ) : = { j : vj 6= 0 } to denote its support . 2 RELATED WORK . In this section , we discuss the works that are mostly related to our paper . Generalization gap between Adam and ( stochastic ) gradient descent . The worse generalization of Adam compared with SGD has also been observed by some recent works and has motivated new variants of neural network training algorithms . Keskar & Socher ( 2017 ) proposed to switch between Adam and SGD to achieve better generalization . Merity et al . ( 2018 ) proposed a variant of the averaged stochastic gradient method to achieve good generalization performance for LSTM language models . Luo et al . ( 2019 ) proposed to use dynamic bounds on learning rates to achieve a smooth transition from adaptive methods to SGD to improve generalization . Our theoretical results for GD and Adam can also provide theoretical insights into the effectiveness of these empirical studies . Optimization and generalization guarantees in deep learning . Our work is also closely related to the recent line of work studying the optimization and generalization guarantees of neural networks . A series of results have shown the convergence ( Jacot et al. , 2018 ; Li & Liang , 2018 ; Du et al. , 2019b ; Allen-Zhu et al. , 2019b ; Du et al. , 2019a ; Zou et al. , 2019 ) and generalization ( Allen-Zhu et al. , 2019c ; a ; Arora et al. , 2019a ; b ; Cao & Gu , 2019 ; Ji & Telgarsky , 2020 ; Chen et al. , 2021 ) guarantees in the so-called “ neural tangent kernel ” ( NTK ) regime , where the neural network function is approximately linear in its parameters . Allen-Zhu & Li ( 2019 ) ; Bai & Lee ( 2019 ) ; Allen-Zhu & Li ( 2020a ) ; Li et al . ( 2020 ) studied the learning of neural networks beyond the NTK regime . Our analysis in this paper is also beyond NTK , and gives a detailed comparison between GD and Adam . Feature learning by neural networks . This paper is also closely related to several recent works that studied how neural networks can learn features . Allen-Zhu & Li ( 2020b ) showed that adversarial training purifies the learned features by removing certain “ dense mixtures ” in the hidden layer weights of the network . Allen-Zhu & Li ( 2020c ) studied how ensemble and knowledge distillation work in deep learning when the data have “ multi-view ” features . This paper studies a different aspect of feature learning by Adam and GD , and shows that GD can learn the features while Adam may fail even with proper regularization . 3 PROBLEM SETUP AND PRELIMINARIES . We consider learning a CNN with Adam and GD based on n independent training examples { ( xi , yi ) } ni=1 generated from a data model D. In the following . we first introduce our data model D , and then explain our neural network model and the details of the training algorithms . Data model . We consider a data model where the data inputs consist of feature and noise patches . Such a data model is motivated by image classification problems where the label of an image usually only depends on part of an image , and the other parts of the image showing random objects , or features that belong to other classes , can be considered as noises . When using CNN to fit the data , the convolution operation is applied to each patch of the data input separately . We claim that our data model is more practical than those considered in Wilson et al . ( 2017 ) ; Reddi et al . ( 2018 ) , which are handcrafted for showing the failure of Adam in term of either convergence or generalization . For simplicity , we only consider the case where the data consists of one feature patch and one noise patch . However , our result can be easily extended to cover the setting where there are multiple feature/noise patches . The detailed definition of our data model is given in Definition 3.1 as follows . Definition 3.1 . Let d = Ω ( n4 ) , each data ( x , y ) with x ∈ R2d and y ∈ { −1 , 1 } is generated as follows , x = [ x > 1 , x > 2 ] > , where one of x1 and x2 denotes the feature patch that consists of a feature vector y · v , which is assumed to be 1-sparse , and the other one denotes the noise patch and consists of a noise vector ξ . Without loss of generality , we assume v = [ 1 , 0 , . . . , 0 ] > . The noise vector ξ is generated according to the following process : 1 . Randomly select s coordinates from [ d ] \ { 1 } with equal probabilities , which is denoted as a vector s ∈ { 0 , 1 } d. 2 . Generate ξ from distributionN ( 0 , σ2pI ) , and then mask off the first coordinate and other d− s− 1 coordinates , i.e. , ξ = ξ s. 3 . Add feature noise to ξ , i.e. , ξ = ξ − αyv , where 0 < α < 1 is the strength of the feature noise . In particular , throughout this paper we set s = Θ ( d1/2 n2 ) , σ2p = Θ ( 1 s·polylog ( n ) ) and α = Θ ( σp · polylog ( n ) ) . The most natural way to think of our data model is to treat x as the output of some intermediate layer of a CNN . In literature , Papyan et al . ( 2017 ) pointed out that the outputs of an intermediate layer of a CNN are usually sparse . Yang ( 2019 ) also discussed the setting where the hidden nodes in such an intermediate layer are sampled independently . This motivates us to study sparse features and entry-wisely independent noises in our model . In this paper , we focus on the case where the feature vector v is 1-sparse and the noise vector is s-sparse for simplicity . However , these sparsity assumptions can be generalized to the settings where the feature and the noises are denser . Moreover , we would like to clarify that the data distribution considered in our paper is an extreme case where we assume there is only one feature vector and all data has a feature noise , since we believe this is the simplest model that captures the fundamental difference between Adam and SGD . With this data model , we aim to show why Adam and SGD perform differently . Our theoretical results and analysis techniques can also be extended to more practical settings where there are multiple feature vectors and multiple patches , each data can either contain a single feature or multiple features , together with pure random noise or feature noise . Two-layer CNN model . We consider a two-layer CNN model F using truncated polynomial activation function σ ( z ) = ( max { 0 , z } ) q and fixed the weights of second layer to be all 1 ’ s , where q ≥ 3 . Mathematically , given the data ( x , y ) , the j-th output of the CNN can be formulated as Fj ( W , x ) = m∑ r=1 [ σ ( 〈wj , r , x1〉 ) + σ ( 〈wj , r , x2〉 ) ] = m∑ r=1 [ σ ( 〈wj , r , y · v〉 ) + σ ( 〈wj , r , ξ〉 ) ] , ( 3.1 ) where m is the width of the network , wj , r ∈ Rd denotes the weight at the r-th neuron , and W is the collection of model weights . Here the use of the polynomial ReLU activation function is for simplifying our analysis . It can be replaced by a smoothed ReLU activation function ( e.g. , the activation function used in Allen-Zhu & Li ( 2020c ) ) . If we assume the input data distribution is Gaussian , we can also deal with ReLU activation function ( Li et al. , 2020 ) . Moreover , we would like to emphasize that x1 and x2 denote two data patches , which are randomly assigned with feature vector or noise vector independently for each data point . The leaner has no knowledge about which one is the feature patch ( or noise patch ) . In this paper we assume the width of the network is polylogarithmic in the training sample size , i.e. , m = polylog ( n ) . We assume j ∈ { −1 , 1 } in order to make the logit index be consistent with the data label . Moreover , we assume that the each weight is initialized from a random draw of Gaussian random variable ∼ N ( 0 , σ20 ) with σ0 = Θ ( d−1/4 ) . Training objective . Given the training data { ( xi , yi ) } i=1 , ... , n , we consider to learn the model parameter W by optimizing the empirical loss function with weight decay regularization L ( W ) = 1 n n∑ i=1 Li ( W ) + λ 2 ‖W‖2F , ( 3.2 ) where Li ( W ) = − log e Fyi ( W , xi ) ∑ j∈ { −1,1 } e Fj ( W , xi ) denotes the individual loss for the data ( xi , yi ) and λ ≥ 0 is the regularization parameter . In particular , the regularization parameter can be arbitrary as long as it satisfies λ ∈ ( 0 , λ0 ) with λ0 = Θ ( 1 d ( q−1 ) /4n·polylog ( n ) ) . We claim that the λ0 is the largest feasible regularization parameter that the training process will not stuck at the origin point ( recall that L ( W ) admits zero gradient at W = 0 . ) Training algorithms . In this paper , we consider gradient descent and Adam with full gradient . In particular , starting from initialization W ( 0 ) = { w ( 0 ) j , r , j = { ±1 } , r ∈ [ m ] } , the gradient descent update rule is w ( t+1 ) j , r = w ( t ) j , r − η · ∇wj , rL ( W ( t ) ) , where η is the learning rate . Meanwhile , Adam store historical gradient information in the momentum m ( t ) and a vector v ( t ) as follows m ( t+1 ) j , r = β1m ( t ) j , r + ( 1− β1 ) · ∇wj , rL ( W ( t ) ) , ( 3.3 ) v ( t+1 ) j , r = β2v ( t ) j , r + ( 1− β2 ) · [ ∇wj , rL ( W ( t ) ) ] 2 , ( 3.4 ) and entry-wisely adjusts the learning rate : w ( t+1 ) j , r = w ( t ) j , r − η ·m ( t ) j , r/ √ v ( t ) j , r , ( 3.5 ) where β1 , β2 are the hyperparameters of Adam ( a popular choice in practice is β1 = 0.9 , and β2 = 0.99 ) , and in ( 3.4 ) and ( 3.5 ) , the square ( · ) 2 , square root √ · , and division ·/· all denote entrywise calculations . We do not consider the initialization bias correction in the original Adam paper ( Kingma & Ba , 2015 ) for the ease of analysis .

Summary: This paper propose a non convex optimization problem where the batch version of Adam has worse generalization than GD. Both methods achieve 0 training loss asymptotically but GD has a zero loss on the test set, while Adam has a large loss.  The paper studies a simplified CNN model with only two patches. The two layers are multiplied with the same weight matrix, fed into the activation function and then summed. A cross entropy loss is used for the classification. The activation function must be of the form $\max(0, x^q)$ for $q \geq 3$. Note that the second layer weights are not learnable and fixed to 1.  The dataset is built as follow: the label is y in {-1, 1}. One of the patch is just $x_1 = [y, 0, 0, ...]$. The second patch is $x_2 = [-\alpha y, \text{sparse gaussian noise}]$ ($0 <\alpha < 1$). The sparsity level is so high that no two examples have the same features.  The idea is that Adam will use a much larger effective learning rate for the sparse noise features than SGD, and grow the coordinates quickly. At some point the dot product with the noise in $x_2$ will be larger than the dot product with $x_1$. Because the activation function is of the form $x^q$ with $q \geq3$, this means the gradient from $x_2$ will overtake the one from $x_1$ and the weight for the first dimension will be negative (while a positive weight is the best thing to classify from $x_1$). So Adam will focus on the weaker signal from $x_2$ (remember $\alpha < 1$).  To summarize further, one patch has the right answer, the second path has the right answer (but with opposite sign) + noise. Adam will focus on the second patch which has lower signal, while GD will focus on the first patch and ignore the noise.  The authors provide experimental verification of this fact.