Understanding the Generalization of Adam in Learning Neural Networks with Proper Regularization
Authors: Difan Zou
Yuan Cao
Yuanzhi Li
Quanquan Gu
Venue: ICLR 2023 poster (2023)
Keywords: generalization
gradient descent
Neural tangent kernel
Adam
nonconvex optimization

Abstract: Adaptive gradient methods such as Adam have gained increasing popularity in deep learning optimization. However, it has been observed in many deep learning applications such as image classification, Adam can converge to a different solution with a worse test error compared to (stochastic) gradient descent, even with a fine-tuned regularization.  In this paper, we provide a theoretical explanation for this phenomenon: we show that in the nonconvex setting of learning over-parameterized two-layer convolutional neural networks starting from the same random initialization, for a class of data distributions (inspired from image data), Adam and gradient descent (GD) can converge to different global solutions of the training objective with provably different generalization errors, even with weight decay regularization. In contrast, we show that if the training objective is convex, and the weight decay regularization is employed, any optimization algorithms including Adam and GD will converge to the same solution if the training is successful. This suggests that the generalization gap between Adam and SGD in the presence of weight decay regularization is closely tied to the nonconvex landscape of deep learning optimization, which cannot be covered by the recent neural tangent kernel (NTK) based analysis. 

Published as a conference paper at ICLR 2023
UNDERSTANDING THE GENERALIZATION OF ADAM IN LEARNING NEURAL NETWORKS WITH PROPER REGU-
LARIZATION
Difan Zou', Yuan Cao², Yuanzhi Li Quanquan Gu
1 Department of Computer Science & Institute of Data Science, The University of Hong Kong dzouřcs hku. hk
2 Department of Statistics & Actuarial Science, The University of Hong Kong yuancaо@hku. hk
8 Machine Learning Department, Carnegie Mellon University
yuanzhi 1 andrew. cmu . edu
4 Department of Computer Science, University of California, Los Angeles qgußcs wwl. edu
ABSTRACT
Adaptive gradient methods such as Adam have gained increasing popularity deep learning optimization. However, it has been observed in many deep learn- ing applications such as image classification, Adam can converge to a different solution with a worse test eπror compared to (stochastic) gradient descent, even with a fine-tuned regularization. In this paper, we provide a theoretical expla- show that in the nonconvex setting of learn-
nation for this phenomenon: ing aver-parameterized two-layer convolutional neural networks starting from the same random initialization, for class of data distributions (inspired from image data), Adam and gradient descent (GD) can converge to different global solutions of the training objective with provably different generalization errors, even with weight decay regularization. In contrast, we show that if the training objective is convex, and the weight decay regularization is employed, any optimization algo- rithms including Adam and GD will converge to the same solution if the training is successful. This suggests that the generalization gap between Adam and SGD in the presence of weight decay regularization is closely tied to the nonconvex landscape of deep learning optimization, which cannot be covered by the recent neural tangent kernel (NTK) based analysis.
INTRODUCTION
Adaptive gradient methods (Duchi et al., 2011: Hinton et al., 2012; Kingma & Ba, 2015; Reddi et al., 2018) such Adam very popular optimizers for training deep neural networks. By adjusting the learning rate coordinate-wisely based on historical gradient information, they are known 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 gra- dient 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 exper- iments to justify this, the results are reported in Table 1). Even with explicit weight decay regu- larization, achieving good test error with adaptive gradient methods seems to be challenging. More- over, have also visualized the first layer of AlexNet trained by Adam and SGD in Figure 1, where
Models
AlexNet
VGG-16
ResNet-18
SGD
75.22
93.25
94.62
Adam
73.08
92.19
92.93
Table Test accuracy (%) comparison between Adam and SGD on the CIFAR-10 dataset.

---

Published as a conference paper at ICLR 2023
can also observe clear difference between Adam and SGD: the model learned by Adam is more "noisy" than that learned by SGD.
Several recent works provided theoretical explanations of this generalization gap between Adam and GD by showing that Adam and GD have different implicit bias. Wilson et al. (2017); Agarwal et al. (2019) considered 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 achieve zero test error. This example in linear regression offers valuable insights into the difference between SGD and Adam. However, there is a gap between their theoretical results and the practical observations, since they consider convex optimization setting, and the difference between Adam and SGD will no longer be observed when adding weight decay regularization. In fact, as we will show in this paper (Theorem 4.2), regularization can successfully correct the different implicit bias and push different algorithms to find the same solution, since the regularized training loss function of convex model becomes strongly convex, which exhibits one unique global optimum. For this reason, we argue that the example in the convex setting cannot fully capture the differences between GD and Adam for training neural networks. 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 weight decay regularization?
Specifically, we study Adam and GD for training neural networks with weight decay regularization on an image-like data model, and demonstrate the different behaviors of Adam and GD based on the notion of feature learning/noise memorization decomposition. Inspired by the experimental observation Figure where Adam tends to overfit the noise component of the data, we consider a model where the data are generated as combination of feature and noise patches, 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 follows.
We establish global convergence guarantees for Adam and GD with 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: when performed on the considered image-like data model, 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 due to the different training behaviors of Adam and GD: Adam is more likely to fit dense noises and output a model that is largely contributed by the noise patches; GD prefers to fit training data using their feature patch and finds a solution that is mainly composed by the true features.
We also show that for convex settings with weight decay regularization, both Adam and gradient descent converge to the same solution and therefore have no test error difference. This suggests that the difference between Adam and GD cannot be fully explained by linear models or neural networks trained in the "almost convex" neural tangent kernel (NTK) regime (Jacot et al., 2018; Allen-Zhn et al., 2019b; Du et al., 2019a; Zou et al., 2019). It also demonstrates that the inferior generalization performance of Adam is closely tied to the nonconvex landscape of deep learning optimization, and cannot be solved by adding regularization.
2 RELATED WORK
In this section, we discuss the works that are closely related to our paper. Generalization gap between Adam and SGD. The worse generalization of Adam compared with SGD has also been observed by some recent works and has motivated new variants of neural net- work 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 adap- tive methods to SGD to improve generalization. Our theoretical results for GD and Adam can also
2

---

Published as a conference paper at ICLR 2023
(a) Adam
(b) SGD
Figure 1: Visualization of the first layer of AlexNet trained by Adam and SGD on the CIFAR- 10 dataset. Both algorithms are run for 100 epochs with weight decay regularization and standard data augmentations. Clearly, the model learned by Adam more "noisy" than that learned by SGD, implying that Adam is more likely to overfit the noise in the training data.
provide theoretical insights into the effectiveness of these empirical studies. Optimization and generalization in deep learning. Our work is also closely related to the recent line of work studying the optimization and generalization guarantees of neural networks in the neu- ral tangent kernel (NTK) regime (Jacot et al., 2018) or lazy training regime (Chizat et al., 2019). In particular, recent works (Du et al., 2019b;a; Allen-Zhu al., 2019b; Zou al., 2019) showed that the optimization only happens within a small neighborhood region around the random initialization and proved the global convergence of GD and SGD when the neural network is sufficiently wide. Moreover, the generalization ability of GD/SGD has been further studied in the same setting (Allen- Zhu et al., 2019a; Arora al., 2019a;b; Ji & Telgarsky, 2020; Chen et al., 2021), which suggests that wide neural network trained by GD/SGD can learn a low-dimensional function class. Moreover, Allen-Zhu & Li (2019); Bai & Lee (2019) initiated the study of learning neural networks beyond the NTK regime as it differs from the practical DNN training. 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 several recent works that studied how neural networks can learn features. Allen-Zhu & Li (2020a) showed that adversar- ial training purifies the learned features by removing certain "dense mixtures" in the hidden layer weights of the network. Allen-Zhu & Li (2020b) studied how ensemble and knowledge distillation work in deep learning when the data have "multi-view" features. Frei al. (2022b) studied the fea- ture learning for two-layer networks, and demonstrated its superior performance than linear models. Shen et al. (2022) explored the benefit of data augmentation by showing its ability to achieve more effective feature learning. This paper studies 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 CNN with Adam and GD based on n. independent training examples {(x;, y:) y:)) 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 data model is motivated by image classification problems where the label of image usually only depends on part of an image, and the other parts of the image showing random objects, or fea- tures that belong to other classes, 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 al. (2018), which are handcrafted for showing the failure of Adam in term of either convergence or generalization (de- tailed illustrations of the data models in these works are deferred to the appendix). 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 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.
be considered noises. When using CNN to fit the data, the
Definition 3.1. Each data point (x. y) with x E 2dd and E {-1,1} is generated as follows:
== x 2 JT where of x, and x2 denotes the feature patch that consists of a feature vector
3

---

