Understanding the Generalization of Adam in Learning Neural Networks with Proper Regularization
Authors: Difan Zou
Yuan Cao
Yuanzhi Li
Quanquan Gu
Venue: ICLR 2022 Submitted (2022)

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 is fundamentally tied to the nonconvex landscape of deep learning optimization, which cannot be covered by the recent neural tangent kernel (NTK) based analysis. 

Under review as a conference paper at ICLR 2022
UNDERSTANDING THE GENERALIZATION OF ADAM IN LEARNING NEURAL NETWORKS WITH PROPER REGU- LARIZATION
Anonymous authors Paper under double-blind review
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 different solution with a worse test error compared to (stochastic) gradient descent, even with a fine-tuned regularization. In this paper, we provide theoretical explanation for this phenomenon: we show that in 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 explicit regularization is fundamentally 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 al., 2011; Hinton et al., 2012; Kingma & Ba, 2015; Reddi et al., 2018) such as Adam are very popular optimizers for training deep neural networks. By adjusting the learning ratethis coordinate-wisel based on historical gradient information, they are known to be able to automatically choose appropriate learning rates to achieve fast convergence in training. Because of this advantage, Adam and its variants are widely used in deep learn- ing. Despite their fast convergence, adaptive gradient methods have been observed to achieve worse generalization performance compared with gradient descent and stochastic gradient de- scent (SGD) (Wilson et al., 2017; Luo et al., 2019; Chen et al., 2020; Zhou et al., 2020) in many deep learning tasks such image classification (we have done some simple deep learn- ing experiments to justify this, the results are reported in Table 1). Even with proper regu- larization, achieving good test eπrоr with adaptive gradient methods seems to be challenging.
Several recent works provided theoretical expla- nations of this generalization gap between Adam and GD. 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 specif- ically designed data, while SGD can learn the linear model to achieve zero test error. This ex-
Models
AlexNet
VGG-16
ResNet-18
SGD Adam
75.22
93.25
94.62
73.08
92.19
92.93
Test accuracy (%) comparison between Adam and SGD on the CIFAR-I0 dataset. ample in linear regression offers valuable insights into the difference between SGD and Adam. However, it is under a convex optimization setting, and as we will show in this paper (Theorem 4.2), the performance difference between Adam and GD can be easily avoided by adding an arbitrarily small regularization term, because the regularized training loss function is strongly convex and all algorithms will converge the same unique global optimum. For this reason, argue that the example in the convex setting cannot capture the fundamental differences between GD and Adam. More recently, Zhou al. (2020) studied the expected escaping time of Adam and SGD from local
Table 1:

---

Under review as a conference paper at ICLR 2022
basin, and utilized this 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 there generalization gap between Adam and gradient descent in learning neural networks, even with proper regularization?
Specifically, we study Adam and GD for training neural networks with weight decay regularization on an image-like data model, and demonstrate the difference between Adam and GD from a feature learning perspective. We consider model where the data generated as combination of feature and noise patches under certain sparsity conditions, and analyze the convergence and generalization of Adam and GD for training two-layer convolutional neural network (CNN). The contributions of this paper are summarized as follows.
We establish global convergence guarantees for Adam and GD with proper weight decay regular- ization. We show that, starting at the same random initialization, Adam and GD can both train a two-layer convolutional neural network to achieve zero training error after polynomially many iterations, despite the nonconvex optimization landscape.
We further show that GD and Adam in fact converge to different global solutions with different generalization performance: GD can achieve nearly zero test error, while the generalization performance of the model found by Adam is no better than a random guess. In particular, we show that the reason for this gap is due to the different training behaviors of Adam and GD: Adam is more likely to fit noises in the data and output a model that is largely contributed by the noise patches of the training data; GD prefers to fit training data based their feature patch and finds a solution that is mainly composed by the true features. We also illustrate such different training processes in Figure 1. where it can be seen that the model trained by Adam is clearly more "noisy" than that trained by SGD.
We also show that for convex settings with weight decay regularization, both Adam and gradient descent converge to the exact same solution and therefore have no test error difference. This suggests that the difference between Adam and GD cannot be fully explained by linear models or neural networks trained in the "almost convex" neural tangent kernel (NTK) regime Jacot et al. (2018); 8); Allen-Zhre et al. (2019b); Du et al. (2019a): Zou et al. (2019); Allen-Zhu et al. (2019a); Arora et al. (2019 9a;b); Cao & Gu (2019); Ji & Telgarsky (2020); Chen et al. (2021). It also demonstrates that the inferior generalization performance of Adam is fundamentally tied to the nonconvex landscape of deep learning optimization, and cannot be solved by adding regularization.
Notation. For a scalar v, we use [x]+ to denote max 0}. For a vector v = (vi...., va)T we
denote by ||v[|2 == ( uz) 1/2 its f2-norm, and use supp(v) = (i vj # 0} to denote its support.
RELATED WORK
In this section, we discuss the works that are mostly related to our paper.
Generalization gap between Adam and (stochastic) gradient descent. The worse generalization of Adam compared with SGD has also been observed by some recent works and has motivated new variants of neural network training algorithms. Keskar & Socher (2017) proposed to switch between Adam and SGD to achieve better generalization. Merity al. (2018) proposed 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 smooth transition from adaptive methods to SGD to improve generalization. Our theoretical results for GD and Adam can also provide theoretical insights into the effectiveness of these empirical studies.
Optimization and generalization guarantees in deep learning. Our work is also closely related to the recent line of work studying the optimization and generalization guarantees of neural networks. A series of results have shown the convergence (Jacot et al., 2018; Li & Liang, 2018; Du et al., 2019b; Allen-Zhr et al., 2019b; Du et al., 201 9a; Zou et al., 2019) and generalization (Allen-Zhh) al., 2019c;a; Arora et al., 201 99;;; Cao & Gu, 2019; Ji & Telgarsky. 2020; Chen al., 2021) guarantees in the so-called "neural tangent kernel" (NTK) regime, where the neural network function approximately linear in its parameters. Allen-Zhu & Li (2019); Bai & Lee (2019); Allen-Zhu & Li (2020a); Li et al. (2020) studied the learning of neural networks beyond the NTK regime. Our analysis in this paper is also beyond NTK, and gives detailed comparison between GD and Adam.
2

---

Under review as a conference paper at ICLR 2022
(a) Adam
(b) SGD
Figure 1: Visualization of the first layer of AlexNet trained by Adam and SGD on the CIFAR-10 10 dataset. Both algorithms are run for 100 epochs with weight decay regularization and standard data augmentations, but without batch normalization. Clearly, the model learned by Adam is more "noisy" than that learned by SGD, implying that Adam is more likely overfit the noise in the training data.
Feature learning by neural networks. This paper also closely related to several recent works that studied how neural networks can learn features. Allen-Zhu & L.i (2020b) showed that adversarial training purifies the learned features by removing certain "dense mixtures" in the hidden layer weights of the network. Allen-Zhu & Li (2020c) studied how ensemble and knowledge distillation work in deep learning when the data have "multi-view" features. This paper studies 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.
PROBLEM SETUP AND PRELIMINARIES
We consider learning a CNN with Adam and GD based on n independent training examples {(x,, generated from data model D. In the following. we first introduce data model DD and then explain our neural network model and the details of the training algorithms.
Data model. We consider data model where the data inputs consist of feature and noise patches. Such a data model motivated by image classification problems where the label of an image usually only depends part of an image, and the other parts of the image showing random objects, or features that belong to other classes, can considered as noises. When using CNN to fit the data, the convolution operation is applied to each patch of the data input separately. We claim that our data model more practical than those considered in Wilson et al. (2017); Reddi et al. (2018), which are handcrafted for showing the failure of Adam in term of either convergence or generalization. For simplicity, we only consider the case where the data consists one feature patch and one noise patch. However, our result can be easily extended to cover the setting where there are multiple feature/noise patches. The detailed definition of our data model given in Definition 3. follows.
Definition 3.1. Let = M(x²), each data (x, y) with x E 2d and y E {-1,1} is generated as follows,
x = [x] x
where one of x and X2 denotes the feature patch that consists of feature vector y v, which is assumed to be 1-sparse, and the other one denotes the noise patch and consists of a noise vector &. Without loss of generality, we assume [1, O] The noise vector & is generated according to the following process:
1. Randomly select s coordinates from [d}{[} with equal probabilities, which is denoted as a vector 8 E {0, 1d.
2. Generate from distribution (V00, I),, and then mask off the first coordinate and other - coordinates, i.e., g = E (
3. Add feature noise &, i.e., & ooy, where 0 < a 1 is the strength of the feature noise.
3

---

Under review as a conference paper at ICLR 2022
In particular, throughout this puper we set 0 ((5 1//22) polylog(n)).
oz +0 (ppollog(()) and a O(0p
12=
The most natural way think of our data model is treat as the output of some intermediate layer of a CNN. In literature, Papyan et al. (2017) pointed out that the outputs of an intermediate layer of a CNN are usually sparse. Yang (2019) also discussed the setting where the hidden nodes in such an intermediate layer are sampled independently. This motivates us to study sparse features and entry-wisell independent noises in model. In this paper, we focus the case where the feature vector v is 1-sparse and the noise vector is s-sparse for simplicity. However, these sparsity assumptions can be generalized to the settings where the feature and the noises are denser.
Moreover, we would like to clarify that the data distribution considered in our paper 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 multiple feature vectors and multiple patches, each data can either contain a single feature or multiple features, together with pure random noise feature noise.
Two-layer CNN model. We consider two-layer CNN model F using truncated polynomial activation function (=)) = (max (0, =)) and fixed the weights second layer to be I's, where q 2 3. Mathematically, given the data (x, y), the j-th output of the CNN can be formulated as
m
m
(WW = Σ (o((w)).:x)) + a x2))] = Σ [o((w y v)) + o(Ws,r {))], (3.1)
where m is the width of the network, Wjr E Rd denotes the weight at the r-th neuron, and W the collection of model weights. Here the use of the polynomial ReLU activation function is for simplifying our analysis. It can be replaced by smoothed ReLU activation function (c.g., the activation function used in Allen-Zhu & Li (2020c)). If assume the input data distribution is Gaussian, we can also deal with ReLU activation function (Li et al., 2020). Moreover, we would like to emphasize that X₃ and X2 denote two data patches, which randomly assigned with feature vector or noise vector independently for each data point. The leaner has no knowledge about which one is the feature patch (or noise patch).
In this paper we assume the width of the network is polylogarithmic in the training sample size, i.e., vn po)ylog(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 N (0, of) with oo O(d-1/4)
Training objective. Given the training data {(x:, pim.............. we consider to learn the model parameter W by optimizing the empirical loss function with weight decay regularization
L(W)=
(W)
(3.2)
n
where (W) = log P. (W.,)
1, W,, denotes the individual loss for the data (xi, ys) and
is the regularization parameter. In particular, the regularization parameter can be arbitrary as long as it satisfies À (0, >o) with Ao O ① ((o=-////oooyyo ). We claim that the Ao is the largest feasible regularization parameter that the training process will not stuck at the origin point (recall that L(W)
admits zero gradient at W = 0.)
Training algorithms. In this paper, we consider gradient descent and Adam with full gradient. In particular, starting from initialization Wi() = (0 = {+1). [ml). , the gradient descent update rule is
((+1)
w j.r V w,, L(W()).
W j.r
where 7y is the learning rate. Meanwhile, Adam store historical gradient information in the momentum m(!) and a vector v1) as follows
((+1)
91 m j.r (1 (1) Vw I(W()))
(1)
(3.3)
jrr
4

---

Moreover, the Under review as a conference paper at ICLR 2022
(++1)
B2v (1 32) DV.,, L(W(0)))
(3.4)
and entry-wisely adjusts the learning rate:
(t) jrr
(2) 3,m
(3.5)
Ar
where B1, B2 are the hyperparameters of Adam (a popular choice in practice is 31 0.9, and 32 = 0.99), and in (3.4) and (3.5), the square (.)². square root v. and division //- all denote entry- wise calculations. We do not consider the initialization bias correction in the original Adam paper (Kingma & Ba, 2015) for the ease of analysis.
4 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 a two-layer CNN model in the following theorem. Recall that in this setting the training objective is nonconvex.
Theorem 4.1 (Nonconvex setting). Consider a two-layer CNN defined in (3.1) with 22n 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 on 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 17 = poly(n) then with probability least "
1 - 0 we can find NN model W W model W Adam also satisfies:
Adam such that || (WL(WW... Adam) I1 S TTy
Training error is zero: E [F (W Adam: x;) S F F-u. (W Adam xx)] = O.
Test error is high: P(x.))~D [F Adams x) S F-y (W* Adam x)] 2 2.
2. Suppose we run gradient descent for T poly(n) iterations with learning rate 1j poly(n) then with probability least 1 1n "), we can find a NN model WeD such that [V
<
Try Moreover, the model Wön also satisfies:
(W x;) S F-9. (W x:)] = 0. Training error is zero: #/2 E [P
Test error is nearly zero: IP ((xx)~DD [F, GWG x) S F-y GD x)] po)y(z)
From the optimization perspective, Theorem 4.1 shows that both Adam and GD can be guaranteed to find a point with very small gradient, which can also achieve zero classification error the training data. Moreover, it can be seen that given the same iteration number T and learning rate 1), Adam can be guaranteed to find point with up 1/(Ti)) gradient norm in & metric, while gradient descent can only be guaranteed to find a point with up to 1/VI'n gradient norm in E2 metric. 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 to point out that there is no contradiction between our result and the recent work (Reddi et al., 2019) showing that Adam can fail to converge, as the counterexample in Reddi et al. (2019) is for the online version of Adam, while we study the full batch Adam.
In terms of the test performance, their generalization abilitics are largely different, even with weight decay regularization. In particular, the output of gradient descent can generalize well and achieve nearly zero test error, while the output of Adam gives nearly 1/2 test error. 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 data patch of lower sparsity and thus tends to fit the noise vectors &. gradient descent prefers the data patch of larger l 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 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 the true feature vector and therefore leads to a test error no better than a random guess. More details about the training behaviors of Adam and gradient descent are given in Section 5.
to converge
5

---

Under review as a conference paper at ICLR 2022
Theorem 4.1 shows that when optimizing nonconvex training objective, Adam and gradient descent will converge to 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 reg ularization parameter A, suppose run Adam and gradient descent for T poly(n) iterations, then with probability at least - n the obtained parameters Adam and WeD satisfy that || (WL Adam)ll S o and [[V Adam)lé S Ty respectively, Moreover, let F(W, x) E R be the output of the convex model with parameter W and input x, holds that:
"
. Training errors are both zero:
1 Σ 1 sggu (F(WGD: # mi] = 0.
1 sgn (F(WA Adam x:)) + vi]
Test errors are nearly the same:
IP ((xx))=D( [sgn (F(WAFaWa Adam: x:)) # v] = ((.a)) (sgg(s(W(aaa (F(W x)) + u] 1/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. Combining this result and Theorem 4.1, 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 abilities 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 {E:]==i...,, have nearly disjoint supports.
Lemma 5.1. Let (xx, (i)):=,...... be the training dataset generated by Definition 3.1. Moreover, recall that x; = [yvv _ (or x, le? "I), let B, supp(E.) (E)) [1} be the support of & except the first coordinate. Then with probability at least 1 - n - B, nB = 0 for all i *
This lemma implics that the optimization of each coordinate of the model parameter W, 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 pattern 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 the patterns learned by these two algorithms are different: Adam tends to fit the noise patch while gradient descent will mainly learn the feature patch. In the second stage, which we call regularization stage, the effect of regularization cannot be neglected, which will regularize the algorithm to converge at some local stationary points. However, due to the nonconvex landscape of the training 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 moving-averaged gradient, normalized by a moving average of the historical gradient squares. As pointed out Balles & Hennig (2018); Bernstein et al. (2018), Adam behaves similarly to sign gradient descent (signGD) when using sufficiently small step size or the moving average parameters 31, are nearly zero. This motivates us to study the optimization behavior of signGI and then extends it to Adam using their similarities. In particular, sign gradient descent updates the model parameter according the

---

(To) Under review as a conference paper at ICLR 2022
following rule:
(t-1) 1)
w 17 sgn V Vw,, (w).-(((o))
(1+1)
Recall that each data has two patches: feature patch and noise patch. By Lemma 5.1 and the data distribution (see Definition 3.1), we know that all noise vectors (Ei)..1....m
supported on disjoint coordinates, except the first one, For data point x,, let B, denote its support, except the first coordinate. In the subsequent analysis, we will always assume that those B, s are disjoint, i.e., B, n B3 = 0 if + j-
Next we will characterize two aspects of the training process: feature learning and noise memoriza- tion. Mathematically, we will focus on two quantities: (.. v) and (w E,). In particular, given the training data (x4; 34) with x; [u,v "T larger (w (r,,,r: V/s v) implies better feature learning and larger (w ye,ro Ei) represents better noise memorization.
Stage I: Learning the pattern. Mathematically, the first stage is defined the iterations that the neural network output is smaller than some constant. In this stage, all training data remains under- fitted and can provide large gradient for model training, and the effect of weight decay regularization can be ignored due to our choice of X. We will show that this stage the inner product (wilrr, Ex) grows much faster than j.rt jv) since feature learning only makes of the first coordinate of the gradient, while noise memorization could take advantage of all the coordinates in B, (note that [B,] >>>>> 1).
Lemma 5.2 (General results in Stage I). Suppose the training data is generated according Defi- nition 1, assume À = - 2pp/n) and v) = 1 1/poly(d), then for any S To with To 0 1 1jsoy
and any i E [n].
(w(tt") (++1)
1j vv S (w v) + O()),
(w (++1) E.) = (w/l), E.) + (11sop))
Since (w f Ex) enjoys much faster increasing rate than that of [w j v), after certain number of iterations, the learning of noise patch will dominate the learning of feature patch. Thus, the model will tend to fit the feature noise in the noise patch, leading to a flipped feature learning phenomenon.
(1)
Lemma 5.3 (Flipping the feature learning). Suppose the training data is generated according to Definition 3.1, a 2 (((0r)) 1-9 V og-1) and do < Ō((sa,))1), then for any E [Tr. To] with T, = ( 1 S To.
(w (1+1) .j v) = ((w (w j. v) - O())
Moreover, it holds that
-gn(() (j) -6(
h ==1,
(T6) J.r
[k]
sgu(E:(4]
52 s0,)) or + Ö(n), E B,, with yı = j.
+ō(n).
otherwise.
From Lemma 5.3 it can be observed that at the iteration To. the sign of the first coordinate of w jrr is different from that of the true feature, i.c., j v. This implies that at the end of the first training stage, the model is biased towards the noise patch to fit the feature noise.
Stage I: 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 a 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.
7

---

Under review as a conference paper at ICLR 2022
Lemma 5.4 (Maintain the pattern). If O and o(A), then let p"
arg mngr (w (2) E.), for any 2 To. i E [n], E [2] and E [m]. it holds that Ei) = @(),
(w 1
Σ
B [+]| 1& [+] @(1),
REB.
E E [m].
(w (0) j.m sgn(s) v) E [-
Lemma 5.4 shows that in the second stage, (w Whir: E.) will always be large while (w 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.
(9)
Lemma 5.5 (Convergence guarantee). If 1) =
then for any it holds that
L(w(++1) - (W))) - -nl VL( + G(n2d).
Lemma 5.5 shows that we can pick a sufficiently small and T = poly(n)/n to ensure that the algorithm can find a point with up (()(TTn) in norm. Then can show that given the results in Lemma 5.4, the formula of the algorithm output W can be precisely characterized, which we show that (w ;;.r: y/: v) < O. This implies that the output model will be biased fit the feature noise -ayv but not the true one v. Then when it comes to 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, also decompose the entire training process into two stages.
Stage I: Learning the pattern. this stage the gradient from training loss function is large and and the effeet of regularization can be ignored. Unlike Adam that is sensitive to the sparsity of the feature vector or noise vector, gradient descent is more focusing on the & norm of them, where the vector (which can be either feature vector or noise vector) with larger & norm is more likely to be discovered and learnt by GD. Note that the feature vector has larger l2 norm than the noise, we can show that, in the following lemma, gradient descent will learn the feature vector very quickly, while barely tend to memorize the noise.
Lemma 5.6. Let (t)
(++1) .j. v), ((t) = max,e[m] (w j.m Ei). and r max,e(m] (w j.r
(() ma j. Let T be the iteration number that A, reaches ((1m)) = ((1), then we have
T, = - for all j E 1-1,1].
Moreover, let To [T then for all t < To it holds that 15 (c.) for all jE (-1,
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 the first stage will remain unchanged.
= Lemma 5.7. If 1 S (gro), it holds that A = ((1) and r = Ō(ao) for all min, Tj.
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 17 = o(1), then for any t 2 0 it holds that
I(W(+1)) (w)))) < - (V
Lemma 5.8 shows that we can pick sufficiently small 17 and T poly(n)/n to ensure that gradient descent can find point with up to ((1 output model of GD can well learn the feature vector while memorizing nearly nothing from the noise vectors, which can therefore achieve nearly zero error.
in &a norm. By Lemma 5.7, it is clear that the
EXPERIMENTS
In this section we perform numerical experiments on the synthetic data generated according to Definition 3.1 verify our main results. In particular, we set the problem dimension d 1000, the
8

---

Under review as a conference paper at ICLR 2022
1.2
2.0
0.8
1.5
Value
Festure Learning: mix (W....M)
Feature Learning: mpx (WL. v)
Noise Memorization: min max (WL 52
1.0
Noise Meroriiatiio:: minn max (w....
0.4
0.5
-0.2
0.0
2000
6000 10000
8000 10000
Iterations
# Iterations
(a) Adam
(b) GD
Visualization of the feature learning (max, (Wi,r, v)) and noise memorization (min; max, (Wi,r, 6.) in the training process.
Figure 2:
training sample size n. 200 (100 positive examples and 100 negative examples), feature vector v v= [1, .,...... noise sparsity & = 0.1d = 100, standard deviation of noise op = 1// 0.1, feature noise strength
0.2, initialization scaling dg 0.01, regularization parameter À = 1 x 10-6, network width m = 20, activation function (2)) = max 2}3, total iteration number T = 1 10t, and the learning rate 7) 5 x 10 for Adam (default choices of 81 and B2 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 2, where the test error calculated on a test dataset of size 10 It 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.1).
Moreover, we also calculate the inner products: max, v) and min; max, (W1.r: representing feature learning and noise memorization respectively. to verify our key lemmas. Here only consider positive ex- amples as the results for negative examples are similar. The results are reported in Figure 2. For Adam, from Figure 2(a), it can be seen that the algorithm will perform feature learning in the first few iterations and then entirely forget the feature (but fit feature noise), i.e., the feature learning flipped, which verifies Lemma 5.3. the meanwhile, the noise memorization 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, (oxx,(W1,> (W,,,, E in the first stage min, max, (exx,(W1.) (W1.r. &s) increases linearly, and in the second stage its increasing speed gradually slows down and min; max, (W1.r. {.) will remain in a constant order. This verifies Lemma 5.2 and Lemma 5.4. For GD, from Figure 2(b), can be seen that the feature learning will dominate the noise memorization: feature learning will increases to a constant in the first stage and then remains in a constant order in the second stage; noise memorization will keep in low level which is nearly the same that the initialization. This verifies Lemmas 5.6 and 5.7.
Algorithm
Adam
SGD
Training error Test error
0.884 0
Table 2:
Training and test errors achieved by GD and Adam.
7 CONCLUSION 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 converge 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, and demonstrates that this difference is tied to the nonconvex optimization landscape of neural networks. Built up on the results in this paper, there are 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 is motivated by the image data, where Adam has been observed to perform worse than SGD in terms of generalization. Studying other types of data such as natural language data, where Adam is often observed to perform better than SGD, is another future work direction.

---

Under review as a conference paper at ICLR 2022
ETHICS STATEMENT
This paper mainly concerns the theoretical understanding of the behaviors of different optimization algorithms in deep learning. We don't see any potential ethical issues in our work.
REPRODUCIBILITY STATEMENT
Our theoretical results can be reproduced according to the assumptions and problem set up stated in Section and the detailed proof provided in the appendix. The experimental results shown in Table 1 and Figure are obtained via standard training (using SGD or Adam) on CIFAR- 10 dataset (Krizhevsky al., 2009). Our experimental results in Figure 2 and Table 2 can also be reproduced according to hyperparameters described in Section 6.
REFERENCES
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Revisiting the generaliza- tion of adaptive gradient methods. 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. What can ResNet learn efficiently, going beyond kemelss? In Advances in Neural Information Processing Systems, 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04411 2020a.
Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. arXiv preprint 2xiv:2005..01190 2020b.
Zeyuan Allen-Zhr and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning, arXiv preprint arXiv:2012.09816, 2020c.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in Neural Information Processing Systems, 2019a.
Zeyuan Ahue Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp. 242-252, 2019b.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In Advances in Neural Information Processing Systems, 2019c.
Sanjeev Arora, Simon Du Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization 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 exact 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 higher-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 Arizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569. PMLR, 2018
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, 2019.
Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang. Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. In International Joint Conferences Artificial Intelligence, 2020.
10

---

Under review as a conference paper at ICLR 2022
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much over- preeraaaamttiiiiiionnn is sufficient to learn deep relu networks? In International Conference on Learning Representations, 2021.
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 wver-parameteriied neural networks. In International Conference Learning Representations, 2019b.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 120ul):2121-2159, 201
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview mini-batch gradient descent, 2012.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances in neural information processing systems, pp. 8571-8580, 2018.
Ziwei and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily 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 sgd. arXiv preprint
2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Yuanzhi Li and Yingyu Liang. Learning overparameterizzed neural networks via stochastic gradient descent on structured data. In Advances Neural Information Processing Systems, pp. 8157-8166, 2018.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Learning over- wverpparmeetrieed 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 10Xiv:1902.09843: 2019.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing Istm language models. In International Conference on Learning Representations, 2018.
Vardan Papyan, Yaniv Romano, and Michael Flad. Convolutional neural networks analyzed via convolutional sparse coding. The Journal of Machine Learning Research, 18(1):2887-2938, 2017.
Sashank J Reddi, Satyen Kale, and Sanjiy 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: 1et.::9999.900777...00 2019.
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.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv: arXiv:1902.000 1902.04760, 2019.
11

---

Under review as a conference paper at ICLR 2022
Pan Zhou, Jiashi Feng, Chao Ma. Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why sgd generalizes better than adam 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

---

Under review as a conference paper at ICLR 2022
A DISCUSSIONS
In this section, we would like to reflect some possible extensions of our work to more general settings.
Mini-batch stochastic gradients. One natural extension of our paper is proving the separation between mini-batch SGD and mini-batch Adam, which we believe is not difficult. In particular, can be shown that the learning behavior of SGD will be largely similar to that of GD since using mini-batch stochastic gradients does not change the rates of feature learning and noise memorization. In contrast, compared to full-batch Adam, the rate of mini-batch Adam in memorizing noise will be scaled down by most n. times (due the preconditioning in Adam), while the of feature learning remains the same. However, the difference between SGD and mini-batch Adam still exists (though the difference might be smaller). This is because, in paper, the separation between Adam and GD is characterized by a 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 (n) (in each step), while the rate of noise memorization in Adam is proportional to the number of nonzero entries, which is in the order of 7 poly(d). In contrast, using mini-batch stochastic gradients can only reduce the difference by at most an ((n) factor, which is still dominated by the poly(d) separation as long as the dimension sufficiently large (overparameterizee)
Learnable second layer. Our results and analyses can also be extended the case where the second layer is trained. The reason behind this is that the difference between Adam and GD comes from the learning of the hidden layer, while the second layer weights will not affect which patch and neuron is more likely to be learned by different optimization algorithms. Consequently, we can still get the result that GD tends learn the feature patch while Adam tends learn the noise patch.
B PROOF OF THEOREM 4. 1: NONCONVEX CASE
In the beginning of the proof we first present the following useful lemma.
B. PRELIMINARIES
We first recall the magnitude of all parameters:
d = poly(n), poly(n)' 8 = @ d1/2
of =0
.ollyloog()) (pooolon))) % =
1
0
n²
(d1/2
= polyloggnn) a= O(a, - polylog(n)), A o
ddg-11/44n polylog(n)
Here poly(n) denotes a polynomial function of n with degree of a sufficiently large constant, poly(n) denotes a polynomial function of log(n) with degree of sufficiently large constant. Based on the parameter configuration, we claim that the following equations hold, which will be frequently used in the subsequent proof.
A
0
0
Op
= =w((0)))
.a=0 (1)))=00=00
00 80
So?
17 (2o800)))))
12
Lemma B.1 (Non-overlap support). Let {(x;, y/)) )im1..... I..... be the training dataset sampled according to Definition 3.1. Moreover, let B; supp ) ) {1} be the support of x; except the first coordinate Then with probability least 1 -22 B, B, = @ for all i,
j E [n].
Proof of Lemma B.I. For any fixed le [n] and j E (upp(ex)) by the model assumption have
(((())) * 0} s/(d 1),
for all i e [m]][ Therefore by the fact that the data samples are independent, we have [1 ss/(d= - s/(d-
P(3i E [n] (+} (Ca), * 0)
Applying a union bound over all &s e [n] and E supp(E (E1) we obtain
P(ak e [n].j e supp(Ex) ((1),, e [x] (E.), * 0) n {1 s/(sssddd-- }) (B.1)
Recall that all data inputs have nonzero first coordinate by Definition 3.1
13

---

Under review as a conference paper at ICLR 2022
By the data distribution assumption we have I vd/(2n²), which clearly implies s/ s/(d-1) - 1/2. Therefore we have
= n expln log s/d-ss/(d]]] n. [1 (1 s/d)"]
I nt 8- exp(n (s/(d-1)
$ n exp(n 4s/d)
S nt 8- (Ans/d)
4n 8² /d
S n -
where the first inequality follows by the inequalities log(1 - 2) 2 -2z for z E [0,1/2]. the second inequality follows by s/(d 1) 2 28/d, the third inequality follows by the inequality 1- exp(-2) S for R. and the last inequality follows by the assumption that S (2n2)n Plugging the bound above into (B.1) finishes the proof.
B.2 PROOF FOR ADAM
Before moving the detailed proof, we first state the update rules of feature learning and noise memorization when the sign gradient is applied.
(w jrr .jv) (w jv) -1 17- {sgg (V w.. L(W(())).jv
(w(+t) (++1)
(ww() (1)
(0) yel [o'([oo'' B(v)) - ( E.) n.w (1)
jv) + jn- sgn
(B.2)
From (B.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 in 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 of the noise vector &, can be described by the following update rule,
where e/0 = 1 y -lggit (F, x;) and (fg Exe(-1.1)
(w((+)) Bear E (w(t) Birrt Ei) 11 (sgu( V Wwy..r
L(W(()))
(w p Ei) 1 Σ
( 1i oll.o'((""
((w
yo.r [.)) (1n1)1111))) Ei(A]
sgn
kEB;
aty: sgn
Σ
yr
(() we.d
(o'((w(())
ao'((w)) (2 E,1) ndw (2) Dirr y(v)) ao'
Bir [1]
(B.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 B.2 (Closeness to SignGD). Recall the update rule Adam, let W(() be the f-th iterate of the Adam algorithm. Suppose that (w))) v), (w. Ei) = @(1) for all j E {+1} and [m]. Then if B2 2 88, we have
. For all & E [d].
(+) m j.r
[+]
< O(1..
v [b]
(4)
14

---

Under review as a conference paper at ICLR 2022
For every l & L (including Is 1) we have either IW,,.. ((W @())
m (t) [k]
sgn V Ww,,. ((W(())]k) (1))
[k]
For every & E B, we have |V Vw,,,- ((W())) < 6mm s0 (0 - ((nso,) or
m j.r [k]
sgn (ww,, L ((W())) @(1).
Proof. First recall that the gradient V w... I(W()) can be calculated as
Vwg.. I(((1)) Σ nelff o' ((w
J ((
(8)
B(v))
o d' (.. s. + AW
More specifically, for the first coordinate of V Vw... L (W(1), we have
V Wyy (W
Σ S ((w j.r v(V)) Σ me o' ((w <i)) + AWW
(B.4)
For any h E B; by Lemma B.1 I we know that the gradient over this coordinate only depends on the training data E., therefore, we have
V Vw Vws.. L(W()) ------- n
1 1 ((w (M) + (2)).
(B.5)
For the remaining coordinates, we have
L(W(0) = Aw [M
(B.6)
Now let us focus on the moving averaged gradient m f and squared gradient v (0) We first show that for all & E [d], it holds that
m (+) [e]l
< O(1).
(B.7)
[c]
By the update rule of m
j.m we have
(t)
(c-1) (--1(1)) [A]] + (1-B:)- Vw.. Vw...(WW ((W(())(1))
= 3
m
j.r
8[ (1)))))) V Vwy.
((W
7=0
Similarly, we also have
11 [
Σ 35 B2) Vw,, L(W((=>))1))
Then by Schwartz inequality we have
[BT - 31)12 Vw,,, (W (Ia+)
(m (4) (1))
2
az
Let a3 [8 (1
(1-A1_a) which forms an exponentially decaying sequence if 3g 2 18. Therefore, we
have Σ o = (1) and the above inequality implies that
(m
(kk)2
(0 [k] e(1).
15

---

Under review as a conference paper at ICLR 2022
which proves (B.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 threshold polylog(n such that Σr= BT B1) have
(t)
polyly 15. Then for each [d] we
[. = Σ
81 -11) Vw,,. L(W( +
Σ 7=r
31 - B1) Vws.. L(W(C-))
m
T=0
= Σ - Vw,,,
+ poly(1-1)*
7=0
where in the last inequality we use the fact that [V Similarly, we can also have the following on (1) Virr
[(W((-)) Õ(1) for all & [d].
O [K] = Σ 85 32) V Wy L(W
[4] poly(n-1)
T=U
Here we slightly abuse the notation by using the same F. Then we have
(6) [A]
Ei=o 37 B1) V Ww,, L(W [[ = poly(a
V (1- - 32) Vw L(W( p)ly(n-
In order to prove the main argument of this lemma, the key is to show that within F iterations, the gradient V w.,. [(W())[[] barely changes. In particular, by (B.7), we have the update of each coordinate in one step is at most ((n). This implies that
(w Jorn v) (w v) ((17),
(w Ex) - -(w <:) S @(1Tsop));
|w [[]]]]]] w/ [+] - ((77).
Then applying the fact that ((w 1.r v)1 S ((1) and [(w (T) &i>ع < @(1), we further have
(W(()))
(r) x:) F3 (t)) x:) < ((mifssa)) = @17T8ap))
where we use the fact that m (11) and sap = (1). Then it holds that
(F,( x.)
5 ExE(-LAI) F Fl((())))))
())(()))))+(((Fo)) + (W =gn((())) e((2))))
where the fact that @(nTso,,) o(1). Similarly, also show that e/3 2 sgu(e 0(1())) which further implies
+0 sgn
for all 7 E [t - 7,t). Note that (0 < 1, then it holds that
E72 00 v\) = sgn((()) e(e) o' ((w (ww...v) v))
S sggn(( ((1e))) o ((w(t) j.r v\) + +(ee(2) -((F))
(t)
16

---

Under review as a conference paper at ICLR 2022
We can also similarly derive the following
( (( (T) v\) 2 sgn((())) 01 ((
for v)) - +(()))) ((7F). (0)
9
j.r (T)
(W jr: Ei) S sgn 0
o'( (0 6i) + + @(TTsop).
(( (i) 2 sgn
E.) - +( (17
Combining the above results, applying (B.4), (B.5), and (B.6), we can show that for the first coordi- nate, we have
V W. (1 -V.... L(w()))) + 0
Σ
( + e(xff);
for any le E Bj. we have
Vw wwrr I(W())) = @V I( ((Wll)]k)) + 2 [e(0]
(1) = +(uuf);;
and for remaining coordinates, we have
[( = (( ((177))
[(W
Now we can plug the above results into the formula of m
(1) j.r and y.r Using the fact that i = 6(1)
A o(1). and le < 1, we have for all k = 1 or k & B, for any i,
[K]
V w.... L L( [K] = O())
L L( (1) + (())
For B3 we have
V W..- L ((W())( = 6 n
= =(An)
[k]
θ VW,., T. [.((0)) θ nag, [0(9
+ @(Ar)
Then, we can conclude that for all k = or k B, for any we have either /w,,, L [(W()))
θ (n) or
m (t) [k]
sgn (Vw,.. I (((((())))) ((1).
[K]
For any h E B,, we have either /V,, (W(+))(k] < (mmmmsaplee + Xn) or
(4) [k]
sgn L(W(0)][))) (1().
j.r [
This completes the proof.
Lemma B.3 (Lemma 5.2, restated). Suppose the training data is generated according to Definition
3.1, assume À o 2-2 /n) and 17 1/poly(d), then for any t < To with To Ö 1787, any E [n]]
and
(w (1+1) .j . v) S (w ( v) + O(n). Jrr
(w(e+1) _Ei (w(!) pr Ei) ((sap))
17

---

Under review as a conference paper at ICLR 2022
Proof. At the initialization, we have
(0) (w j.r v) ((ro),
[( (0) Gi)l = 0(s)/2 0p00 + +a) (5s1/22 ,pou).
w opo)))))ou)
Ō(oo).
which also imply that 5 (1). Besides, note that e//3 j=y. - logit, (F10, (F,0),,);,
xz), we have
sgm (u sgn(j):
where recall that j {-1,1). Therefore, given that À olog- ) o(1), S Ö(1) and assume p(t) j. = (1) (which will be verified later),
sgm Σ
yi (1) 5 ((w (1) j.r v.v) a Σ vel 1 j.r (i)) - ndw n2w....1)
sgn [ @ j. @(an( ( +
sgn(j).
Since v is 1-sparse, then by Lemma B.2, the following inequality naturally holds,
(t) (w (++1) J v) I (w). j.m j v) 7 m j.r f.r: .v S (w. v) + (1))
Additionally, in terms of the memorization noise, we first consider the iterate the initialization. By the condition that = o(1/d) = o(1/(s0,)) and note that for sufficiently large fraction of
E B; (e.g., 0.99), we have [A] 2 (a,) 2 1mm so, (((0)) and thus
sgn V Wxy. L ((W(0))[4)) = sgn
(t)) v1 o ((w (1)
Ve.ro [ - nAw (0 [4]
sgm
θ ((d) 00000))) op sgn (g:(E((k:)) + 100000100000 = sgn(( (k)(
(B.8)
Therefore, by Lemma B.2 we have the following according to (B.3),
(w(()
Besrs {i) (w Wor Ex) 77 m
(t) f/r
Mesrs &
N (w (0)
barr Ei) + ((n)
O(nso)) - ((na)
REB
= (w E.) + ((nsop))
where in the first inequality the term (1780)) represents the coordinates that [c]]]]] < ((pp) (so that we cannot use the sign information of V (w)))(((())) but directly bound it by @(1)) and the last inequality due to the fact that [B.] 2 5 - and = o(1). For general 6, we will consider the
following induction hypothesis:
(w(1)) (.) (w(()
minn {.) + (()sap))
(B.9)
which has already been verified for t = 0. By Hypothesis (B.9), the following holds at time t,
Haaro (.) = (w(0) (0) :: + + (tnssa,) = (22 (20000 + tmsop))
In the meanwhile, we have the following upper bound for (t)(
j.r [k]l-
w j.r [+]| < [ (.
(
[+]] + nl sign( Vw,, L(W())) S |w j.r
[&]] t t 5 + tn))
(B.10)
Besides, it is also easy to verify that for any < To
ã (
6 (())_)) we have
(wl)) m. &i), (wih whr v) < (11mm) and thus 10 = O(1) Then similar to (B.8), we have
(1)
sgn (W [())))(())
= sgn ((0)
Vi.s /0 ((w(()
Wiiro E.))
nAw
(0)
[k]
18

---

Under review as a conference paper at ICLR 2022
-gn ( (s 0pou + tmsap)* )0-1)) op.sgn(E)(()) = o op (ou + tn))]
= (()))(
(B.11)
This further implies that
(wft+) .Et) 2 (w (t)
{.) O(1)
Σ
(sgn (V W., L( L(W()))))))) &.-)) 0(123007) (
BEB.
(w(f)) (t)
Nirr E,) + (nsop))
where the term 80)))
) is contributed by the gradient coordinates that smaller than p(sop))
This verifies Hypothesis (B.9) at time and thus completes the proof.
From Lemma B.3, note that sap (11,, then it can be seen that (w j.ri v) increases much faster (1)
than (w).r v). By looking at the update rule of (w W.rn J. v) (see (B.2)), it will keeps increasing (0) E.)). Since (.jr Ex) increases much faster than (w (w.. j.r.j v), can be anticipated after a certain number of iterations, (w v) will start to decrease. In the following lemma, we provide an upper bound on the iteration number such that this decreasing occurs.
only when, roughly speaking, o' d(ww J. v)) > ar' ((w
Lemma B.4 (Lemma B.4, restated). Suppose the training data is generated according to Definition 3.1, a 2 (((OO))2- V a)= and on < (((sop))) -1), then for any E [Tr, Tol with T, =
O 1180 (180p1//1-11) < To.
(1+1) .j. v) = (w -
Moreover, it holds that
-sgn(j) 52
== O2( x0p or Öm)) e B, with y; = j.
(T)) [k] =
sgn(E:(4])
Jur
+Õ(1),
otherwise.
Proof. Recall from Lemma B.3 that for any / S To have
(t+1) (w fir . v) < (w j.r j v) + (n)) S (w [W]r v) + O(m)),
(w(11) E.) = (w!f Ex) (1sap)) S (w (0)
2 E.) (bpssop).
Besides, by Lemma B.2 we also have |w f.r ||] S w w [k] + ((n)). Then it can be verified that for some Ö( (1 we have for all E [n]] and E To
ad (od((wwl giar (.)) 2 C. (d((( v)) Xr [0]]
for some constant C. This further implies that
sgn V (e,,,,LWW [(W())
= sgn Σ
Ne (1) 3, (
fim ((v))
Σ
Mi 32
((w (...))
j.m E.))
ndw now (0111)
1=1
1=1
= sgn - a
Σ
((w(0) E.)
sgn(s),
where we use the fact that sgn(y; sgn(f) for all (n). Then by Lemma B.2 and (B.2), have for all E To].
(++1) (w jr
.j.v] v) = (w j.n j. v) - -0()) ssn(() sgn (V ((W()) (w ( v) -(n).
19

---

Under review as a conference paper at ICLR 2022
Then at iteration To, for the first coordinate we have
(TO) [
w 5.r []] + sgn(s) (T,,)) - sgn(s) (asggg))(((()) O ((TT-((TT)))))))) T,)n) 2 sgn(j) σ
1
sspp)
For any & B, with ys = j. have either the coordinate will increase at rate of (1) fall into As a consequence we have either w) ",r [k] [-@ {()) or
(T))
(0) []] (g)( (Tinn) 2 (E])) ? To [|] w
p0p
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 Top- = w(oo). By Lemma B.2 is clear that
w j.r
(2b)
[k] [-)(n), (n))
Lemma B.5 (Lemma 5.4, restated). If a = 0 and 7 (A), then let r*
arg max (w (t) Ni Ei), for any To, i [n]. E [2] and [m]. it holds that
(w(t) (1) Wert ( 001);
(())) se. [B]] [E] 0(1)
kEB.
Vr E [m],
(w (t)
sgn(i) v) E 1-Ō
) 0(1~11))-
sap
Proof. The proof will be relying on the following three induction hypothesis:
(w/()) (1)
Ve.r. E.) 5(11,
(B.
w wir [k]] [A:] = 6(1),
(B.13)
Vr E [m], (w
f.r sgn(j) v) E - Ō nic
-
0(A-11)
(B.14)
so?
which we assume they hold for all < and E [m]. E [n]]. and j E [2]. It is clear that all hypothesis hold when f. To according Lemma B.4.
Verifying Hypothesis (B.. We first verify Hypothesis (B.12). Recall that the update rule for (w Bur: E() is given as follows,
(w(++), Birr (x)
(w(!) E.) - .7- (ml) / Vi Ei) Birri
(w(t) Whrr: E.) - -(()) (sgn (V Wyy L(W(")), - 20
(w())
Birri E.) (())
sgn
1 1 ((m.. [.))) (1) [0.11b]] E
kEB
n Σ
On) sgn
Σ
jort Biv)) ('
y Eunefto (000'(((((0 ji {i)) - ndw (111
0(n²s² 2
(B.15)
Note that for any a and we have sgn(a b) a 2 |al 2((b]. Then follows that
(1)
War [M)) n.w (0. E.
Σ
[&]]
2n w
vi.r [k]]
(t)
sgn
p())
(w (t)
EE)
kgB.
KEB.
@(sap) o
nd
p(t)
v.. op
20

---

Under review as a conference paper at ICLR 2022
where the last inequality follows from Hypothesis (B.12) and (B.13). Further recall that À 3/000"0p/10 sp/,,, plugging the above inequality to (B.15) gives
(w (++1) B:. .a) 2 (w s... E.) + ((nsa,) - - mnd
)
81)23000)
p(()
1000 8-2
2 (w]) Bcrt E.) + @(nsop) ((a)) -6
(B.16)
(t)
Then it is clear that (wher, Ei) will increase by ((ssap) if g(t)
yis is larger than some constant of We will first show that as soon as there is a iterate w satisfying (
order ( 2 O
08
sap
e(T) 0
for some
then must hold that dr2 will also be smaller than some constant ua.
in the order of 0 nd for all r' [T, (+, 1/. To prove this, we first note that if p(t)
2:3 reaches some
constant in the order of Ō nd so? ), we have for all r [m]] by (B.16)
( EE) 2 (w mo {,) (nsopp)
(w (wiri Ei) S (w (1)
(w-...r {.) + O(an),
(ww(tt)) (e+1)
vv)l S (w j.r v) + O(n).
(B.17)
Therefore, we have
e n (Wi(11) xx))
p(t+1) 21.5
EjE(-1113 (W(c)))))
1 exp |Em [a( (o(((w)))) N:.t
(++1) v)) + ol (w Eii) a(w ((+1)
-seins v)) - o((w ((+1)
3arr {t))]
1
exp
[Em (o((w porr v\) ((w (1) Ei)) o ((w wirrs v\) o a(ww -n
wi.r {]) (msoa?)
[o((w purty v)) 02 {\) a(ww -umm v)) o ((w(() -u +]]]] |E
exp
(())
where inequality follows from (B.17). Therefore, this implies that as long as M larger than some constant = Ō nd 20 then the adam algorithm will prevent it from further increasing. Besides, since
mnoo = o(1), then we must have (t++1) E [0.5e(() ya.i: 20(1 ] As consequence, we can deduce that p(t) cannot be larger than 2b, since otherwise there must exists iterate W with S such that 'y.3 should decreases (r) M., 2 b. < = 5 ((n)) for some S f, then we have
2e.s E [b. 2b) and p(r+1) 31.3 Therefore, we can claim that if p(r)
2 ((T)
"u..e which contradicts the fact that p(T)
S δ nd
(B.18)
9i. 80?
for all 7' [T, 1].
Then further note that
& (W(() xx.) xx.)
20(++1) vid
m
2 exp
(o((w(() p vvv)) ((w((() f {i)
2 exp ( (- (m max o ((w(() turr {,))))
(B.19)
re(m
21

---

Under review as a conference paper at ICLR 2022
where in the last inequality we use Hypothesis (B.14). Then by the fact that ((++1) < 0 n1
o(1) and m = (()) it is clear that exp O maxre(m) m) a a((w oer (()) o(1) so that mm (W ((+1)
(4+1)
E = (2).). This verifies Hypothesis (B.12).
Verifying Hypothesis (B.13). Now we will verify Hypothesis (B.13). First, note that we have already shown that (w yurt E() (1) so it holds that
(1+1)
Σ
[]]]
[E][ + a (1+2) Burr [T]] (w
Bor E.) (1)).
(1+1)
REB;
By Hypothesis (B.14), we have | 1.7 [4]] S yur [L]]+ + (L). Besides, since each coordinate
in &, is a Gaussian random variable, then maX
[[]]]]] (op)). This immediately implies that [E][ = 5(1).
Σ REB:
|w W [k]]
Then we will prove the upper bound of E ERB. any le B, such that V Vwu..- ((W(.))()))))))))) O(n mnso ( we have
(++11 [A]]] [K]]]]. Recall that by Lemma B.2, for Wait
((tH)) Mear [k]
[k] + (()) sgn p(0) Mir
o
Mer (.) - nAw( (1)
2e.r [k]]
(T)) [k] sgn((:((k)) 0 ( or Note that by Lemma B.4, for every k B,, we have either w W
KOp
Iw yo. [A:] S n.- Then during the training process after To. we have either sgn( (w (t) yer [k)) = sgn ( [k])
(Ta)
or sgn( E [k]) w (t)
2 (n) since if for some iteration number e' that we have sgn( (W [b))
sgu & [K]) but sgn( We.r (K6) sgu(E: [A]). then after F Õ()) steps (see the proof of Lemma B.2 for the definition of F) in the constant number of steps the gradient will must be in the same
direction of [Kk, which will push w,..r [k] back to zero or become positive along the direction of
[k). Therefore, based on this property we have the following regarding the inner product (w Ei),
(wi(t) Ex)
Σ
w
[K]
[A)
KEBIU(L)
Σ
( [A]
[ Ō(1)
|c] (1)
Σ
(EB,U(1)
Σ
[K]] [c(]]] (17sop)
kEB
where the second inequality follows from the fact that the entry w (t) [kc] that has different sign of &: [A]
satisfies (w(t), [k]] S Ö()). Then let ((()) ΣjeR((()) [w [2.[[k] 1[ (t [k]| 2 Ö(n)) | [k]].
which satisfies TT) = @(1) by Lemma B.4. Then assume B(()) keeps increasing and reaches some
value in the order of @ log(dns)
it holds that according to the inequality above
(w()) Weur E,) + (log(ddnn-)) ((soo,) O log(dnn-)1))
where we use the condition that y = 0((0r,) Then by Hypothesis (B. and (B. we know (9) that kw j.rt v) = o[1) (w p..r. E.) = 2(1), and (w r.. E.iL = (dd) + (w (t) then similar to 19), holds that
m. v\ = o(1)
(W x.)
(1)
(c) ΣSE(-Ln) (w()))) I exp (xpp--(((((w)))) ( {i)))) < poly(d in ,n)
BL.i
Therefore, at this time we have for all & E B,,
(1) 1 ((w b {.)K)]]]] < poly(d-1 n -.,,) .n) O (oo0ogg11 (dnn-)))))) -(o,) s nnm.
22

---

Under review as a conference paper at ICLR 2022
Then for all |W p.. [A]] 2 Ö(1),
the sign of the gradient satisfies
sgn VW,,, VS ((W(())A))
L (w
e
ndw ph. [k:
-gm
= sgn(nnn w (2))))
sgn(w!!) oe.r [A]]]-
Then note that Vw (W]))&]] = e(IAw s [R]]) 2 nnn
vid + Xn). by the update rule of w w [R] and Lemma B.2, we know the sign gradient will dominate the update process.
Then we have [w goar
[ (n)) sgn (.r [x)) S |w [k]] which implies that sirr [+]] 2 {n)) decreases so that B(") also decreases. Therefore, we can conclude
( [k]
we. [k]] 1(w
(1)
that B will exceed @ log(dnn Then combining the results for all E [n] gives
Ew
W [K]]- [E] - ((") + Ö(maa)) S 6 (logg(dnn + 0(1) = @().
kEB;
where in the first inequality we again the condition that 17 o(1/d) (((0,) - This verifies Hypothesis (B.13). Notably, this also implies that (w
(w yrr { maX,g|m) (w pi.r, <.) - (())
Verifying Hypothesis (B.14). In order to verify Hypothesis (B.14), let us first recall the update rule of (w (9) jart v):
(4) J.r
(++1) v) = (ww j.rt v) 17
m
Then by Lemma B.2, we know that if VV I [(W())][]] S (n). then m /.// V v j.r otherwise
(1)
S @1(1) and
sgn
Σ
niele (e) 2
(olloo((wW (c)
jar ((v)))))))
Uc d j.n Ei)) ndw nn O().
Without loss of generality we assume j = 1. then by Lemma B.4 we know that w w (Tb) [1] = 50
In the remaining proof, we will show that either w wir ((+)((1) (o, wir ((+1((1) [1] E
Ō n0 0).
First we will show that wi. [1] E [o, e()-1n)) for all r. Note that in the beginning of this stage,
(1++)(
we have w (To) [T] < 0. In order to make the sign of w [ flip, we must have, in some iteration
C' S'St l that satisfies w (t') [1] [0, (,))(((1))) therefore
-n wi. (W())) Σ vie 3. a o'((w (c)
J. vvv)) Σ yaelt') 0 ((w] Ei) n2 (1
< <n[( [1)) --))===~~ >] w 12 [1] ((n) 10
where the second inequality holds since 17
Note that V L L(W()) ē()).
then by Lemma B.2 know that Adam is similar to sign gradient descent and thus w wi.
(((+1)
w (t')
wi. [1] (n) which starts to decrease. This implies that if w/, (++1) is positive, then it cannot [1]
exceed @ o(1).
Then we can prove that if w ((+1) [T] is negative, then |wi.r
((+1) [4]]
Ō
na ). In this case we have for so??
all <
v L L(W(O)
Σ
ye (l') (1)
1, d'((w[[f])
i. (vv)) - Σ
yol p(
1,8 ('(WW") (.)))))) nAw (c)
23

---

Under review as a conference paper at ICLR 2022
((a) + n m/ f1ll
14.1 /w 1f
((a + n2 [4].
where in the inequality we use Hypothesis (B. and (B. 14) to get that (&, [k]] + all (t))
(w
w
(1) ye.
[k]]
max REB;
wurr v) 6(1)
Recall from (8.18) that have 1/ Ō
m therefore have if w [2[11] is smaller than some
value in the order of - na sog)
polylog(d). (d), then
V (W -
an
nd na so
+ θ
polylog(d) 2 O(nn),
50
(1) which by Lemma B.2 implies that w jir [1] will increase. Therefore, we can conclude that (1++))
0) in this case, which verifies Hypothesis (B.14).
sor
Lemma B.6 (Lemma 5.5, restated). If the step size satisfies O(d-1/2), then for any t it holds that
((Wi++))-LVC - - - VLL[ (()"d).
Proof. Let AFJ, = (3(W x.) F3( (W()) x:). Then regarding the loss function
(W.
L:WW) = - log E, ef(xx) Fyy, (W, x;) + log Σ De
It is clear that the function L,(W (W) is 1-smooth with respect the vector [F x,), F1 (W, Then based on the definition of A we have
OL. (W()) OF OF x.)
L (.((+))) L, ( Σ
AFF Σ(AF)A)²
(B.20)
Moreover, note that
(,(() x;) = Σ [o((w(()) Jort
(.. v(v)) + (((w E.)]
By the results that (w/r (t)
jurt v) & ((1) and (w E) s ((1), for any
we have
(w (1+1) Jor
(w j.m v) + n & @(1) (w/t+") . S (w j. {.) + @(ns 1/2) < @(1)
(t)
(t) which implies that the smoothness parameter of the functions a (ww!!) j.m vav)) and o (o(oo((w))) 3.r Ei) are at most @(1) for any w in the path between w and (((+1), Then we can apply first Taylor expansion (2)
on a(w j.m v)) and ((w j.r (.)) and bound the second-order error follows,
lo [o(( jrr vvv) - ((w S. v.v)) (Vw,.- [Vw,.-o((W uv)) (0+)) w [
(t+1)
< < (++1) w (1111 6(π2d).
(B.21)
where the last inequality is due Lemma B.2 that
2 I {()dd).
W (1+1) Jur
w
(1/
j.r
n
j.r
2
24

---

Under review as a conference paper at ICLR 2022
Similarly, we can also show that
(o((W() (t+1) Ar Ei)) al ((w
10 j.rt Ei) (Vw,,. (( (6) J.r E.)). w frr
(1+1)
w & @(rid).
(B.22)
Combining the above bounds on the second-order errors, we have
F (Vw FV (w ( xx),, w(++)) w s d) (()²d),
(B.23)
where the last equation is due to our assumption that m = (1). Besides, by (B.21) and (B. (B.22) the convexity property of the function (2x), we also have
|o((w]) (t ,vv))) - a( (0) j.rt u/ v())) - (V ws. o jor U.))), wr ( - +0)++6017
= @ (1+1)
jor M(v))l -Ivllll)) + ((7dd)
+(n+
(((W((t))) o(( {.)) (o((( jr
(()))))))))) S w(V,,,- σ (( j.r (.)). w/
(t+1) w
( (nnlo((ww)
((wj.r (++1) Ei)) (
Ö(nsop + rddd..
These bounds further imply that
F31l S (m (n80, + r/d)) + (118ap + +7²d).
(B.24)
Now we can plug (B.23) and (B.24) into (B.20) and get
(L((W())
(w((-))) - L. (W(())
Σ
AFGs + >'(AEA (AF))))
OF, (F x,)
(W()
x:), OF, (W(,, x²) (VwI)(W(0) xx), -
+ (() + +0((180, + +1/d)33
= (VL, w(+1) w(1) + +O(ndd)
(B.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 = 0ss-1/2) -1/2) so that O((nso, + 1)² (n)
Now we are ready to characterize the behavior on the entire training objective L(W) n E (W) A||Wl Note that W|| is 2X-smoothness, where X = o(1). Then applying
(B.25) for all E [n] gives
1 Σ L. (W(+1) [W()) + X(I I
L L(W ((W())
(VL(W()), w(++1) w + @(12d),
where the second equation uses the fact that [[(++1) w =6(n" d). Recall that we have
(++1)
Jar
w jir
-1)
j.r
Then by Lemma B.2, we know that m m/ large. Then we have
[k] is close to sign gradient if VL(W())[()) is j.r
Vw, L(W())
m
L(W
-(d 3(d-n) -n) - -(ns (ss,))
25

---

Under review as a conference paper at ICLR 2022
e V ws.. l(w())l)) O(dn),
where the second and last terms the R.H.S. of the first inequality contributed by the small gradient coordinates & B and & E = B, respectively, and the last inequality is by the fact that ns²o, (d). Therefore, based on this fact (B.25) further leads to
L(W - L(W v -nl VL((())) + d()2d).
which completes the proof.
Lemma B.7 (Generalization Performance of Adam). Let
W argmin
||V((W]]]-
W w())
Then for all training data, we have
Σ
[ox (W", x,) S S (W",
(W", x.)] =0.
n
Moreover, terms of the test data (x,y) ~ D, we have
P(xy)) (W", x) FF, (w', x)] 2 2"
Proof. By Lemma B.6, we know that the algorithm will converge to point with very small gradient (up to (rjd) in & norm). Then in terms of noise vector &,, we have
w L(W"))))))) [k]] O(dd).
(B.26)
KEB.
Note that
nV W... L(W [[K]
Ex : o' ((w W [.)))
nAw" or.r [Kc],
where V1, logit,, (F", xx). Then by triangle inequality and (B.26), we have for any r E [m].
(owllo(((" Cnniloo ((Wiers Wear [ nd
| year
[k]]]
S
V
L(W*) [T]||| - (wndd).
REB;
KEB
REB
Then by Lemma B.5, let r" arg max (w" wr Ex),
have (W (Wy) @(1) and
ΣLEB, (we.. [K]] [A]] = ē(1). Note that [k] = Ö(op), we have EnEB, WW [k]] 2
((1/op). Then according to the inequality above, holds that
1 ((sa,) 2 ē nd Σ W W..... - mand 2 θ nd
Op
AEB.
where the second inequality is due our choice of n). This further implies that 16 18 θ a. by combining the above results with (B.18). Then let us move to the gradient with respect to the first coordinate. In particular, since || VV((")) Ill < O(1d), we have
Σ Vil vvv) Σv v V ('w", Ei) ndw;.. [1]
s (nnd).
(B.27)
Then note that sgn(u: (z.)) = sgn(i), it is clear that w [1] < 0 since otherwise
V L(W >
Σ f |0' (dw(w ((w)..,E:))
E.) o' ((w), [v]]]
2
> A(ndd)
so?
26

---

Under review as a conference paper at ICLR 2022
which contradicts (B.27). Therefore, using the fact that w [1]-
O, we have
V w.... [T] = a Σ
o'
Σ
Yo j. ((w [ID]] n [1]
A:gh =
Then applying (B.27)and using the fact that / - 1f--v
θ
n2 30
for all i E [n]]. it clear that
[w [4]] 2 θ
a A
na
+
na
807
where the second equality is due to our choice of op and a. Then combining with Lemma B.5 and the fact that w [1] j < O, we have
j.r- [1] -
na sop
Now we are ready to evaluate the training error and test error. In terms of training error, it is clear that by Lemma B.5, we have (Wa..r,, E.) 2 (L), (w E.) 2 -o(1), and (( ii v) o(1),
((w" pe. 6.I = o(1). Then we have for any training data (x,,
(w, (x,) Σ
(o DI.ri B(V)) + (w or. Ei)) Ō(1),
F-g. W" xx) = [o((w- wu. vVv + ((w wu.r {I) = o(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 = [), which generated according to Definition 3. Note that for each neural, its weight wi,,- be decomposed into two parts: the first coordinate coordinates. As previously discussed, for any j [2] and sgn()) w []] S -((n)/(sa7)) and sgn(j) w [1] S @(>-1n) for r + r*. Therefore, using the fact that θ (nnx/(so²)) = w>>-'n) and Lemma B.5, given the test data (x, y), we have
and the rest d
r", have
Fy xx) =
(o(( p.rt vv)) + o(w p.rt E))
+ Carr ) so
Σ
F F-g(W", x)) =
( (w"
y.r. vv)) + o ((w w.r: (.:r:)))) {D)
0 | n.
[1]| + 15 15- y [4]
na
9
+ + 12))
so
where the random variables fy.r and Gy.r are symmetric and independent of v. Besides, note that o(1), it can be clearly shown that na/(so3)) < na/(sof?) Therefore, if the random noise Sy,r and G-y.r are dominated by the feature noise term (w--,r-: vv),, we can directly get that (W"" x) < F-v(W",x)) (recall that m = @(1)) 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 Cy and C-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) S F (W*, x), which implies that the output of Adam achieves 1/2 test error.
27

---

Under review as a conference paper at ICLR 2022
B.3 PROOF FOR GRADIENT DESCENT
Recall the feature learning and noise memorization of gradient descent can be formulated by
(W (1+1) J. v) (1- nd): (w (e)
j.rtJ vv)
+ 7 j Σ Wil o' (' j.m
((w v.))
Σ Di of
j.r Ei))
(wift)) = - .))) (w Blirs E:) +
p(0) Vi,
o' sooo(wWl
Ber Ei) Ei(1
kEB,
na
Dins o
(( mi
Ex))
Σ
Us es o' ((w per (o (B.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.c., the stage where all outputs Fj (W(()) x;) remain in the constant level for all and i.
Lemma B.S. [Lemma 5.6, restated] Suppose the training data is generated according to Definition (w
3.1, assume À = o 2o op/n). Let A = max (w]tt (t+1) .j v), r(t) J, = max,e(m² (w) J.ct Ei),
and 1 = max.y j 1 Then let Tj be the iteration number that (()) reaches ((1/m), we have
T = Ō (n)) for all j E (-1,1)
Moreover, let To max, } then for all < To it holds that r Ö(oo) for all jE 1..
We first provide the following useful lemma. Lemma B.9. Let (It, V1 h=...... be two positive sequences that satisfy
T++1 2 F4 + Azt
3 S yt + Byd 1-
for some A O(1) and B = o(1. Then for any g 2 3 and suppose yo -20o) and 17 < O we have for every C E [0o, 0()). let T₂ be the first iteration such that oı 2 C, then we have T2n = (z 2-) and
yr. S (zo).
Proof. By Claim C.20 in Allen-Zhu & Li (2020c), we have T.n = Tn 00zz-9). Then we will show
yt S 2zg
for all t < Tz. In particular, let Tan = C' 2-4 for some absolute constant C" and assume C' B22-1 < 1 (this is true since o())), we first made the following induction hypothesis on ye for all < To
yt S yo + Bn (2x00)
Note that for any t < To, this hypothesis clearly implies that
20 Then we are able to verify the hypothesis at time + based on the recursive upper bound of Uts i.e.,
yr S yo + Ban 9-1 S xp CB22-1 2- x g-1 <
2
Ve+1 S Ve + Byy
S yo + tn) 9-1 +1 Byi
yo + (t + 1))0(220)*
Therefore, we can conclude that yt < 2ng for all < T.. This completes the proof.
Now are ready to complete the proof of Lemma B.8.
28

---

Under review as a conference paper at ICLR 2022
Proof of Lemma B.8. Note that the initialization, have (w (0) 0,0 Then it can be shown that
(0) v) = and kw Ex)
FW x2) Σ [o((w))
(0) j.r (ir) +l(
(. Ei)]] = o(1)
for all j {-1,1). Then we have
14 e (i(W x.)
ef O().
Then we will consider the training period where [ for all j. i, and t. Besides, note that sgn(y, (7)) j. Therefore, let r' arg max, (w 1)-1)_
.j v). (B.28) implies that
1 2 (w) .j. vv)
17A) (w) j.r j v)
n
+ "
J.t ((w). (1-0) ((v))) Σ 10 (t-1) (.
n
2 (1 nd) (w
(t-1) 1) j v) - O(1) [o'(( Jir- ij v)) ao (T))
2 (1 n 1) n 0 ((A (t- (-)))-) n (
(B.29)
Similarly, let = arg max, (w
(1) g. E.). we also have the following according to (B.28)
(1-1) yrr- E.) + θ ?
r(t) n.t (w/ :. E,) S (1 m))
(t)
o' d'w (W E()
+ θ na²
((w (1-1)
ye.r" Es))
1 (1(--1)) p
((--))
+0
Vick
θ
ya²
(1)))))11))
24.
Wis
n
Then by our definition of (())) maxxejj] ((t)
we further get the following for all j E (-1,1],
p (T 3 + nso? + nnaa
(r)) (19--1)3-11 = -
110
(1(6-1))
+ θ
n
n
(B.30)
where the last equation is by our assumption that a = Ō0s00//))
Then we will prove the main argument for general 4, which based on the following two induction hypothesis
5 2 A. +n- ((A D1)))))))--
(t-1)
(B.31)
noo n
(r (116-110-1)
r50 < <
(B.32)
Note that when 0, we have already verified this two hypothesis in (B.29) and (B.30), where we use the fact that À = olog 2-2 2 op/n) < (1))))-2 and a = o(1). Then at time 4, based on Hypothesis (B.31) and (B.32) for all r < t, we have
r < O(A5))
as so so2/n o(1) and A increases faster than 19) Besides, we can also show that Ar < (1())))) which has been verified at time = 0, since 186)) keeps increasing. Therefore, (B.29) implies
A Altt) 2 (1 n))Aaa +n ((( - n ((a(T!)))-)
29

---

Under review as a conference paper at ICLR 2022
24 + ((A
which verifies Hypothesis (B.31) at + 1. Additionally, (B.30) implies
1889 (1))))-1)
r - ( + +
which verifies Hypothesis (B.32) at + 1. Then by Lemma B.9, we have that A.7 = Ö(1) for
all < To = (((A)))2-9/))
9(000%/)) Moreover, Lemma B.9 also shows that (
O(A) = (pn)) This completes the proof.
Lemma B.10 (Off- diagonia correlations). For any data (x,, yt) and for any & T-y. it holds that (w()) (1)
Mart Ei) S ((a).
Proof.
By the update form of GD, we have for any ls E Bi.
(0) w - [.( = (1 - nA) [k. & ik] + 2
((+1))
Σ
((0)
-
wirr Ei) E:(A)
KEB,
which keeps decreasing. Therefore. for all and i, we have
(w Burt Ei) S --.rr
(t)
[1]]++
Σ
w (0) - [K]]
(
KEB.
< ((a)= + @(anops1/2)
@(a):
(1) where the second inequality follows from the fact that / v) S @(1) for all t s TJ- This completes the proof.
Note that for different j. the iteration numbers when A reaches ((1/m) are different. Without loss of generality, we can assume Ti I T-1- Lemma B.8 has provided clear understanding about how A varies within the iteration range [o, T. However, it remains unclear how r(t) varies within the iteration range [Ti. since in this period we no longer have |e(01 = ((1) and the effect of gradient descent on the feature learning (i.e., increase of (Wj.r. v)) becomes weaker. In the
following lemma we give a characterization of (!) for every E [Ti,T-1].
Lemma B.11 (Stage of GD: part II). Without loss of generality assuming Ti T-1- Then it holds
that (t) = O(1) for all E [-iT--]].
Proof. Recall from (B.29) that we have the following general lower bound for the increase A
A(t11)= 2 - (1-1))- (w (w.
v)
7
Σ
p(0) lo (ow ((W B())) a
el0)
(looo(((w))) j.r.
n
1=1
2 (1 - ()))) + θ
(A-0)=+-((an) O(am) (r v ((a)) (B.33)
where the last inequality by Lemma 1. Note that by Lemma B.8, we have r Ō(oo) for all t t T-, and. Then the above inequality leads to
1 (0+) 2 (1 (1-))) +
+0(2)
(0) 12
Σ Py
(B.34)
where we use the fact that a = w(ao). The the remaining proof consists of two parts: (1) proving A 2 0(1/m) Õ(1) and (2) A O(log(1/A))
30

---

Under review as a conference paper at ICLR 2022
Without loss of generality consider 1. Regarding the first part, we first note that Lemma B.8
implies that A (fi)
2 (1/m). Then we consider the case when A < ((1og((/a)/m) holds
that for all V/ 1
(-1( x.)
10) 1.3
E
21E(-1,1) xi)
exp (0)
Σ
[a( (ow (Im o(w =.)] - Σ [o((w (iv a((wf).. (.)1
2 exp
( (²)
2 (xxp(-0log(11/000
-(a).
Then (B.34) implies that if ( < (llog(1/o0)/m) we have
1(+11) (1 A (ma) 1 O(a"n) A + O(pa) A 1
(1)
where the second inequality is due to À o(a). This implies that A will keep increases in this case so that it impossible that A ( which completes the proof of the first part. For the second part, (B.28) implies that
(++1 S - -))A +
? n
Σ
4
(A
(B.35)
Consider the case when 1 2 O(log(d)), then for all yı = 1,
cF-1(W(),x) x()
((0)
Eje(-1.1ef55) (W
(0
exp
[o((w (0)
-1,, vtv)) + ((w (1)
-1,: {.)]
Σ (o((w o((w(?).
I exp (-0(A
< expppx exp(-O((og((/A)
(Apo)
Then (B.35) further implies that
A (++1) < (1 (1)))) + (poii)a)) (Af9)--1
< A (1) - - - (A-poly(A)
(A(9) < 11
which implies that (9) will decrease. As a result, we can conclude that X will not exceed Ollo((1/)), this completes the proof of the second part.
Lemma B.12 (Lemma 5.7, restated). If 1 s O(ou), it holds that Ant (1) and r Ō for all / [T-1T]).
Proof. We will prove the desired argument based on the following three induction hypothesis:
1 2 (1 1
+
Σ
O 1 Σ
(
5.rb
(B.36)
T Ō(oo).
(B.37)
31

---

Under review as a conference paper at ICLR 2022
A -
(B.38)
In terms of Hypothesis (B.36), we can apply Hypothesis (B.37) and (B.38) (B.33) and get that
1(t11) 2 (1 (1- +
Σ
9- O(an) --(n)- (T.)) v ((a)) 9-1 1 Σ 14
n
((ann) 1
21 (n- θ
where the last inequality we use the fact that 2 0o- This verifies Hypothesis (B.36),
In order to verify Hypothesis (B.37), have the following according (B.36),
Σ
1ltt1)= 2 xn))
Σ
1
θ
7
(aa
71
(E(-1.1
jEl 1,1)
(1 - X
Σ
1
)E(-1.1)
where the last equality holds since o(1) Recursively applying the above inequality from T-1 to gives
1
E
A (4)
(1-
Σ
A A +
(
Σ
xn)
Σ
(t-
11-11-1)
n
JE(-1.1)
5E(-1.1)
i=1
Then by Hypothesis (B.38) we have
(-T-1-1
n
7 n
Σ
(1 (1- xn)
Σ
|(|---1)
S 6(1)
T=0
1=1
Now let us look at the rate of memorizing noises. By (B.28) and use the fact that a² < (s0 we have
3 s - (()-D))) +
n
(1(4-1))0-1
n
( (1-))))((()))
+
Real
nn80300
T-T-11
< <1 (T-1) + 3
Σ
(1 An)
Σ
10 (0-1-1)
7=
1=1
3(00 + sofa8-11
< {((o),
which verifies Hypothesis (B.37).
Given Hypothesis (B.36) and (B.37), the verification of (B.38) straightforward by applying the same proof technique of Lemma B.1 and thus we omit it here.
L
Lemma B.13 (Lemma 5.8, restated). If the step size satisfies, then for any t 2 T-, it holds that
I(W(++)) - ((w()) < - |I
Proof. The proof of this lemma is similar to that of Lemma B.6, which is basically relying the smoothness property of the loss function L(W) given certain constraints on the inner products (w,,r, v) and (Wir, E:).
32

---

Under review as a conference paper at ICLR 2022
Let 3F F, (W((++)) .x)) - Fy (W()) xx),, can get the following Taylor expansion the loss function (w((+))))
- S E OF x.)
( AF + E(APAN
(B.39)
In particular, by Lemma B. we know that (w j.r viv) S (11) and (w Si) - O(oo) S
(1). Then similar to (B.21), we can apply first-order Taylor expansion to P2 (W(+1), x:), which requires to characterize the second-order error of the Taylor expansions on (( jrr vvv)) and a(w ((+1) fr
(t+1)
E.)).
a (o((w)) 1++1) (vv))) - a((w ((w vbv)) - (V Wy.. ((w)) (1) vv))), wor w (
(1+1)
s (1 (1+ -w (..113 6(1 ww,,- (i.PlI (1
(o(((W))) (++1) .E)) o(w
(1) Ei) - (V w,, o ((w (2)
J.n E.)), Wer
(e+1)
w
s (11w - (0111) 0(1²' v w,,- L13
(1+1)
(B.40)
Then combining the above bounds for every r [m], we can get the following bound for AF,,;
|AF - (Vw F)((0, x;), w(++1) w n² Σ VV
I(W(3))
rE(m]
(B.41)
Moreover, since (.r v.v) < <0(1) and (w E.) Ō(1) and (-) is convex, then have ((+1) (++1) lo((w(tt) jr .vv))) a(ww (1) vvv))l S max (10' ([w] ((+1) ((v)) ((w (0)
jo wv)))) 1(v, (1+1)
s @((lw (++1) w (.
Similarly we also have
(t+1) [o(( Ar
.E)))= - ((w
Ji +1)i - 6([ fo - (f
Combining the above inequalities for every r [m]. we have
(1
w (.hllla) 1) < =
|A 2 I Ō =6(
Σ
(w(t))) j.r
rE(m)
(B.42)
Now we can plug (B.41) and (B.42) into (B.39), which gives
L(((+1)) L. (W())) s (L.(W(O)
OF, (WO,,, x;) AFF Σ(AP),)
-V((), w(+1) - w()) 6
(B.43)
Taking sum over i E [n] and applying the smoothness property of the regularization function X|| WI we can get
L L( ( L(W())= 1 n n
Σ
L (w(++1) (.())) (( -
s (VL(W()), - - W +
- (n -n- 0(r²)) ||
S " (
where the last inequality is due to our choice of step size 1) = o(1) so that gives - -5(1?) 2 n/2. This completes the proof.
1
33

---

Under review as a conference paper at ICLR 2022
Lemma B.14 (Generalization Performance of GD). Let
W arg
min W W(O) [|VL()())]l
Then for all training data, we have
FF xx) S F- (W",x)) = 0.
Moreover, in terms of the test data (x D, we have
((.,))DD F (W" x) S F (W", x) o())
Proof. By Lemma B.12 it clear that all training data can be correctly classified so that the training error zero. Besides, for test data (x, y) with [yv T : _ it is clear that with high probability (w vrr yv) = O1) and [( +) O(oo). then
Fy (W", x)
Σ [o((w v.r: vv)) + o((w p.r: E)]] 2 ((1).
ww-yor If = -y we have with probability at least 1/poly(n), (w" yor vv) I and [w" y.r E) S Ōa)) which leads
F (_y(W", x)
+ or oo {)> S =(aa)) o(1)
This implies that GD also achieve nearly at most 1/ 1/)oly(n) test error. This completes the proof.
C PROOF OF THEOREM 4.2: CONVEX CASE
Theorem C.1 (Convex setting, restated), Assume the model is over- wer-parameterized. 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 1 - n the obtained
"
parameters Adam and WeD satisfy that I[VI. Addam 111 S 1 and [|VL(W* Addm)) 11 S T7 respectively. Moreover, holds that:
In
Training errors are both zero:
1 |sgn (F(W Adam x:)) + vsl
Σ
1
[sgu[ (FW x:)) # #vdl = 0.
==1
Test errors are nearly the same:
((xaa)>D sgn (F(W Adam xx)) # 8) (x.))) (sgn x)) Fv o11)
Proof. The proof is straightforward by applying the same proof technique used for Lemmas B.6 and B. 13, where we only need to use the smoothness property 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.e., point with sufficiently small gradient. This further imply that the points found by Adam and GD must be exactly same and thus GD and Adam must have nearly same training and test performance.
Besides, that the problem is also sufficiently over- parameterized, thus with proper regularization (feasibly small), we can still guarantee zero training errors.
34


Reviews:
[{"review_id": "c7LiMRSoZKb", "reviewer": "Reviewer_9CZz", "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": "Strengths:\n1. The proposed analysis of the comparison between GD and Adam is, as far as I am concerned, novel.\n2. The results look interesting and promising, the dataset proposed by the authors is interesting and could be useful in other analyses as well.\n\n\nWeakness:\nAlthough I am not an expert on the NN analysis literature, I happen to be familiar with some of the papers the authors have mentioned, especially those related to NTK. I have to admit that the settings in this paper are so novel that it looks a bit strange to me. Although I understand that it is hard to analyze neural networks, especially when the authors are trying to compare GD with Adam, which is notoriously known to be unnecessarily fancy, I still have the following questions/concerns.\n\n1. First of all, the dataset looks interesting to me, but it seems that it is particularly designed for Adam on purpose. That is to say, it is already known to the public that Adam has convergence issues (Reddi et al, 2018), then it may not be too hard to construct such a dataset where Adam performs badly. Also, such a specific construction process raises the question of whether it is also possible to construct a case where GD performs pretty badly but Adam performs fairly well. I would sincerely encourage the authors to give more intuitions about their dataset, and maybe relax the assumptions (at least for GD)\n\n2. Second, the neural network looks a little strange to me. The authors are using a truncated polynomial activation function, instead of ReLU. Is this on purpose? I don't see such kinds of activation functions very often in the literature. Is it possible to replace it with, say, sigmoid, tanh, or maybe ReLU so that the analysis is more similar to the real training process? I do not see where the technical difficulty is. It would be helpful if the authors can explain it a little.\n\n3. If I understand it correctly, all the algorithms in this paper use full-batch gradients as their updates, instead of mini-batch gradients. I wonder whether it is possible to generalize the results to the stochastic mini-batch version.\n\n\nSashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.\n\n", "summary_of_the_paper": "This paper proposes a new perspective to explain the difference in generalization ability between Adam and SGD. The main contribution of this paper is that it analyzes the behavior of Adam and SGD training on a two-layer network and a very special dataset, and the authors claim that Adam generalizes poorly while SGD generalizes well. The authors also provide analysis for convex functions and some experimental results to validate their claims.", "summary_of_the_review": "Although I find this paper to be novel and interesting, I have to admit that some of the settings in the paper look strange to me. However, I would like to encourage such explorations and vote for acceptance because the authors are not using NTK, and I could imagine how difficult it is when trying to analyze Adam on neural networks. However, I would also like to encourage the authors to further generalize and improve their analysis.", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "T4XKWWHcVyR", "reviewer": "Reviewer_ahKf", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "Using artificial datasets to simplify the analysis of complex optimization dynamics seems to be a highly promising approach. It has the potential to provide a theoretical yet intuitive understanding of the interactions between data, models, and optimizers. While I believe this paper has made some progress in this direction, many of its assumptions are too contrived.\n\nFirstly, the same dimensions of different patches of the constructed data are reserved for either feature or feature noise, and the norm of the noise vector is specifically chosen to be much smaller than that of the feature vector, both of which are not well motivated. It would be interesting to see if some of the results still hold with weaker assumptions. Secondly, it seems the convolution part of the two-layer model is only introduced to impose weight sharing, such that the feature patch and the noise patch have to compete for the same set of weights, which is not necessarily the case in practice. Furthermore, the choice of truncated polynomial activation function is unconventional, and the second layer of the model simply sums up all the activations of the first layer without learned weights. Thirdly, in practice, the generalization gap is mostly observed between Adam and SGD rather than GD, and it is well known that SGD usually generalizes better than GD; therefore making directly comparisons between Adam and SGD would be more convincing.", "summary_of_the_paper": "This paper aims to explain why there is a generalization gap between Adam and gradient descent in learning neural networks, even with proper regularization. To this end, the authors construct an artificial dataset in attempt to capture the basic properties of real image datasets that can lead to the generalization gap. Based on this dataset, the authors prove that Adam is more likely to fit noise in data, while GD tends to fit real features. The theoretical results are verified by experiments on the proposed dataset.", "summary_of_the_review": "This paper presents a promising approach to explaining the generalization gap between Adam and gradient descent, but many of the assumptions on its data, model, and optimizers are far from realistic.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "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": "KcdVMKAo9cO", "reviewer": "Reviewer_yf6t", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "Strength:\n- experimental verification helps believe the result is true, as the proof is complex.\n- a simplified version of the result is presented assuming SignGD is used instead of Adam and can be intuitively understood.\n- food for thought paper on the limitations of Adam even if on a very specific example.\n\nWeaknesses:\n- not really a two layer network because the second layer is not learnable. If it were learnable, then setting a negative weight on the second patch would remove the issue that the dimension with the answer has different signs in the two patches.\n- The behavior only exists because q>1 (amplification of the gradient from the dot product with the noise in the second patch). Otherwise, the gradient from the first patch for the coordinate 1 would always win.\n- Behavior only exist for d > n^4.\n- The level of sparsity is quite strong ($\\sqrt{d} / n^2$), which seems a bit unlikely for a neural network (the citation (Papyan et al. 2017) does not just really justify such as assumption).\n- The result is for batch GD and Adam. Notice that for the stochastic version, we wouldn't have such a strong difference between GD and Adam, because most of the time the gradient would be zero for the noise coordinates, and when they are seen, Adam would rescale by 1/sqrt(n) rather than 1/n with the batch version.\n\n\nQuestions:\n- Each noise coordinate is present only once. So its gradient will be of the order of 1/n. Adam will rescale the gradient by a factor of n. What if you would do GD but with a learning rate n times bigger for the sparse features ?\n\nMinor remarks:\n- Lemma numbering is not aligned between main paper and appendix, while the Lemmas are the same. Either use same numbering or mention the Lemma number in the appendix with the proof.\n- mention that you assumme d>>n^4, this is not expicitely the case.\n- definition of Adam is missing the initialization bias correction, mention you drop it for simplicity.\n\nupdate: after discussing with the authors on the activation function, I believe the authors reply on the smooth ReLU, while technically correct, is unrealistic. In particular, the authors rely on the scaling down of the pre-activation outputs to a region where the smooth ReLU is polynomial, in order to obtain again the amplification that allows the bit flipping. Neural network initialization is specifically made so as to avoid such small pre-activation outputs, and I believe this makes this theoretical example too far off current practices to be considered an illustration of the generalization issues faced by Adam.", "summary_of_the_paper": "This paper propose a non convex optimization problem where the batch version of Adam has worse generalization than GD. Both methods achieve 0 training loss asymptotically but GD has a zero loss on the test set, while Adam has a large loss.\n\nThe paper studies a simplified CNN model with only two patches. The two layers are multiplied with the same weight matrix, fed into the activation function and then summed. A cross entropy loss is used for the classification. The activation function must be of the form $\\max(0, x^q)$ for $q \\geq 3$. Note that the second layer weights are not learnable and fixed to 1.\n\nThe dataset is built as follow: the label is y in {-1, 1}. One of the patch is just $x_1 = [y, 0, 0, ...]$. The second patch is $x_2 = [-\\alpha y, \\text{sparse gaussian noise}]$ ($0 <\\alpha < 1$). The sparsity level is so high that no two examples have the same features.\n\nThe idea is that Adam will use a much larger effective learning rate for the sparse noise features than SGD, and grow the coordinates quickly. At some point the dot product with the noise in $x_2$ will be larger than the dot product with $x_1$. Because the activation function is of the form $x^q$ with $q \\geq3$, this means the gradient from $x_2$ will overtake the one from $x_1$ and the weight for the first dimension will be negative (while a positive weight is the best thing to classify from $x_1$).\nSo Adam will focus on the weaker signal from $x_2$ (remember $\\alpha < 1$).\n\nTo summarize further, one patch has the right answer, the second path has the right answer (but with opposite sign) + noise. Adam will focus on the second patch which has lower signal, while GD will focus on the first patch and ignore the noise.\n\nThe authors provide experimental verification of this fact.\n", "summary_of_the_review": "Interesting example where Adam generalizes worse than GD. The assumptions are very far from a real network, and removing any of them make this example incorrect: second layer is not learnable, very high level of sparsity, and activation function has an amplification behavior.\n\nAs far as I can say, claims are correct but the proof is quite hard to verify.\n\nTo summarize, I'm recommending weak acceptance, because the example is interesting but the assumptions are unrealistic.", "strength_and_weaknesses": "", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "-QRejmaVj5P", "reviewer": "Reviewer_6gZd", "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": "The paper makes interesting observations about generalization properties of GD vs ADAM on a few settings.\n\nStrengths\n- Global convergence results for full batch ADAM and GD in non-convex setting when proper regularization is applied.\n- Both Thrm 4.1 and 4.2 explores interesting differences/non-differences of ADAM and GD with respect to their generalization abilities.\n- Based on a carefully constructed data distribution, the paper makes the argument that the inferior generalization performance of ADAM is closely tied to non-convex landscape of deep learning.\n\nWeakness\n- While all their observations are interesting in their own right, to me, the main weakness of the paper is its limited scope (2 layer conv net, full batch ADAM instead of online setting), some of which the authors also acknowledge in Section 7.\n- Since ADAM performs *better* than SGD in when used with transformer/lstms in language setting, a more complete analysis would be to also consider a data distribution/model where that can be explained as well.", "summary_of_the_paper": "This paper studies, in more detail than predecessors, the generalization gap between Adam and gradient descent for image-like datasets.", "summary_of_the_review": "This paper makes interesting and novel observations about generalization gap between adam and GD. Their observations are somewhat limited in scope. ", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper considers the difference between GD and ADAM in terms of implicit bias. It considers a specific distribution and architecture where the two algorithms converge to different solutions while perfectly fitting the training data. The authors highlight the fact that this happens while adding regularization, which does not happen in the linear case.
The reviewers found some of the insights and analysis interesting. However, they also had reservations about the impact of the results given that it is known that GD and ADAM have different implicit biases, and that the distribution appears specifically crafted towards  showing this effect for the architecture studied.
In future versions, the authors are encouraged to better motivate the chosen distribution, use more standard neural architecture (e.g., standard relu), and provide more explanation about the role of regularization in their result.
Decision: Reject