Published as a conference paper at ICLR 2023
3 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 = 0, (1. OJ". The noise vector g is generated according the following process:
Randomly select s coordinates from [d} uniformly, denoted as a vector s E {0, 13a
Generate from distribution (0, Tf and then mask off the first coordinate and other d - coordinates, i.e., 5 =
Add feature noise to &, i.e., ç = - where a < 1 is the strength of the feature noise.
In particular, throughout this paper we set d = 2(n*), & = 0 1/13) n²
of = + (spooyylgttn) and
(0p polylog(n)).
The most natural way to think of our data model is to treat x as the output of some intermediate layer of CNN. In literature, Papyan 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 sampled independently. This motivates us study sparse features and tittyy 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, as long as the sparsity gap between feature and noises exists.
Note that in Definition 3.1, each data input consists of two patches: a feature patch yv that is positively correlated with the label, and a noise patch & which contains the "feature noise" -ayv as well as random Gaussian noises. Importantly, the feature noise -ayyv in the noise patch plays a pivotal role both the training and test processes, which connects the noise overfitting in the training process and the inferior generalization ability in the test process.
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 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 mul- tiple feature vectors and multiple patches, each data can either contain single feature multiple features, together with pure random noise or feature noise.
Two-layer CNN model. We consider two-layer CNN model F using truncated polynomial acti- vation function (z) = (max/0, 2))² and fix the weights of second layer to be all I's, where q 2 3. Given the data (x, the j-th output of the CNN can be formulated as
F,W, Σ [o((w (((W))))))) Σ [ollWsrr v)) ((),, (3.1)
where m the width of the network, Wj.r E Fdd denotes the r-th CNN filter, and W is the collection of model weights. For the ease of analysis, we set the output layer as all I's. Our analyses and results can still be applied we use random second layer weights.
Besides, the motivation of using polynomial ReLU activation function is to guarantee that the loss function is (locally) smooth and the amplification ability of pattern learning. can be replaced by a smoothed ReLU activation function (c.g., the activation function used in Allen-Zhu & Li (2020b)). we assume the input data distribution is Gaussian, we can also deal with ReLU activation function (Li et al., 2020). A set of similar smoothed ReLU-type activation functions have also been widely considered to study the generalization performance of two-layer neural networks from different aspects (Frei et al., 2022a; Cao et al., 2022; Shen et al., 2022; Chen et al., 2022). Moreover, would like to emphasize that x; and X2 denote two data patches, which are randomly assigned with feature vectors or noise vectors independently for each data point. The leaner has 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 po).log(n). We assume j E { {-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 (O, of) with og (dd11/1)
4

---

Published as a conference paper at ICLR 2023
Training objective. Given the training data points {(x,, [+)] 1e consider to learn the model parameter W by optimizing the empirical loss function with weight decay regularization
L(W)= = L. (W)
2 WT
(3.2)
where L((W) - -log 1) ( denotes the individual loss for the data point (x.,Bt) and
(W)))))
À 0 is the regularization parameter. In particular, the regularization parameter can be arbitrary as long as it satisfies À E (0,)o) with λn = O (/e polf--ooyyoogaa>> We claim that the Ag is the largest feasible regularization parameter that the training process will not stuck at the origin point (recall that L LW) admits zero gradient at W = =.)
Training algorithms. In this paper, consider full-batch gradient descent and Adam', In particu- lar, starting from W(0) {w for {+1), ),,,), r E [m]]. the gradient descent update rule
((+1) w W., L(W()).
j.r
where 77 is the learning rate. Meanwhile, Adam store historical gradient information in the momen- tum (() and a vector y(t) as follows
(++1) B1 m (t) + +1 31) Vw Moy.-L(W' L(W()).
(3.3)
jrr (++1) B2v (t) +1 - 82) |V Ww... L I(W(1))3
(3.4)
j,r
and entry-wisely adjusts the learning rate:
(++1) w
(t) 77-
(2)
(3.5)
where 31, B2 are the hyperparameters of Adam (a popular choice in practice is 81 0.9, and 82 - 0.99), which are considered as constants in our paper, and in (3.4) and (3.5), the square ()) square root V and division ./- all denote entry-wise calculations. We would like to clarify the original Adam paper (Kingma & Ba, 2015) considers to normalize the gradient m j.r via
(t) [v(()
j.r +e]///2
while the small bias term is ignored in paper. In practice, tuning can help improve the generalization ability of Adam (Choi et al., 2019), as it allows to make a trade-off between the normalized gradient update and gradient update. We remark that considering tunable e is beyond the focus of this paper. For the ease of analysis, we do not consider the initialization bias correction in the original Adam paper either and set m
(0) V W L(W(0)) and Vj..
V w L(W(0)))2
MAIN RESULTS
In this section we will state the main theorems in this paper. We first provide the learning guarantees of Adam and Gradient descent for training two-layer CNN model in the following theorem. Recall that in this setting the training objective is nonconvex.
Theorem 4.1 (Nonconvex setting). Consider two-layer CNN defined (3.1) with d 2(n and regularized training objective (3.2) with regularization parameter À >>>> 0, suppose the network width m polylog(n) and the data distribution follows Definition 3.1, then we have the following guarantees the training and test errors for the models trained by Adam and Gradient descent:
1. Suppose we run Adam for T poly(n) iterations with 1) pnly(n)> then with probability least
1 - O)n=)), we can find a NN model W Adam such that (W Adam
11 Tp
Moreover, the
model W Adam also satisfies:
Training emoor is zero: EE 11 [Fy, (W* Adam: x;) S F-g; W
Adam: x;) =0
9.
Test error is high: P P(xx))~D Fy (W+ Adam x) F-y (W Adam= x)
2-
2. Suppose we run gradient descent for polyín) iterations with learning rate v/ poly(n) then with probability at least 1 (0 we can find a NN model W such that (VVL(W..)ll
Tv Moreover, the model W also satisfies:
Our theory can still hold for mini-batch stochastic gradient descent, which we will discuss in Appendix.
5

---

Published as a conference paper at ICLR 2023
Training enror is zero: Ei=1 1 Fy (W CD: xi) S -y: (W GD x,i] = 0. Test error is nearly zero: P P(x.w)>D [F,, ( (W, x) S F-y(W x)] [oly(n)
From the optimization perspective, Theorem 4.1 shows that both Adam and GD can be guaranteed to find a point with a very small gradient, which can also achieve zero classification error on the training data, Moreover, it can be seen that given the same iteration number T and learning rate 7, Adam can be guaranteed to find point with up 1/ (T')) gradient norm in & metric, while gradient descent can only be guaranteed to find a point with up to I/vTn gradient norm in Ez metric. More specifically, let e be sufficiently small quantity and ignoring other problem parameters, we can set 17 = O(e) for Adam and 1j = O(²²) for GD, then Adam and GD will need T = O(e²) and T (Et)) to find first-order e-stationary point. This suggests that Adam could enjoy faster convergence rate compared to SGD in the training process, which is consistent with the practice findings. We would also like point out that there is no contradiction between recent work (Reddi et al., 2019) showing that Adam can fail to converge, as the counterexample in Reddi et al. (2019) for the online version of Adam, while we study the full batch Adam.
result and the
In terms of the test performance, their generalization abilities are largely different, even with weight decay regularization.
particular. the output of gradient descent can generalize well and achieve nearly zero test eπror, while the output of Adam gives nearly 1/2 test eπror. In fact, this gap is due to two major aspects of the training process: (1) At the early stage of training where weight decay exhibits negligible effect, Adam and GD behave very differently. In particular, Adam prefers the denser and thus tends to fit the noise vectors &. gradient descent prefers the data patch of larger & norm and thus will learn the feature patch; (2) At the late stage of training where the weight decay regularization cannot be ignored, both Adam and gradient descent will be enforced to converge to a local minimum of the regularized objective, which maintains the pattern learned in the early stage. Consequently, the model learned by Adam will be biased towards the noise patch to fit the feature noise vector -oyv, which is opposite in direction to the true feature vector and therefore leads to test error no better than a random guess. More details about the training behaviors of Adam and GD are given in Section 5. Experimental justification are provided in Appendix.
Theorem 4.1 shows that when optimizing a nonconvex training objective, Adam and gradient descent will converge different global solutions with different generalization errors, even with weight decay regularization. In comparison, the following theorem gives the learning guarantees of Adam and gradient descent when optimizing convex and smooth training objectives (e.g., linear model F(w.x) = w x with logistic loss).
Theorem 4.2 (Convex setting). For any convex and smooth training objective with positive regu- larization parameter X, suppose we run Adam and gradient descent for T pollii) iterations, then with probability at least 1 - -- the obtained parameters W and W satisfy that (VL.aad)))ll
"
Adam
1l1 S Tn and || VLLWA Adam)) < Ty respectively. Moreover, let F(W, x) E R be the output of the convex model with parameter W and input x, it holds that:
Training errors are the same, 1 Σi= I [sgn( (WW Adam x4)) F sgn (F(W GD x,] = 0.
7 Test errors are nearly the same: P ((xx))>> [aggi/sssssnng (F(WA Adam x,)) # sgm [F(W &D x,)] S poly(n)
Theorem 4.2 shows that when optimizing a convex and smooth training objective (e.g., a linear model with logistic loss) with weight decay regularization, both Adam and gradient can converge to almost the same solution and enjoy very similar generalization performance. The proof will be relying on the strong convexity of the training objective and the convergence (to the first-order stationary) guarantee of Adam (Défossez et al., 2020) and GD. Combining this result and Theorem I. it is clear that the inferior generalization performance is closely tied to the nonconvex landscape of deep learning, and cannot be understood by standard weight decay regularization.
PROOF OUTLINE OF THE MAIN RESULTS
In this section we provide the proof sketch of Theorem 4.1 and explain the different generalization abilitics of the models found by gradient descent and Adam.
Before moving to the proof of main results, we first give the following lemma which shows that for data generated from the data distribution D in Definition 3.1, with high probability all noise vectors {&i j=l.... have nearly disjoint supports.

---

Published as a conference paper at ICLR 2023
Lemma 5.1. Let {(x,, ye) be the training dataset generated by Definition 3.1. Moreover, recall that x, lyvv I (or x; "yv let B, -supp { be the support of E except the first coordinate. Then with probability at least 1 - 22-2, B, n Bj = 0 for all i + j.
This lemma implies that the optimization of each coordinate of the model parameter except for the first one, is mostly determined by only one training data. Technically, this lemma can greatly simplify the analysis for Adam so that we can better illustrate its optimization behavior and explain the generalization performance gap between Adam and gradient descent.
Proof outline. For both Adam and gradient descent, we will show that the training process can be decomposed into two stages. In the first stage, which we call pattem learning stage, the weight decay regularization will be less important and can be ignored, while the algorithms tend to learn the pattern from the training data. In particular, we will show that in the pattern learning stage, the optimization algorithms have different algorithmic bias: Adam tends to fit the noise patch while gradient descent will mainly learn the feature patch. In the second stage, which we call it regular- ization stage, the effect of regularization cannot be neglected, which will regularize the algorithm to converge some local stationary points. However, due to the nonconvex landscape of the train- ing objective, the pattern learned in the first stage will remain unchanged. even when running an infinitely number of iterations.
5.1 PROOF SKETCH FOR ADAM
Recall that in each iteration of Adam, the model weight is updated by using a noving-averaaed gradient, normalized by a moving average of the historical gradient squares. As pointed out in Balles & Hennig (2018); Bernstein et al. (2018), Adam behaves similarly to sign gradient descent (signGD) when using sufficiently small step size the moving average parameters 31, .3 are nearly zero, which is also justified in our Lemma C.2. Specifically, we show that when considering constant B1 and B2g the Adam update on the coordinates with large gradient (e.g.. (V((W)) can be well approximated by the signGD update (i.e., sign(V (WL(((()))))) understand the optimization behavior of signGD and then extends it to Adam using their similarities. In particular, signGD updates the model parameter according to the following rule:
> 7) This motivates us to
(t-1) (1+1) j,r 77 sgn( V W (Vw, L(W())).
jur
Recall that each data has a feature patch and a noise patch. By Lemma 5.1 and the data distribution (see Definition 3.1), all noise vectors {E:)i=....,, are supported on disjoint coordinates except the first coordinate. For x,, let B, denote its support excluding the first coordinate. In the subsequent analysis, we will always assume that those B; are disjoint, i.e., B, n Bj- = 0 if i + j.
Next we will characterize two aspects of the training process: feature learning and noise memoriza- tion. Mathematically, we focus on two quantities: (w.)) v) and (w Berrs E.). In particular, given
the training data point (x,, v4) with x; [y.vT IT, larger (w!2r, V/ v) implies better feature learning and larger (w E.) represents better noise memorization. Then regarding the feature vector that only has nonzero entry at the first coordinate, we have the following for signGD:
((41) .jv) = Jv) - 17- <sgn (Vws, ((.(((())))) (w jrr
(5.1)
n yee
=(w((h + in- sgn
[ (W.... v(v)) [o'((
ao'
J.r {,))] ((w
where 5/ = Ly.=y- (F, (F) and 1F,it,(F,x,) Exe(- W.J.
From (5.1) we can observe three terms in the signed gradient. Specifically, the first term represents the gradient over the feature patch, the second term stems from the feature noise term in the noise patch (see Definition 1), and the last term is the gradient of the weight decay regularization. On the other hand, the memorization of the noise vector &, can be described as follows.
(++1) E.) (w (t) g..r. Es) = -1 (sgg(( V Ww, (.....(())) E.)
(5.2)
(1) Σ sgn
(1) [kk) (2w [.]
=n 1
kEB,U(1)
7

---

(2b) Published as a conference paper at ICLR 2023
Throughout the proof, we will show that the training process of Adam can be decomposed into two stages: pattern learning stage and regularization stage. In the first stage, the algorithm learns the pattern of training data quickly, without being affected by the regularization term. In the second stage, the training data has already been correctly classified since the pattern has been well captured, the regularization will play an important role in the training process and guide the model to converge. Stage I: Learning the pattern. At the beginning of training, the neural network output is smaller than some constant for all data, and therefore all training data remain under-fitted and can provide large gradient for model training. We specify this stage of training as Stage I. In this stage, the effect of weight decay regularization can be ignored due to our choice of X. We will show that this stage the inner product (wo..r, E:) grows much faster than jirr jv) since feature learning only makes use of the first coordinate of the gradient, while noise memorization could take advantage of all the coordinates in B; (see (5.2), note that |B,] $ >> 1).
Lemma 5.2 (General results Stage D. Suppose the training data is generated according to Defi- nition 3.1, assume À = ol 22 2p//n) and v = 1/poly(d), then for any L < To with To = =0( 1yso =
and any i [n].
j.r (1+1) .j v) S (w jrr v) + O()), (w (1+1)
..) = (w!)) (4)
wi.rt Ei> + ((nso,))
j,r j v), after certain number of iterations, the learning of noise patch will dominate the learning of feature patch (i.e., ad' (( j.r (.)
Since (w(0) Ei) enjoys much faster increasing rate than that of (w j.r.
o' d(((!" j.r. y.v)). Thus, by (5.1), the model will tend to fit the feature noise in the noise patch (i.e., v), leading to flipped feature learning phenomenon.
Lemma 5.3 (Flipping the feature learning). Suppose the training data is generated according to Definition 3.1, 2 )((so,)² 1-9 V vg--1) and op < =((so,)))) then for any E [or with
T, = peapaa/t--1) S To.
(++1) .j v) v)= =(
(w j. v) (
Moreover, it holds that (1) W
(2b) [k] (E))))()) 52 (d)))) or (Tb) [1] sgg(j) 52( (2) wor
w j.r
(TTb [k] = =n() for E B, with ys = j: and (3) w (Tb)
j.r [k] (n) otherwise.
From Lemma 5.3 it can be observed that the iteration To. the sign of the first coordinate of w j.r is different from that of the true feature. i.e.. j v. This implies that at the end of the first training stage, the model is biased towards the noise patch fit the feature noise.
Stage II: Regularizing the model. In this stage, as the neural network output becomes larger, part of training data starts to be well fitted and gives smaller gradient. As consequence, the feature learning and noise memorization processes will be slowed down and the weight decay regularization term cannot be ignored. However, although weight decay regularization can prevent the model weight from being too large, it will maintain the pattern learned in Stage and cannot push the model back to "forget" the noise and learn the feature and stops at some local stationary points. We summarize these results in the following lemma.
Lemma 5.4 (Maintain the pattern). If a = (s073/n) and = o(A). then let r Ua.rs Ei), for any 1 2 To. i E [n]. E {+1} and r E [m]] it holds that
arg maxre(m] (w]) (t)
(w (1) wart (,) = ((), Σ
w (t) Ba.
[ [+]]|
|&, kb] = 5(1), (w j.r sgn(() v) E (0o(1), In)]
KEB
Lemma 5.4 shows that in the second stage. (w Ex) will always be large while (w we 81 v) is still negative, or positive but extremely small. Next we will show that within polynomial steps, the algorithm can be guaranteed to find point with small gradient.
(1)
(t)
Lemma 5.5 (Convergence guarantee). If 1 = ((d-1/2), then for any it holds that
L(W(++1) - ((W()) < -nl -//VV + d))
Lemma 5.5 shows that algorithm can find a point with up to (()(TTn) in l norm. Then we can show that given the results in Lemma 5.4, the formula of the algorithm output W can be precisely characterized, which
pick sufficiently small n and T = poly(n)/n to ensure that the
8

---

Published as a conference paper at ICLR 2023
can show that gt.r+ v) < O. This implies that the output model will be biased to fit the feature noise -ayv but not the true one v. Then when it comes to a fresh test example the model will fail to recognize its true feature. Also note that the noise in the test data is nearly independent of the noise in training data. Consequently, the model will not be able to identify the label of the test data and therefore cannot be better than random guess.
5.2 PROOF SKETCH FOR GRADIENT DESCENT
Similar to the proof for Adam, we also decompose the entire training process into two stages. Stage I: Learning the pattern. In this stage the gradient from training loss function is large and and the effect of regularization can be ignored. Unlike Adam that is sensitive to the sparsity of the feature vector noise vector, gradient descent is more focusing vector (which can be either feature or noise) with larger & norm is more likely to be discovered and learned by GD. Note that the feature vector has larger l norm than the noise, gradient descent will learn the feature vector very quickly, while barely memorize the noise.
the lo norm of them, where the
show that
Lemma 5.6. Let A
max,e[m] (W (++1) j.r .j v), r(t)) = max.e[m] (W Ei), and r
ya j.s Let T be the iteration number that A, reaches ((1/m) = @(1), then we have T = ( 2-9) for all j E 1-1,1). Moreover, let Tu = max then for all t I To it holds that 1 Ö(7o) for all E 1-1,
((t)
Stage II: Regularizing the model. Similar to Lemma 5.4, we show that in the second stage at which the impact of weight decay regularization cannot be ignored, the pattern of the training data learned in the first stage will remain unchanged.
Lemma 5.7. If 1 & O it holds that A = ((1) and r = Ö(ao) for all miny Ty.
The following lemma further shows that within polynomial steps, gradient descent is guaranteed to find a point with small gradient.
Lemma 5.8. If the learning rate satisfies 1 = =(1), then for any t 2 O it holds that L(W(++1) - L(W())) < - "
Lemma 5.8 shows that we can pick sufficiently small rj and T = poly(n)/1 to ensure that gradient descent can find point with up to ((1/(Tn))* 1/2) in l norm. By Lemma 5.7, it is clear that the output model of GD can well learn the feature vector while memorizing nearly nothing from the noise vectors, which can therefore achieve nearly zero test error.
Experiments. We perform experiments on synthetic data (generated according to Definition 3.1) to validate our theoretical findings: Adam performs stronger noise memorization than feature learning while GD performs stronger feature learning than noise memorization, when conducted on the data distribution constructed in Definition 3.1. We consider both the two-layer CNN model studied in this paper and a 5-layer CNN model for further justification. Experimental setup and results are deferred Appendix A due the page limit.
CONCLUDING REMARKS AND FUTURE WORK
In this paper, we study the generalization of Adam and compare it with gradient descent. We show that when training neural networks, Adam and GD starting from the same initialization can con- verge to different global solutions of the training objective with significantly different generalization errors, even with proper regularization. Our analysis reveals the fundamental difference between Adam and GD in learning features or noise, and demonstrates that this difference is closely tied to the nonconvex landscape of neural networks.
We would also like to remark several important research directions. First, our current result is for two-layer networks. Extending the results to deep networks could be an important next step. where we will not only look at the input data but also consider the output of each intermediate layer as "input", Second, our current data model motivated by the image data (i.c., sparse feature and denser noise), where Adam has been observed to perform worse than SGD in terms of generalization. In fact, our theoretical analysis can lead to an opposite conclusion on the generalization comparison between Adam and GD if the noise is sparse and the feature is denser. Therefore, it would also be interesting to explore whether this is the case in other machine learning tasks such as natural language processing, where Adam is often observed perform better than SGD.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We thank the anonymous reviewers and area chair for their helpful comments. YL is supported by the National Science Foundation CCF-2145703. QG is supported in part by the National Science Foundation CAREER Award 1906169, IIS-2008981 and the Sloan Research Fellowship.
REFERENCES
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Revisiting the general- ization of adaptive gradient methods. 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. What can ResNet learn efficiently, going beyond kernels? In Advances in Neural Information Processing Systems, 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. arXiv preprint arXiv 2020a.
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816 2020b.
Zeyuan Allen Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization overparameter- ized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. International Conference on Machine Learning, pp. 242-252, 2019b.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li. Ruslan Salakhutdinov, and Ruosong Wang. On ex- act computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, 2019b.
Yu Bai and Jason D Lee. Beyond linearization: On quadratic and bigher-order approximation of wide neural networks. In International Conference on Learning Representations, 2019.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference on Machine Learning, pp. 404 413. PMLR, 2018
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569, PMLR, 2018.
Yuan Cao, Zixiang Chen, Mikhail Belkin, and Quanquan Gu. Benign overfitting in two-layer con- volutional neural networks. arXiv preprint arXiv:2202.06526, 2022.
Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang. Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods training deep neural networks. International Joint Conferences Artificial Intelligence, 2020.
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much partmttriizzii is suf- ficient to learn deep relu networks? In International Conference on Learning Representations, 2021.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding mixture of experts in deep learning. arXiv preprint arXiv:2208.0281. 2022.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems, 2019.
Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint 19Xiv:1910.05444 2019.
10

---

Published as a conference paper at ICLR 2023
Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv preprint arXiv:2003.02395, 2020.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pp- 1675- 1685, 2019a.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations, 2019b.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(ju):2121-2159) 2011. 1.
Spencer Frei, Niladri $ Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pp. 2668-2703, PMLR, 2022a.
Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202.07626 2022b.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning
lecture 6a overview of mini-batch gradient descent, 2012.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemnl:: Convergence and gen- eralization in neural networks. In Advances in neural information processing systems, pp. 8571 1- 8580, 2018
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve ar- bitrarily small test error with shallow relu networks. In International Conference on Learning Representations, 2020.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv: 171::077,, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over prramertaerr two-layer neural networks beyond ntk. In Conference on Learning Theory, pp. 2613-2682. PMLR, 2020.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv: 190.09990..09833, 2019.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing 1stm lan- guage models. In International Conference on Learning Representations, 2018.
Vardan Papyan, Yaniv Romano, and Michael Elad. Convolutional neural networks analyzed via con-
volutional sparse coding. The Journal of Machine Learning Research, 18(1):2887-2938 2017.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv
preprint arXiv: 009337, 2019.
Ruoqi Shen, Sebastien Bubeck, and Suriya Gunasekar. Data augmentation feature manipulation. In International Conference on Machine Learning, pp. 19773-19808 PMLR, 2022.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, pp. 4151-4161, 2017.

---

Published as a conference paper at ICLR 2023
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760, 2019.
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoreti- cally understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems. 33, 2020.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep ReLU networks. Machine Learning, Oct 2019.
12

---

Published as a conference paper at ICLR 2023
A EXPERIMENTS
A.1 EXPERIMENT DETAILS FOR FIGURE
The experiments in Figure 1 are performed by running (stochastic gradient) Adam and SGD in training AlexNet on the CIFAR-10 dataset. Specifically, the first layer of AlexNet is set as: ker- nel size=11, stride=4, padding=2, in order to match the size of CIFAR-10 image. In terms of the input data, we use standard random crop and horizontal flip data augmentations. In terms of the model training, we set the batch size 64, the epoch number as 100, the regularization parameter as À = 5 x 10-4. Besides, we set the learning rate 7 = 0.01 for SGD and v = 1 x 10-6 for Adam, where B1 and 32 are set as their default values in pytorch.
A.2 NUMERICAL, EXPERIMENTS ON SYNTHETIC DATA
In this section we perform numerical experiments on the synthetic data generated according to Def- inition 3.1 to verify main results. In particular,
the problem dimension d 1000, the training sample size n = 200 (100 positive examples and 100 negative examples), fea- ture vector
[1.0..... OoT, noise sparsity 8 0.1d 100, standard deviation of noise o, = 1/21/2 1/2 0.1, feature noise strength a 0.2, initialization scaling dg 0.01, regularization parameter À = 1 x 10-5
Two-layer CNN model. We first consider the exact two-layer CNN model studied in the paper. We set the network width m 20, activation function (2)) max 2}3, total iteration number = 1 x 10*, and the learning rate 5 x 10-3 for Adam (default choices of 31 and B3 in pytorch),
? 0.02 for GD.
We first report the training error and test error achieved by the solutions found by SGD and Adam in Table where the test error is calculated on test dataset of size 10a. It is clear that both Adam and SGD can achieve zero training error, while they have entirely different results on the test data: SGD generalizes well and achieve zero test error; Adam generalizes worse than SGD and gives > 0.5 test error, which verifies our main result (Theorem 4.
Algorithm
Adam
SGD
Training error Test error
0.884
O
Table 2: Training and test errors achieved by GD and Adam.
Moreover, we also calculate the inner products: max, (wix,(Wi,r v) and min; max, (Wi.r, E.). repre- senting feature learning and noise memorization respectively, to verify our key lemmas. Here we only consider positive examples as the results for negative examples are similar. The results reported in Figure 2. For Adam, from Figure 2(a), it can be seen that the algorithm will perform fea- ture learning in the first few iterations and then entirely forget the feature (but fit feature noise), i.e., the feature learning is flipped, which verifies Lemma 5.3. In the meanwhile, the noise memoriza- tion happens in the entire training process and enjoys much faster rate than feature learning, which verifies Lemma 5.2. In addition, we can also observe that there are two stages for the increasing of min, max, (wi.r. E:): in the first stage min; max, (W1rr> {.) increases linearly, and in the second stage its increasing speed gradually slows down and min; max, (W1.r. E.) will remain in a constant order. This verifies Lemma 5.2 and Lemma 5.4. For GD, from Figure 2(b), it can be seen that the feature learning will dominate the noise memorization: feature learning will increase to a constant in the first stage and then remain in a constant order in the second stage; noise memorization will keep in low level which is nearly the same that at the initialization. This verifies Lemmas 5.6 and 5.7.
5-Layer CNN model. We further perform numerical simulations for the deep neural network mod- els. particular, we consider a 5-layer CNN model: the first layer exactly the same as the two- layer CNN model, followed by a 4-layer ML. with ReL.U activation. The number of neurons is set
13

---

Published as a conference paper at ICLR 2023
1.2
2.0
1.0
0.8
1.5
Value
Feature Learning: max (Wi...
Feature Learning: max (Wi.r V)
Noise Memorization: min max (Wi.
1.0
Naise Memorizatian: min max (Wi, Ei
0.2
0.0
-0.2
0.0
2000
4000
6000
8000
10000
0
2000
4000
6000
8000
10000
# Iterations
# Iterations
(a) Adam
(b) GD
Figure 2:
Visualization of the feature learning (max,, (max,((1,r) v)) and noise memorization (min, max, (Wi,,,.)) for training the two-layer CNN model.
0.6
1.75
0.5
1.50
1.25
Valuue
0.4
Feature Learning: max (Ws, V)
0.3
0.75
Noise Memorization: max (W1.n. E/
0.50
Feature Leamiing max (Wa,r, V)
Noise Memorization: min max (Wa. EA
0.25
0.00
2000
4000 Iterations
6000
8000 10000
2000
4000 Iterations
6000
8000
10000
(a) Adam
(b) GD
Figure 3: Visualization of the feature learning (max,((1,,) v)) and noise memorization
(min, max, (W1,r, E.)) for training the 5-layer CNN model.
as m = 20 for all layers. The total iteration number = 1 x 104, the learning rate is rj = 5 x -5 for Adam and 7 0.1 for GD.
We then show the feature learning and noise memorization of the first layer of this neural network model by calculating the inner products: max,(w1,r. v) and min; max,(W1.rs E.), notes the weights of the r-th neuron in the first layer. The results are shown in Figure 3. It can be observed that, when applied on the data distribution in Definition 3.1, Adam tends perform stronger noise memorization than feature learning, while GD performs stronger feature learning and nearly negligible noise memorization. Moreover, we also visualize the first layer of the 5-layer CNN trained by Adam and GD in Figure 4. It can be seen that the CNN model found by Adam is clearly more "noisy" than that found by GD. This consistent with our theoretical findings and empirical observation on the real-world dataset (i.e., Figure 1).
(Wi.s
where Wi,r de-
B EXTENSIONS TO MINI-BATCH STOCHASTIC GRADIENTS
One natural extension of our paper is proving the separation between mini-batch SGD (without replacement) and mini-batch Adam, which we believe is not difficult. particular, let I, of size B be the set of indices of the mini-batch data used in the t-th iteration, the update rule of SGD is
w ((+1) (t 17 B
V Wy, L (w(()) - (! j.r
EET.
14

---

Published as a conference paper at ICLR 2023
Adam
(b) GD
Figure 4: Visualization of the first layer of the 5-layer CNN model on the synthetic dataset.
The update rule of mini-batch Adam is
m j.r
(++1)
=B,m
(1) j.r + (1- B1)
1
W... L:(((()) Aw (1)
B
v frr = 32 + (1 (1)
((-1))
Aw (2) Aw V s.r (((())))
B
SEI
and
w (1+1) jrr w/
Jrr +c.
Here the bias term is set as e = (Aoo). We claim that this parameter introduced to guarantee that the regularization term will not dominate the training process when using stochastic gradients in Adam.
Then will take deeper look at the speeds of feature learning and noise learning for mini-batch (t)
SGD and Adam, where we focus on the period that (w j.rt v) / j.r & o(1) for all j. i, and r
(i.e., the pattern learning stage). This further implies that | = 0.5 + o(1) for all j. i, and t. Thus
in the following, we will assume that all (e(t) has nearly the same quantity. 2,1
Feature Learning. First, according to Definition 3.1, we know that the feature vector v and feature noise are the same for all data, which implies that the learning pattern of the feature coordinate will be largely the same as that of full-batch algorithms. In particular, for mini-batch Adam, we can show that the update of the first coordinate (i.e., feature coordinate) is similar sign-Gr GDD when using sufficiently small learning rate y 1 1/poly(d) since all stochastic gradients V (W()) have the same component in this coordinate. Then using the fact that |e I's are nearly the same for all i, have
(w (++1) .jv) ~ jv) + in sgn Σ Eueef lo (w vvv) ad' ao'(( {X)] n.w (ul)
which is the same as full-batch Adam (see (5.1)). For SGD, using the fact that 10%11' 's are nearly the same for all we can get that
(w (++1) jv) (1 mA) (w J.
+ 77 "
Σ yo 2
uv)))
Σ i=1
viel 2 ((w (( Eil)
which is also the same that of GD (see (C.28)).
15

---

Published as a conference paper at ICLR 2023
(6) Noise Memorization. Note that due to the normalization term j.r in the Adam update, all coor- dinates will be updated with nearly the same amount. Therefore, we only need to count the number of coordinates that are updated by full-batch Adam and mini-batch Adam.
Recall that we have shown that using mini-batch gradients will not affect the feature learning. How- ever, the noise memorization will be slightly different, since in each iteration, full-batch Adam can update (ns) coordinates while mini-batch Adam can only update @( coordinates. To show this, we note that for any coordinate l + 1, the gradient momentum of full-batch Adam is
(1) jur
[+]
(1 B1)
1 n
Σ xE|m
Vwg, L (w(( Aw
(0--))010 [b)
==0
while for mini-batch Adam,
m M
Σ (8(1-1)) ==0
Σ
V Vwg.. Le [)] + Aw (1 [b)
B
5,r
VEI.--
where we only maintain the recent F polylog(n) gradients since for S t- f- the decaying terms (31)² S (3,)² becomes negligible. Therefore, by comparing the above two equations and applying Definition 3.1, it is clear that for full-batch Adam can update all noise coordinates, i.e., E Use[m] B,, which is of size (nn)). In contrast, mini-bateh Adam can only update a subset of noise coordinates, i.e., & Ute[e] _ B;, which is of size 7 FBs = (Bs)) This further implies that in each epoch (one pass of the data, ((n/ B) steps), the noise coordinates in B, will be updated by mini-batch Adam in at most F = ((1) steps, while within the same amount of iterations, the noise coordinates in B; will be updated by full-buich Adam for (n B) steps, suggesting that mini-batch Adam admits slower rate of noise memorization by O(n/ B) factor.
For SGD, it is easy to show that the rate of noise memorization will still be nearly the same as that of GD. In particular, during each training epoch (0(n/ B) steps), SGD will learn the noise vector & in only one step with the mini-batch gradient 1 B VT.(((rr) for some in this epoch, while within the same amount of steps. GD will Iearn the noise vector & in all ((n/ B) steps but with strength I (VL(((()) giving the same total learning ability. This suggests that SGD admits nearly the same rate of noise memorization compared to GD.
R
Overall, we are able to deliver the following lemmas that characterize the feature learning and noise memorization of SGD and stochastic gradient Adam in training Stage I.
Lemma B.1 (SGD, Informal). Suppose the training data is generated according to Definition 3.1, then given proper configurations of X and e and sufficiently small learning rate 77. define A maxye[m) (w j.rtj v) and r maxye[m] (w
j.ri Ei), for any satisfying A,))), r((a)
o(1/polylog(n)), we have the following one-epoch update of feature learning and noise memoriza- tion for SGD
) (otln/B))
2 Ataa) +n e n B (A.
(to)
(A//o)3-11)
((o+)/BB tt) + n- θ
B
Lemma B.2 (Stochastic gradient Adam, Informal). Suppose the training data is generated according to Definition 3.1, then given proper configurations of X and e, for any to S To with To Ö(,, Bs and any E [m], we have the following one-epoch update of feature learning and noise memorization (w In+BB
n
< (w!t) (to) .j v) + O(nn/ B);
(w yeur
fto (oooot//B))) Ln/BB) E (t) (*a) wt.r: &.) ((ssa,)
To sum up. we have shown that (1) mini-batch SGD and
Adam will not change the leamn ing speed of feature vector v compared to their full-batch counterparts, i.c., Lemma C.3 and (C.31) (needs to covert to |n/B B iterations); (2) mini-batch Adam reduces the noise memorization speed of full-batch Adam by a ((n B)) factor, while mini-batch SGD has nearly the same noise memo- rization speed compared to full-batch GD, by comparing to Lemma C.3 and (C.32)). Additionally,
16

---

Published as a conference paper at ICLR 2023
recall that in paper, the separation between Adam and GD is characterized by poly(d) factor: the speed of feature learning in Adam and GD, and the rate of noise memorization in GD are both in the order of O(n) (in each step). while the rate of noise memorization in Adam is proportional number of nonzero entries, which is in the order of poly(d). Therefore, the separation between mini-hatch SGD and mini-batch Adam in terms of the generalization error still hold under stronger over- weer-parameterization condition sop (dd [tdd///((polylo/))) ((pplyylog(n)) w(n/ B) (in contrast, the over-parameetrrizaion condition for full-batch Adam is s0p ((1).
the
C PROOF THEOREM 4. 1: NONCONVEX CASE
In the beginning of the proof we first present the following useful lemma.
C.I PRELIMINARIES
We first recall the magnitude of all parameters:
d poly(n), = poly(n)* 8 0
d//2
oz = 0
( poloylogg(m) of
n²
(11/2
m = polylog(n), a = polylog(n)), x 0 ddu-11/4n polylog(n)
Here (nly(n) denotes a polynomial function of n with degree of sufficiently large constant, poly(n) denotes a polynomial function of log(n) with degree sufficiently large constant. Based on the parameter configuration, we claim that the following equations hold, which will be frequently used in the subsequent proof.
Jp
00 - o 1
(1))))
80 n
À =
0
a =
11
71
80
Lemma C.1 (Non-overlapping support). Let {xx, 1/4 |= be the training dataset sampled ac- cording to Definition 3.1. Moreover, let B, = supp(E, {,) [13 be the support of x; except the first coordinate?. Then with probability at least n B; n @ for all i,j E [n].
Proof of Lemma C.I. For any fixed k E [n] and j E supp(Ex) (1)), by the model assumption we have
P((()) + 0} s/ (d
for all i [n]][k]. Therefore by the fact that the data samples are independent, we have
P(3 E [n] [k) (Ei), + 0) 1- s/dd 1)]" Applying a union bound over all & E [n]] and supp(E:)\(1), we obtain
P(3k [m].j E supp(Ex)\(1), E [m] (E:), + 0) 8 {1 (1- 1)]"). (C.1) By the data distribution assumption we have S vd/(2n²) which clearly implies s/dd-1) & 1/2. Therefore we have
[1 (1- s/d)"] = n {1 log|nlog(1 s/ s/d--1))) 1))1}
nn.s I -exp(n 2s/(d 1))]
S n
exp(n 4s/d)]
< nn- 8- (4ns/d)
4n² 5²/2 < <n 2
where the first inequality follows by the inequalities log(1 - 2 -22 for [0. 1/2], the second inequality follows by s/dd11 > 2s/ dd. the third inequality follows by the inequality 1 -exp(-2)) S for E R, and the last inequality follows by the assumption that vd/(2n²). Plugging the bound above into (C.1) finishes the proof.
RReaall that all data inputs have nonzero first coordinate by Definition 3.1
17

---

Published as a conference paper at ICLR 2023
C.2 PROOF FOR ADAM
Before moving to the detailed proof, we first state the update rules of feature learning and noise memorization when the sign gradient is applied.
(1+1) (w)) Jr .jv) (w/..vv) - 17. -ssg Vw,,r (ew,,,,, L(W()))).
(w(+))
n
(V) jn sgn
2 yel (' jir biv)) - ad ((w E.) ad'(ww
nAw (.
(C.2)
where e// = llggtt (F.x,) and logit,(F.x.)) = E xg)
From (C.2) we can
observe three terms in the signed gradient. Specifically, the first term represents the gradient over the feature patch, the second term stems from the feature noise term the noise patch (see Definition 3.1), and the last term is the gradient of the weight decay regularization. On the other hand, the memorization the noise vector & can be described by the following update rule.
(w (1+1) Berr
Et)
(w W1.r Ei) 17. (sgn( V Wyyl ((W())).
= (wll)) Ex) +
Weirt
sgn
Miit o (' (1)
We.rt [.ME
nAwp).rl1) Ei(A]]
kEB;
03/77 sgn
Σ yie (1
[o'( wi
yirro gir - ao ((wl wirr Ei))] ndw (2)
vi. [1]
(C.3)
In this subsection we first provide the following lemma that shows for most of the coordinate (with slightly large gradient), the Adam update is similar to signGD update (up to some constant factors). In the remaining proof for Adam, we will largely apply this lemma to get a signGD-like result for Adam (similar to the technical lemmas in Section 5). Besides, the proofs for all lemmas Section can be viewed as simplified version of the proofs for technical lemmas for Adam, thus are omitted in the paper.
Lemma C.2 (Closeness to SignGD). Recall the update rule of Adam, let W be the f-th iterate of the Adam algorithm. Suppose that (with j.r v), (w!") j.r: Ex) = @(1) for all j {+1) and E [m]. Then
if B2 2 B8. we have
For all l E [d].
(t) [k]
S O(1).
(0)
For every &s & UP_1B (including k- 1) have either Ww,,, I...L [W(())[]] θ()))
mjrr [k]
sgn Vw,,. L(W(0) ((1).
For every & e B,, we have /W WV,..L(W(())) [WW()]]]] < 6(m 1 (.p1e(11) < ((nso,) or
[
sgn (Vw,. L(W()) -(1))
Proof. First recall that the gradient V Vw Vw,..((((t) L(W())) can be calculated as
V Wi. L(W())
Σ el 0.10
(8) v(vi) v
ja o (fw (9)
j.r vrE,)) .g. + +Aw
18

---

Published as a conference paper at ICLR 2023
More specifically, for the first coordinate of Vws,, L(W()), have
R yi o ((W firt vvv))
Vw,,. ((W(0
we d'(w Eil)
(0)
(C.4)
For any l E B;, by Lemma C.I we know that the gradient over this coordinate only depends on the training data Es, therefore, we have
1 do n 3. (( (.)) [.
V L [(
(C.5)
For the remaining coordinates, we have
V W L (W(()) = Aw [+].
(C.6)
Now let us focus on the moving averaged gradient m and squared gradient jir We first show that for all [d]. holds that
I (1).
(C.7)
V [c]
(0) jir we have
By the update rule of m
[
B,m (t-1) 1) [k] + (1-3)) V.,,. [(W()))
Σ BT - 31) Vwy, L(W 0=
Similarly, we also have
[kj
Σ 35 - B2) Vw5,. L(W(( (Ap
T=U
Then by Cauchy- Saach-SSchwarzz.ee inequality we have
(m (B)) <
[BT 3111
V Vw L(W(1-7)) (APP)) (Σ
Let a3 (1-))=1 which forms exponentially decaying sequence if 33 2 B7. Therefore, (1-)
have Σ o? +11) and the above inequality implies that
mm [k)) [k]
which proves (C.7).
Now we are going to prove the main argument of this lemma. Note that m m..- which is a weighted average of all historical gradients, where the weights decay exponentially fast, then we can take on
(t)
threshold polylog(n such that D 3[ B1) poly(o- Ty. Then for each ks [d] we have
[Kk
Σ 01 - (1)) Vw,.. L(W + Σ अ5 B1) Vw,,. ((W((-T))
T=0
37 - B1) V Vw L(W( + Σ
poly (1) 1)
19

---

Published as a conference paper at ICLR 2023
where in the last equality we use the fact that VV IVw., W...- ((W((- [[]]] (1) for all ls [d]. Simi- larly, we can also have the following on v (t)
[kj
Σ
By (1 - (2) Vw.... Lo.,L [A] 1)
Here we slightly abuse the notation by using the same F. Then we have
o BT BI) V Vw,,- L(W [k] + poly(1
V (5 - 32) Vw L(W [ + poly(n-1)
In order prove the main argument of this lemma, the key is to show that within iterations, the gradient V Vw,.- L [k] barely changes. In particular, by (C.7), we have the update of each coordinate in one step is at most ((n). This implies that
kw v) (w v) = @(17),
[(w J.r {t) - -(w Et) S @(1T8op);
|w] [k]] w - ((77).
Then applying the fact that (w (T)
Jor v) S (1) and (w S. Et)ع S ((1) we further have
(3((()) x:) F3 (W())) (t) xx) < (mmps sa,) (1758pp):
where we use the fact that m (1) and sap w (1). Then it holds that
(w(()) x.)
ENE(-1.1) (FL((())).)))
(: (((tsar,) + F ((W()):))
= sgn(e())) 01
where we use the fact that @(nFso,)) o(1). Similarly, we can also show that e5 2 sgn(( e(14))) which further implies
3,8 = sgn S
for all 7 E [t 7, r,t). Note that |c(771 < 1, then it holds that
('
(1) v)) = sgn (10 - v)) v))
<gnn 0(10 ((w /o v)) + 0((e()) Ō(17).
We can also similarly derive the following
(r) v\) 2 sgn( e d'(ww v)) e((ef)) 6
(W 2 E() sgn(e))) 0 o' @ 41) + e(le()) (1Trsop))
2 p El) sgn(e))) e
fop &)) e((e @1jFsop).
Combining the above results, applying (C.4), (C.5), and (C.6), we can show that for the first coordi- nate, we have
V Wi. L(W e ((W)), + 0
Σ
- + O(x7);;
1=1
20

---

Published as a conference paper at ICLR 2023
for any lo E Bi, we have
[000 2
V W,,r [(W = @VV Ww,.. (w [()
0 ((Nn));
n
and for remaining coordinates, we have
V W [(W(r))[ = (V,
(
(t) Now we can plug the above results into the formula of mj. and j.r Using the fact that 7 = 6(1)
À o(1,, and
gett Jull I 1, we have for all k = 1 or k & B, for any i,
(9
L(W [R] + (n)) L(W [1]) + (())
For & E Bi we have
Vw,,- L (( = 3 g00ple
= {(Ar)
n
[
2 LV. +
1 1ef0ll
+ @(A))
Then, conclude that for all 1 B, for any have either Vw T [.W()))(]] θ (n) or
[k]
sgn V W,,,, L(W(0)]k) ((1).
For any le E B,, we have either Vw [(W())( S 1 Xn) or
m [k]
sgn T. 0(1)
This completes the proof.
Lemma C.3 (Lemma 5.2, restated). Suppose the training data is generated according to Defini- tion 3.1, assume À = o(a 2-2 app/) and 17 (d), then for any s To with To = Ö (19 and any i [n].
(w(t+")) j v) - (w v) + O(n). (w(tt)) .Ei) = (w p..r + + ((7sop)
Proof. At the initialization, we have
v) @(oo),
[(w])) J.r Ed) = 0(s +a) =0(31/2 [0 =((o))
(0)
which also imply that (0(0) J. ((1). Then recalling that = olo0 (1), we have
/ 2
= o(1), s2/20,
sgn Σ
vi (0) j.8
d'(w
j. vvv)) uv))) - a
Σ yi 0'(( ('(((
(0) j.r {i)) now 200
i=1
sgn
[j @(nog ) j ((an(s 00,00) +
sgn(j).
21

---

Published as a conference paper at ICLR 2023
Since is 1-sparse, then by Lemma C.2, have
j v) S (w jont v) n (0)
jond v) S (w j.r: J v) + (
Now suppose that the first inequality holds for iterations o,...,; 1. Then have j.r: O(n To) S O(oo).
Besides, note that 1 = j= -logit, (F(1), xx)) we have
sgn (Ur (vre = sgn(j);
where we recall that j E {1,, Therefore, given that À = oloo oo//)), Ö(1), and assume e (11) (which will be verified later).
o a = o(1), s1120,
sgm Σ vi We (( (2) uv)) Eueli ((w (.m, - mdw (.w])
sgn
[ θ(ng) j (0
) a ap)
sgn(j).
Since v is 1-sparse, then by Lemma C.2, the following inequality naturally holds,
(++1) (w jr .j v) S (w j. v) 77 m //2 j.r:j v) (w 19 j v) + O
(1)
Additionally, in terms of the memorization of noise, we first consider the iterate in the initialization. By the condition that 7) o(1/d) o(1/(sap)) and note that for a sufficiently large fraction of
k E B; (e.g., 0.99), we have |&: [b]| > (ap)) 2 (mmmm'sa (0)1)
and thus
sgn V Was L [(W())()) = sgn
(00) 0' (( (0) Vi.d
Wi.ro [ - ndw (0.
sgn
θ 00000)
(1)002)000))) - op sgn (g:(<:()))) + o60000100 -sgn(EEs
(C.8)
Therefore, by Lemma C.2 we have the following according to (C.3),
(1)
Berrs Ei) = (w (0)
Bert
E - 77
m
j.r
(0)
yeirs E:
(w) Ei) + (n)) Σ
(E)) O(nsa,)) (nja
ye.r
kEB.
(w E.) + ((sso)),
where in the first inequality the term O(sST)) represents the coordinates that [5/ S O(gp) (so that we cannot use the sign information of V ",.r (.) but directly bound it by (1))) and the last inequality is due to the fact that [B,| 2 s - and = -(1). For general t, we will consider the
following induction hypothesis:
(w((t-)) (
(w(() (t)
Bi.ro {.) + ((ssap);
(C.9)
which has already been verified for /. = 0. By Hypothesis (C.9), the following holds at time 1,
Hearo = (0) r: < + (tnssa,) = 025' + tmssp))
In the meanwhile, have the following upper bound for | (2)
[k]] [k]l
w ( [k]] < |w!r (t-1) [k]] + nl sign (V Vi
((W((-)))
(0) |w j.r [k]] + tn = O(op + tr). (C.10)
22

---

Published as a conference paper at ICLR 2023
Besides, it is also easy to verify that for any t 1 To
θ s we have
(w Ei), willi. v) < (1/m)) and thus | ( (1) Then similar to (C.8), we have
sgn (V wwy.. ((W())
(
o' (2)
bi.r [i))s n.w ve.r [..([k]
= sgn Ni.t
= sgn @
0000 0000 + tnsop))- sgu(E,[b]) + oo (0o + tt])]]]
-gg((;(())) (k)).
(C.11)
This further implies that (t++)) yur (xi) 2 (w yirrs
Ei) O(n)
Σ
(sgu( Ww,,, - ((( 02 - O(ja)
KEB;
= (w (t) {i) + (msop))
where the term 62 is contributed by the gradient coordinates that are smaller than (n)8pp). This verifies Hypothesis (C.9) at time and thus completes the proof.
L
From Lemma C.3, note that sa, = (1), then it can be seen that (w j.rr J v) increases much faster
than (w J v). By looking at the update rule of j.r: v) (see (C.2)), will keeps increasing only when, roughly speaking, o' ('( J.. v)) > aoo'(w j.r E.)). Since (..+.)))
(t)
E.) increases much faster than (w!t) j. v). can be anticipated after a certain number of iterations, j.r. v) will start to decrease. In the following lemma, we provide an upper bound on the iteration number such that this decreasing occurs.
Lemma C.4 (Lemma C.4, restated), Suppose the training data is generated according to Defini- tion 3.1, 2 @((00p)²~ v o)-1) and (((s0))=)), then for any / Tb) with do T, = ( al//01//--))) I To.
(w (1+1) . v) (w j v) O(n).
Moreover, it holds that
sgg(() ④
k= == 1,
20p
(Tb) w
[ Jir
sgn(Er[A]) Ω1 )0. or + Omm) k E B, with = j, otherwise.
+Õ(1).
Proof. Recall from Lemma C.3 that for any < To we have
(w jrr . v) < (w J.r v) + (T) s (wjrrj v) + O(t)),
(t+1)
(w+1) {x) = (w(t) (t) v..r {x) + + (1sp S (w B Ee) ((tsoo)) Dert
(0)
Besides, by Lemma C,2 we also have [w]] 1.2 [k]]] A |w / [k]] + (tr)). Then it can be verified that for
some T, Õ( n
on (/-1)) we have for all E [n] and E T, Tol
ad' ((w g {.)) 2 C (
[o'((w j.n. v)) + A [1]]
for some constant C. This further implies that
sgn (V,,. L(W L(W())
= -gnn Σ ye J,s
0 ((w
j.r vvv)) - Σ
p(t) 1'o''/www ((w jort E.)) - ndw n1w
yi
Vil (lll() 0 ((w y.ro {x))
sggn
sgn(j),
23

---

Published as a conference paper at ICLR 2023
where we use the fact that sgn(yi ( sgn(f) for all [n]. Then by Lemma C.2 and (C.2),
have for all E [Tr,, To].
(++1) . v) = (w jnn v) (n) : sgn(s) sgn ( (Vw,,. (()))) (w.).(((())) (w) Jrr
(w (t) v) (n))
Then at iteration To, for the first coordinate we have
w j.r [4] + sgm(j) (TT,)) sgu(i) (T)T))))))) 2 sjn(
S0p
For any &e E B, with = j. we have either the coordinate will increase at a rate of (1) or fall into As consequence have either w j.r
(4bb) [k] [-0(1), (())
(Tu) j.r
[ = w W [k] sgm(E:(4)) (Ton) 2 (E]((]
For the remaining coordinate, its update will be determined by the regularization term, which will finally fall into the region around zero since we have Ton w(oo). By Lemma C.2 it is clear that
w (Tb)
[k] (()) 6()))
j.r
Lemma C.5 (Lemma 5.4, restated). If = - and "f o()). then let r*
arg [m]] (w (t)
parr E,), for any 1 2 To. i E [n]. E [2] and E [m]. it holds that
(t) Weir Ei) = =
Σ
ww()) Wear [B]] 1 [4: @(1).
KEB
Vr E [m], (4)
(w. sgn(j) v) E (-ō
na so?
O(A
Proof. The proof will be relying on the following three induction hypothesis:
(w())
Mirr - {:) 5(1);
(C.12)
(++
[&]] [e]] 6(1),
(C.13)
KEB
Vr E [m], (w jim sg) v) E - Ō na
0(1-1n)) m)
(C.14)
sop
which we assume they hold for all < t and [m]. i E [n], and j E [2]. It is clear that all hypothesis hold when f = Ty according to Lemma C.4.
Verifying Hypothesis (C.12). We first verify Hypothesis (C.12). Recall that the update rule for (w (1) barr: &i) is given as follows,
(++1) -Ei)
(1) (w(t) bi.r: Ex) -.- (mmmm) ( Vinir Ei)
(w())
Di.r: E.) -O(n) (sgn (V Wy.. L(W(1)), <.) - -612)303)
(w(t) : E,) + +()) Σ
th 1 si.r (i)) nAw (2.1166] EE(
sgn
kEB;
n
-oy, sgn
Σ
vae j.s ((w v(v)) Σ
yi
viefo oure ('( Ei) - (111111111
- -
(C.15)
Note that for any a and b we have sgn(a - 6) lal -2/61 Then it follows that
sgn V..t 1 i [.)) nAw ([][1k]) [E((A)
(1)
Σ KEB.
2nA w (4)
W..r [k]]
(
( E3)
(0)
kEB.
@(8ap) o
nd
(())
y. Tp
24

---

Published as a conference paper at ICLR 2023
where the last inequality follows from Hypothesis (C.12) and (C.13), Further recall that À - op/n), plugging the above inequality to (C.15) gives
(w (w(+ Ei) 2 (w ys.r E.) + @(nsap) - mnd
p(t)
2 (wit) ( Mr.r: E.) + ((nsop) O -
1100
(C.16)
p())
Then it is clear that (wiers Ei) will increase by ((sso,)
e(t)
is larger than some constant of
order Q2( aA) p(T)
We will first show that as soon as there is iterate w(F) satisfying o say
<
n. for some
S L, then it must hold that elr²) 21 will also be smaller than some constant
in the order of Ö ni
for all T' E [T,t+ 1]. To prove this, we first note that if p(0) 9:
reaches some
constant in the order of 0 ad we have for all [m]] by (C.16)
(wlltt (E)) 2 (w (t)
Miir .Ei
Werr Ei + (nsop))
(w ((+2) Ei) I (w m. 8 + O(an),
((w 3or
vvll s (w jrr v) O(n).
(C.17)
Therefore, we have
e 2. (Wit+)) xx:)
((+11) Wis
Σ E e (Wict)))
1 exp
EEm [o((w Dirr v)) + or
(++1)
(ww
(++1)
.Ei))
o(w ((+)))
verrs v)) - (((w (++1)
Baar Et))
1
[o((w])) pi.r. v\) ((w :ir: Ei) o ((w(())
Wirr v\) o (( Virr E1) (mnso?)
exp
1
1 exp
m Dm
[o((w pi.r.
v)) ((w (3) (w. p. E.)] o(ww!" (1) Bier v)) a(w -yt.r [,]]]]
where inequality follows from (C.17). Therefore, this implies that as long as (()) is larger than some constant ( nd then the adam algorithm will prevent from further increasing. Besides, since so
mnoo = o1), then we must have g(++1) E [0.5e(() yn.i 20 As a consequence, we can deduce that ((t) Mi. cannot be larger than 2b, since otherwise there must exists a iterate (r)) with S such that (T)
WEs E [b, 2b) and ((++1) 9:.
((T) which contradicts the fact that ((T)) should decreases if p(T) 9.d 2 b. "9.
9:3
Therefore, we can claim that if p(T)
S Ō( for some S t, then we have
an
elr S δ
nd So?
(C.18)
for E [|T 1). Then further note that
F v. (W()) xx))
20(4+1)
Eje(-1.1) (W(O).
m
2 exp Σ
[o((w(t) wh. v.v) + + (((w() aer [i)]
exp
(- Omm max o ((w(())
pi.r E.))))
(C.
rE[m]
25

---

Published as a conference paper at ICLR 2023
where in the last inequality we use Hypothesis (C.14). Then by the fact that cl+t1) < 0 s
o(1) and m = (1). it is clear that exp O [a(m a em (Weer E.) = [2()). This verifies Hypothesis (C.12).
a(w Mir 6()) o(1) so that
(4+1)
Verifying Hypothesis (C.13). Now we will verify Hypothesis (C.13). First, note that we have (f+1)
already shown that (w yurt E (1(1) so it holds that
[3] [ + a w
(1+1) Mir
[1]] 2 (W
(H+1) yr E.) = (1)).
REB;
By Hypothesis (C.14), we have |w (1+ [1]] S w Des [1] +1 = (1). Besides, since each coordinate
in & is a Gaussian random variable, then max
[E]]]]] = Ö(o,). This immediately implies that
Σ REB.
|w w
[K]]
[E:[A]] 50(1.
Then we will prove the upper bound of E KEB, Woier any ls B, such that V Ww,,. L(W( 'nnnnnso p
(1+1) [[:]] | [K]] Recall that by Lemma C.2, for
p(t) w. we have
wit BuT
[k] + O(7) sgn ( ((w!) We EMEE - (2. yo.
Note that by Lemma C.4, for every k E B; we have either (To) Wi. [k] sgn(E;(<t) O
STT or [ oaa [k]]] $ 17. Then during the training process after To, we have either sgn p..r [k))
(To)
sgn(4 (4: [ or (g)( Wair 2 -O(n) since if for some iteration number t' that we have [R)] but sgn( sgn(w ( (<
sgn(E;(k))., then after Õ(1) steps (see the proof of Lemma C.2 for the definition of 7) in the constant number of steps the gradient will
sgn(wg (w 9 [k]) = sgn((
(t')
must be in the same direction of [k], which will push w back to zero or become positive
along the direction of [A]. Therefore, based on this property we have the following regarding the
inner product (w werr: Ei)
(w(() Merr Ei)
Σ KEB,UL)
w Mhir [k] [A]
Σ
(4)
[b]] [s( Ö(1)
[E: kEB,U(1)
Σ
9
(EB
Σ
[M]+ Ö(1spp),
LEB
(t) where the second inequality follows from the fact that the entry w your [k] that has different sign of [i[k] satisfies w (t)
[k]] < Ö()). Then let Bi"h) ΣjeB,u()) [ W. [k] 11 1[( g..r [4]] Ö(n))
[A]],, which satisfies BiTa) ē(1) by Lemma C.4. Then assume BS") keeps increasing and
reaches some value in the order of O log(dnu)- it holds that according to the inequality above
(w()) Werr 5 = θ [ogg(dn 1)) - ((soo)) = O log(dnn-)))
where we use the condition that y = ((so,)) Then by Hypothesis (C.12) and (C.14) we know (t) that (w jor v) = o(1), (w k..r. {.) (11,, and (w -..... E.>) = Ö(dr) +a (w - v) = o(1)
then similar to (C.19), it holds that
F- (W()) x.)
Σ Eje(-1n1 (wi) I exp (- {,)))) s ( ,n
91.5
Therefore, at this time we have for all & E B,,
(1) 1 ((w (0) bur {()))]]]] < poly(d-1 nn -n..)) .n) O (oo (doon""(dnnn1)) @(o,) s ndnn
26

---

Published as a conference paper at ICLR 2023
Then for all | pi [A]]] 2 Ö(1),
the sign of the gradient satisfies
sgn Vww.. L. [[))
sgm
((0) Bi,
ndw(()
Miir
[R:)
sgg(nnnn w w(lr
Then note that V Ww L ((W()) = @( ( [+]) 2 @n= + An), by the update rule of Wit [k] and Lemma C.2, we know the sign gradient will dominate the update process. 0. [k] (nn) sgn( (en(wwl. 1 [..r+())) & w. [k]],, which implies that wae [k] 1( [ 2 O(r)) decreases so that B(C) also decreases. Therefore, we can conclude
Then we have Wagtr
[k]] [+++D((]]
that B.d will not exceed @ log(dnn) Then combining the results for all E [n] gives
Σ
w (C)
[+] [b]] & ((t) + Ösngp) & e (log(ddm) 1) + 0 = 6(1)
Mart
REB,
where in the first inequality we again use the condition that v) = o(1/d) = -((80))))) This verifies Hypothesis (C.13). Notably, this also implies that (w "u.r" E,) maxre/m] (w ( 5.) S (1).
Verifying Hypothesis (C.14), In order to verify Hypothesis (C.14), let us first recall the update rule of (w!0 jart v):
(w(1)) Ao
v) = (w j.rt v)
Then by Lemma C.2, we know that if V W,,,. L - ((n). then Im j.j.// V V.rl I (1) and otherwise
(t)
m
Jur
sgn
Σ
viell 33
[()) 0' ((w
j.r viv)
ee 23 o' ('w jorr Ei) nAw (011 9(1).
(T) [1] =
Without loss of generality we assume j 1, then by Lemma C.4 we know that w
-((. ). In the remaining proof, we will show that either w w (t=1) [1] E [0, 0(x-)))))))))))) or w 1, [1] E
Ó no ))
First we will show that w [1] [0, o()-1n)) for all r. Note that in the beginning of this stage,
we have [T) < 0. In order to make the sign of w (f) flip, we must have, in some iteration
t' < that satisfies w (t') [1] [o, ((x--1))) therefore
Vn wi.,. I(W())) [1] = Σ yi 3x 0 (o(((w (1)
J. , Σ
yaelt') 0 ((w Ei) now (1
n ( 8-2 --2-
>1 w (1) (nnn) 20
where the second inequality holds since " o(A o(A0-11)(0-2)) Note that Ww.,, L(W()) > (n), then by Lemma C.2 we know that Adam is similar to sign gradient descent and thus w ( w (1 O())) which starts to decrease. This implies that if w (e+1) positive, then it cannot exceed @(~-1,)) o11)
27

---

Published as a conference paper at ICLR 2023
Then we can prove that if Wir [1] [ for all l' S L
(++1) negative, then [wi,
[1] [.7
Ō sdp
In this case have
v L(W(
= Σ Vi 1n 0' ((w) (.11... - Σ viel of ((w (1) Ei) - nAw 1
((a) +
((a) w
=
where in the inequality we use Hypothesis (C.13) and (C. 14) to get that
6 (w Wir
S
[+]] max
& |+] + a (l( (t') surt
v) @(1)
KEB
REB
Recall from (8.18) that we have
(
therefore we have if w (12 [1] is smaller than some
so
value in the order of - - so?
(d), then
Vv (W()) -
am?2> so?
nd na + θ
polylog(d) O(nn),
p
which by Lemma C.2 implies that w jir
[1] will increase. Therefore, we can conclude that w (1++)
[- 0 no sa ), 0) in this case, which verifies Hypothesis (C.14).
L
Lemma C.6 (Lemma 5.5, restated). If the step size satisfies Odd-1/),, then for any it holds that
L(W(++1) - L(W()) S - VL((()))) + (()"d).
Proof. Let F,, F5(W(+) x.) FJ( (W()) xx). Then regarding the loss function
(W.))
((W) = - log E, efiiwx)))
(W, + log Σ Ee
It is clear that the function L; (W) is 1-smooth with respect to the vector [F [F Fi Then based on the definition of A AFG we have
(i,,,,;))
OL. (W()) (WO)) APJ + Σ (AP
L L(((+)) - L. (w Σ
(C.20)
Moreover, note that
F3 FW (0) x;) = Σ [a((w ( v:v)) + +a((w 0 (.)))
(1)
By the results that (t) jor v) S ((1) and (w(0)
j.r E) < ((1), for any (d-1/2), we have ((+1)
(1+1) (w j.r v) + s @(1), (w!--" . S j.r E.) + ((ns 1/2 - 6(1),
which implies that the smoothness parameter of the functions a (( j.m v and ( j.r: {x) are at most @(1) for any w in the path between w and j.r
(t+1) Then we can apply first Taylor expansion
on a a(w j.m v(v and ((w J.r {i)) and bound the second-order error as follows,
(t)
(1+1) o (o(wWf (e+1) .biv)) - a(w [.7 v(v)) - (Vw,,-0((W()) v(v)), w j.r w w
< @((l jrr w111111) ((72d),
(t+1)
(C.21)
28

---

Published as a conference paper at ICLR 2023
where the last inequality is due to Lemma C.2 that
[w (++1)
(t) n?²
S dl
3.r 2
Similarly, we can also show that
[w((W] (+t)) E:)) a(w 1. (.) (Vws,, o ((w (1) Jir E.)), wrr
((+1)
w S d()
(C.22)
jrr
Combining the above bounds on the second-order errors, we have
AF (V F xx),
- -wll - =)mm)d) ((1))).
(C.23)
where the last equation is due to our assumption that m @(1) Besides, by (C.21) and (C.22) the convexity property of the function a (zr), we also have
(a(( f " ,uvv)) - a(w (1) 3.r svvi S WV,,,, o ((w
jr uv)),, W
(t+1) w w ((r)d))
jr
[ ((w jrr
((+1)
vv)))l -Ivll1) + ((7²d)
θ(1 + +f2d);
|o((w))t+ (1+1) . - o(( j.r: {.))
S 1(V w,, o (w j.r Gi)), (++1) - -
w
= e(T)o' (w (t+1) -i))l [k]l] + 6(77d)
= Ō(nso, + 1))d.
These bounds further imply that
AF S 6(m (nso, + r²d)) ã (nsop + -2dd..
(C.24)
Now we can plug (C.23) and (C.24) into (C.20) and get
(W("))
L((W((1)) - L.((W))) (1) = OF,(W()) xx) AF + E(AF),))
(W())
$ (W x:) (WW0, x;), w(t
+ + + +((SSO + +)dd)
(W(),, W(++) W()) + @(n²d),
(C.25)
where in the second inequality we use the fact that L. (W) is 1-Lipschitz with respect to F, x;) and the last equation is due to our assumption that op = O(s (0s-1(() so that G((nsa, + + = Ö()2)).
Now we are ready to characterize the behavior on the entire training objective L(() = En i= (W) 1l Note that W|| is 2A-smoothness, where X = o(1). Then ap- plying (C.25) for all i E [n] gives
1 T n
Σ [L (W(+t)) L; (W())] + +> X -
< (VL(W()), w(++1) - W()) + +(()d))
where the second equation uses the fact that [ - - = (v)2 Recall that we have
(4)
(1+1)
w 1r
1)-
v j.r
29

---

Published as a conference paper at ICLR 2023
Then by Lemma C.2, we know that m large. Then we have
m 1/14// / (..l [k] is close to sign gradient if V [L(w(0)]A) is
(0)
Vw, Vwy.. L(W()))
m
0 V W.,-- L(W()))))) ((d- 1)) O mssop)
e V w,,. (( Vw,..
-(dn),
where the second and last terms on the R.H.S. of the first inequality are contributed by the small gradient coordinates & & i and & E U B3 respectively, and the last inequality is by the fact that ns² o (d). Therefore, based on this fact (C.25) further leads to
L(W(++1)- - I(W())= & -nl VL((())]] + d))
which completes the proof.
Lemma C.7 (Generalization Performance of Adam). Let
W
argmin RE(w").....o..
[|VL(W)]lI
Then for all training data, we have
(W' x. S (W", (W",,)) = 0.
Moreover, in terms of the test data (x,y) v) ~ D, we have
P(x B,( x) S F-,
x)] 2
Proof. By Lemma C.6, we know that the algorithm will converge to point with very small gradient (up to (nd) in l norm). Then in terms of noise vector &,, we have
Wwy. L(W*) [k]]| s (nd).
(C.26)
REB.
Note that
nV. w... L(W [k] oo ((W oi.r (.)) nAw Ik.
where " = logit,y, (F*, x:). Then by triangle inequality and (C.26), we have for any r E [m].
L lo ((w werr [e))] - nd Σ
|w (w
S n
Σ V L(W [b]||| - Onndd)
REB,
REB.
REB.
Then by Lemma C.5, let r* arg ma (W...r Ei), we have (Wyir Ex) ē(1) and
ENER, |w W.. [k]] |&: [b]] = @(1). Note that [k]] = Ö(ap). we have ΣxEB, /w [A:] >
((1/op). Then according to the inequality above, holds that
1 ((sa)) 2 Ō nd
W W nndd 2 3
ndd
nd
op
AEB
where the second inequality due to our choice 1j. This further implies that 1 = g by combining the above results with (C.18). Then let us move the gradient with respect to the first coordinate. In particular, since |]V((*)) < ((jd), we have
ZtA
12
In w L(W [[]
Σ yil ((W)) fm v(v) Σ y 2nly ((w Ei) ndw3..
$ O(nd).
(C.27)
30

---

Published as a conference paper at ICLR 2023
Then note that egn(y:().) = sgn(j), it is clear that w < 0 since otherwise
2 ( A
V L [W)) 2
a Σ G.s [o ((Wj,r, Ei)) o' ((w
2 (nnd)).
sos
==1
which contradicts (C.27). Therefore, using the fact that w [1] J S O, we have
V L(W [D]] = α
Σ
Vi o' ((i, Ei))
Σ
yo Vi (( G. [1]]] - -nnnn/w,,,, [1]]
Then applying (C.27)and using the fact that (fgel that
A mm for all i [n], it is clear
[ [1]] 2 T
na
na
so?
30
where the second equality is due to our choice of op and a. Then combining with Lemma C.5 and the fact that w j.r- [1] j < O, we have
j.r 11:55 -5
na so?
Now we are ready evaluate the training error and test error. terms of training error, it is clear that by Lemma C.5, we have ( 2 ((1) (wh..n E.) 2 -o(1), and (w v) o(1),
((w 22. S:I = o(1). Then we have for any training data (x;,
Fg. (W" xx)
(ww w.. m..r U:)) + (w Ei)) = =(1),
F-9. W xx) =
Σ [o -.rr. vvv)) + wa.. Ei)]] = o(1),
r=1
which directly implies that the NN model W can correctly classify all training data and thus achieve zero training error.
In terms of the test data (x, y) where x = [yv, E). which is generated according to Definition 3.1. Note that for each neural, its weight wj.- can be decomposed into two parts: the first coordinate and the rest d - 1 coordinates. As previously discussed, for any j E [2] and r = r", we have sgn(l)- [i.r S -((n)((so)) and sgn(j) [ < G(A in) for # r". Therefore, using the fact that é (nna(s80?)) = w(X- and Lemma C.5, given the test data (x, y). we have
Fg ,)
(o((w g.r. vv)) + o ((w g.r. E))
$ Σ T ((.
+ Coir 19) ).)
na
so?
m
F- F-y(W",x)
Σ
[a((w y.r yv)) + <(W p.r: E)]
0 |
[1] 9 + 15- [i]
2
2
na
q
1 17
$a?
where the random variables Cyrr and Sy.r are symmetric and independent of v. Besides, note that a = -o[1, it can be clearly shown that a na/(so?) < naa(soo)) Therefore, if the random noise Gy.r and (-g.r are dominated by the feature noise term (w- yV), we can directly get that F2, (W", < F=,("W xx) (recall that m = (())), which implies that the model has been biased by the feature noise and the true feature vector in the test dataset will not give any "positive" effect to the classification. Also note that ñy and (-v are also independent of v, which implies that if the random noise dominates the feature noise term, the model W will give at least 0.5 error on test data. In sum, we can conclude that with probability at least 1/2 it holds that F,, (W*, x) < F F,,(", which implies that the output of Adam achieves 1/2 test error.
(W*, x),
31

---

Published as a conference paper at ICLR 2023
C.3 PROOF FOR GRADIENT DESCENT
Recall the feature learning and noise memorization of gradient descent can be formulated by
(++1)
(1- nd): (w (1)
jirt .v)
+ 2
Σ Vel o' ((w j.n
((w v(v))
Σ Mr
e o' jor Eil)
((( Eii (1 n7) (w Beirs E.) +
Σ
wi o' (( Bears E.) E,]
(0)
kEB;
na
Vins o
((w Dirr Ex))
Σ Us e0 o' (' ber Bo))) (C.28)
Then similar to the analysis for Adam, we decompose the gradient descent process into multiple stages and characterize the algorithmic behaviors separately. The following lemma characterizes the first training stage, i.e., the stage where all outputs F, (W()), xx) remain in the constant level for all j and i.
Lemma C.8. [Lemma 5.6, restated] Suppose the training data is generated according to Definition
3.1 and X o(oo 2 20, Let A maxye[m] (w (t+1) .j v), r = max+e(m] (w jart Es), and
jir
1 max: (1) Then let Tÿ be the iteration number that A/) reaches ((1/m), we have
T = ē
2/20//0 /v) for all j E (-1
Moreover, let To = min {Ty}. then for all t s To it holds that 1.00 = Ō(oo) for all j e {-1,1}.
We first provide the following useful lemma.
Lemma C.9. Let {æt, Ve h= be two positive sequences that satisfy
J++1 2 Ft + 1) Azt
Ue+1 S yt + Byy 2-1
for A O(1) and B o(). Then for any q 2 3 and suppose yo O(ro) and v < (2r)) we have for every C' E [zo. (1)], let Te be the first iteration such that re 2 C, then we have
T2n = (1) 2-9 9) and
VT. S Ozzo).
Proof. By Claim C.20 in Allen-Zhu & Li (2020b), have Tun O(23-9). Then we will show
S 2r0
for all Tz. In particular, let Tel C²23-9 2-9 for some absolute constant C' and assume C' B22-1 < (this is true since B = o(1)), we first made the following induction hypothesis on Ye for all & < Ta,
y: S 90 + tm (2x0)
Note that for any s To, this hypothesis clearly implies that
Vt S yo + T.mB'2r-12 120-1 < xp CB29-1 122- 1.299.1 ng
2np-
Then we are able to verify the hypothesis at time /. + based on the recursive upper bound of yts i.e.,
Vt+1 V 7- By:
< y0 + (2x0 18 + +1 Byy
$ yu + (t + 1))x(22))*-1
Therefore, conclude that y 2ng for all < Tz. This completes the proof.
32

---

Published as a conference paper at ICLR 2023
Now are ready to complete the proof of Lemma C.8.
Proof of Lemma C.8. Note that the initialization, we have (w j.r v) = ((o)) and 3w Ex)
(0)
(0)
0(5 Then based the parameter scaling summarized in Appendix C.I, we have
F FW x.)
Σ
(o Sor Biv)) + ((w 3,r E.))] o(1)
for all j {-1,1}. Then we have
cs xx.)
e Fwi xx))
min Σ, ef
(().
Ej (W()) x.)
Then we will consider the training period where |/ = (((1) for all j. i, and t. Besides, note that sgn( yi ((t) 2,1 = j. Therefore, let = arg max, (w (t-1)
j. .j - v). (C.28) implies that
4,90 (w/)t/") j v)
(1 -n2) (w) (t-1) 1) J v)
(1-1) (t- ((w) j.r" ,((v)) 1 10' ((W
+
(t-1) 2 (1 - (A- (w j.r J v) - O(1) 1o' ((w j,r .j v)) ao (T
(t-1)
2 (1 (1-nA)A 11 + n 0 ((4 (-))) -n ((a(r(--))--1
(C.29)
Similarly, let = arg max, (w
(t) (Wrrrrs yi.r, Ei), we also have the following according to (C.28)
p(t)
(w (1)
; Ei)
11802
na²
n
s (1 (1-nA)(we Mer E θ
0' (W E
A
Σ
o' d'llw ((Wtt-1)
" Ex))
(t-1)
n
n
nso (T(t))(--))) < (((-1) y0.
na²
a Σ
9-1 (1)
wi.
n
n
Then by our definition of ( = max (1) j,8- we further get the following for all j E (1,
p (T((=1) + n + anaa 2
(1(6- (05--1)0-1)) -- + A nso P
(T (1(6-11)0-1)
(C.30)
where the last equation is by our assumption that a = O(s02/n).
Then we will prove the main argument for general 4, which is based on the following two induction hypotheses
A 2 1.t (t-1) +
(C.31)
!os (1((-1)0-11
p < 3
(t-1)
(C.32)
n
Note that when f we have already verified these hypotheses in (C.29) and (C.30), where we use the fact that À = o(ag-2 (0o003ap//)) S (A (A 9-2 and a = o(1). Suppose that (C.29) and (C.30)
hold for iterations 7 < At time + 1, for all < t, we have
r < ((A)))
33

---

Published as a conference paper at ICLR 2023
So² /n (11) and A increases faster than r.) Besides, we can also show that Ar(()) S )r))) which has been verified time = since r5) keeps increasing. Therefore, we have 5 S (
- {0(())))))
and hence (C.29) implies
A 1++1)
2 2(1 1)) +1 ((AA)))) )
2 +1 ((AA)))0-))
which verifies Hypothesis (C.31) at + 1. Additionally, (C.30) implies
re n()) + e (nsop
(r) (1))))))
which verifies Hypothesis (C.32) at t + 1. Then by Lemma C.9, we have that A. = Ö(1) for all < To = = = ( 9/n). Moreover, Lemma C.9 also shows that r(t+1)
0(4,)) = (o))) This completes the proof.
Lemma C.10. For all i E [n] and S T it holds that (w (t) B.rr r. E1) S ((a).
Proof. First of all, for E [+1}, by the definition of Tj. we have
(w j v) O(1.
Moreover, with the same proof as Lemma C.8, it is clear that - (W j. v) is decreasing in t for I s Tj. Therefore, by the fact that (w jur v) < (1), we have (w
(0)
kw Jir v) < ē(1)
(C.33)
for all t S Ty.
Now by the update form of GD, we have for any Bi
(++1) w.r [k] [M] = (1 - m/ w [k) E
[k] +
Σ
(0) o' (( -n E.) E.
KEB.
Note that e0) of ((w -i.d
wert {x) < O, which implies that w
[k] [:] is decreasing in t. There-
fore, for all and we have
Ei) = w -..rr [ +
Σ
[k]
REB
S [1] [5:(1 +
[K]
REB
S ((()) (t)
[
[1]l +
Σ
(0)
[
KEB.
S @(a) + +
= Ō(a).
where the third inequality follows by (C.33). This completes the proof.
Note that for differentt j. the iteration numbers when ( reaches ((1/m) are different. Without loss of generality, we can assume Ti < T-1. Lemma C.8 has provided clear understanding about how A varies within the iteration range [o, T;). However, it remains unclear how
r(()) varies within the iteration range (Ti, T-1] since in this period we no longer have |e(0) = @(1) and the effect of j.a
gradient descent on the feature learning (i.c., increase of (W j v)) becomes weaker. In the
following lemma we give characterization of A A() for every E [Ti, T-1).
34

---

Published as a conference paper at ICLR 2023
Lemma C.11 (Stage of GD: part II). Without loss of generality assuming T1 T-1- Then it holds that ( = (1) for all E [-i].
Proof.
Recall from (C.29) that we have the following general lower bound for the increase of A
A(+11 (1 1/a) (w (w v)
(1) 9u5 lo ((w)) v(v)) Σ 2all ((w)
n
(1 - -))A + x02))
Σ (
- (T/)) v v((a))* (C.34)
where the last inequality by Lemma C.10. Note that by Lemma C.8, we have y Ö for
all t T-1 and. Then the above inequality leads to
1 ((+2) 2 (1- + θ
Σ (4 EnL=S
(C.35)
where we use the fact that w(o)). The the remaining proof consists of two parts: (1) proving A 2 (L/m)) = (11) and (2) A < (log(1/A)).
Without loss of generality consider 1. Regarding the first part, we first note that Lemma C.8 implies that A (T) 2 @11mm). Then we consider the case when A s (llog(1/a)/mg holds
that for all y/: 1,
FF-I x.)
E 25E(-1,1) e (w(,))
m
m
(a((w" (8)
(((w-I,.:. o(w E\)
Σ
(((w a((w()) {.,,
(t)
exp
2 exp (- ((m)("))
2 exp(- ((lo0(((((((/)))
O(a).
Then (C.35) implies that if [ < @(log(1/og) (/mm))/m) we have
A (1+1) 2 (1 - + ((ja) A (() 2 A + O(yo) A 2 (9
where the second inequality is due to À o(a). This implies that A will keep increases this case that it is impossible that A (1/m). which completes the proof of the first part. For the second part, (C.28) implics that
A (++) < - 1 99)
+
1
Σ
42 (A1))-11
(C.36)
Consider the case when ((() > O(log(d)), then for yı = 1,
e F-(((()) x.)
E E e5 (w(),.))
exp 0
(o((w (8)
-(,,, (( &.))
Σ
(o o((w/)) ((w))...))))
I exp ( (- ((A
< expl( exp(-O((og((/))
((poly(A)).
35

---

Published as a conference paper at ICLR 2023
Then (C.36) further implies that
A S (1- (1)1)))ff) + ((poolya) (4f9)
s 1 ((1A()) ( (A-poly(A) (4 S (0)
which implies that A will decrease. As a result, we can conclude that ( will not exceed ((log[I/A)) this completes the proof of the second part.
Lemma C.12 (Lemma 5.7, restated), if & O(oo), it holds that Al9 @(1) and r50 = Ō( for all E [T.
Proof. We will prove the desired argument based on the following three induction hypothesis:
1 1(11 2 (1 - 1 +
+0
2 n1
Σ
-(a
(C.37)
jab
p = Ō(oo).
(C.38)
A = O(1)
(C.39)
In terms of Hypothesis (C.37), we can apply Hypothesis (C.38) and (C.39) to (C.34) and get that
1ft11 2 (1 - n2)A50 +
m
Σ
2. (A/0)=-1 - O(an) (r V ((a))) 1
2 (1- (1 (0)
+
((a 1
jrl
where the last inequality we use the fact that 2 og- This verifies Hypothesis (C.37).
In order to verify Hypothesis (C.38), we have the following according to (C.37),
Σ
1 1 (1 - An)
Σ
A (8)
@()²)) 1 Σ 10
71 1=1
1E1-
1,1-1,1)
= (1 Xn)
Σ
1000
θ
(
-E(-1.11
where the last equality holds since o(1). Recursively applying the above inequality from T-1 to gives
4
Σ 1500 (1)
Σ
A
+0(2)
Σ xn) Σ
1
JE(-111)
JE(-1.1
Then by Hypothesis (C.39) we have
T- Σ
θ
(1-
xn) Σ
plfe 1-T) < x(11)
n
-=0
1=1
36

---

Published as a conference paper at ICLR 2023
Now let us look at the rate of memorizing noises. By (C.28) and use the fact that a² s (0so0p/n) we have
1 (1 1A)r + +
leall
(TI--
n
9-1
(1 (1-r -1-11 +
Re
n
nso 100300
t- 1-T-111 Σ
(T 1) + 0
(1 -nn)²
Σ
1 p6- 10-1-
r=0
3 + s0 20208-1)
((o))
which verifies Hypothesis (C.38).
Given Hypothesis (C.37) and (C.38), the verification of (C.39) is straightforward by applying the same proof technique of Lemma C.11 and thus omit it here.
Lemma C.13 (Lemma 5.8, restated). If the step size satisfies, then for any / T-1 it holds that
I(W(++)) L(W()) s 7 (
Proof. The proof of this lemma similar to that of Lemma C.6, which is basically relying the smoothness property of the loss function L(W) given certain constraints the inner products (W,rr, v) and (W,rr,
Let AFJ = F,((++)) xx) F((W()),x)) we can get the following Taylor expansion on the loss function L. (W((L))
L((((t)) ((W(t))-L(((())) < Σ (L((W(9)
AFJ. + Σ(DE,,)
(C.40)
(W x.)
In particular, by Lemma C.12, we know that (w Then similar to (C.21), we can apply first-order Taylor expansion F2 (W((11), x;),, which re- (++1) quires to characterize the second-order error of the Taylor expansions on al a((w vvv) and (t+1) a a(w jr
j.r viv) S 6(1) and (w (w (. Ei) s Ō(oo) S ē(1).
E.)).
a (( (t+1) y(v)) - ((W.. - (Vw,, ((w (V)) w ((+1)
w w(())
j.r
< é (ll Jir
(1+11 w (12111)) V 1, (11
0 (a for
(++1) _E))) - ((w (.. E1) (V W,,,, or ((w (2)
for Ei)). w j.r
((+1)
< <(([w (++1) w (11113) 6(12 (III,,,.L((())) (13
(C.41)
Then combining the above bounds for every r [m]] we can get the following bound for AF,
AF (VWF, x:), w(++1) - w(t <
Vw,, Ww,.. L(W("))
rE|m]
5(n² (V
(C.42)
Moreover, since (w j.r y.v) & 6(1) and (w E.) s ((1) and ((-) convex, then we have ((+1) la((w((t") jrr .v(v)) a ((w(() ((w vi))) & max (lo'((w ((+1) ,y.r))l lo (( (0) fm v(vi)) ((v, ((+1)
-w([))
< Si(Iw (0+1) (111)))
37

---

Published as a conference paper at ICLR 2023
Similarly we also have
[o((w), Ar E.) - ol (( jirs (.ro +1) S ((1 oo
(++1)
((tl)
- 111)
Combining the above inequalities for every [ml], we have
(
[w( - w.f.hlle 1) <
A 2 5
Σ
rE(m]
(C.43)
Now we can plug (C.42) and (C.43) into (C.40), which gives
OL, (W(")
- - L. (W(()) s E
OF, (WO,, x;) AFJ + I(AF,,)²
= (W(), w(+1) W + 1() V (C.44)
Taking sum over i [n]] and applying the smoothness property of the regularization function \/| W[ we can get
L( L(W())) 1 Σ (w((t)) -( + A A -|
S (VL(W()), w((+1) W(o) + ((
- (1 - -((2)) I|VL(
s 2 (
2
where the last inequality is due to our choice of step size = o(1) so that gives lj - -O(17) 2 n/2. This completes the proof.
0
Lemma C.14 (Generalization Performance of GD). Let
W arg
min (w.).... W())
||
Then for all training data, we have
II (W*, xx) F- (W", x;)] 0.
Moreover, in terms of the test data (x. D, have
P(x (W* x) < (W" x)] o(1).
Proof. By Lemma C.12 it is clear that all training data can be correctly classified so that the training error is zero. Besides, for test data (x, y) with x = lyv _ T it is clear that with high probability (w" n.r yv) @1)) and [(w", g.r: E) < Ö(ao). then
F2 (W', x) Σ [o((w" y. vv)) + o(ww r {)) ((1).
If j -y, have with probability at least 1/poly(n), (w" y. vv) I 0 and [w" [w",,r:: E)I+ s Ö(a), which leads to
F-y (W" x) [o((w"g.r: vv)) + o ((w g.ri (.rr+)))))) {\)
= < =(maa))=(aa))))) =(1))
This implies that GD can also achieve nearly at most 1/poly(n) test error. This completes the proof.
38

---

Published as a conference paper at ICLR 2023
D PROOF OF THEOREM 4.2: CONVEX CASE
Theorem D.1 (Convex setting, restated), Assume the model is werrpaarameterizd. Then for any convex and smooth training objective with positive regularization parameter X, suppose we run Adam and gradient descent for T = poly(n) iterations, then with probability at least n the obtained parameters W Adam and W CD satisfy that ( Ty and (FLmm Ty respectively. Moreover, it holds that:
Training errors are the same:
1 sgn (WW Adam xi)) # y1 = Σ
(s # v.].
π
Test errors are nearly the same:
(xxx)>DD ssgn (FWAAdaan Adam xx)) # [ IP (x (sgn( (F(W GD- x))
fu + o(1)
Proof. The proof is straightforward by applying the same proof technique used for Lemmas C.6 and C.13, where we only need to use the smoothness property of the loss function. Then is clear that both Adam and GD can provably find point with sufficiently small gradient. Note that the training objective becomes strongly convex when adding weight decay regularization, implying that the entire training objective only has one stationary point, i.c.. point with sufficiently small gradient. This further implies that the points found by Adam and GD must be exactly the same and thus GD and Adam must have nearly the same training and test performance.
Besides, when the problem is sufficiently over- parameterredeed with proper regularization (feasibly small), we can still guarantee zero training errors.
E DISCUSSION ON THE DATA MODELS IN WILSON ET AL. (2017); REDDI ET AL. (2018)
Data model in Wilson et al. (2017). In particular, given the binary label E {-1, 1}, the feature vector x; is set as
yi
j =
1,
j = 2,3
x.
1, j = 4 5(i - 1), 4+ 5(i 1) + 2(1 yi)
0, otherwise.
Data model in Reddi et al. (2018). In particular, Reddi al. (2018) considers one-dimensional optimization objective. Besides, in each iteration Adam, the stochastic gradient taken based on the function (xx) defined as follows:
fe(x) = Cr. mod 3 =
-x, otherwise.
Then it can be seen that in these two prior works, each coordinate of the feature vector (or the objective function) is hard coded. In contrast, our data model allows randomness in the data gen- eration process. This implies that our theory can hold for the data points generated from a certain distribution, while these prior works can only cover one particular data or optimization objective.
39


Reviews:
[{"review_id": "wuzC5h_q1-", "reviewer": "Reviewer_zveC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "On the specified data distribution, this submission shows that Adam and GD, starting from the same initialization, can converge to different solutions with significantly various generalization errors, even with proper regularization. ", "summary_of_the_review": "Although there are some limitations and issues in the theoretical part, this submission attempts to explain an interesting experimental observation from the theoretical perspective, so I have a positive score on it.", "strength_and_weaknesses": "####  Strength\n- This submission is well-written and the key motivation is well-discussed.\n- It reveals the difference w.r.t. generalization and convergence speed between Adam and GD on a non-convex problem under the non-asymptotic setting, which was an open problem before.\n\n\n\n#### Weaknesses\n- The submission does not consider the de-bias term in Adam; hence, it should specify the initialize condition for $m_{j,r}$ and $v_{j,r}$ in Eq. (3.3) and Eq. (3.4), respectively.\n- The noise vector's norm is much smaller than that of the feature part.\n- The stochastic gradient complexity to find an \u000f$\\epsilon$-approximate first-order stationary point on the general non-convex problem (non-stochastic) for GD is $\\epsilon^{-2}$. Namely, $T = O(\\epsilon^{-2})$ in Thm.4.1 for GD, which is right. Not the complexity's lower bound for the first-order optimation methods on the non-stochastic non-convex problem is $\\epsilon^{-1.75}$. The complexity for Adam in Thm.4.1 is $\\epsilon^{-1}$, which obviously violates the known lower bound. I know some assumptions exist on the data and model while the low bound is proved in the worst case, but the author should also discuss the part against a common guarantee.\n- It is strange that the results for Adam have no connection to the moving average coefficients $\\beta_1$ and $\\beta_2$. \n- As the author also claimed, Adam is only similar to signGD when  $\\beta_1$ and $\\beta_2$ are small. However, in practice and as pointed out by the author, we usually set  $\\beta_1=0.9$ and $\\beta_2=0.99$, which distinguish signGD and Adam. Moreover, small  $\\beta$'s may make Adam diverge on some examples, but large $\\beta$'s will not. Hence I quite doubt the proof routine that extends the results on signGD to Adam.\n- For Lemma 5.2, $\\eta s \\sigma_p$ seems much smaller than $\\eta$ due to $s = O(1)$ and $\\sigma_p = O((poly log (n))^{-1})$, why to claim that $\\eta s \\sigma_p$ increases. much faster?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "r1w3BbcTor", "reviewer": "Reviewer_52yW", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proves under some particular situation, e.g. the input vector is the concatenation of a 1-sparse feature vector and a noise vector, under certain assumptions, the two layer CNN (whose second layer is assumed to be fixed as all 1’s), converges to “bad” local optima with high probability when trained using ADAM. However under the same setting, when trained using gradient descent, with high probability the final convergence is a relatively “good” local optimum. The authors also show for the convex settings such as linear models + logistic loss, both ADAM and gradient descent converge to “good” local optima with high probability.  \n", "summary_of_the_review": "I have mixed feelings about this paper. On one hand it is a solid move towards understanding the generalization capability of ADAM and GD. On the other hand, the assumptions and simplifications made in the paper have made me less excited. The models and data discussed in the draft are too far away from the real models and data used in applications, so that I am hesitating to use the insights from the paper in any real applications.  \n", "strength_and_weaknesses": "Strength: It looks like the first work to show the finite sample generalization guarantee on the local optima of a CNN trained using ADAM as compared to GD. I guess one reason we haven’t seen such results is that in general it is not easy to give rigorous finite sample guarantees on ADAM’s generalization errors for the non-convex neural networks. In that sense, this paper has its unique value. \n\nWeakness: Even though I understand without the assumptions it is hard to provide rigorous proof, I would still say the assumptions and simplifications made in the paper are too much. Those include\n1) two layer CNN but the second layer is fixed to all 1s\n2) the simplifications on the input vectors seem weird. 1-sparse looks over simplified to me, and in general the noise comes together with the input signal instead of coming from a totally different dimension.  \n3) the width of the network is assumed to be polylogarithmic in the training sample size. Even though over-parameterization is used a lot, usually people do not assume the model’s width grows with the sample size in applications\n4) The work only gives results on full batch ADAM and full batch GD. However in real applications full batch optimization is rarely used. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "CEz0fIURB-", "reviewer": "Reviewer_oAui", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work mainly focuses on the generalization comparison between Adam and GD. Specifically,   for a well-designed two-layer network,  Adam often converges to a bad minimum while GD can converge to a good solution in terms of generalization error, even with weight decay regularization. For convex problem, with the weight decay, Adam and GD converge to the same solution. ", "summary_of_the_review": "Overall, this work provides some good initial results. But its assumptions may not be consistent with the real settings.", "strength_and_weaknesses": "Strengths: \nThere are main three contributions in this work. \n1) On a well-designed two-layer network,  Adam often converges to a bad minimum while GD can converge to a good solution in terms of generalization error.   \n\n\n2) For convex problem, with the weight decay, Adam and GD converge to the same solution.  \n\n3)The above two results show the different behaviors of Adam and GD and may provide some insights for their deep learning application. \n\nWeaknesses:\n1)The analysis in this work heavily depends on the well-designed network and data assumptions. For network, this work designs a special two-layered network which is shallow and also uses a rarely used  polynomial RELU. This network has a big gap with the real network. For data, it assumes the feature is 1 sparse while the noise is also s-sparse. For feature, it is too sparse to hold in practice. Though the authors claim their results still hold when feature is k-sparse, as long as the sparsity gap between feature and noises exit. However, they did not explain more how sparsity gap is required and whether it is reasonable in practice setting.  So it is a big concern whether this analysis with these well-designed network and data can really reflect the performance for network training in practical applications.    \n\n\n2)In deep learning, one often SGD and stochastic Adam to train network. But this work analyzes GD and  deterministic Adam which yields a gap. Though the authors also say their results can be extended to stochastic setting but did not give many details. It is better to directly put the results under the stochastic setting in the paper, if stochastic setting can be well handled, since one often cares more about the real settings. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3oo_UsOzzK", "reviewer": "Reviewer_1Tav", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the generalisation properties of Adam compared to gradient descent (GD). In order to simplify the analysis and isolate some distinguishing hypothesised differences between the two optimisation variants, the authors make use of a data model consisting of simple, sparse feature vectors and noise. The paper then provides a theoretical analysis in terms of convergence, generalisation, feature learning and noise memorisation, using a 2-layer neural network trained with weight decay regularisation. The main results are that both Adam and SGD behave similar in the convex case, but in the convex case they converge to different global solutions, Adam tending to fit the noise in the data.", "summary_of_the_review": "My overall impression of this paper is positive, since it is a clearly written, solid analysis of relevant optimisation algorithms such as gradient descent and Adam. While I think the set of assumptions to carry out the analysis (shallow network and simple data model) limit the relevance of the results, I still think that the contributions in this paper advance our understanding of Adam and gradient descent.", "strength_and_weaknesses": "### Strengths\n\nAs main strength, I highlight that this paper is technically solid, mostly well-written and easy to follow. I value that introduction is clearly written, it includes motivations for the work and a concise review of the relevant related work. Further, both the methodology and the main results are discussed with sufficient detail, and I value that the significance of the mathematical results are unpacked nicely in the text, for example in Section 5. Finally, I also appreciate that, in the last section of the paper, some of the limitations of the paper are transparently mentioned as future work, for instance the limitation of the analysis to a two-layer neural network.\n\n### Weaknesses\n\nIn my opinion, the main weakness of this paper is that the limitations adopted to enable a solid and conclusive analysis may impact the relevance and generalisation of the conclusions for practical applications. In particular, I identify two main limitations: \n\n* The use of a highly synthetic data model (almost as simple as it can get), crafted to highlight the disadvantages of Adam with respect to gradient descent.\n* The dependence of the results on a two-layer neural network.\n\nThe first limitation is directly connected with the data. Data sets used in practical applications vastly differ from the data used to derive the theoretical results in this paper, which consists of the combination of a 1-sparse feature vector and noise vector. This choice highly simplifies the analysis, but it leaves doubts to whether the conclusions can be extended to more complex data. As a matter of fact, the authors point out that their \"theoretical analysis can lead to an opposite conclusion on the generalization comparison between Adam and GD if the noise is sparse and feature is denser\"\n\nThe second limitation is also important, as it restricts the relevance of the conclusions to shallow models (2-layer neural networks), whereas Adam was born as an optimisation algorithm for deep neural networks and it is in this realm where it has been found to be superior--- according to certain aspects such as convergence rate, lower dependence on hyper-parameter optimisation, as well as generalisation in some cases---to gradient descent.\n\nWhile I value the contribution of an analysis with a set of limitations to enable detailed conclusions, I think it is fair to highlight that the significance of the results are limited by such assumptions.\n\nSome more minor weaknesses I would like to point out are the following:\n\n* One of the main conclusions highlighted by the authors in the paper is that Adam and SGD converge to the same solution in the convex case, but differ in the non-convex case. That both algorithms are able to find the global optimum in convex optimisation (with weight decay) is intuitive (the differences lie in the convergence rate), and that the two algorithms are most different with a non-convex objective is also expected. Therefore, it is unclear how much technical novelty this conclusion is able to offer.\n* Figure 1 shows visualisations of the features learnt by the first layer of AlexNet trained with Adam and SGD on CIFAR-10. However, no further details are provided beyond those mentioned in the capture of the figure. This visualisation, in and on itself without further details and analysis, seems to me like a rather superficial result, disconnected from the rest of the paper which is focused on theoretical results.\n* In Section 3, the authors claim that the data model used in the paper is \"more practical\" than that in Wilson et al. (2017) and Reddi  et al. (2018). In this regard, I missed a more detailed discussion of the justification of the data model used in this paper, as well as the differences with models used in previous works analysing Adam. Furthermore, it is unclear to me how the data model in this paper can be \"more practical\", while it is an extreme case where the feature feature is a 1-sparse vector.\n* The last item in the review of related work (Feature learning by neural networks) lacks diversity, since only works from the same group of authors is cited. As a matter of fact, citations of Allen-Zhu & Li are over-represented in the list of references.\n* While the paper is well-written in general, some aspects could be improved in my opinion:\n    * Some sentences contain long list of citations, which makes them unreadable. For instance, see the paragraph \"Optimization and generalization guarantees in deep learning\"\n    * Typos or writing issues:\n        * \"In this section, we discuss the works that are mostly related to our paper\": the word \"mostly\" sounds strange to me in this sentence.\n        * The word \"data\" is used several times in singular, where I would instead use \"data point\". For instance: Definition 3.1, after Equation 3.2, second paragraph in Section 5.1\n        * patter -> pattern (in last paragraph (Proof outline) of the introduction of Section 5)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster