Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks
Authors: Xuan Tang
Han Zhang
Yuan Cao
Difan Zou
Venue: NeurIPS 2025 poster ()
Keywords: Generalization Error
Implicit Bias
Stochastic Gradient
Adam
Primary area: theory

Abstract: Adam is a popular and widely used adaptive gradient method in deep learning, which has also received tremendous focus in theoretical research. However, most existing theoretical work primarily analyzes its full-batch version, which differs fundamentally from the stochastic variant used in practice. Unlike SGD, stochastic Adam does not converge to its full-batch counterpart even with infinitesimal learning rates. We present the first theoretical characterization of how batch size affects Adam's generalization, analyzing two-layer over-parameterized CNNs on image data. Our results reveal that while both Adam and AdamW with proper weight decay $\lambda$ converge to poor test error solutions, their mini-batch variants can achieve near-zero test error. We further prove Adam has a strictly smaller effective weight decay bound than AdamW, theoretically explaining why Adam requires more sensitive $\lambda$ tuning. Extensive experiments validate our findings, demonstrating the critical role of batch size and weight decay in Adam's generalization performance.

Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks
Xuan Tang
Han Zhang' Yuan Cao
Difan Zou! 1.2
School of Computing and Data Science, The University of Hong Kong Institute of Data Science, The University of Hong Kong (xuantang8, hzhang23) @connect. hku hk, fyuancao, Hakuu hk
Abstract
Adam is a popular and widely used adaptive gradient method in deep learning, which has also received tremendous focus in theoretical research. However, most existing theoretical work primarily analyzes its full-batch version, which differs fundamentally from the stochastic variant used in practice. Unlike SGD, stochastic Adam does not converge its full-batch counterpart even with infinitesimal learn- ing rates. We present the first theoretical characterization of how batch size affects Adam's generalization, analyzing two-layer over wer-parameterized CNNs on image data. Our results reveal that while both Adam and AdamW with proper weight decay À converge to poor test error solutions, their mini-batch variants can achieve near-zero test error. We further prove Adam has strictly smaller effective weight decay bound than Adam W, theoretically explaining why Adam requires more sensitive À tuning. Extensive experiments validate our findings, demonstrating the critical role batch size and weight decay in Adam's 's generalization performance.
Introduction
Adaptive gradient methods, such as Adam (Kingma and Ba 2015) and its variant AdamW (Loshchlloy and Hutter 2019). have emerged as widely adopted optimizers for training deep learning models across diverse tasks (He et al. 2016 Ma and Hovy] 2016). More recently. Adam and variants have also been used to train large language models (LLMs) like GPT (Brown et al. (020), LLaMA (Touvron et al., 2023). and Deepseek (Bi et al., 2024). In practice, Adam known for its fast convergence during training, yet its generalization performance varies significantly depending on the task. Despite its empirical success, the theoretical understanding Adam remains incomplete, especially regarding its generalization performance.
Recent theoretical work has sought analyze the task-dependent behavior of Adam and compare with other optimizers like gradient descent (GD). For instance. Wilson et al. (2017) demonstrated that adaptive methods like Adam exhibit poor generalization on linear models, while GD and stochastic gradient descent (SGD) can achieve zero test error. Further, Zhou Thot et al. (2200) theoretically characterized the generalization gap between SGD and Adam through local convergence analysis, though their work did not account for neural network architectures or test error behavior. Other studies have focused on the implicit bias of adaptive methods: Wang et al. (2022) analyzed momentum's role in generalization, proving that GD with momentum and its adaptive variants converge to the l2 max-margin solution; Xie and Li (2024) showed that full-batch AdamW converges only to KKT point under an foo norm constraint; and Zhang al. (2024) established Adam's 's convergence to a maximum foo -margin classifier in linear logistic regression with separable data. In nonconvex settings. Zou et al. (2023b) revealed that full-batch Adam and GD converge to distinct solutions with differing generalization performance, which shows that even with weight decay, Adam fails to achieve low test error in overparameterized CNNs. Following the nonconvex analysis of Adam vs.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

GD by Zou et al. (2023b) Li et al. (2025) show that Sign Gradient Descent sign-only surrogate for Adam (Balles and Hennig 2018 Bernstein et al. 2018 achieves fast convergence but poor generalization when training two-layer Transtormers.
While existing theoretical analyses have provided valuable insights into the behavior of full-batch Adam, these results may not fully capture the characteristics of stochastic gradient Adam commonly used in practice. Notably, although stochastic gradient descent (SGD) and full-batch GD exhibit similar training dynamics in expectation (Bottou 2012), stochastic gradient Adam demonstrates fundamentally different behavior from its full-batch counterpart a distinction that persists even with vanishingly small learning rates. This gap raises important questions about how stochastic gradient Adam, particularly with small batch sizes, affects model generalization, an aspect that remains largely unexplored in current literature.
Motivated by this, in this paper, we investigate how the generalization of mini-batch Adam and Adam W differs from that of large-batch Adam. We analyze the convergence and generalization of Adam (and AdamW) with different batch sizes on two-layer wer-parameterized convolutional neural networks (CNNs) for an image data model. This analysis follows the settings outlined in the recent study of full-batch Adam in Zou et al. (203bb) We also compare the sensitivity of the weight decay parameters À for effective weight decay in Adam and AdamW.
The main contributions of this paper are summarized follows.
Theorem and rigorously prove that in the large-batch regime, both Adam and A converge to solutions with poor test error in nonconvex settings, even with proper weight decay. This extends prior results for full-batch Adam to AdamW, W. showing that adaptive methods inherently overfit noise in low-stochasticity training. Real-world data experiments in Figure [ demonstrate that large-batch Adam and Adam W suffer drastic test error increases, while synthetic experiments in Appendix D confirm this failure stems from noise-dominated solutions.
For mini-batch training, theorem 4.2 and 4.3 prove that stochastic Adam and AdamW achieve near- zero test error in nonconvex settings with appropriate weight decay. The key mechanism is twofold: (i) stochastic gradients implicitly regularize the optimization trajectory by slowing noise fitting while preserving feature learning dynamics, preventing Adam from overfitting noise patches; (ii) weight decay explicitly suppresses residual noise components. This synergy ensures convergence to solutions dominated by true features. Real-world data experiments in Figure [ demonstrate that mini-batch Adam and A significantly improve test performance, with synthetic-data experiments in Appendix [ further validating our theoretical insights. Moreover, under constant 31.82 hyperparameters, we prove stochastic Adam and Adam can be rigorously approximated by SignSGD (Bernstein et al. 2018) and SignSGDW (with decoupled decay) respectively. This extends the known full-batch Adam MiinnS correspondence to stochastic regimes- a crucial advancement given mini-batch noise fundamentally modifies approximation dynamics. Our analysis in Appendix reveals this approximation holds precisely when gradient magnitudes dominate optimization noise (e.g., 19.9. [K]] 2 θ()) where is the learning rate).
4.3 and 4.6 derive distinct theoretical upper bounds for weight decay parameters in nonconvex settings: Adam permits strictly smaller maximum effective than AdamW. This arises because Adam's is adaptive gradient normalization amplifies the effective impact of weight decay, causing excessive regularization to destabilize updates. In contrast, AdamW's decoupled weight decay mechanism avoids this issue. Experiments in Figure 2 validate that exceeding Adam's 's upper bound (e.g., A= > 0.05) leads to catastrophic test error increases, while Adam W tolerates much larger À values (e.g., À = 0.5) without significant performance degradation. This demonstrates that the interplay between batch size and weight decay is critical: mini-batce training enables effective regularization, but Adam's 's narrow tolerance demands precise À calibration.
Corollary
The rest of paper is organized as follows. Section 2 discusses the works that are most closely related to this paper. Section ( describes the problem settings. Section presents the main results of this paper. Section[5] proovides the proof outline of stochastic gradient Adam. Section 6 concludes this paper and discusses future research directions. Additional experiments and all experimental details can be found in Appendix D All proofs are provided in the remaining appendices (Appendix A
Notation. Scalars are denoted by lowercase letters 2; ....., vectors by bold lowercase letters x, and matrices by bold uppercase letters A.... For any integer d 2 1, denote the set [d] {1, ..., d).

---

For x E R. define [x]+ = max (z. 0} and sgn(x) = x/|0| for Z * 0, sgn(0) = 0. For x = (01,...,;dd) Rd define |X||p = (2: 12,l1)//p (p 2 1) and supp(x) = (i E [d] : x; # 0). For real sequences {an], }.n denote a, =Obb) if there exist C, > 0, s.t. [an] S C/b.), Vn 2 N; denote On = (b(bb) if bn = (aa)); on - (bn) if O(b,)) and (bb,,) hold; denote an (b,)) if for any C > 0, there exist N > 0, s.t. |an| < (/ba), Vn 2 N: and denote On = w(bx) if bn = o(an).
We write ((), ())). ( to suppress logarithmic factors, an = poly(b,) an = (()))) for some D > 0, and an = polylog(b,) (bn) if on = poly(log bn).
2 Related Work
Adaptive Optimization Methods. There are a series of papers on adaptive gradient methods, including AdaGrad (Duchi et al. 2011), Adam (Kingma and Ba] 2015). Adam (Loshchilov and Hutter 2019). and second-order information methods (Yao et al.; 2021 Liu et al.j 2024). The convergence of Adam and related methods has been analyzed in a line of papers under various conditions (Chen et al. 2019 Guo et al. 202T] Défossez et al. 2022). However, some work presented the possible case where Adam fails to converge to an optimal solution even in simple one-dimensional convex settings (Reddi et al., 2018). The generalization performance of Adam has been investigated
and compared with that of gradient descent in Wilson et al. (2017); Zhou et al. 2020); Zou et al.
(2023b). To better understand the performance of Adam, Bernstein et al. (2018 2019): y Kunstner ct (2023) analyzed its similarity with signGD. Similar works have also been done for AdamW (Xie and Li 2024p. Loshchilov and Hutter (019) demonstrated that improper use of weight decay in Adam could lead poor generalization performance, and proposed the AdamW that improves generalization in comparison to Adam. Recent work has also highlighted the role of weight decay in modern deep learning setups, showing its impact on optimization dynamics and generalization ( I'nggo et al., 2024). While L2 regularization and weight decay are equivalent for standard SGD) and GD (with rescaled by learning rate), that is not the case for adaptive methods like stochastic gradient Adam and full-batch Adam (Loshchilov and Hutter 2019 Zhang et al., 2019) /Zuuan et al., 2022). However, the true reason why Adam with weight decay fails to improve the generalization remains unclear. Therefore, the current understanding of how batch size and weight decay influence the generalization performance of Adam still relatively limited.
Implicit bias. Implicit bias refers to the tendency of machine learning algorithms to favor certain solutions. This phenomenon has also been studied in neural networks theoretically to understand how they generalize and converge to solutions. Lyu and Li (2019) and Ji and Telgarsky (2020) studied the
implicit bias of gradient descent on the homogeneous neural networks. Kunin et al. (2023) extended the results to a wider class of networks with varying degree of homogeneity. Cai et al. (2024) focused on the large stepsize gradient descent on two-layer non -onooooeeeeuus networks. [Frei et al. (2022)
analyzed the implicit bias of gradient flow in two-layer fully-connected neural networks with leaky ReLU activations for netroo-orthogonal data. Kou et al. (2024) extended this results and analyzed the implicit bias of gradient descent on similar settings. For Adam and Adam W. the implicit bias of Adam have been analyzed in Wang et al. (2021) 2022); Zhang et al. (2024). and the implicit bias of Adam W has been analyzed in (Xie and Li 2024).. Recently, Cattanco et al. (024)) showed that
Adam penalizes the l n-norm of perturbed gradients, favoring flat minima. Our work complements this view by analyzing, in discrete-time feature learning setting, how batch size and weight decay jointly regulate noise suppression and generalization.
Feature learning. There are a series of papers that studied the feature learning theory in neural networks. Allen-Zhu and Li (2020) investigated the feature learning of ensemble methods and knowledge distillation in deep learning when applied to data with multi-view features. Cao al. (2022) examined the benign overfitting in the supervised learning of two-layer convolutional neural networks, and proved that under certain conditions on signal-to-noise ratio (SNR), arbitrary small training and test loss can be achieved. Zou et al. (2023b) compared the feature learning of full-batch Adam and GD on two-layer convolutional neural networks. It demonstrated that GD learns the features, but full-batch Adam, even with proper regularization, may still fail. Some works have studied the feature learning of contrastive learning method (Zhang and Cao| 2024), federated learning (Huang et al. 2024b) on two-layer convolutional neural networks, and multi-modal contrastive learning on single-layer ReLU networks (Huang et al., 2024a). Additionally, some papers have
analyzed feature learning on other architectures, such as transformers (Jelassi et al. 2022 Li ct al.
3

---

2025), and diffusion models (Han et al. 2025a b); and other training configurations (Zou et al. 2023a Lu al. 2024). Unlike the aforementioned works, this paper focuses on the feature learning Adam and Adam W algorithms with different batch sizes on the two-layer convolutional neural networks.
3 Problem Setup
In this paper, we train the two-layer convolutional neural network (CNN) with Adam and AdamW on the training dataset S : {(x,, y/) )=) of size n,, which is generated from data model D. In this section, we introduce the data model D, the two-layer CNN model, and the training details of two algorithms (Adam and AdamW) analyzed in this paper.
Data Model. We adopt the feature-noise patch concatenation framework from Definition 3.1 aligning with previous studies (Allen-Zhu and Li 2020 Cao et al. 2022 Jelassi et al. 2022 Zou
ct al. 2023b Huang et al. 2024blat Zhang and Cao 2024 Li et al. 2025 Han et al.) (025a))
Definition 3.1. Let each data point (x, y) consist of feature vector x E Rdd and a label y E (11,
The data generated as follows:
x x2
where X1 and X2 represent two distinct feature patches. One of these patches corresponds to the signal patch and consists of feature vector y where R is assumed to be sparse vector, specifically 1-sparse. The other patch represents the noise patch and is noise vector denoted by &. Without loss of generality, assume = [1, (1.0......0 O The data is generated from the following distribution D:
The label y is generated as a Rademacher random variable with y E (-1,+1}.
2. Randomly select s coordinates from the set [dM[]} with equal probability. This selection is represented by a binary vector s E 10. Then generate & from the Gaussian distribution (O La) and apply the masking operation such that & = & @ where @ denotes element-wise multiplication. Finally, add feature noise to the vector & by updating it as & = oryv, where (0,1) 1) controls the strength of the feature noise.
One of the two patches X1, X2 is randomly selected and assigned as y v, representing the signal patch, while the other patch is assigned as &, representing the noise patch.
We set s θ (d2/2 of
θ (a, polylog(n)) in this paper.
The data model formalizes image classification dynamics where localized label-relevant features coexist with global noise-alignn noise-aligning with CNN behaviors: sparse mid-layer activations (Papyan et al.) 2017] vs. non-informative regions as independent noise (Yang) 2019). By isolating 1-sparse feature and s-sparse noise patches, we distill the feature learning vs. noise memorization interplay. Though our analysis uses simplified single feature/noise patch model for clarity, the results can be extended to broader settings (e.g., multi-patch or denser features/noises) by assuming sub-Gaussian noise and using concentration inequalities (e.g., Bernstein bounds) to control overlapping or structured perturbations, with similar qualitative behavior expected as long as the total noise remains controlled.
Two-layer CNN model. We define the two-layer CNN considered in this paper as follows.
Definition 3.2. Given the data (x. 4) ~ D and the activation function (xx) [<]]] with 2 3, the j-th output of the neural network F with width m is
m
F,( = Σ
x2)) Σ (o((W,r.. + ((W E))],
where War is the weight at the r-th neuron and initialized from Gaussian distribution N (O o Ia). In this paper, we assume E {+1} for clarity, ensuring the logit index matches the data label. Additionally, also assume polylog(n) and do = 1/11

---

Training algorithm. We investigate the behavior of stochastic Adam and Adam starting from same initializations and training on the same dataset S {x,, y:)) The loss function for each data point (x,, y:) is denoted (W) log
e (W.x:)
Ezet 1.11 e "
For stochastic Adam and AdamW, the CNN model is trained by minimizing the empirical loss function
(Adam) L(W)
L W + 2
W
(3.1)
==1
Σ ((W))) (AdamW) L(W)= n
(3.2)
where || denotes the Frobenius norm, À the weight decay regularization of Adam. Therefore, (t)
the stochastic gradient gr.j. be calculated
(Adam)
gij.r
(t)
1
w to
(W()))
Aw
(AdamW) gij.r B IEI
(t)
V ( (.(()))
where the subscript of Pr.j.r represents the batch Ij at the t-th iteration and the superscript t of Pt..r represents the model W(1) at the t-th iteration. Herein, we emphasize a fundamental distinction: Adam's stochastic gradients inherently incorporate weight decay regularization, whereas deliberate design choice to prevent momentum
(t)
AdamW W's gradients remain ingularization-free based normalization from destabilizing regularization effects (Loshchilov and Hutter 2019). This architectural distinction, also analytically demonstrated in our proof, crucially impacts the training
process. The momentum estimates m mjr Vj.r of Adam/ Adam//damW are updated as follows
(t) (t)
j.r = m,m j.r + (1 - 31) Pijir:
(t-1)
(3.3)
(t+1) B2v Jir j.r + - 52) [g 19..j.r
(3.4)
where B1, f2 are the hyperparameters of Adam/Adami W and we initialize m
(0)
= =v 0. Finally,
(0)
the update rule of stochastic Adam/AdamW for model W can be formulated as
(Adam)
(++1) ja
w. Wir
17
m j.r
(3.5)
+
(AdamW)
w (t+1) = (1 md)w
j.
(3.6)
+
j.r
where vj is the learning rate, 6 = ((ry) is stability constant and À the decoupled weight decay parameter of AdamW. In particular, in (3.4). (3.5) and (3.6), the square (-)2, square root vs and division //- all denote entry-wise calculations. The details of gradient calculation and its expansion can be found in Appendix A
4 Main Results
In this section, we present the main results of our study. We begin by introducing the primary metric used evaluate generalization performance: the classification error rate.
Given training dataset S = {(x;, " generated from data model D in Definition 3.1 We define the training enror crr (W) and test error err (W) of model W follows,
errs (W)= Ex (W, x) F-, (W, F-,WW,WWW
5

---

32
VGG16
26
VGG16
ResNet18
ResNet18
24
28-
Test Eroo 189
22
25
22
2 20
18
16
14
10
10
10
104
10a
10
10
104
Batch Size
Batch Size
(a) Test error vs. batch size under Adam
(b)
Test error vs. batch size under Adam W
Figure 1: Test error vs. batch size for VGG16 and ResNet18 on CIFAR-10.
erTD W)= ((xx))/DD I (W S F-,(W,)))
While theoretical analyses often prioritize mathematically tractable surrogate losses (e.g., cross- entropy, hinge loss), classification error rate remains the most direct and practical performance metric. Unlike continuously approximated surrogate losses, error rate directly quantifies discrete misclassification events, better reflecting models' true decision-making ability in classification tasks.
4.1
Theoretical Results for Adam
The following Theorem 4.1 characterizes the behavior of Adam in the targe-batch regime.
Theorem 4.1 (Large-batch Adam). Suppose 7)
poly(n) and À satisfies 0 < X o( o we
train our CNN model in Definition 3.2 on loss function B.1 for poly(n)
" epochs using Adam 3.5 B (1). Then with probability at leas! n we have
with batch size B satisfies
n
The training error is zero: errs (()) =0
The test error high: errp (W (T) 2 o(1)
Theorem 4.1 extends Zou et al.
2023b''s 's full-batch analysis to large-batch regimes (B = (n))..
Basically, it states that under the nearly same data model in Zou et al. (2023b), large-batch Adam cannot effectively learn the feature vector from the training dataset, and finally attains nearly 0.5 test accuracy, despite its perfect fitting on the training data points.
In stark contrast, we further provide Theorem 4.2 which proves that stochastic gradient Adam with smaller batch size can achieve good generalization performance.
Theorem 4.2 (Mini-batch Adam). Suppose pnly(n) and <\= o we train our CNN
model in Definition 3.2 on loss (3.1) for = polyin) epochs using stochastic Adam (3.5) with batch
size B satisfies 2 (loge where is the hyperparameter of Adam. Then with high probability
least we have
The training error is zero: ars(W(()) = 0.
The test error near-zero: errp (T) o(1)
Our theoretical results demonstrate that mini-batch Adam achieves near-perfect test accuracy when the ratio n/B is large, significantly outperforming its large-batch counterpart which exhibits only random- random-guessing guessing performance. This advantage can be attributed to three fundamental properties of stochastic Adam optimization: First, since the feature vector is shared across all data points, its learning remains robust regardless of batch size. In contrast, noise vectors are data-specific and vary across different samples. When using mini-batches, only subset of noise vectors exposed during each update, creating inherent asymmetry in learning dynamics. More importantly, Adam's coordinate- wise normalization amplifies this effect: maintains consistent learning rates for shared features while substantially slowing down noise memorization. This selective suppression of

---

noise learning explains the superior generalization performance of mini-batch Adam compared to large-batch implementations.
Besides the results on the generalization performance, we further deliver the following corollary, which states the feasible range of the weight decay in Adam. Theorems 4.1 and directly yield:
Corollary 4.3 (Effective weight decay in Adam). Suppose the same conditions as in Theorem 4.1
and 4.2
If X = w(ag then with probability at least - n2 training stuck at the initialization.
This corollary provides a theoretical upper bound on the effective weight decay that allows Adam to successfully train models, aligning well with previous empirical observations that Adam typically performs better with small weight decay values compared to AdamW (Loshchilov and Hutter 2019). This sensitivity arises because weight decay regularization implicitly entangled with the normalization step in Adam, i.c., when the gradient of weight decay is greater than that of the cross-entropy loss, it will fully dominate the Adam update. In the next subsection, we will show that weight decay will exhibit different behavior in AdamW, leading to different feasible range for X.
4.2 Theoretical Results for A
We first establish the theoretical results of Adamw in Theorem 4.4 under large-batch training.
Theorem 4.4 (Large-hatch AdamW). Suppose v poly(n) À
12 B²
A1) and À = Õ(1), we train
our CNN model in Definition 3.2 on loss function for T poly(n) epochs using AdamW (36))
"
with batch size B satisfies B ((1) or B p( Then with probability at least 1 n we have
The training error is zero: errs (W()) = 0.
The test error is high: errp W (T) 2 2 o(1)
The results for large-batch Adam W closely resemble those of large-batch Adam in Theorem 4.1 the learned model consistently exhibit test errors of at least 1/2 o(1), performing no better than random guessing. This phenomenon arises from the training dynamics in the early stages, where the influence of weight decay is minimal. As a result, both Adam and AdamW exhibit similar behavior, tending to fit noise. By the time the decoupled weight decay in AdamW begins to take effect, the model has already overfit to the feature noise -ayv. The weight decay then guides the model toward nearby local minima, effectively preserving the previously memorized noise. Then at test time, this overfitting to feature noise causes the model to predict labels that are systematically misaligned with the true labels, leading to test performance that is no better than random guessing.
In contrast to the results for large-batch Theorem 4.4 the following Theorem characterizes the generalization ability of mini-batch AdamW.
Theorem 4.5 (Mini-batch AdamW). Suppose
1) and À Ö(1), train (32 for T poly(m) epochs using stochastic > Ologg 1) and B w(Sar V 12)2). Then with
poly(n) X π2( 8
our CNN model in Definition 3.2 on loss function
AdamW (36)) with batch size B satisfies
probability at least 1 -1 we have
The training error is zero: errs (W()()())) = 0.
The test error is near-zero: errp errp(w(()) (W(T)) o(1)
Under mini-bated training, AdamW achieves near-zero test error with partial similarity to Adam. The extended iterations per epoch slow early-stage noise overfitting (notably for s-sparse noise). However, AdamW's W's decoupled weight decay penalizes weights independently of gradients, exerting significant regularization only in later phases converge toward generalizable minima. Thus, AdamW can leverage a much larger than Adam, which is shown as follows:
Corollary 4.6. Regarding the effective weight decay coefficients of Adam and AdamW W for achieving good generalization performance, we have AAdam ~ og B² A AAdauW-
Corollary reveals fundamental gap between the effective weight decay regimes of Adam and AdamW, consistent with empirical observations. For Adam, the admissible weight decay AAdam
7

---

Adam
Adam
24
20
AdamW
AdamW
Test Error (%)
22
Test Error (%)
19
20
18-
17
18-
16
16
15-
14
10
10-3
10 -4 10
10
10-1
14
10
6
10-5
10
10
10
10-
Weight Decay Adam
Weight Decay Weighl Decay 1/25 for AdamW
Weight. Detay 1/25 for AdamW
(a) Training VGG16
(b) Training ResNet18
Figure 2: Test error weight decay (batch size 16), comparing Adam and AdamW each model.
extremely small, bounded above by an initialization deppedddn term of Consequently, even moderate weight decay values can destabilize training due to the entanglement between gradient updates and weight decay, leading to suboptimal generalization. In contrast, A decouples weight decay from the gradient update, applying regularization directly on the weights. As result, it requires a larger weight decay to effectively suppress noise overfitting and exhibits greater robustness. The lower bound 52( n? A 1) serves as sufficient condition ensuring that weight decay is strong enough prevent noise amplification. This robustness is reflected in its much broader effective range-from S2( n A 1) up to Ó(1) repprenpnnnngning a large, constant-order window independent of initialization. This theoretical separation explains the empirical fact that Adam requires delicate tuning and highly sensitive to weight decay, whereas AdamW is considerably more robust and casier to tune Loshchilov and Hutter (2019). Our experiments (Figure 2 further corroborate this prediction.
Experiments. We train VGG16 and ResNet18 CIFAR-10 10 with Adam (A = x 10 and AdamW (A= x 10 2), selecting the optimal learning rate from f5x10- 1x 10-4, 1x10010- 63 and varying the batch size measure test error (Figure TT. Both optimizers exhibit sharp performance degradation once the batch size exceeds a critical threshold, in line with theoretical predictions of fixed batch size of 16 (Figure 2. we find that Adam's ' error spikes for X 0.05, whereas Adam remains robust even up to 0.5, highlighting the benefit of decoupled weight decay in adaptive optimizers. This observation consistent with our theoretical analysis (Corollary 443 and Additional experiments, including the dynamics of feature learning and noise memorization (Figures 3 4 sensitivity weight decay
large-batch generalization collapse (Theorem 4.1 and 4.4). Separately,
(Figures 5 6). error bars across random seeds and momentum parameters (Figures 7 g 10 and
large-scale vision experiments with ResNet-50 on ImageNet- IK (Figures 12 all corroborate our theoretical findings. Complete experimental details and results are provided in Appendix
Proof Outline of the Main Results
In this section, we mainly outline the proof sketch for the theorem 4.2 in Section Proof sketches for remaining theorems are deferred to Appendix B Following the two-stage analysis framework of Cao al. 2022); Zou et al. (2023b), we decompose the proof into two distinct stages:
Stage I: Pattern Learning. During the initial phase of training, the effect of regularization is negligible. The model operates an underfitting regime, where it rapidly learns dominant patterns in the training data, leading to improved empirical performance on test error.
Stage II: Regularization. As training progresses, the model's classification accuracy the training set approaches convergence, resulting in diminished gradient magnitudes. At this stage, regularization dominates the optimization dynamics, steering the model converge to local minima. Due to the nonconvex nature of the loss landscape, the model retains the patterns acquired during the pattern learning stage.
Furthermore, motivated by the behavioral similarity between Adam and SignGI when the learning rate is sufficiently small or 31, 8, approach zero (Balles and Hennig 2018 Bernstein al. 2018).

---

we present results for SignSGD. We subsequently extend these results to stochastic Adam, which provided in Appendix C
The update rules for SignSGD are given as follows:
(SignSGD)
[r+1) Wir
") sgn (g...))).))
(5.1)
where 9r. jor (.1) denotes the stochastic gradient of 3.1) The detailed update rules of Adam with the SignSGD approximation are provided in Eqs. (B.3) and (3.4) while those of AdamW are given in Eqs. B35 and (B.6).
Next, following the framework of feature learning (Allen-Zhu and Li. 2020; Cao et al. 2022 Zou al. 2023b; Han et al. 2025a), we primarily focus on two key quantities: 1) Feature Learning jv)::.vv: This term captures the alignment between the learned weight vector Wj,r and the true feature direction jv, reflecting the model's ability to extract meaningful latent structures from the data. 2) Noise Memorization (w w...r:si Ei): This term measures the correlation between w y,rr and the noise patch & of individual samples, characterizing the extent to which the model overfits to stochastic perturbations or idiosynerasies in the training set. This decomposition allows us to separately analyze the model's generalization behavior (driven by feature learning) and its memorization capacity (influenced by noise fitting).
We first clarify that, although the sketch appears straightforward, the underlying process is non-trivial Numerous intricate and interesting details arise, which we elaborate on in the proof presented in Appendix C
We begin by characterizing the dynamics feature learning and noise memorization under large-batch training, facilitate a comparative analysis with mini-batch regimes, as formalized in Lemma 5.1
Large-batch Adam. We consider
B = (1), which is the large-batch batch setting.
Lemma 5.1. Given the training dataset S, if B = (1), = 1/poly(d) and 0 < X = =000a=0000 op/n).
then for any S To with To Ō "j00
and any i [n].
(1+1) (w j.r (.v) < (w (1) jv) + e()), (w(t=") E.) = (w g..r E.) + (1)sop))
We observe that Lemma 5.1 is identical to Lemma 5.2 in (Zou et al. 2023b), allowing us to directly extend their full-batch Adam results the large-batch setting. The remainder of the proof follows identically, as the underlying theoretical machinery remains unchanged under this batch size scaling a (1). We observe that under large-batch setting, the optimization dynamics of Adam closely resemble those of the full-batch setting. This similarity arises because the algorithm traverses the entire dataset within few iterations, resulting in nearly identical momentum estimates and, consequently, comparable training dynamics between large-batch and full-batch regimes.
We next consider the mini- batch setting, which yields conclusions that differ fundamentally from those in the large-batch setting.
Mini-batch Adam. We consider E O(loge which is the mini-batch setting.
Lemma 5.2 (Stage I). Given the training dataset S, if BB Z O(loge e()oge-)), n = 1/poly(d) and an/n),, then for any To with To O noa
0 x olog "000-app/oi
and any E [n].
(w ((+))) jor
n ) v).))= = (w j 0
1D 10 E1) < @(1sop))
Compared to Lemma 5.1 Lemma 5.2 reveals fundamentally different optimization dynamics. Specif- ically, feature learning progressively increases throughout Stage I, whereas noise memorization remains suppressed near the initialization. The key distinction from the batch setting lies in the fact that, under the mini-batch regime, traversing the entire dataset requires many iterations. Since noise is sparse and uncorrelated across samples while features are dense and shared, feature learning can proceed effectively during the early training phase without being hindered by weight decay regularization. In contrast, noise memorization is strongly suppressed by weight decay due to its sparsity. As training progresses, the momentum estimates in Adam gradually forget the gradient contributions from noise, allowing weight decay to dominate. As result, recently acquired noise

---

memorization is continually erased, keeping noise-related parameters close to their initialization throughout Stage
In the following lemma 5.3 we show that the patterns learned by the model during Stage are preserved in Stage II, due to the nature of our non-convex optimization landscape.
Lemma 5.3 (Stage II). Suppose the same conditions hold as in Lemma 5.2 For >> To, j e
{#1}, [m], i E [n], let
argmax,e|m) (w jv). then j.::iv) (1) and
Mur: E.) (1sop))
Given Lemma 5.2 and Lemma 5.3 we can characterize the convergence rate of Adam as follows.
Lemma 5.4 (Convergence), Suppose the same conditions hold as Lemmal 5.2 and 5.3 if the step size 1j - (d 3), then for any 6,
E ((((W())))
L(W())) --]]]V(((())]] -
Combine Lemma 5.3 and Lemma 5.4 we observe that the model successfully learns the true features
and eventually converges to local minimum, retaining strong generalization performance.
Conclusion and Limitation
In this work, we theoretically and empirically analyze the impact varying batch sizes and weight decay parameters the generalization of Adam and AdamW when learning two-layer CNNs. Our results demonstrate that large-batch Adam and Adam inherently overfit noise-dominater solutions with weight decay, while their mini-batch counterparts achieve strong generalization through the synergy of implicit stochastic gradient regularization and explicit weight decay. Furthermore, we establish that Adam's 's adaptive gradient normalization imposes stricter constraints weight decay parameters compared to Adam necessitating precise calibration for stable optimization.
While our theoretical framework provides insights into the interplay between batch size, weight decay, and generalization, several limitations highlight critical directions for future research. First, the current analysis is restricted to two-layer networks. Extending the results to deeper architectures and investigating how batch size influences the dynamics of hierarchical feature learning presents promising direction. Second, extend the analysis to domains with fundamentally different data structures, such as NLP, where batch size and weight decay may impact model performance through different mechanisms. Finally, other critical hyperparameters, such as momentum, learning rate schedules, and gradient clipping, are considered in our analysis, and some modern vision architectures succeed with large batches (Liu jtt al. 2022 2023 Chen et al. 2024) despite our theoretical predictions, suggesting that additional factors like architectural design and normalization may play significant role.
work focuses image data, and important direction is to
Acknowledgments
We would like to thank the anonymous reviewers and area chairs for their helpful comments. Xuan Tang and Difan Zou acknowledge the support from NSFC 62306252, Hong Kong ECS award 27309624, Guangdong NSF 2024 1515012444, and the central fund from HKU IDS. Yuan Cao is partially supported by NSFC 12301657 and Hong Kong ECS award 27308624.
References
ALLEN-ZIIU, Z. and LI, Y. (2020). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816.
BALLES, and HENNIG, (2018). Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference Machine Learning. PMLR.
BERNSTEIN, J., WANG, Y.-X., AZIZZADENESHELL, K. and ANANDKUMAR, A. (2018). signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning. PMLR.
10

---

BERNSTEIN, J., ZHAO, J., AZIZZADENESHELI, K. and ANANDKUMAR, A. (2019). signSGD with majority vote is communication efficient and fault tolerant. In International Conference on Learning Representations.
BI, X., CHEN, D., CHEN, G., CHEN, S., DAI, D., DENG, C., DING, H., DONG, K., DU, Q. FU, Z. ET AL. (2024). Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint WwX::i241.0025555
BOTTOU, L. (2012). Neural Networks: Tricks of the Trade: Second Edition. Springer Berlin Heidelberg.
BROWN, T., MANN, B., RYDER, N., SUBBIAHI, M., KAPLAN, J. D., DHARIWAL, P., NEELAKAN-
TAN, A., SHYAM, P. SASTRY, G., ASKELL, A. ET AL. (2020). Language models are few-shot learners. Advances in neurul information processing systems 33 1877-1901.
CAL, Y., WU, J., MEI, S., LINDSEY, M. and BARTLETT, P. (2024). Large stepsize gradient descent for con-oooogeeeous two-layer networks: Margin improvement and fast optimization. Advances in Neural Information Processing Systems 37 71306-7135/
CAO, Y., CHEN, Z., BELKIN, M. and GU, Q. (2022). Benign overfitting in two-layer convolutional neural networks. Advances in neural information processing systems 35 25237-25250.
CATTANEO, M. D., KLUSOWSKI, J. M. and SIIGIDA, B. (2024), On the implicit bias of adam. In Proceedings of the 41st International Conference on Machine Learning. vol. 235 of Proceedings of Machine Learning Research. PMLR.
CHEN, H., CHU, X., REN, Y., ZHAO, X. and HUANG, (2024). Pelk: Parameter-efficien large kernel convnets with peripheral convolution. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition.
CHEEN, X., LIU, S., SUN, R. and HONG, M. (2019). On the convergence of class of adam- type algorithms for non-convex optimization. In 7th International Conference on Learning Representations, ICLR 2019.
D' ANGELO, F., ANDRIUSHCHENKO M., VARRE, A. V. and FLAMMARION, N. (2024). Why do we need weight decay in modern deep learning? Advances in Neural Information Processing Systems 37 23191-23223.
DÉFOSSEZ, A., BOTTOU, L., BACII, F. R. and USUNIER, N. (2022). A simple convergence proof of adam and adagrad. Trans. Mach. Learn. Res.
DUCHI, J., HAZAN, E. and SINGER, Y. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research 12.
FREI, S., VARDI, G., BARTLETT, P., SREBRO, N. and HU, W. (2022). Implicit bias in leaky relu networks trained on high-dimensional data. In The Eleventh International Conference on Learning Representations.
GUO, Z., XU, Y., YIN, W., JIN, R. and YANG, T. (2021). A novel convergence analysis for algorithms of the adam family. arXiv preprint arX0v:2112.0345.
HAN, A., HUANG, W., CAO, Y. and Zou, D. (2025a). On the feature learning diffusion models. In The Thirteenth International Conference on Learning Representations.
HAN, Y., HAN, A., HCANG, W., LU, C. and Zou, D. (2025b). Can diffusion models learn hidden inter-feature rules behind images? In Forty-second International Conference on Machine Learning.
HE, K., ZHANG, X., REN, S. and SUN, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition.
HUANG, W., HAN, A., CHEN, Y., CAO, Y., ZHIQIANG XU and SUZUKI, T. (2024a). On the comparison between multi-modal and single-modal contrastive learning. In The Thirty-eighth Annual Conference Neural Information Processing Systems.

---

HUANG, W., SHI, Y., CAI, Z. and SUZUKI, T. (2024b). Understanding convergence and generaliza- tion federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations.
JELASSI, S., SANDER, M. and LI, Y. (2022). Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems (S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho and A. Oh, eds.), vol. 35, Curran Associates, Inc.
J1, Z. and TELGARSKY, M. (2020). Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems 33 17176-17186.
KINGMA, D. and BA, J. (2015). Adam: method for stochastic optimization. International Conference on Leurning Representations.
KOU, Y., CHEE,, Z. and GU, Q. (2024). Implicit bias of gradient descent for two-layer relu and leaky relu networks on nearly-orthogonal data. Advances in Neural Information Processing Systems 36.
KUNIN, D., YAMAMURA, A., MA, C. and GANGULI, S. (2023). The asymmetric maximum margin bias of guugi-homoeeneous neural networks. In The Eleventh International Conference on Learning Representations.
KUNSTNER, F., CIIE,, J., LAVINGTON, J. W. and SCHMIDT, M. (2023). Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In The Eleventh International Conference on Learning Representations.
LI, B., HUANG, W., HAN, A., ZHOU, Z., SUZUKI, T., ZHU, J. and CHEN, J. (2025). On the optimization and generalization of two-layer transformers with sign gradient descent. In The Thirteenth International Conference Learning Representations.
LIU, H., LI, Z., HALL, D. L. W., LIANG, and MA, T. (2024). Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Conference on Learning Representations.
LIU, S., CHEN, T., CHEN, X., CHEN, X., XIAO, Q., WU, B., KÄRKKÄINEN, T., PECHENIZKIY, M., MOCANU, D. C. and WANG, Z. (2023). More convnets in the 2020s: Scaling up kernels beyond 51x51 using sparsity. In ICLR.
LIU, Z., MAO, H., WU, C.-Y., FEICHTENNIOFER C., DARRELL, T. and XIE, S. (2022). A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition.
LOSHCHILOV, I. and HUTTER, F. (2019). Decoupled weight decay regularization. In International Conference on Learning Representations.
LU, M., WU, B., YANG, X. and Zou, D. (2024). Benign oscillation of stochastic gradient descent with large learning rates. In The Twelfth International Conference on Learning Representations Vienna).
(ICLRR107/05/222411
LYU, K. and LI, J. (2019). Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 1906::990000599
MA, X. and Hovy, E. (2016). End-to-end sequence labeling via bi-directional lstm-enns-erf. arXiv preprint arXiv: 1603.01603.01354
PAPYAN, V., ROMANO, Y. and ELAD, M. (2017). Convolutional neural networks analyzed via convolutional sparse coding. Journal of Machine Learning Research 18 1-52.
REDDI, S. J., KALE, S. and KUMAR, S. (2018). On the convergence of adam and beyond. In International Conference on Learning Representations.
TOUVRON, H., LAVRIL, T., IZACARD, G., MARTINET, X., LACHAUX, M.-A., LACROIX, T.
ROZIÈRE, B., GOYAL, N., HAMBRO, E., AZHAR, ET AL. (2023). I.LMM: Open and efficient foundation language models. arXiv preprint 233.
12

---

WANG, B., MENG, Q., CHEN, W. and LIU, T.-Y. (2021). The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning. PMLR.
WANG, B., MENG, Q. ZIIANG, H., SUN, R., CHIEN, W. MA, Z.-M. and LIU, T.-Y. (2022). Does momentum change the implicit regularization on separable data? Advances in Neural Information Processing Systems 35 16764-267766
WILSON, A. C., ROELOFS, R., STERN, M., SREBRO, N. and RECHt, B. (2017). The marginal value of adaptive gradient methods in machine learning, Advances in neural information processing systems 30.
XIE, S. and LI, Z. (2024). Implicit bias of adamw: -nomm constrained optimization. In Forty-first International Conference on Machine Learning.
YANG, G. (2019). Scaling limits of wide neural networks with weight sharing: Gaussian pro- cess behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760.
YAO, Z., GHOLAMI, A., SHEN, S., MUSTAFA, M., KEUTZER, K. and MAHONEY, M. (2021). Adahessian: An adaptive second order optimizer for machine learning, In proceedings of the AAAI conference on artificial intelligence, vol. 35.
ZHANG, C., Zou, D. and CAO, Y. (2024). The implicit bias of adam on separable data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
ZHANG, G., WANG, C., XU, B. and GROSSE, R. (2019). Three mechanisms of weight decay regularization. In International Conference on Learning Representations.
ZuANG, H. and CAO, Y. (2024). Understanding the benefits simclr pre-training in two-layer convolutional neural networks. arXiv preprint arXiv:2409, 118224999..8...... 18685
ZHOU, P., FENG, J., MA, C., XIONG, C., HOI, S. C. H. and E. W. (2020). Towards theoretically understanding why sgd generalizes better than adam in deep learning. In Advances in Neural Information Processing Systems, vol. 33.
ZHUANG, 7.,, LIU, M., CUTKOSKY, A. and ORABONA, F. (2022). Understanding adamw through proximal methods and scale-freeness. Transactions on Machine Learning Research.
Zou, D., CAO, Y., LI, Y. and Gu, Q. (2023a). The benefits of mixup for feature learning. In International Conference on Machine Learning. PMLR.
ZOU, D., CAO, Y., LI, Y. and GU, Q. (2023b). Understanding the generalization of adam in learning neural networks with proper regularization. In International Conference on Learning Representations.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses the limitations of the work in the conclusion section, which are also our future research directions.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: We provide the main results in the main paper, and state all the assumptions and the complete proof in Appendix. We also provide proof outline in Section 5 in the main paper.
Guidelines:
The answer NA means that the paper does include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide all the configuration of the experiments in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We use VGGI6, ResNer18, ResNet50 models and the CIFAR-10, ImageNet- IK datasets, which are easily available on the Internet. All the experimental details are provided Appendix D
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc
quuiddes/CoeSSumissiiiii
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
ssiocoolicyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide all the experimental details in Appendix D
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We report the error bars in Figures g
9 TO
Appendix
D
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide all the experimental details in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http://eerrri...coooo/o/..aaaaaa
Answer:
[Yes]
Justification:
The research follows the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: This paper a theoretical analysis paper on the generalization of stochastic gradient Adam (AdamW) and does not have potential societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: This paper does not use existing assets.
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
18

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy hatt::///eeerrssssse////ffferers for what should or should not be described.
20

---

A Preliminaries
In this section, we give the asymptotic equations for all the parameters use, useful lemmas, and the gradient and weight update equations.
A.1 Asymptotic Equations
First, we give the asymptotic equations for all the parameters we use in the proof. Condition A.1. Suppose that the following conditions on data model 3. hold,
1. The dimension d satisfies d = poly(n).
2. The number of noise coordinates satisfies θ d1/2
The variance parameter o of the noise vector satisfies oz = s-polylog(n)
The feature noise strength a satisfies = e (ap polylog(n))-
Condition A.2. Suppose that the following conditions on hyper-parameters hold,
The initialization variance of the model weights og satisfies oî
(xt²).
2. The width of the network satisfies m = polylog(n).
3. The learning rate satisfies 1| poly(n)
The purameter Stochustic Adam and Stochustic AdamW satisfies e = @(An).
Based on the parameter configuration, we claim that the following equations hold, which will be frequently used in the subsequent proof.
a- (
s02
a = w ((
Go = 0
= o (No8ot)
(A.1)
80pp
The following Lemma A.3 describes the initialization.
Lemma A.3. At the initialization, for vjj {=1),r
[m], [n], we have
(0) j.r v\ ((
((oo),
kw (0)
E.) 0sst1220000 00 aga) ( R) [0 = =Ō(g)),
Proof. By Definition 3.2 we have
d.r NO, La)-
For data (x;, y/) D, by Definition 3.1
let B, =
we have
(1,0,...,.)),
[|[ = -ayi
& [k] ~ NO, of)
(for E B,)
Therefore, condition on the training dataset S, we have
(w j.r v) ~ N()
By standard Gaussian tails, we get
(0) v\ 6(oo))
[k] 6(Go).
jr
For (w) j,r E.). we have
(w j.r E.) = -aaw
(0)
[1]
Fi +
W.r [ [M).
(0)
ReB.
21

---

where
[. ~ NNO,,
5. Ei(K] N 0,00 (E sIKIP) ))
KEB.
KEB.
~ NO, (0 20,00)
since [i(k] ~N(00, of). Therefore, we have
(0) (w &.) 225 0 aoo) -(5 o,,0).
A.2 Preliminary Lemmas
The following lemma studies non-overlap support property of noise patch in data model D in
Definition 3
Lemma A.4 oven-overlap support, Lemma C.I in Zou et al (023b)). Let {(x,, 4/ i be
the training dataset sampled according to Definition 3.7 Moreover let B, supp {.) {1} be the
support of x4 except the first coordinate. Then with probability at least 1 n B, n B3 = 0 for all i,j [n].
A.3
Gradients and Updates
We first calculate the gradient of the individual loss function I2 with respect w
Lemma A.5. Consider the CNN model defined in Eq. 3.2 Let (x,, y4 be a data generated from data
(W.
model D in Definition 3.1 The gradient of T. (W) log Eje(-1a) Sww with respect War is:
V Li (W) Ww,r Sy..F (, x;) -zji
- (yelg,o' ((wjr. y(v)) v + ((W))))))) Ei),
where Eax y== - logit, (F:x:) and logit, (F, Exe(- 1.1 eki..
Proof of Lemma
A.5 For j Vi+
Ww,, (W)
xw,, ee (Wx)) (WWx)) (Wx). Vw,, FWW, (W + (W.)))
(W.))
(W))
Ww,. F(W, (.()))) + (Wx)
(I- x.))
Vw,, W.,,,.. F,W,
(F)( + eF-)(Wxx.)
- F,(,, x:) bgi
-yilg,o Suull,o'((WW.r)) ((W viV)) + (3,00((W Ei) Ex)
For j # Vi.
(WW) e (w) + +e (PF (L( Pw...F(WW:
(W.x(W
(V) + (W.x))
(()))))
FyW, F.(W.)) + (W.)
= - F, (0-
(W.))
(F..(x))) ( FF
22

---

Ww,, (,(,,, (.,.).))..... Eg, (u1e),o'((W3.. (WWrr) v + fg,00 (W),.:E)))) Ei)
Based on the definition of Ejá we have useful lemma follow:
Lemma A.6. Given data (xi, yi) generated from data model 3.1 define b -l,git (P, x;)
and , = Eve(-1.1) (W we have
Proof of Lemma A.6 For j = yi-
sgn(yils,)) sgn(j) = sgn(
sgn(1 - logit, x:))
(logit, (0, 1))
1.
For # yt.
sgn((eez,) sgn(j) sgn(L),s) = sgn(- (ologott, , x:))
sgn(logit, y x:))
(F, E (0, 1))
Now we calculate the gradient of loss 33.1 and loss responding to stochastic Adam and stochastic AdamW respectively. Here we slightly abuse the notation. We use g represent the (t)
stochastic gradient with respect to w!!)) at the t-th iteration. The subscript of Pij.r represents the
batch the t-th iteration and the superscript of g...rr represents the weight matrix W(t) at the t-th iteration.
Lemma A.7 (Gradient of Stochastic Adam). Consider the CNN model in Definition 3.2 Let
be the training dataset generated from data model in Definition 3.1 Using stochastic Adam to train the neural network, at the 1-th iteration with batch data index set I, of size B, the
(xx,y:))E= y:
stochastic gradient of the loss 3.1 with respect to w (8)
jar is as follows:
1 (m v
Σ
C. (' (i))))))))))))))))))) & + +w
LiEI,
SEI
More specific, for the k-th coordinate, we have
k
92..
[4]
1
Σ
yr 1 ((w (())) a
Σ 31 d' Eil) + +xw
B
LiEZ
IEI
EE B. i
gtg.r [k|
of jar (.) EdAA] + AW
B (0)
. k + and & & Uiez By:
Ii.j.r [k] (91))
23

---

Proof of Lemma A.7
The loss of stochastic Adam at the t-th iteration with batch data index set I, is
L(W())) = 1 Σ (.((()))
1
+ A w
B
VEI,
By Lemma
A.5
we have
(t)
10 L(W(1))
98.j.
1
Wit
La(((()))) + Aw (0)
B
w
3.r
Uie 0' ((w Jut
Jor Biv)) v Σ
(10)
3it o'
((wl((wl))
(2)
J.r Si)) &
+ +w j.r
B
LiEI,
EEI
For the k-th coordinate, if k 1, we know v v= =1,,0,..... O0T and E: -oys.. So we have
(2) - [1]
Σ
Vi elt) ((w J.r vvv1) a Σ vill o' ((w(()))) jor Eii)
Aw (2) [1]
B
LiEZ
EEI.
If k E B, E I, by Lemma A.4
we know B; n3 = @ for i + j. So we have
[k]
= B
(0) nu
(
(t) fir: Ei) & [k) + Aw 31
If k + 1 and ls & Usez, B, it is obvious that
g [k]
= Aw [b]]..
Lemma A.8 (Gradient of Stochastic AdamW). Consider the CNN model defined in Eq. 3.2 Let
{(x;, : E [n]} be the training dataset generated from data model Use stochastic Adam W training the neural network, at the t-th iteration with batch data index set Ia of size B, the stochastic
gradient of the Loss defined in Eq. 3.2 with respect to w
(1) jor is as follows:
(4) St.grr
1
Σ
veeff of o' ((w Ulv)
Σ
2 C0 Ex))
B
LiET.
More specific, for the k-th coordinate, we have
k ===1 1:
(0) 11
1
Σ yiell ((w
j.m B(v))
Σ
snell
fir E.1))
B
SEI.
E B. i
[. =
( [k]
1 (1)
o
(( Jir (:) [.]...
B
k + and k & B,;
(0) 9 [k] 0.
Proof of Lemma
A.8 The loss of stochastic AdamW the t-th iteration with batch data index set Ie
L(W())) Σ L, (W(!)).
TEI.
24

---

By Lemma A.5
we have
(e)
V
w ((Wl))
Wo (i(())
Yil 2 3 5, 0' ((w ((w
j.r Biv)) vvv)
.v+
- Σ
p o'
((w))
Jirt E:) E
B
iET
For the k-thh coordinate, if
1, know [1,0,...,,0 and [: -tfi So have
1
vel ye ('
(1) j.r. a Σ vee d' (( (t) Ex)
B
LiET,
iET,
If k E B,, E I, by Lemma
A.4
we know B, n B3 10 for + j. So we have
1 p() 3.1 o' o'
(1) Ei) [A].
Pt.. =
If l # and le & Uiez B,, it is obvious that
9 [k] =
B Proof Sketch
In this section, we mainly outline the proof sketch for the main results in Section Following the two-stage analysis framework Cao et al. (2022): Zou et al. (2023b), we decompose the proof into two distinct stages:
Stage I: Pattern Learning, During the initial phase of training, the effect of regularization is negligible. The model operates an underfitting regime, where it rapidly learns dominant patterns in the training data, leading to improved empirical performance on test error.
Stage II: Regularization. As training progresses, the model's classification accuracy the training set approaches convergence, resulting in diminished gradient magnitudes. At this stage, regularization dominates the optimization dynamics, steering the model converge local minima. Due the nonconvex nature of the loss landscape, the model retains the patterns acquired during the pattern learning stage.
Furthermore, motivated by the behavioral similarity between Adam and SignGD when the learning rate is sufficiently small or B1, 82 approach zero (Balles and Hennig 2018: Bernstein al. (018). we present results for SignSGD and SignSGDW (SignSGD with decoupled weight decay). We subsequently extend these results to stochastic Adam and AdamW, which provided in Appendix C
The update rules for SignSGD are given as follows:
(SignSGD)
(t+1) =
J.r 1) sgn(( (g...).
(B.1)
jr
where g:g.r in B.I is stochastic gradient of 33.1 The updata rules for SignSGDW are given as follows:
(SignSGDW)
(t+1)
Xn) J.P 1 sgn(g).).))
(B.2)
j.r
where gi.j.r B.2 is stochastic gradient of 3.2 D. and À is the weight decay parameter.
Next. following the framework of feature learning (Allen-Zhu and Li 2020 Cao et al. 2022 Zou al. 2023b Han et al. 2025a), we primarily focus on two key quantities: I) Feature Learning jv): (v).:jvv: T his term captures the alignment between the learned weight vector Wj.r and the true feature direction jv, reflecting the model's ability extract meaningful latent structures from the data,
2) Noise Memorization (W (w..r:&i) This term measures the correlation between w wrrr and the noise patch &, individual samples, characterizing the extent to which the model overfits to stochastic perturbations or idiosyncrasies in the training dataset. This decomposition allows to separately analyze the model's generalization behavior (driven by feature learning) and its memorization capacity (influenced by noise fitting).
25

---

B.1 Proof Sketch for Stochastic Adam
(2) We present the dynamics of feature learning (w [v) and noise memorization (Warrr E.) for SignSGD as follows. The details of calculation are provided in Appendix A
(++1) jr .iv)
(w Jun 3 -11 (s8 sggu
(1.j.r) .iv)
(t)
-w//2 (anniv) jv) + +jn sgm
Σ 1 2.
('w(( B(v)) ao' ao'(( (())) BAW (B.3)
((+ oour ;5:)
(we Berr Es)
(2) Pobierr .E:)
sgn
= (w(t) Merrs 5
+
Sgn
(10
Ui.o Vi.00 ((wl))
Meir (.)) - ndw
weir (bl)) 5 {x(k)
KEB
Σ
wr (o((w thor U/v)) ao aow((( yerr {n))]
BXw 11. (B.4)
at/: sgn
JET
where (1) = - logit,(( (F,x;)
and logit,(F,xx) Cx
(F)(W.x))
B.1.1 Proof Sketch for Theorem 4.1
In this section, we present the proof sketch for Theorem 4.1 We consider B O which the large-batch setting.
Lemma B.1. Given the truining dataset S, if - (()) 1/poly(d) and 0 < A = =0 20p/n).
then any S To with To Ō (3 and any i [n].
(w.r (++1) .jv) S (W O()), (w (w/ Ex) (w {.) + ((ssop))
We note that the above lemma is equivalent to Lemma 5.2 in (Zou et al. 2023b), which enables us to directly extend their results for full-batch Adam the large-batch regime. The remainder of the proof proceeds in the same way, as the core theoretical framework remains invariant under the batch size scaling 1/2 (1). Recall that the condition So, (11) implies that noise memorization outpaces feature learning and we have fj.r @(1) throughout Stage since the outputs are small. As result, after a certain number of iterations, the direction of feature learning is reversed, as indicated by the update rule in Equation
B.3 Specifically, the noise-driven term aa' ((w ca''(ww j.r {.) becomes dominant, Ex) d'((w j.r. B.v)) + n2 n (f [1]] By the end of Stage I, the model's s feature learning direction has been inverted, while noise memorization reaches a quusi-stationna state. In the subsequent regularization phase, weight decay drives the model toward convergence. However, lacks the capacity to eliminate the memorized noise. Consequently. the model fits the feature noise -vyv and converges to a local minimum that preserves the patterns acquired in Stage ultimately leading to poor generalization performance.
satisfying ao' ((w (t)
(1)
j.r
We observe that under large-batch setting, the optimization dynamics of Adam closely resemble those of the full-batch setting. This similarity arises because the algorithm traverses the entire dataset within few iterations, resulting in nearly identical momentum estimates and, consequently, comparable training dynamics between large-batch and full-batch regimes.
We next consider the mini-batch setting, which yields conclusions that differ fundamentally from those in the large-batch setting.
B.1.2 Proof Sketch for Theorem 4.2
In this section, we present the proof sketch for Theorem 4.2 the mini-batch setting.
We consider
n
Ollgg -1), which is
B
26

---

Lemma B.2 (Stage 1). Given the training dataset S. if
z @ "agee")). )) 1 = 1/poly(d) and and any i [n]]
olog (o o,//),, then for any S To with To 0 my
(w frr
((+1)) #)) .. v)
(1: (W jrr
(w j,rt Ei) s (1pso,)) 3 v) + (n B
n
Compared to Lemma B. Lemma B.2 reveals fundamentally different optimization dynamics in the mini-batch regime. In Stage I, feature learning advances memorization remains at initialization scale. This divergence arises because mini-batch sampling requires many more iterations to traverse the dataset: dense, shared features receive consistent gradient updates and resist weight decay, whereas sparse, uncorrelated noise is continuously attenuated. As features strengthen, network outputs grow, the loss gradients diminish, and weight decay takes over, marking the transition into Stage II. We now show that the structures acquired in Stage persist throughout this regularization phase.
for (11)-wiile noise
Lemma B.3 (Stage II). Suppose the same conditions hold as in Lemma B.2 For t t> To, J E
(+1), E [m], i E [n]], let r* argmax,e(m) (w) (t) (9) WHiT E.) (Issop))
jor jv), then (w ja j.r:: jv) = O(1) and
This lemma follows because, once feature learning has increased accuracy and reduced gradients, weight decay takes effect but cannot reverse the established feature alignment; finally the model converges to a local minimum that preserves the patterns learned in Stage I.
Lemma B.4 (Convergence). Suppose the same conditions hold in Lemma B.2 and size 1j = Od ), then for any i,
B.3 if the step
E ((W(+)))-((W') s <=
Combining Lemma B.3 and B.4 we observe that the model successfully learns the true features and eventually converges local minimum with infinitesimal learning rate 7) and = poly(n)/n retaining strong generalization performance.
B.1.3 Proof Sketch for Corollary 43
We next show that if the weight decay parameter satisfies À \= w== then the learning dynamics of Adam effectively suppressed. This implies that the effective weight decay for Adam is of the order - which is significantly smaller than that required for AdamW, will be discussed later.
Corollary formalizes this observation by showing that if the Adam weight decay parameter satisfies = w(ag-²). the training process becomes stagnant and remains near the initialization. This corollary follows directly from the proof sketches of both harge-batch and mini-batch Adam. By Lemma
A.3 we know that at initialization:
[(w = ((ao).
kw j.rt & 0(s 12 12000 + aoo) = =e(s ((po))))o0o)
J.r [k] = ((ao).
Then, from the update rules given in Equations (B.3 and (B.4. we observe that the updates are dominated by the weight decay term, i.e.,
A/w
(0) []] a' ((w) jir yev)) + ad a ((w)) E.)),
and
A/w (0) [k] o' ((w (0 (.)) [.),,,
due the condition À 3) As result, the learning dynamics are overwhelmed by the regularization term, preventing meaningful updates. Consequently, the model parameters remain close to their initialization throughout training. We formalize this Lemma B.5
Lemma B.5. Suppose the same conditions hold as in Lemma
and
if A w(o 2), then
1 v) < Ō(og).
E2)
S 0(5'122,00).
27

---

B.2 Proof Sketch for Stochastic AdamW
Motivated by the similarity between Adam and SignSGD, and to better illustrate the core idea, we present the dynamics of feature learning (.r+JV) and noise memorization (w Ei) under SignSGDW. The detailed derivation of update formula is deferred Appendix A However, we emphasize that there are key differences between AdamW W and SignSGDW.
In SignSGDW, due to the presence of the sign operator, the weight decay affects (w j.rivv only after it grows beyond certain threshold, and similarly for ( In contrast, for AdamW, (6) weight decay becomes effective once (.r: jv) or (w))) E.) reaches a level where the gradient magnitudes become sufficiently small. At this point, the update is normalized by the stability constant
(t) c and dominated by the weight decay term, which causes both (w [.:::vv and (warr Ei) to cease increasing. Besides, as the lemmas in this section are simplified instances of those presented in Section C.2
we omit them for brevity. For more details, refer Section C.2
(((t)) .iv) jr
= - An))
-1-
w (.m:.v) -1 (sgu (gl (hiser) .5v)
(1)
= (1 - Xn)
(w/)) (. jv) + im sgn
Σ
lo' (ww (r. an uod((w <i>)
(B.5)
NEI
(1+1)
yirr 5.
(1 - An)
(1)
6 -17- sggn
Piy.rr EEi
(1)
Ve.r
(1 Xn))
(t)
+ : E:
sgn (.. (oo'((wll yirr (.)))())))))))))) .E((k)
Vi.ro
kEB
n Σ
03y.7 sgn
yi eell
widd
O ('(wll)) yerrs v(v)) - a(w(( wir {1)1
(B.6)
where ((0) = -
-o:-looge logit, (F,x;) and logit, , Exe(-ran cFei
BBy (W.))
Moreover, should note that the duration of Stage I under SignSGDW differs markedly from that under SignSGD, owing to the decoupled weight decay mechanism. During Stage I, model parameters grow unchecked by gradient-based regularization, allowing features accumulate strength until the decoupled weight decay term begins to exert significant influence. Once this threshold is reached, training transitions into Stage II, in which weight decay counteracts further parameter growth and stabilizes the weight norms.
B.2.1 Proof Sketch for Theorem 4.4
In this section, present the proof sketch for Theorem 4.4 We consider = (1) or o(sap). which is the latch-batch setting.
The following Lemma B.6 characterizes the duration of Stage I in the large-batch SignSGDW setting and provides upper bounds on feature learning and noise memorization.
Lemma B.6 (Stage I, pattern learning). Given the training dataset S, if 7 O(1) or #/ = o(sop). (d)), A = δ B² A 1) and X = Ö(1), then for any t S To with To Ō
B
Xnm
((++)) 1)
. v) S (w (w!!:²) .j v) + O n
(e:4)
(w jr
B
(w (0+1)- 8) .Ei) = (w {2) + (1)op))
(1.
Since in the latch-batch regime B lates faster than feature learning. At the beginning of Stage I, feature gradients dominate because
n of p(8ap), Lemma B.6 implies that noise memorization accumu-
28

---

d'(W B(v)) > ao'((w)) yu.r, E.)), given a o(1) and negligible weight decay influence. After certain number of epochs, the noise term grows until ax ao'((w g. > a' d' (ir..r yvv)), at which
(0)
point feature learning reverses and eventually flips direction. Lemma B. below provides precise description of this transition.
Lemma B.7 (Stage I, fitting feature noise). Suppose the same conditions hold as in Lemma B.6 if
2 θ (1 (a,p)= then for any L E [Te, To] with T, 0 n00 17
S To.
(1: ((+1) B) .. v) = (w jr j v) (11
Y ))
(Wir jr
and at epoch To. we have (a) w (To [1] sg)(() ()) (b) w
(oo
( [R]] sgn(E,(R:)
jr
jrr
Q
B BA
or + OO)) for k E B, with yi j: (c) w
(To 8) [k] = =Ō()) otherwise.
j.r
Lemma B.7 implies that, by the end of Stage 1, the model has fitted the training noise. The
following Lemma shows that these pattern persist throughout Stage II, ultimately leading to poor generalization.
Lemma B.8 (Stage II. preserve the noise). Suppose the same conditions hold as in Lemma B. B.. B.7 for 1 > To, j E {+13, r E [m], i E [n]] lel r" = argmax,e(m) ( (8) yorro Eed. then {w (t)
jrr j v)
(((/AA) and (w!... "irr+
6( ((axe))
The following Lemma B.9 prove the convergence under certain conditions.
Lemma B.9 (Convergence). Suppose the same conditions hold as in Lemma B.6 B.7 and B.8 if the step size satisfies 1j = 1/2), thenffor any t, L(W())
< -
E
L( (+ L(W(++1)
<- + =()
we observe that, with an infinitesimal learning rate 7 and T poly(n)/n, the model ultimately fits the feature noise and converges to a local minimum, resulting in poor generalization performance.
Combining Lemmas B.8 and
B.9
B.2.2 Proof Sketch for Theorem
In this section, we present the proof sketch for Theorem 4.5 and n = w(srp), which is the mini- mini-batch setting.
We consider 21 O(n v log E 11)
The following Lemma B. 1.100 characterizes the duration of Stage I in the mini-batch SignSGDW setting and provides upper bounds on feature learning and noise memorization.
Lemma B.10 (Stage I). Given the training dataset S, if + Z ( log and मे w(so)).
77 l/poly(d), A σ B² A 1) and X Ö(1), then for any t S To with To Ö B
Xmm
(1+1)))) .jv) = (w) (1-H) jjv) + O "", (w)) j.r
((++)) (Weer
(+*#) W
E.
E.) + ((ssop).
In the mini-batch regime, b = w<) so feature learning outpaces noise memorization-unlik in the large-batch case. Consequently, noise cannot reverse the feature learning; instead, features are learned continuously until decoupled weight decay intervenes. Noise memorization also grows until this point, but because noise is both sparse and independent, accrues only during a few iterations per epoch and is concurrently suppressed by weight decay. Hence, both feature learning and noise memorization reach their peak at the end of Stage I, after which weight decay governs Stage II. The next Lemma B. formalizes this behavior.
Lemma B.11 (Stage II). Suppose the same conditions hold in Lemma B.10 for To, j {+1}, E [m]. E [n], let r- = argmax,e(m (0) jv),,,, then (wl()) (v. = G(1/)) and (w & < O
(adgy) nd
The following lemma establishes convergence under the specified conditions.
Lemma B.12 (Convergence). Suppose the same conditions hold as in Lemma B.10 and B.ll 11 if the step size satisfies 1) = Odd-1/2), then.for any t, I(W(1)] < L- 5(r2d).
E [[(((+)))) (())))))-(((())
29

---

B.2.3
Proof Sketch for Corollary 4.6
By Conditions
A.1 and A.2 along with Definition
3.2
we know that d = poly(n), and hence
-- θ
with 2 3.
d//-2)/4
This directly implies that the effective weight decay parameter for Adam satisfies
{
B²
Addmm -og²
min
n 111
Adamm.
This completes the proof.
C
Proofs
First we give a general upper bound of the moving average in stochastic Adam and stochastic AdamW. (t)
(8) Lemma C.I. Let j be the first momentun estimate, ji be the second momentum estimate at the t-th iterate in the update rule of stochastic Adam stochastic AdamW. Then for all j E {=1}, r E [m] and & & [d], if 32 > 8i 31, B2 E [0, 1), we have
20()).
+e
Proof of Lemma
C.1
Let us expand the moment estimates
(9) [k] = Bmm J.r
(t-1) [k] + (1 B1) [t.j.r
Σ BT BT
B1) 9 - [k]:
[ = B2
(t-1) [b]] + B22 Di... [k]
t-1
Σ (1 32)) 9 ( [k]².
Let {= = --3-0 [810--1)] P)(1-Pa)
be convergent series since 32 > BT. Then we have
jur [k]
2
BT - (1) -: [k]
2
8T B1)
(1-7) [K]]
2r
[B] (1 - B1)1
((--7) [k|2
Σ 2
z7
B7 (1 B2) Ir - I 14/2 O(1)
(()), where the first incquality we use Cauchy-Schwarr inequality and the third equality we use the fact (1-8)_)) is convergent series. So we have
that 23
9511 3g)
(4)
[k] j.r
m (2)
< (().
[k]] +
v 19
30

---

C.1 Proof of Stochastic Adam
First, we try to approximate the update stochastic Adam to sign update since the similar perfor- mance between Adam and SignGD (Bernstein et al. 2018 Balles and Hennig 2018 Zou et al. 2023b Xie and Li 2024, Li al., 2025).
Lemma C.2. Consider the update of stochastic Adam in (3.5. Let W be the weight the t-th iteration. Suppose that (w jor yov), (w j.r E.) ((1) for all j E {=1}, r E [m]. i E [n] and
8 < 32. We have the approximate update rule for each coordinate weight as follows:
For k- = 1, we have either 10 [1]] & ((n) or
sgn (0. [k))) (())
10 [k] +
For every k B,, E It-Ti T Ta = (ro tro i represents the number of iterations away from the current iteration coordinate is affected by &: sampled at the iteration To since the moving average, and we define F = ((log(An)
In/B 1
To B). where To
If h (1), for any To B. we have either
Di- [k]] <ē = insoo 16 + /AW (0(1))
or
m [[[[k]] sgn f (1 O().
[.211 +
; 2 O(log(An)) 6(1), for Tu (1) such that 3T0 (1),, we have either
mo.iir [k]]
< B 175 (6(t) + 1Aw (1) [+]]
or
m [0
sgn (1) (1 O()). 3.00 +
For To 2 O(log(An)-)) such that 3 7
< ((1)), we have either
w [11k1kk] s 20(1)
or
m.r [
sgn (w)))))))) 1(1)
G. +e
For the remaining coordinates + and & B,, i IT- To E [0) U (7), where To represents the number of iterations away from the current iteration I, coordinate k. is affected by &, sampled at the iteration To since the moving average, and we define 7 = log(>n) -1, Then we have either
19.09 19 [k]] < (Arr)
(0) m 200
[k]
sgn (1)).111) O(1).
+E
31

---

Proof of Lemma
C.2
Let first focus on the first momentum estimate,
[
3 3,
(t-1)
[k] + (1 - B1) 1.3
Σ 81 -B1) -
Σ (1 Bi) o [M] + Σ
8T - 1-)33:)).. 91--7.10 [M]
+=
(I - Bi) - [+] + Ö(A))
0
where the last equality we select 7 = (log (n))
such that Eier+1 BT - B1) = -()))) and
lo(--3)) [+]] = Õ(1) for all l E [d] by Lemma A.7 since the facts that (w v). (w (w/ E.) = j.r
j.ro
O(1). Similarly, for the second momentum estimate,
5.r [
Σ
35 32) -
Σ
85 (1 32) - [k]]
Σ
B5 (22
( - [2]]
Σ 81 32) s [k)² x = Ö(An).
Here we use the same 7 because we can always reselect i of smaller one to larger one, and the absolute value of the tail will not increase. Then we have
(9 [
E1 (T - B1) -2 (L-T) [b]] + Ō(Ar)
(C.1)
++
35 3 32) (1- [K|2 + Ö(Ar)
since (Ar)). Now we want to use sign update to approximate (C.I). First, we should note that once the signs of - [k] for [0, F] aligned, (C.I) can be approximated as sgn(g [A]) -(1). since
85 32) -
32 (1 B2) | [k]
T=0
35 B2) 1
33 (1 11--30 1,)+++++ |
7
Recall the gradient of stochastic Adam given Lemma A.7 we want approximate - [k]] to (())) g [k], such that we can use the current weight to approximate sign update. By Lemma C.1 the
upper bound of each coordinate in one step is (()). Then for + E [t - 7,4). we have
ww (t) BeV - (w wwff..uv vvv)
(k) (w( (k+1) .yiv) ( (w krr Bvv)
< (17T).
(C.2)
Similarly, we have
(ww() Es) - (w(T) Ei) S O(1Fsoap),
(C.3)
32

---

[2 [. - ((17).
(C.4)
Then recall the predict function
F,(W x2) Σ 0 (ww Sir vir)) + a (i.....
We have
B3 -x)- F,(Wi, xx.|
(1w (vv) - (1w557 J.r miv))
-
+ Σ> 0 (ww
rm1
$ @(1) , (w (Ev) - (w/77 vivi)
+ Σ 0(1)
m
Ei) - - -(w/?? Es)
- (mmy pap) + ((m))))
= (17Tsop);
(C.5)
where the second inequality we use the convexity of (-) and the facts that (w (2) ( joro 1,) = 6(1) The last inequality we use m - @(1) and sap = w(1).
J.r Mov) = @(1) and
Then we can approximate
g(T) pl)) in the gradient A.7
j.r jir
(-,(W( x1)
ExE(-1.1) (W(()))
F (x),,,))
(x + (W
= sgn(e()) D
(y,- = j)
(W(())))
2
EKET-LI) (W(O))))
sgn(e))) 0((e))))
(yi + i)
where we use the fact that @(nFso,, = o(1) and C.5 So we have
D
sgn
for all E [t-7, 7, f). Further, by Ö(1), recall (z)) = max (,, 2, we have
and the facts that (( f. vvv) = Ö(1) and Kw Ez)
(w
277 of (T) uv))
< 2 of d'
for v(v)) + 15771 ((FF) -(())
(1)
=gg((())) e((())) o' d''(w
j.m vv)) + + -((F),
2 (w (( uv))
2 247 of fr v(v)
((w))).)))) - 1657711 -((7))
sgn(())) e((())) (( vv)) -v0)))((())))))) -((F).
33

---

So we conclude that
o' ( (( vvv) sgn(()))) - o' ((w = +(0))) @(7F).
(C.6)
Similarly, we have
o'
(())) (.) sgn(
e d'((W)
(0) Jar {.) + ((e))))) (n)
(C.7)
((-) Now, have all the tools need to approximate It- - [k] to - [k]. Recall Lemma
A.7
substitute C.. (C.6) and (C.7 into ( [k},
we have
For =1, 1,
9 9
((-T)
θ (1) = (1 Σ 14%1)
0(1)) = O
(C.8)
1ET
For all k E B, E II-T.
g [k] = 1 (8 [bl) + O
Ö( + O(AnF).
(C.9)
For + and 1o & B, E I+-T.
b [k] = θ (
(1) (f + (Am)).
(C.10)
Plugging C.8). C.9 and (C.10) into CC.I) with facts that j = 6(1). x o(1), ((1)) (.)
6.)
= ē(1).
(w (t) Jm vvv)l = (11. 1 = ( c ((2)) and Lemma A.6 we have
For k = 1,
m j.r []
(+)
j.r +
Σr=o B[ B1) -
(t-r) [1] + Ö()))
Et= 85(1 (2) - [1]2 + Ō(An)
EF BT (1)) e ( (11) 0 +o2 EieT.. (511) Ö(TF) + =((A))) =((A))
1 El=o 35(1 - 32) (0 ( [11) 0 0 0 EiEn., (0921)
Ō(17) = = + Ō()m)
ΣF (T B1) ( 0 9 [1] Ō(1F) + +((2n)
Σl 35 32) (0 ( 0(m))) Ō(()))
(019.11 + ((T)) + Ō(An)
(11)).11)) 6(m))) 2
+ Ō(Am)
( + ((n)
o (()) + θ())
34

---

For k E B, i E To E [F), where to represents the number of iterations away from the current iteration coordinate k is affected by & sampled at the iteration t To since the moving average. We note that if the number of iteration in one epoch u. is less than 7. the moving average B
will use some sample multiply times. We denote Tx == tio B : E [0}U] n 10--1pp 11. S g) the timestamp set involved using noise & (i.e., i - for any T&), and Bi, To S B
If B F, in this case we have To = {т0} for any l E Bi, and & was used in iteration 70- Then we have
(4)
Jir [| +E
Ei= 87 - B1) -1- 9t- [k: + Ō(v))
V VE== 35(1 (2) - (1-7) [K]2 = =((A))
ΣFET. BT(1 -Bi) 9r- r.j. [k]] + +ErE (1
(1-T)
- B1) 1 + Ō(Am)
V Σt (1 B2) -r- - [b]] + Ö(A))
EFET. BT 7o [k]]) + (0(1) -IEET 37) - + Ō(A))
Er+ 13 - (B2) -h-= + Ö(vr)
3/ [k]]
+ θ
Ö((Tsop) + (AmT)) @(A (..11111)) + +
Eu= 35 (1 B2) (1-T)
- |1|2 + Ō(Ar) (2)
where the third equality we denote (t-T) [b] = O(o J. [k]] + @(AnF) = @(Aw]]]16) + ((2TF) as for E (0) U [+]. For the denominator, we have
Σ 35 32) (1--) [R]2
0=0
+
Σ
B (1 32) | [k]]
(((t)
3
(
96 1 [ki] + +
(t)
Ō (1mpsa,) + ((w))) O([Aw + +(())
Σ
35(1 1 32) -/ [k]]2
T=0
B3 (1 32) -
(e(t)
1
[kt]) +
Ö(1mss0,))) + 0(77)) - 6((AW (111)) + +0(0))
B
So we have
m kk]
(4)
[. [ +e
870
@ (1/1 (1 [b]
+ θ
n80p le!t) B
+ (((n)) (Aw - (0 + @(An)
(0(1ef) [6ll)
+ t
5(()) - O(IA (((Am(11111)) @(1Aw (o1]] + O(An)
13
35

---

where we use the fact that = =((n)). Now we handle 310 and 32"
with more care. First we have
$1 32 (c) 1o.. [k]| (1) and To
n B
Then if / = (1), we have 37 = 32 = (1), then
m [
[k] [k] +E
yoap (e(t))
θ (1 + θ
+ ((A)) + ((AW
10 mo [kll]
meople
((An) = A(IAW [2])
13
6 (
1 c(0
(( jirt {.)) E((A]
-
yeap 1
= = = ((A (Aww)..01)) [+]
B ji
B
θ (1
B 3 (
1 p(t)
((w
j.rr (, [k]
(0)
7/80pl
(
((h)) + @(IAW [.(]
B
If g 2 ((log(An) = 5(1) such that B2 < O(An), then for mp O(1)
m m
[0 [[]]+++]] +
e ( (- B j o' ((w j.rt E.) s(())
d(0
+ 0 180 (c())
+ (((n) ((w j.r (B)(
B
θ (1 B _ell 0' ((w jims E.) E [k:]
(0)
+ θ
yeo 1
- ((1n)) + ([12
B
For To < 12
B. To 2 O(log(An) - such that 332 < (Ar)), we have
[K]]
@(Aw (2) [k) + ((An)
6([Aw
3.r [+]] + ((A))
1
[k] +e
Ö(1), Xy o1) and nsop nsop= o()) We claim that the intersection (gap) of ((log(An) ) ) and @(1) very small for O(log()n)~)). that is, considering the intersection part c <log(An) (n)) for sufficiently small constant >> 0, the impact of the intersection (gap) is very small, since e log(An)) = o(log(An)-)) Therefore, for most of To
since _0 1 [K]|
m
(4)
@(Aw
1. [k]) + ((n))
1000 [k] +E O(IAw
4 5.r [ (Aan)
For l: * and lo & B E I E [0, F). recall (2
g [k] = Xw (2) [k), we have
(2
[ +e
oF - 31) 9 - [k] + Ō(An)
1 V 83 - (2) - [k]] + Ō()))
((-T)
Σt [[ - 31) (o g/th ( + O(ArF)) + OA())
Dimo o 83 - B2) @
gr [. [b]] + =(()))))
(t)
+ O(())
36

---

(41 6(A))
((0 ((A))
((w = O())
((w j.r [+]] ((1)
This completes the proof.
C.1.1 Proof of Theorem
Lemma C.3 (Stage D). Given the training dutaset S, if B @ = 1/poly(d) and
olog- o op)n), then for any < To with To 0 nooo and any i E [n].
(w jrr
((+1) ..v S (w j.r: + (()),
(wit+!) E.)
(w-(t) Birr Ex) @(ssap))
Proof of Lemma
C.3
By Lemma C.I we have
m (E)
(++1) (w) j.r
.jv) S (w j.n: jv) 17
jv
3ir +
< (w/ jv) + (
Then, prove (witt
(++1) .E)) (W W E() + θ @(nsor)) by induction. By Lemma A.3
we have
(w jar v) = ((oo), ( &.) = 0/2 Oppo)) w j.r [k] @(oo).
which imply that 1 (1). Additionally, we have 178 o(0g and 0 < A ol0g 9-2 0p/ pp/n).
For a sufficiently large fraction of k B; (e.g., 0.99), we have B B-1 e [+]] 2 On 80, (000
A/ A/w [/w![[]]] [k]) for i E I.. Therefore, by Lemma
C.2 and A.6 we have
sgn
(0)
Mhes o ((w
10 E.) [kk]]
sgn
(0)
Be.r [xk]
sgn & [k]).
(C.11)
Recall (1)) By Lemma C.2 we have the following update according to B.4 Lemma
(C.1 and
(w (1) ve.r: Ei)
(U) m (0)
(ww(0) vorr Ex) -1
j.r
j.v
j.r
2 (w (0)
Birrs {i) + O
(gg Ei) - O(ja) (118,,
keB,
(w (0)
Wers t + p))8op))
For general
Top assuming (. { 2 (w (t-1) {i) + (()sapp) Then we have
(w()) (1)
anrr Ev) = (w Ex) (I1sa,,) @ 0,00 + (1180)) S O(1).
By Lemma C.I we have
[b]] < (w [+]| O()) < Ō(aa bn) < @(1)
37

---

So we have 15,7 (1). Besides, we still establish the condition (w [s >
Ó(1 s0p 1/ + Al f [k] +t7l) since 0 < À =(oo0= op/n). Then we have (C.II for Follow
the same proof above, we have
(w (1+1) Mer
.E,)
(w Blrrs Ei) -1'
m jr
jv
Vv
2 (w) (t)
Werrs {t) + O())
Σ (E] - O(ja) - (1s50)
kEB,
(wl)) Werr b,) + (nso
The term ((nSop)) in the above inequality arises because, for coordinates that [[:[A]] < (( cannot exploit sign information. Instead, we directly apply Lemma C.I This completes the
proof.
Lemma coincides with Lemma A.3 of (Zou et al. 2023b). allowing to transfer their full-batch Adam analysis to the large-batch case under 7 = O(1) Therefore, the remaining proofs are omitted for brevity, they coineide with those in Zou al. 2023b).
C.1.2 Proof of Theorem 4.2
Lemma C.4 (Stagel, nearly zero noise memorization). Given the training dataset S, if (log ()), = 1/poly(d) and 0 < x olog (log)2a//n), then for any t S To with To = (
B 2
and any i [n]. suppose Wir Ui v) -O(o)), then
(. E, S @(Voopo0 + a)-
Proof of Lemma C.4
By Lemma
A.3
at initialization
@(V5 + aao) < @((3o,no))
since a = o(1). Lemma C.2
ensures that stochastic updates slow down noise memorization-alllvinn
(1) (w Ed) to grow for only
o(log(An))))
iterations after & is sampled while the remaining
o(log(An) > ollog(An))))
iterations, weight decay dominates. In particular, whenever |w (0) [b] 2 ①()) we have
w (1+1) w (t) j.r kk]
[k]
sgn (w) ( [k]) O()).
Concretely, if &; sampled at iteration ti the first epoch, then
(w (1+ {.) (m) E.) + +(nso,)) s @(Vāapon + +s7,) S @(VSopoo),
where we directly bound the update by ((1) according to Lemma C.I Over the next o(log(Nn)))) iterations the noise memorization (w!(.)
(c) j.r: E.) increases by at most
o(log(An) (mssa)),
and thereafter weight decay decreases it in each of the remaining o(log(An)~1) steps. Hence, we can calculate the maximum value of the noise memorization
(w
(n/
{i) I max ((VVs @(Vso,00 + + O(nso,)) (ollg(A))))))) B B.
ji
s <ã(VSopo0 + aa),
since n/ B ((log(An)) and [1 = -cy. This is true in every epoch. So we complete the proof.
D
38

---

Lemma C.5 (Stage I, feature learning). Given the training dataset S, B 2 (loge 1 1 1/poly(d) and 0 < x = ol08 2 op/n), then for any S To with To x(+) and E {+1}, we have
(w)) (1+))
jr .jv)
Jv) +
Proof of Lemma C.5
by Lemma A.7 we have
sgn (90.,
(0)
sgn
Σ
n o' ((w ( B(v))
(0)
Σ
vi o ((
(0) Ei) (Aw
IE
IEIn
Then with Lemma A.3 and facts that 1/0 (1)) by Lemma A.6 we have
Si o' (' (0)
j.r vvv)) = sgn(j) (0))
o (0) 0' d'(w jor 6,) sgn(j) O(a(( -
Aw (0) [1] = =o(oj 'o)).
Substituting them into jr with a = d1). 0, = Ö(1), X = oloo-= 2 2opp/), we get
(0)
sgn (0)
)= sgn(g
-snn(j -(00) = ss)
By Lemma C.2 and t) we have
(0)
(w (1) (w = (w (w)rriv) -
j
j.r +e
(w (wW + sgn(j) (11)
=(w g..jv (n))
Now suppose that the equality holds for iterations O,....f. f. Then (w (w/0 v) = Õ(1), (w.).. (nssa)) = O1). Therefore,
g(t) 3.4 = @(1) By Lemma C.4 we have
vel 10 aye j.r Ex) S sgn(j)
( ((w
j.r v() sgn(s) @((o + 1n)
(0)
0p00 + 1n)o-1),
(1) J.r [1] =o(0g= 20,(00 + tn))).
Substituting them into (c)
Iig.r- with o(1), s s1/20, = Ö(1), x = olog- 30-00ppp100/ 2op/n), we get
sgu(g)))r) -sg((j Ō(0g ssg((i).
By Lemma C.2 and 1 o(a we have
(w (1+1) .jv) (w (2) (. -1
m jir
.
(1)
Vir +e
(W Jv) +j.ssn(j) O(1) (w(() jv) + On)).
This completes the proof.
Lemma C.6 (Stage I, general dynamics). Given the training datasel S, if 3 2 O(log= " =
(dd and 0 < A (008-0pp//) then for any < To with To O and any i E [n].
(w (t+1) . v) (w (t) [ v) + O "
(w E,) < @(Veo aa).
39

---

Proof of Lemma C.6 We prove the claim by induction on E, using Lemma C.4 and C.5
At = 0, by
Lemma A.3 initialization we have
(0) (w Jor jv)l _ ((o))
and hence
(w.rr jvv - (aa))
(0)
Therefore Lemma holds at = 0. Suppose Lemma C.4 and C.5 for some f 2 0, then
(w (1+1) jvv 2 - (
by exactly the same proof used in the proof of Lemma C.5 This lower bound remains valid step t ++ 1. Consequently, Lemma C.4 continues to hold, and the induction carries through all iterations.
This completes the proof.
Lemma C.7 (Stage II). Given the training dataset S. if n Ollog e () n 1/poly(d) and
ola o(og-3a (oog)pp/n)) op/n), then for any 1. > To j E {=1], r E |m], [n], let ((1) and (w (1)
(6) a (w j.m jv), then (w]?)
(t)
j.r: Ei) S Onso,, + aa).
Proof of Lemmaa C.7
We begin by establishing the bound (w E.) I (nsoo a). According to Lemma C.4, during the first To epochs, the number of iterations in which weight decay dominates the update dynamics is at least
(0)
n2 (log(An))))) To (
In each such iteration, the contribution from weight decay is lower bounded by @(jso,)), leading to cumulative effect of
3 To
n (== (log(An)-1)) msso,) =
(nS0pp
B
This term is asymptotically larger than @(v i.e., nsop/ w( To). Therefore, conclude that over the first To epochs, the weight decay effectively suppresses noise memorization, ensuring that (w()) Ex) I j.r aayvv + ((nsop) @nsoo + aa) holds.
Next, we focus (w (1) j.) @(1) for r* argmax/e(m| (w.... By Lemma C.6 we
know (w (w/PP) .j v) @(1) and ir
(()) For > To. We show if (w (w).r. v) <
( log 1))= then for (x; vs) with y/ j.
c ,(w()xx) x.)
3.
Σ (-1(-111) x.)
1 exp
exp
E a ((w
o ((w jimo Ei)) a "(ww!) jv) ((w").r EX1)
2
Em ((w j.n
++000 exp
1 exp m (1 log ((2) 1 1))
=A)).
where the inequality we use (w Ex) < @(nsop + a) and n o(1), a a= o(1) Then we have
Σ nel o 3.
o (( v(v)) Σ
yeel (( (w Ei)) BAw59214)
sgn
IEL
IEI
sgn(j) > log ((2)-1 1)
+ (
log ((2) (0))
sgn
m
40

---

= sgn(j),
where we use Lemma A.6 (w) j.r E.) - @(sso, + a) and a o(1). So we have
(w :j v) (wf.... v) + (n))
If ((w))) (t)
jrm: v) log(A= :), then e o()), then for (x+, yi) with yi j.
eF-,(W(()) xx))
E F (W()).))
1 exp
E (v)) + o (W {.) o( ((w")))) o(w E.))
1
$
mml - a) (log(A ( }) 9
+ exp
1
(og(>>+5) )
exp
x
A
(log(>-+)) 9-2 +)))
Then we have
sgn
Σ
yi dit o ('o((((())
Jims Biv)) - a Σ
sielt)
(i)) (i))))
BAw
IET,
iEI,
sgn(j) (10g()-))) 9-1 -0 (
A
+ > Aog(A>2)
sgn
8-2 ((Oo((A
sgu(f) A (1og(A+t) + Alog(A= -xb)
sgn
sgn(-ww) (1(((1))
= (ggg(
So we have
(w ((+1) j vv (w()) (1)
j. v) O()).
Therefore, - v) 5(1) for > To This completes the proof.
Lemma C.8 (Convergence). Suppose the same conditions hold as in Lemma C.6 and C.7
if the step
size 1 = O(d ), then any t,
E [(W((++))
L(W())) < - -
We aim to prove the convergence of the objective function under the Adam optimization algorithm in non-convex setting, Recall the loss function for each data point i is
Proof of Lemma C.8
L (W) log (1+ exp (Fy, (W, xx) F-, -. (W. (W.,)))
41

---

where W represents the parameter matrix. is the imput data, ys is the true label, and F, (,W,,, are the logits for class j. The total objective is:
L(W)- Σ ==1
with x = ((1) as a small regularization parameter.
Since (W) is non-convex, we exploit its smoothness with respect to the logits [F, (W, x;) Specifically, (:W) is 1-smooth in [F,WW x;)lı due to the properties of the cross-entropy loss. Define:
A = F3 (W((())) x,) FJ ( x.).
Using the smoothness property, we apply a second-order Taylor- like expansion around W
(W(") S E (F,(WO) x:) AAF + E(SEE)
L(
(C.12)
This upper bound arises because the second derivative of L, with respect to the logits is bounded by 1, a standard result for cross-entropy loss. The logits are defined as: Fj(W(C x;) =
(l + ((. (1) E.)). where ( the r-th neuron in j-th output of (' or))
[E] a smooth activation function (e.g., with 2 3). By Lemma jari
C.7 and C.I we have
(w())
(1) v) I @(1)) and (w!y) (.) s ((1), ensuring the local smoothness of remains ((1) be- (1+1) yov) and (w (i vev) (similar for ( E.)). Then with Taylor expansion, we have tween (w (w((t)),
((w jr :vv)) ((
(e) Jam Bev)) - (V ((w (t) Jor v6v)).
(t+1) w (t) j.
jr
< <6(1) [w
(1+1)
(1)
w
j.r
2
(+)
= =(1)
m
).r
17
(2)
jor
< d))
(C.13)
where the last inequality we use Lemma C.I Similarly, we have
[a(wW)t (+1)) EE))
a j.r (.) (V Ww,. ((w (/ Ei)), 35r w
(t)
(1)
(1+1)
< <0(1)
(((t))
(4)
j.r
(t)
ē(1)
m
j.r
(t)
jrr
+
< =
(C.14)
Summing over r (with m = ((1)) we get
AF < VV F Fw((), x:), - + 6(72d).
(C.15)
Additionally, Ww (W(1)x;))lP - 5(1) since m Õ(1), (w/) (vv) s @(1), (w s (1). So we have
AF Fal (nso, + nja + + 122d) Önsoo + r7²d)
(C.16)
Substitute C.15) and
C.16 into C.12):
L (((++1)
L, (W()) < (WWL. _ -
(C.17)
42

---

For the full objective:
L(w - L(W())=
X w -
- L.(W)))]
(C.18)
Since A W is 2X-smooth and À o(). the regularization term contributes:
X(I W l| || 2X(W(), w - -w(") (0
(C.19)
where the quadratic term is absorbed into (1)"d). Substitute CC.17 and C..19 into (C.88) we have
L(W(++)) L(W()) & (VL(W()), w w(0) + (()
(C.20)
Take expectation for the stochastic gradient both side in (C.20).
E [(( T(W())]
s E (VL(W())),
- -y(tt)--((0))) + +
S E
Σ
9t.
(
+ + ((dd-1²²) + ((ns.7'so,) + (
)E(=1) [m)
Σ
Σ
(0)
+ (
-
9h.s
)E(=1) rE[m]
/ VV +
where we use Lemma C.2 that the update aligns with the gradient's sign for large gradient and the fact that ns² o, = (d) and Jensen's inequality. This completes the proof.
Lemma C.9 (Generalization Performance of Stochastic Adam). Suppose the same conditions hold as in Lemma C.8 We have the following results for T poly(n) with training dataset S
The training error is zero: crrs (W()) = 0.
The test error near-zero: errp (W(T) o(1) Proof of Lemma
C.9
By Lemma C.7 we have
E.) = =Ō(noo + a). jv) = ((1), (w)rr (7)
Recall (,, x) in Definition 3.2 with 1750, o(1), o(1), we directly have
errs(W(()) E s 1 (T(W(() x) < F- (W()) xx) 0,
since Fy, (W(2), xs) = πi()) while F-w (W(T),) x) z θ(1so, + a). Besides, for test data (x, D with _ TT, it is clear that with high probability (w (w( +(1)
and [( ( + @(nso, + at), then similar as training error, we have
(T() x) z o ((w g. vv)) 5i)),
while
F_,(W",x) = m
[o((w (.r.. + o((w"-.rr (.x-----)))) {)) s Ö(nsap + a).
Therefore, we have
errp (W)()) Bx
(,(W(), ( F-, (W(T) x) o(1)
This implics that Adam can achieve nearly zero test error. This completes the proof.
43

---

C.1.3 Proof of Corollary 4.3
Corollary 4.3 follows directly from Lemma C.10
Lemma C.10. Suppose the same conditions hold as in Lemuna C.3 and C.6
if X = -wlo then
(w jv) S O(ro),
(0) (w Jan E1) 2 ((50))000)))
Proof of Lemma C.10
This corollary is an immediate consequence of Lemmas
C.3
and C.6
In
particular, Lemma A.3
guarantees that at 0
kw (0)
((oo),
(( j.r E.) @(s)/2 op oo),
(0)
J.0 [k] O(oo).
Since À w(a "), Lemma A.7 implies that, at initialization, the weight decay regularization term
overwhelmingly dominates the gradient:
A >/w]/ (0) []]
o' o'((w), (0) j.r: vvv)) + ao ((w
(( E)))
A w St.r
(0) [+]
>> d'( yort E1) [((A].
Hence, by Lemma C22 the updates remain in the regularization-dominae regime, and no coordinate
ever grows beyond its (oo) scale throughout training. This completes the proof.
C.2 Proof of Stochastic Adam
Lemma C.11. Consider the update of stochastic AdamW in (3.6). Let W(() be the weight at the V.V), (w j.r+s)) = ((1) for all E {+13, E [m], i E [n] and I-th iteration. Suppose that (w j,r
(9)
8 < 82. We have the approximate update rule for each coordinate weight as follows:
For lo we have either (1) I.]],, S θ(n) or
sgn sgn((gg
(0.. O(1)
+E
For every E B₁, i I.-. = (To : ( To < }. where To represents the number of iterations away from the current iteration t, coordinate ls is affected by E sampled at the iteration To since the moving average, and we define F O(log(An))
1f S (1)), any To B. we have either
206 B BBono n ((0)
@
or
m (1
sgn (0) (11 ((1)
[.] +
If 2 O(log(An)-1) =(()),
(0(1).fol for To 10(lo((>-10) -sspp) such that 8 2 O ((,2)) have
either
( [k]]
s < (B- p1
or
( O())
sgn
(t) +
For TD 2 (((ogg(A)) such that 3 < (()21))) we have
=0()m)= o(1)
( [k] +e
44

---

For the remaining coordinates k + and k & Bi, i II-To E U [F]. where represents the number of iterations away from the current iteration t. coordinate affected by & sampled at the iteration i since the moving average, and we define + log(>n) -1 Then we have
m 1
=Ō())) o().
M+ +
Proof of Lemma C.11] The proof is similar to Lemma
C.2 We select F = (log((An)-)) such that
Ei (1 B1) ()221))) and Er (T (1)
92- - [k] O(22172)
(4) R j.r
Et=o 37 - 31) 01 (- [k] + 0(221)²)
VI 35 (1 82) 96- - [6]2 + + ((221)))
+
Recall the gradient of stochastic Adam given in Lemma A.8 we want to approximate 9 -
to (t) sk-r such that we can use the current weight to approximate sign update. By Lemma C.
the upper bound of the normalized moving average of each coordinate in one step is ( since
A O()) Then for E [L- F, 4), we have
(w (1) j.r vov) (w viV)
(b) Wrr .vv) - ((W (w).r.BVV) BIV)
(k+1)
(k)
(8) nn w 3.r viv) 1
(k)
j.r
+
S (17)).
(C.21)
The last inequality we use the fact that (w j.r B.v) = ē(1), X Ö(1) and Lemma C.1 Similarly, we have
(w E.) - (w/77 Ei) S @11Ts50p).
(C.22)
Then recall the predict function
(2((()) (9) x4) = Σ 0 (iv)) +o (Iw E.
We have
F (W(() - - FF((()), xxi)
(w - ( for (.. )iv))
((w55)
+ Σ / (w - 0 (ww (ww(2,)))) si)
m
θ(1).
(w (t) BiV) (w//7) uivi
+ Σ 5(1) (w/ Ei) - -w/f7 Es)
< (ma) + (mm)))
(nFssop)
(C.23)
where the second inequality we use the convexity of (-)) and the facts that (w vvv) 6(1) and (0) + = @(1), the last inequality we use m @(1 and sop (1).
Then we can approximate el (()
frr in the gradient A.8
F-)(W()) x.))
87
ExE(-1.1) x.
45

---

F-,(( x))+
sgn(e))) 10
(yx = j)
50 ExE(-1.1) (W(r))))
(0(00 + + ,(W0,))))
= =g)( D0(
(ys * j)
where we use the fact that (1)Fso = o(1) and C.23 So we have
(5 sgn (1) ((
for all T - (8- f, Further, by C.21). CC22 and the facts that (( ( B(v) Ö(1) and (t)
(( Am 4 = 0()) recall (xx) 20) we have
5770'' ((
s 4 ((w J.r v(v)) + |471 ((1F)
= sgg((())) e((5))) o' d'w
(0) vv)) + + @(mF),
o7 (w
4 o'w (9) )(v)) - (057 ((77)
sgn (9) 0((5)))) ''((w
for v1v))
(4)
0(e(50 -((F).
So we conclude that of (w.r. vvv) ((w)r. = sgn(( - ((W vv)) + + ((7)).
(C.24)
Similarly, have
O ((w S.w.eeee Ei) sgn(e @ e(e o' ((w
(1) j.r Ex) + O 22 @(FFsop).
(C.25)
Now, have all the tools need to approximate gi 3 [k]] to (e) - [k]. Recall Lemma A.8
((-))
substitute (C.24) and C.25 into [ we have
For k = 1,
((-)) [1] θ (1 + 1 1
(
1 Σ (2
Ö(17).
(C.26)
SET.
For all k E B, i E I--..
[20
g o [k] = 0 (1)
(gg + θ
Ö(fFso
(C.27)
B
For k and lo & B, i E In-r.
(t-T)
-.- [k] 0.
(t) r.J
(C.28)
J.r
Plugging C.26
C.27
and CC28 into C.11 with facts that 7 = θ(1), A = o(1), 1(
j.r E.)
(0)
((1), kw in uv)) (1), @(1), θ(Ar) and Lemma A.6 we have
(0)
46

---

For k- = 1,
" +e
Σz= (1
(11-3)) (1) -2- [1] + (())))))
Σ 35 -32) - - [1|2 + = ((221)2)
(1-T)
0 - -31) @0 (0) (1l + 0 B Eiez, 109)) Ō(TF)
+ O((29)9)
1 E 35 - (22) 00 (1 4 1 Eier,,, 1691)
-((F) (1))),
+ Ō(2219)
Ex-o B1 31) (0 9i ( 0(1F))
(4)
+ ((22)7)
Er= 35 - 82) (0 (i-T) (1) + 1(()))
2
+ + Ō(19Tp)
1 (4) + (1FF) + =((277))
(3....111)) - 6(1F) + 0(1777))
e ((0 = (n)
O (( + (())
For k E B,, E I To (0) [7], where to represents the number of iterations away from the current iteration t, coordinate k is affected by & sampled at the iteration t To since the moving average. We note that if the number of iteration in one epoch 2 is less than F, the moving average will use some sample multiply times. We denote T : {To timestamp set involved using noise & (i.e., i E IT-, for any T E Tk), and E B,, To S B.
# i 1 1]} as the
If ; > F, in this case we have To == { for any k E B,, and & was used in iteration t - To- Then we have
[
j.r [k]]
+e
Et=o BT - Bi) ( [K] + 0(177))
ETzo B3 (22 9: (1-T) [M| + c Ō(127)
ENeT, BT - B1) 91 - [k] - Ō(1717)
EIET, (8 - B2) D1 -1- [1]2 + E + 0(22)p2
8T O(g jo.j. [k]) Ö(12772
(1-70)
3 (gi (1-70) [k]22) + + 0(27))
870
0 g
( [k1) ã
(t)
=(()))))
B
82"
0
19 1K11) + T
ye
+ 0(2))))))
Now we handle 8 and 3 with more care. First we have 870h < 3 and _9 [k]] = Ō(1).
47

---

Then if 2 (1), then 810 = (1) and $2 = (1) since TD <
n
In this case, we have
0 ( [kt]
meap (
m mm
B
V te
(199 [k]] O
(t)
noople
(1)
1
If TØ 2 @(loge=)) @(log(An)) @(1) such that 8.3h S ((2(12),
then for TO
(log(A sp))) such that 37 2 0 psop
(pOpp))
(( say ) we have
(0 1kll 5
(1)
180, 1
(2)
m
1. [K]
+E 2
[k]]
+
B
For To B. To 2 ((log(An) -1) such that $12² O(2217?), we have
m 11
=0()21²)
= =0(2)) = o(1),
1 + O(27)))
since = =(\n))
. For l + 1 and l & B, E I--ITTE E [0, F). recall 96 [k] = 0, we have
[. +e
Et=o B7 (1 - B1) ( [K] + Ō(2)))) V V B1 (1- B2) 9 - [k|² + E + Ō(27)2)
1Õ(1777)
+ 0((2)))
= +0(2m)
o(1).
This completes the proof.
C.2.1 Proof of Theorem
4.4
Lemma C.12 (Stage pattern learning). Given the training dataset S, if 2B O(1) or 2 of ssp))
7 1/pollydd)) A = ÿ B² A 1) and X = Ö(1), then for any 1 S To with To Ō maa,
((+1)) B) .. v) S (w (1.))) .).) + O B
(w jr
(++1) B) (x) (wyer (4) (1ssp))
(t:))
yur
Proof of Lemma C.12 We prove this Lemma by induction. By Lemma m
C.I
n 11)
B) (w) jrr . vv =1 (1- (B-1) v).v)
j.
.v
18-1)
v,r
+c
s <11-11) Xw <w1 (1-0((1)))))) | v).v) + e(n)
(#-1)
48

---

3-1
s (1 Xn) B (w (0)
jrr v) + (
Σ (1
(1-\
km0
(W v) ( =
(w (0)
jrrtJ v) O m
B)
B
(w j.r. v) O B)
where the second equality we use Taylor expansion and An B o(1), and the last equality we have (0)
(w j.r . v) = ( by Lemma A.3 Now suppose the inequality holds for t = O,..., to with
to S To. We have
(fo-8) 8) . v) S (w j.r.J v) + O to) s 50 + Bsap )= = Ö(1)
(w jrr
B
since 3 = o(50p). Fort t0
(1o++)) 111)
(16+1))+) (W) jr
(11+1)) --1)
m mjr
.j v) = (1 (1-An)(w)) (w fr
5 v) -
.v
((o++)) 2 11) +e
(10+1))
I (1 (w) 3rr
v..v) +
#-1
$ (1 An) B (w
foo #
.j v) + +())
Σ (1 Xn)
jrr
k=0
(w.rr vj.v) 0(Am
(1o-H)
n
(w (1o
.j.v) v) + O(1 n
B)
B
(W5,r (lo:#)
n B
.j.v v) O(1
Hence, we have (w j.r:J j v) Ō1). Then, we prove (w yi.r
((+1)-
+) ..) = (w bur Ei) + ((nsa,)
(1.
by induction. By Lemma A.3 we have
(0) (( j.r c.) = //22220000)) w w = ((oo),
which imply that 50 = @()) Assume that sample (x,; v)) is in batch I, in the first epoch. Then we have
(w (T)
Wirs 4:) 2 (1 Xn)((
wir ,,))
(.) - O(1a)
2 (w 00 Ei) - O(A7on na)
= O(oo),
since Xn o(1), 17 o(oo), o(1) and 8 op Ō(1) Additionally, have n/s o(oğ and [k]] 2 @(a,)) with high probability. Then B |B-log [x]1 2 O so ( for i E I,.
Therefore, by Lemma C.Il and A.6 we have
sgn
(0) Nh
o
(( (0)
Nr.r [.))&
B
sgm
100 o' ((w (0)
Bi.r [
=
(C.29)
Then, by Lemma C.1l we have the following update according to B.6.. C.29 and Lemma C.I
(w (r+1) Viar EE)
(T)
= =(1 -11-
Wor Ei) (-xn)(((7)
E:
(T) \
yur +e
49

---

(w E. + (())
(s))(E)(()).)) -AArooo) O(pa) O(1sop)
kEB;
5.) + θ (nsap)
At the end of the first epoch, we have
(B)
pi-ro Ei)
(1 (w) Wiir E.) (na)
(8-1)
(7+1) Et) - (0)))) sap) - -(na)
{.) (1sop))
This completes the base case for t 1. For general t to with to To, assuming (w!!?) ()#t) (1-1)-
EEi)
(Wyir
.Ei) + ((nsopp). Then we have
(b:B) (w
(1-1)-
Dest
Ei)
(Wair Et) + ((nsop)
(w!!))
wi.r Ei) + (b11SO7))
@(st/2apm)) k1s0p)
@(1)
By Lemma C.I
we have
(1.B) [w jir
(r 1) [k:] + O(1)
[k] S wrr
W w.r [k]] + +(n B
(0)
Ō(oo +
Bsop
((1)
So we have 1 (1). Follow the same proof above with / = l0 + 1, assuming that sample (xe, Mi) is in batch n/ B+r in the t-th epoch. Then we have
(w #to
(E) 2 (1 (1-An) (to: #+r=)+r-1) (EI) e(pa)
(fo-i) it)
(Weerr Ei) -0(1)) maa)
(1o Ei), (Wair
since v = o(on) and a o(1). Additionally, we have v|s = o(ag and [g[k]] 2 Ō(o,) with
high probability. Then B
(w (0)
[k] (
6(nB-1sopp15) (0
Ô(nB- s
for i E It B+r-
Therefore, by Lemma C.Il and A.6 we have
(to-B B+T)
(to B-T)
sgn
Wisk
a (((W
[x))
(1o vi.d
(to-
sgm
0' ((ww)
[.)
(k))).
(C.30)
Then, by Lemma C.Il we have the following update according to B.6 (w (Na = +r+1) (W Ee)
(C.30) and Lemma C.I
(to Bo +r) Migrr
=(1->n)(ww.. -w (1o +T)
.E,) n.
E:
(fo B +r) Vyne
+
50

---

(to (w ye.r
.E) (
Σ
(sgn(E:(6)), E.) -Ō(Am) - ((na) (nssp))
(to 8) (w
(Way
E.) ((soop)
At the end of this epoch, we have
(1 Merr
-ii
(1: -11
(1
.t: - -(na)
yi.t
(6- #6+++
E.) Ö((n) - O(na)
(10:8)
.E:) + ((180)))
yerr
This completes the proof.
Because in the large-batch regime we have n/Z o(sap). Lemma C.12 tells us that noise memorization outpaces feature learning. Early in Stage I, feature gradients predominate since (!) E.)), given that a o(1) and weight decay effect is negligible. After a certain number of epochs, however, the noise component grows until avo' yor: 6 ((ww.rr
((wwerr 3(v) > a :(:))..::.)))
(t)
(9)
(t) v.v)), at which point feature learning slows, then reverses direction entirely. Lemma C.13
below characterizes this transition in detail.
Lemma C.13 (Stage I, fitting feature noise). Given the training dataset S, if B = (1) or BB2
o(s0p). 1) 1/poly(d), A 22
B² n 1) and X Ö(), then if a 2 θ (( B sap) for any
te [Tr, Tol with T, Ō 1s (-))) I To.
((+1)) B)
(t: jj v) = jo j..v -0(1 n
(wir
B
At epoch Tp. we have (a) w jr
(oo 3) [1] = sgn(()) δ( n
Bea ): (b)
w
j.r
To [)] = (g)((((4))
σ (dd) or + for B, with y j: (c) w jrr
(To #) [k] (n) otherwise.
Proof of Lemma C.13 By Lemma C.12 we have
ao'
((
(wiir
(Tr $)
< 2 a (ō(s 000000000)0000 + @((sop)) 9-1
a (ō(s 0p00) + 0 a1/1--1) 2 -((0 )
9-1
d' (( (T.. -(v))) - (δ(0o) + +T, O(1 n. )))²² 9-1
9-1
((oo) noo
Bsa,at/(a-1)
< <ō(00-1)
Hence, there exists some constant C' > for t E [T, 3 To E1.
coo (( ((w))) ((w(ti.):)) k:)) 2 C . o' ((w v.))
Then by Lemma A.8 C.I. and A.6 we have
(e)
w (++1) .iv) (1 (1xxn)( (w -17
m
jir +e ..v
j.r
51

---

(() 3. ((n) sgn Σ
ayrelio o' ((wfiin(s))
VEZ,
(w)))
jv
((1)) -j sgn(j)
w jv (1)).
So we conclude that
((+1):
(t: .j -j.v) - ) .j.v) - (w
n
jr
Moreover, at the end of epoch To.
(To #)
jiv ((00) + O(7 B)
n
(To --To-TTT
T,) O(n "
= - 11 Bso
Multiply j on both side, we get
(To B
[1] [[g]]]]nn(] O n
j.r
Bsap
(To: #) [k]. where l E B, and E [n]. Lemma C.12 shows that it increases by ①(7) in the (0)
For W,r
direction of [k]( if Mars Ei) > 0, that is.
(To #) [k]
(0) [K] (g
(E:(::) To O(n)
1
sgn (E))( θ
s0p
(To B) Otherwise, weight decay drives w w. [k] toward zero it is initially negative, in this case jrr (-)(1), @(n)]. For the remaining coordinates, Lemma A.8 implies the gradients are zero, so the
[A] E
(To +) updates dominated by weight decay. Given the fact that Top w(o)), we have w Wjr [- (n)).. This completes the proof.
[
Lemma C.13 implies that, by the end of Stage I, the model has fitted the feature noise -ayv. The following Lemma 1.14 C.14 poor generalization.
shows that these pattern persist throughout Stage II, ultimately leading to
Lemma C.14 (Stage II, preserve the noise). Suppose the same conditions hold as in Lemma C.12 and C.13 for t > To 3 E {+1}, [m], i E [n]. let r* = arginaxre(m) ( ycrr Ei), then (t) j.m v) - B - and (w!() (1)
(t)
(t)
(w
Burr, Ei) = O(1)
Proof of Lemma C.14 By Lemma
C.12
and B.6 we have (w Ei) E
(-O(nsa,), ((o))]. Because if (w Burt Ex) 2 (Op)) then we have
(++ wir E.) < (w (9)
-..r a @(1sop))
while if (w (t) We.r Ei) < O, we have
(w (c+1) -Birr
E:) 2 (1 (w Ei) + O
Now, suppose (w()) ;,r,r,,)) Ex) (2 log ((m)) 1)) à then for (x,, vi) with y: j.
(1)
(=)((((),))) xi)
Σ Σ e((()))).))
52

---

+ exp - alw (( jmn) (( jom Ei)) o(w ((w"))m))v)) o(w (0) -j.m {11
+ exp E ((w Ei)
2
+ exp m ( log ((xr)) -1 1)) 6.
(Ar),
where the inequality we use (w (t)
Let jv) < O. Then, in epoch Ta which contains iteration E, it follows
from Lemmas C.11 C.12 and (B.6) that for all la E [t, Ta + # we have
(w (t))
Veire Ei) 2 (1 - Aw) (4) wi. E + (11so,)
2 (w!l) Mar- {,) + (nsop))
If (w (t)
yi Ei) 2 (log ((1n)- 1)) then for 11) with y/: J.
of ,(W(()) xx)
0/1
Σ Ef (W((..))
1 + exp E al ((w jv)) + ((w jor Ei) o(( o((w -.. E.1)
1 + exp
00 ((( E:)
((w
<
(log (2 1))
1 exp
((22)))),
where the inequality we use (w
-j.r -j.mjv jv) I 0 Bsa, ) n o(sop), (w Ei) s @(ou). Then
by Lemma C.II C.12 and (B.6).
have
(w ((+1) yurr Ei) S (1 >nn) Mirt Ei) + θ(An 780p)
S ((w (t)
r. Ei),
since 1180, o(1). For (w jv) - B ) the same proof applies, since sa, o(1) and
B o(s0r)) If (w/()) ".r,,r Ei) (log ((An) -2 - -1) then for (x1,3i) with y; j.
1/ < 77)
Then by Lemma
C.I1 C.12 and B.5 , we have
(w.Ct ((+1) .jv)
(1 Xn) (.: O(M 11)
(0 (1) jv) A(A) Bsop ) O(n n))
(w (1) jv).
since 17 o( -(_))) This completes the proof.
Lemma C.15 (Convergence). Suppose the same conditions hold as in Lemma C.12 if the step size satisfies 1 = O(d-1/2) then for any I,
C. 13 and C.14
E L(W(++1) -((W")) s -
53

---

Proof of Lemma C.75 The proof is similar to Lemma C.8 We aim to prove the convergence of the objective function under the AdamW optimization algorithm non-convex setting. Recall the loss function for each data point is
L (W) = log 1
exp (Fy. (W,x) F-w. (W, (W,x)))
where W represents the parameter matrix, is the input data, y/ is the true label, and F3 (W,, x:) are the logits for class j. The total objective is:
L(W) 1
n Σ L(W)
Since [;(W) (W) is non-convex, exploit its smoothness with respect to the logits [F, (W, x;)];- Specifically, (:W is I-smooth in [P, xx)l; due to the properties of the cross-entropy loss. Define:
A (W xx)) FJ (W x.).
Using the smoothness property, we apply a second-order Taylor- like expansion around W(
( < Σ (L((W(O) AFF, + Σ(AE).).
(C.31)
(F x.)
This upper bound arises because the second derivative of L, with respect to the logits is bounded by 1, standard result for cross-entropy loss. The logits are defined as: Fj(W(o) xxi)
jom B.v)) + ((w j.rt E.))]. where ((t)) the r-th neuron in j-th output of (()), (()) =
(f)
[2]] is smooth activation function (e.g., with q 2 3). By Lemma C.14 and C.1 we have
(w
J.r. v) < @(1) and (w Ei) I @(1). ensuring the local smoothness of remains Õ(1) be-
(1)
(t) tween (w j.r v.v) and (w j.r y&v) (similar for (w E.)). Then with Taylor expansion, we have
(++1)
o(w jr
(1+1)
vvv))) o ((w j. B(v) w,, a(w So vov)), Wir
((+1) w (
< O(1)
(w (1+1)
w w.jr 1
(1)
m (t)
= ē(1)
An (t)
for
j.r
+ e
< < d).
(C.32)
where the last inequality we use Lemma C.1 and || 10011 < (dd) by Lemma C.13 and C.14 Similarly, we have
(++1) (o((w jr
EE)) (( Ei) ww,, ((w Ei)), w (141) w!/
< Ō((1) (
w 12/1
< @(1'd).
(C.33)
Summing over r (with m = Ō(1)) we get
AF < (WWF) x.), w(L
+ 6(722))
(C.34)
Additionally, WV (W()) x;)le= & ē(1) since m Ō(1), (w (//.ir s ē(1), (w E.) < (11. So we have
A F < Õ(ssa, + na + rdd))= < Ō(nsa, + +
(C.35)
54

---

Substitute (C.4 and
C.35
into C331
L La
L
(V (V (ww,)))) w - - (C.36)
For the full objective:
( = Σ
( - L.W(O))..
(C.37)
1=1
Substitute (C.36) into C.37 we have
L(Wi++)) L(W()) S (VL(W()) w(+h)) ( -
(C.38)
Take expectation for the stochastic gradient both side in (C.38).
L ((w(++1) [[(())))]]))])
E
(VL(W())), w(+ -w + (()²d)
16. 1,
-n
Σ
Σ
xn (VL((())))) w + G(d 17²) + r)² ((n²d)
(jE(+1)
rE(m)
-
Σ
Σ
E
(8)
+ 1
Ann Bso
|V + -
rE(
S V
|11 (1)²d),
where we use Lemma C.I1 that the update aligns with the gradient's 's sign for large gradient and the fact that ns o, (d)) and Jensen's inequality, Hölder's 's inequality and o(s0pp) |W
θ B by Lemma C.13 and C.14 This completes the proof.
Lemma C.16 (Generalization of Stochastic Adam W. large-batch). Suppose the same conditions hold
as in Lemma 1.15
We have the following results for T poly(n), with training dataset S
9
The training error is zero: crrs (W()) =.
The test error high: errp (W (2)) 1 o(1)
Proof of Lemma C.16 By Lemma C.14 have
jrr jv)
(T)
O
(w W...rr. {i) = (w Ei) s (
sop
Recall F, (W, x) in Definition 3.2 with nso, o(1), a= o(1), o we directly have
errs (W(T) T2 (x Il (y(W(T) x) < (_,(((T) x)) =0, since F2, (W(T), x;) 5(1), while F-g. (W(T), x) θ केμ + σn) and sop (1). Besides, for test data (x,3) D with x = T JT. it is clear that with high probability (w vv) = é( sa, then similar as training error, have
8
(W(((), (T) x) Σ @a(ww (( + o(w (. = Σ
+o
a
+
Sy.r
sop
r=1
while
1
F F-y (W()) xx Σ uv)) + o((w" y.r. E) = Σ
1 +
IS ,.
sop
r=1
55

---

Here, Syrr and 5-g.r are independent and symmetric random variables. Therefore, if the term dominates Sy.r and 5-g.r, then it is immediate that (W()), x) < F-y - (W((), x), since o (1). This implies that large-batch Adam yields high test error. On the other hand, is dominated by both Sy,r and S-y.r. then with probability at least 1/2 - o(1), we have F (W()) .x) < P- (-v((()), x), as Sy.r and (-y.r are independent of v. In this case, large-batch AdamW incurs at least 1/2 o(1) test error. Therefore, we conclude:
errp (W (T) Exx 1 (F(W s (W 2 -o()..
This completes the proof.
C.2.2 Proof of Theorem 4.5
Lemma C.17 (Stage l). Given the training dataset 8, B 2 (n / log - and n w(80p).
1/polydd), À = 9d n² A 1) and X Ó(1), then for any / S To with To Ō km
((+1)- B)) .jv) = =(
n } (t: 3) .jv) + θ
(w jr
((++1) (Woer
(1))
Eii = (w ( (1sop))
((+1))))
Proof of Lemma C.17 We prove this Lemma by induction. First, we prove (Wyar (1:#)
EEi)
(Waer
.E)) + (nsopp) ) It is same as Lemma C.12 By Lemma A.3 we have
(( j.r E.) 0(s opor), Wj.r [k] (cro),
(0)
which imply that 1000 = @(1). Assume that sample (x:, ve) is in batch I, in the first epoch. Then we have
(T) Wirs Et) 2 (1 (n)
Wir {t) O(1a)
2 (w) Ei) O(Noo n7)
= ((oo):
since Xnj = o(1), 1 o(oo), d(1) and 3²/²o, Ö(). Additionally, we have 178 o(o00
and <: [K]]]]] 2 ((a,) with high probability. Then |B- 1g- 2 @(nB-iso sl (00 Therefore, by Lemma C.1l and A.6 we have
for i E IT.
sgn
(0) Di.
o
((w (0) [k))
B
sgn
(10) vo.s
berr [x)
sgn(gg
(C.39)
Then, by Lemma C.II we have the following update according to B.6 CC29 and Lemma
C.I
(w (r+1)
barr Ex)
(r)
= = (1-An))(w()) Wirr Ex) -m7 17.
E,
yeir +
2 (w( {i) + O(n)
(sgn(E,(k)), Et) - ((Aooo))) - O(1a) O(730,)
Dirr.
keB
(w(0) (0) {i) + ((sop))
At the end of the first epoch, we have
W EA)
56

---

(1 (w (8-1)
Nir (.) O
(w (1 Ei) - so,) - ((ya)
(0) Diirt &.} (nsop))
(1##)
This completes the base case for t 1. For general S to with to S To, assuming (w .E:) (1-1)- -.) + @(ssop)). Then we have
(W yur
(1: ( Best
(1-1)-8)
Ei) (Wour Ei) + ((nsop)
(with) hir Ei) ((1sso,)
O(s!/2 0p00 tm8Spp)
Ō(1)
By Lemma C.1
we have
(1##) [k:]
(1.4-1)
w j.r
- |w j.r
[ki] + +())
S (0) [k]] + O t. "
B
(1).
It's also obvious that (w) j.r .jv) = Õ(1) So we have (0:;3) = O(1). Follow the same proof above with t= to + 1, assuming that sample (x,, m1) is in batch It. in the I-th epoch. Then we have
(1:B)
B+r) WWat
(la B+r-1
(.) (1 An)
(War."
E.) O(ja)
(to +)
Ei) O naa)
yeir
(fo: (Wiir
Ei),
since 7 o(0o) and a o(1). Additionally, have rss o and [ 2 Ō(o,) with (w")
high probability. Then B Therefore, by Lemma and A.6
ir kk] tmsser,))*11
(0)
[ 2 G(n s0 (0000) 1) for i Itg B+r
we have
1
(to
#+T)
(to #-T) [x
sgn
d'((Worr
vi.
(fo-#+7)
(to-#+T)
sgm
o' ((W
[.)
(A)))
(C.40)
Then, by Lemma C.1l we have the following update according to B. (C.30) and Lemma C.I (to # +++1)
(wg.r
Ex)
(fo: +r) myar
(fo +T)
=1
>n)) =(1->n)(w
Eii n7-
:
(1o (1e +T] Vini
+
2 (W
(to:
E.) O())
Σ (sg)(E;(6)), Ex) -Ō(m)) - O(na) - (1soa))
kEB
(w tto 8) y..
Ei) + ((ssp))
At the end of this epoch, we have
(1.%)
(w Wi.r Ei)
57

---

(1- (1 An)(W War (.)) - ((aa
(Warr
(++4)))1)
Si) ÖÖ(A)) - O(pa)
(1):8)
war
.Ei ((ssap).
((+1))
+1) .jv) (w () jjv + inn $). By Lemma A.3
Next, we prove (w), jrr
A.8 and
C.11
we
have
o' d'((W) ((w])).)))) jv)) - ao (w
(0) j.r E,) 2 6 )
since a = o(1), and we have 17 00000 ) By Lemma A.6 we have
sgn(9n 10) = sgn 1-ssgg (j)) sgn(j).
Apply Lemma C.I1
we get
(w/r J.r J v) (1 Xn) (w J.r.j j v) O()) (sgn(0)).r).vv)
= (w (0) j. v) (1))
since (w (w v) = -O(o0) and x = Ō(1. We have
(1+1)) B) Ei) N (wa.r ..) ((sspp))
(1.2)
(w yhr
So we have (wig.r. Ei) S Ō(s 12 0000 + msS for E [0. B1. Thus, for t 10, R. we have
d' ((w jv) > ao' ((w (1)
1.m E.),
since a = o(1) By Lemma A.8 and A.6 we have
sgn (ighr = sgn(-sgn(j)) = -sgn(j).
Apply Lemma C.11 for & [0, 21 we get
(w (w J.r v) = (1 Xn) .j v) ((n) (s
2 (w
(0) jo - v) + 0 1),
since (W j v) = ((oo) and x = Ō(1). So we have for t = 0,
(w ((++)) 3 j v) (w
(t: .j v) + O(1 T
Now suppose the equality holds for 0, 0...., to with to S To. We have (0)
(w (to-B) B) .j. v) (w/ j. v) + O B
1o) < Õ(1)
jr
Since n w(s We have
(w jr ... v) = 5(oo + 1 n
(to B)
1o),
B
(to:
(Weer
Ei) = 0(5'112,,00 00- + nsoo to),
(10:#)
(lo
(W
.j .j.v) 2 (weet E.).
Therefore, for L E [[4o + 1) 3, (lo + 2) Z1.
d d'((w jv) x ao' ((w Ei).
Apply Lemma C.Il for t = to + 1, we have
(w
((+1)) #)
.j v) = (1 - (n)) (1(11) x-1) . v). - -(())
(((+)) -11 (1t(( jvv
j.r
2 (w/ft (t:
n
,j v) + (v 3
since (w/) jv) = O(1) and x = Ō(). This completes the proof.
58

---

Lemma C.18 (Stage II). Given the training dataset S, if 2 0(n V loge 1) and B w(sp))
(t) ? = (d). x = Q B² A 1) and À Ō(1), then for any > Tu. j E {+1], r E [m]. E [n].
let r" argmax, (w (t)
jar jv), then {w. jv) (1) and (Wairs E,) S 0 =
(0)
Proof of Lemma C.18
By Lemma C.17 we know that (w jv) increases at a faster rate than
(w yur: 5.) since n
(1)
w(xap)-
We also have (w (t) -w- E,) E [- ((o))] following
Lemma C.14
Now suppose that (w j,r+, 2 (log ((An)-2- - 1)) $ then for (x))3)) with yi j.
( xx))
52
Σ (-(1111) (W()),))
E o( (c) jv)) o ((w jir {.) a((w ((w o ((w"")r EL))
+ exp
<
(( [v))
+ exp
(log ((Ar)- - 1))
+ exp
(A2r))),
where the inequality we use (w (1) and (B.5). we have
-grrr 0, (w (....rr E.) < <Ō(oo). Then by Lemma C.11 C.17
(++1) 15v) S (1 bn) 10 jv) + 2(. m)
S (w.....v)
since 1) o(1). Similarly, if (w E,) @(oo)
(1+1) (w bir
(t) E.) S (1 - Ar) (w wirrt Ei) + 5(x) mso,)
I (w bi.r E2)
Otherwise, (woorrs Ei) E [-o(0o). (((r)) and satisfies (w (Warr Ex) 0 Bagg
If (w (1)
jrr jv) < s log ((1r)) 1))= 9, then for (x,, vi) with y/:
(w x1) xx)
Σ E e (W(()..)
1 1+ exp Em ((w( + o (w/()) Ei)) ((w (v))).ivv) "(w")).::)))
2
+ exp 2m a((W((...)))
1
2
+ exp [2m (en log (1 7.
= =>()),
where the inequality we use (w (t) (v) 0, (w () E.) @(oo). Then by Lemma C.11 C.17
and (B5)). we have
((+1) (w .jv) 2 (1 wm jv) ((n)
59

---

2 (w jor,
since A Ö(). This completes the proof.
Lemma C.19 (Convergence). Suppose the same conditions hold as in Lemma C.17 and C.18 if the step size satisfies fj = Od1-1/2), then for any t,
E L(W(++)) - <- 6())d).
Proof of Lemma C.19 The proof is same as Lemma C.15 We aim to prove the convergence of the objective function under the AdamW optimization algorithm non-convex setting. Recall the loss function for each data point is
(W) log 1 exp (Fy. (W, xx) F- (W, x;))
where W represents the parameter matrix, x; is the input data, y/1 is the true label, and (W, x;) are the logits for class j. The total objective is:
L(W) Σ L(WW)
Since I2 (W) is non-convex, we exploit its smoothness with respect to the logits (F, Specifically, (WW is I-smooth in (V,,,)) due to the properties of the cross-entropy loss. Define:
(W, xx)l;
A (f(((+)), x;) F2( (W() x:).
Using the smoothness property, we apply a second-order Taylor-like expansion around W(O);
(L((W(9) L (w((+))) I, (W())) (W x,) AF3. + E(AR)))..
(C.41)
This upper bound arises because the second derivative of L, with respect to the logits is bounded
by 1, a standard result for cross-entropy loss. The logits are defined as: F5 (W(0), ,x) =
So (() vv)) + ((w j.r. E.))]. where w the r-th neuron in j-th output of (()), a(2)
[a]2 is smooth activation function (e.g., with q 2 3). By Lemma C.18 and C.I we have (w J.r v) < @(1) and (w E.) < @(1). ensuring the local smoothness of remains Õ(1) be- tween (w/ (w/(t"), yov) and (w))), yov) (similar for (w("). E.)). Then with Taylor expansion, have jo
a(i jrr
(++1) 11+1)
((+1) - w ( .bvV)) o (( B(v)) (V w, (ow([((wW) ( ((w Bov)), Wir
(2)
(8)
< <0(1)
w
(1+1)
w 1
J,r
= =(1)
Sm7- ( +n
(t)
m
< <0(ndd)
(C.42)
where the last inequality we use Lemma C.1 and |w < O(d) by C.18 Similarly, we have
|o(w (++1) Ei)) ((w
(t) j.r
(x)) (Vw,, ((w Ei)), w (t+1) w j.rl
<0(1)- (((tt)) w 1 |
< (())d))
(C.43)
60

---

Summing over r (with m = =(1)), we get
A - WV (W(l), x,), w(( - + 6(77d).
(C.44)
Additionally, VVw (W()) xi)ll= 6(1) since =0(1, (w fn viv) - 6(1, (w (W E.) 5 (1). So we have
A 3 (nsop + mja 71 d) (nso, + 1/dd).
(C.45)
Substitute CC.44 and C.45 into C.41):
L (w((++))) 1)- L; (W $ (V - - (C.46)
For the full objective:
L(W(++))) L(W
(C.47)
Substitute (C.66) into CC47). J we have
L(W(++) ( - -(17d).
(C.48)
Take expectation for the stochastic gradient of both side in (C.48)
L L( [((())
$ E (VL(W())), w(+1)) - -Wo)) + (()2d)
$ S-n.B E
Σ
Σ
Σ
1 1
Am (VV(VZ((()), w()) + O(d x73) + Ō(ns - :/88o)) + ②()²d)
jE(Ll)
rm|m]
S Σ Σ
+ +0(2)) Mejer 11,
(V + -
rE(
~-n VL
+
where we use Lemma C.IT that the update aligns with the gradient's sign for large gradient and the fact
that ns (d), and Jensen's inequality, Hölder's 's inequality and Ō(1), ( Ilo I 6(1)
by Lemma C.18 This completes the proof.
Lemma C.20 (Generalization of Stochastic Adam mini-batch). Suppose the same conditions hold
as in Lemma C.19
We have the following results for T poly(n) with training dataset S
9
The training error is zero: errs (W(() = O.
The test error is near-zero: erT (W()(((T)) o(1).
Proof of Lemma
C.20
By Lemma
C.18 we have
(T) E.) < θ BROpy, (w (. Ei) S @(ou).
jv) @(),
(w mi
Recall F,, in Definition 3.2 with a = o(1), we directly have
errs(W(?) Ex 11 (F(W(TT x) < F-(((()) x)] = 0,
since Fy. (W(T), x;) (()) while F -w: (W < ① + a:).
For test data (x,y) 9) D with x [yv T T it clear that with high probability (w yv) = (1). Let B = supp(E). w ExEr wyrr [k]. Sy.r ΣKEB [k.r(k]-E(K -N(0, \w then we have
(w E) S Cy.r.
61

---

Now we calculate the upper bound of w By Lemma A.4 we know (E..E) [n]. Then let
for i * i.j
E Σ E S spanffi.... En..
We have
(T)
Psw,, (7) jr +r+ rISUIM
where c c= ((1). By Lemma
C.18 we have
x+ (Σ) P| (
12 S jir
(Paw
1 E(Psw
(T) jir Σ (T)
Ex)² Ō 0
B38 1
Since ||<: |1 ogx and w(log n), we have ||::1 (( with high probability. Hence,
mnin (E) = min; | = (( With a little abuse of notation, we have
B
IPs 1P12 - θ
no
By Lemma C.11
C2 B
A 1) and Any (1), we have
||r11 - O(1²d).
So the upper bound of |/ is
B²
[w S PPs (2)
+ ||r|1 < O
fr
no
Finally, with high probability
B²
Su.r S (
σp) = o(1),
no?
since 3 = w22 The same result holds for C-y.r. Then, we have
Fg(W(T) x) 2 (W vv3) = (2)),
while
F-y Fw =
(o(( (W (T)
y.r. vvv) + oo(wo"()))) (T)
-y.r { = m (a 5- = o(1)
Therefore, we have
errp (W (x,))= 1 (w((), x) F F-,(W()), x)]
o(1).
This implies that mini-batch AdamW can achieve nearly zero test error. This completes the proof.
C.2.3 Proof of Corollary 4.6
By Conditions
A.I and
A.2
along with Definition 3.2
we know that d = poly (n), and hence
0g
---
with 9 2 3.
2//4
2//4
This directly implies that the effective weight decay parameter for Adam satisfies
B² n _11
Addam ~ o--²
min
AAdamW.
This completes the proof.
62

---

D Experimental details and results
This section presents the complete details of our experiments.
D.1 Experimental Details for Real-world Data
For the real-world experiments Figures T and 2 we use the CIFAR-10 dataset, VGG16 and ResNet18 architectures. and the Adam and Adam W optimizers, all implemented in PyTorch. We do not use data augmentation in order to avoid any additional regularization effects.
In Figure 1 we report the test error as a function of batch size. The batch sizes considered are {16, 32, 64, 256, 1024, 4096, 8192). with training conducted for 100 epochs. The weight decay is set to 5 x 10-4 for Adam and x 10-2 for AdamW; the momentum parameters are fixed at (31, B2) = (0.9, 0.99) for both optimizers. Each configuration is evaluated with three learning rates: (5 x 10-1, 1x 10-1, 1 10-5), and we report the best test performance for each batch size. All experiments can be run within one hour on a single RTX 4090 GPU. The only exception is training ResNet18 with batch size of 8192, which requires three GPUs due to memory constraints.
Figure ((a) presents the test error versus batch size for Adam with VGGI6 and ResNet18, while Figure (bb shows the corresponding results for AdamW: W. Both demonstrate that test performance
degrades as batch size increases, which is consistent with our theoretical findings in Section 4 showing that small-batch Adam and A outperform their large-batch counterparts.
In Figure 2 we report the test error as a function of weight decay À for Adam and AdamW, using VGG16 (Figure 2(a) and ResNet18 (Figure ((b). We fix the batch size to 16, the learning rate to
1 10-4, and set (B1, 82) (0.9, 0.99). The weight decay values for Adam are {1 10-1 5 x
10-2, x 10 2 x 10 3 1 10 5 x 10
1x 10
5 x 10
1x 10 5
5 10 1x
10 5 10 J. and for AdamW are (5 10 1x 10 5 x 10
1x 10 5 10
1x
10-3, 10 4. 1 10 1). All models are trained for 100 epochs.
Figure 2(a) shows results for training VGG16, and Figure 2(b) for ResNet18, both using Adam and A. For a fair comparison, we scale the weight decay of AdamW by a factor of 1/25. The results show that Adam suffers from poor generalization under large weight decay values (e.g., À > 0.05), while Adam W maintains stable performance even with larger weight decays (e.g.. 0.5), which aligns with our theoretical results in Section 4
D.2 Experimental Details for Synthetic Data
For the data model defined in Definition 3.I1 we set the input dimension to = 1000 and the number of training samples to 200, consisting of 100 positive and 100 negative samples. The sparsity level is set to s = 0.1d = 100, and the noise strength is Op = 1/vs = 0.1. The feature noise strength is set The network, defined in Definition 3.2 has width m = 20.
0.2, and the model weights are initialized with standard deviation du 0.01.
All synthetic experiments are trained for T = 10a epochs with learning rate of 1 = 5 x 10 and evaluated on test dataset size 1 For Adam and AdamW optimizers, we adopt the default momentum hyperparameters 31 0.9 and B3 0.999.
We primarily focus on the following metrics:
Training error: errs (W).
Test error: errp(W)
Feature learning: maxye[m) (Wy...vv).
Noise memorization: min maxre[m) (W (Wj.., E,) or maxxe(n]uw max (W (W
Et)-
Large-batch Adam vs. Mini-batch Adam. We set X (batch size B = 100) and mini-batch Adam (batch size B = 2). Table test errors of the solutions obtained by the two training methods. Although both larce-batch and mini-batch Adam achieve zero training error, their generalization performance differs significantly. Specifically, latch-batch Adam suffers from high test error (greater than 0.5), while mini-batch Adam achieves zero test error. This observation verifies Theorems 4. and and
10 for both large-batch Adam I presents the training and
63

---

Table 1: Training and test errors of Adam with large (B = 100) and mini-batch (B = 2) settings.
Batch size
B
100
B = 2
Training error
0
0
Test error
0.9545
0
2.25
1.25
Feature Leamling: ]=1
Feature Learning: j=1
Value
0.75
Feature Learning: j--1
Value
1.35
Feature Learning: j--1
Noise Memorization: j-1
Noise Memorization: j-1
0.90
Noise Memorization: j=
Noise Memorization: j=-1
0.25
0.45
0.00
0.00
-0.25
-0.55
2000
6000
8000
10000
2000
6000
8000
10000
Epoch
Epech
(a) Large-batch Adam 100)
(b) Mini-batch Adam (B
Figure 3: Feature learning and noise memorization of Adam in the training.
Moreover, Figure ((a) illustrates the dynamics of feature learning, measured by max+e[m] (Wj.r. jv). and noise memorization, measured by minlee max,glm] (Wjrrs fx), under large-batch Adam. The results are consistent with Figure 2 in Zou et al. (2023b). Figure 3(b) shows the corresponding dynamics for mini-batch Adam, where feature learning maxye[m]] (W,,))W,,, J..r: increases steadily, while
noise memorization max;g(n):w= max,e (m) j.r.: E.) remains suppressed the end of Pattern Learning Stage. In the subsequent Regularization Stage, feature learning saturates at a stable threshold and stops increasing. This behavior is consistent with Lemma C.7
Large-batch A vs. Mini-batch AdamW. We set À = 0.01 for both large-batch AdamW (batch size 100) and mini-bated Adam W (batch size = 2). T the training and test errors for both training methods. Although both large-batch and mini-batch Adam achieve zero training error, their test performance differs significantly: large-batch AdamW suffers from high test error (exceeding 0.5), while mini-batch AdamW attains zero test error. This observation supports Theorems 4.4 and 4.5
Figure 4(a) illustrates the dynamics of feature learning, measured by max, m (W..r. jv), and noise memorization, measured by min;g|n):a: max,e[m) (Wj.r, Es), under large-batch AdamW. Initially, feature learning increases, but it is eventually flipped by noise memorization, which grows faster rate. As result, the model begins fitting to the feature noise, which is negatively aligned with the true feature direction. Specifically, noise memorization increases rapidly during the Pattern Learning Stage and saturates at a logarithmic rate in the Regularization Stage. These behaviors are consistent with Lemmas
C.12 C.13.
and
Figure 4kb) shows the corresponding dynamics for mini-batch AdamW. Feature learning increases steadily and remains unaffected by noise memorization during the Pattern Learning Stage. In the Regularization Stage, feature learning saturates at a stable threshold, which causes the gradient to become small and consequently suppresses further growth of noise memorization (recall that [x(1] = ogy:). This behavior consistent with Lemmas C.17 and C.18
Large weight decay regularization hinders Adam training. We repeat the experiments from Large-batch Adam vs. Mini-batch Adam using larger weight decay parameter À = 0.05, and
Table Training and test errors of AdamW with large (B 100) and mini-batch (B settings.
Batch size
B
100
B
Training error
0
0
Test error
0.5485
0
64

---

1.25
1.25
1.00
1.00
Value
0.75
Feature Learning: j=-1
Feature Learning: j-1
Volue
0.75
Noise Memorization: j=1
0.50
Feature Learning: j-1
Noise Memorization: jm-1
0.25
0.25
Feature Learning: j=-1
Naise Mernorization: j=1
0.00
Noise Memorization: j=
0.25
2000
6000
8000
10000
0.25
2000
6000
8000
10000
Epach
Epech
(a) Large-batch AdamW (B 100)
(b) Mini-batch AdamW (B 2)
Figure 4: Feature learning and noise memorization of AdamW in the training.
train error
test
0.8
0.8
Error
Error
0.4
0.4
0.2
train error
test error
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Epoch
Epoch
(a) Large-batch Adam 100)
(b) Mini-batch Adam
Figure 5: Training error and test enror epochs of Adam training with À 0.05.
0.4
train error
0.5
test
0.4
Ener
train error test error
Error
0.2
0.2
0.1
0.0
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Epoch
Epoch
(a) Large-batch Adam (B 100)
(b) Mini-batch Adam (B 2)
Figure 6: Training and test over epochs of AdamW training with À 0.5.
those from Large-batch AdamW vs. Mini-batch AdamW with = 0.5. Figure reports the training accuracy over epochs for Adam, while Figure 6 shows the same for AdamW. It can be observed that Adam fails to train under large weight decay. contrast, AdamW remains robust and achieves results consistent with those in Large-hatch AdamW vs. Mini-batch Adam even with a larger X = 0.5. These results support Corollaries 4.3 and 4.6
D.3
Additional Experimental Results
Error bars across random seeds, 7 We provide additional results support our theoretical findings. To assess statistical significance, we repeat the CIFAR-10 10 experiments from Figures D and 2 with five random seeds (0-4), using the same settings as in Section D.I Figures and 8 report the results, with error bars denoting the standard deviation across runs. The results confirm that both Adam and Adam W degrade performance as the batch size increases, and that
Adam is more sensitive to weight decay than AdamW.
65

---

VGG16
28
VGG16
30.0
ResNet18
ResNet18
26
Test Erron (%)
24
25.0
22
22.5
1est
20.0
15.0
10a
10
10
10
10a
10*
Batch Size
Batch Size
(a) Test emor vs. size under Adam
(b) Test error vs. batch size under AdamW
Figure 7: Error bars across seeds: Test error batch size for VGGI6 and ResNet18 on CIFAR- 10.
26
Adam
Adam
24
AdamW
22
AdamW
Error (%6)
22
Test Error (%)
20
20
18
18
16
16
14
1-5
10
6
10-5
10 -4 10
10
2
10-1
10
6
10-5
10
10
10
10
Weight Decay for Adam
Weight Decay for Adam
Weight Deray 1/25 for AdamW
Weight Decay 1/25 for AdamW
(a)
VGG16
(h) ResNet18
Figure 8: Error bars across seeds: Test error weight decay (batch size 16), comparing Adam and
AdamW.
VGG16
VGG16
ResNNtt18
40-
ResNet18
40
Error (9%
30-
25
20
20-
10a
10a
10*
10a
10a
10*
Batch Size
Batch Size
(a) Test error vs. batch size under Adam
(b) Test error vs. batch size under AdamW
Figure Error bars across (B1, 82): Test error vs. batch size for VGG16 and ResNet18 on CIFAR-10.
Sensitivity momentum parameters (81, 82), Figures and 10
We further study the sensitivity of Adam and Adam W to the momentum parameters (31, (2), which are treated as constants in our theory. We sweep over $1
0.5, 0.9} and 3, E (0.5. 0.9, 0.95), yielding 8 valid combinations under B² < 32. plus the standard setting (B,, B2) = (0.9, 0.99), for a total of configurations. Figures Seuuresgandee [O report the results, with error bars showing the standard deviation across 9 runs. The findings again confirm that both Adam and Adam W suffer performance degradation as batch size increases, and that Adam more sensitive to weight decay than AdamW.
Large-scale vision experiments with ResNet-50 on ImageNet-1KK subset, Figures I1 and 12
further validate our theory, we conduct large-scale experiments I K. We construct subset by randomly sampling 100 training images per class (seed=0), ensuring controlled large- batch regime ( ((L)) while keeping computation feasible. ResNet-50 is trained for 90 epochs
66

---

30
Adam
Adam
AdamW
22
AdamW
Test Error (%)
25
Test Error (%)
20
20
18-
16
15
10
10-5
10
10
10
10-1
10
10-5
10
10
10
10-1
Weight Decay
Weight Decay
Weight. Detay 1/25 for AdamW
Weighl Decay 1/25 for AdamW
(a) VGG16
(h) ResNet18
Figure 10: Error bars across B2): Test error weight decay (batch size 16), comparing Adam
and AdamW.
Adam Top-5 Error
80
AdamW Top-5 Error
Top-5 Validation Error (%)
70-
5000
10000
15000
20000
25000
30000
Batch Size
Figure 11: ImageNet-IK subset: Top-5 validation error vs. batch size for Adam and AdamW with ResNet-50.
with standard ImageNet preprocessing. We report top-5 validation error against batch size (Figure) II and weight decay (Figure 12). comparing Adam and AdamW.
For Figure decay A
we learning rate 1j 1
10-4, (01,32) (0.9, 0.99), and weight
1 x 10 for Adam and x = 1 x 10 for AdamW. Batch sizes are (64, 128, 256, 512, 1024, 2048, 3072, 4096, 8192, 16384, 32768].
For Figure 12
we fix = =4, 6 10-3 and (B1, (2) = (0.9, 0.99). Weight decay values for
10-3 5 x and for Adam are {1 x 10 4,5 10 11 10 10-3 10-3 1 10-2, 10 1,1 10-1,5 10-1).
10-100 Adam [5 x 1, 10-6 10 10-6,11 1 10-5 5 x 10-5 1 10 55xx 10
10 10-3
10 10-2.5 5 10 1,-2,x100,,111011111
The results again confirm that both optimizers degrade as batch size increases, and Adam is more sensitive to weight decay than Adam W.
67

---

100
Adam Top-S Error
AdamW Top-5
Validation Error (%)
10
10
1
10
10-1
Weight Decay (Adam)
Weight Decay + 25 (AdamW)
Figure 12: ImageNet-Iee IK subset: Top-5 validation error vs. weight decay for Adam and Adam with ResNet-50 (batch size 64).
68


Reviews:
[{"review_id": "31Zog4jYNj", "reviewer": "Reviewer_dSx3", "summary": "The paper studies the generalization error of a classification task with a two-layer neural network for Adam and AdamW. The findings are: 1) the two algorithms both perform better when the batch size is small, and 2) AdamW allows a larger regularization parameter $\\\\lambda$ while Adam's parameter selection is restricted to a small $\\\\lambda$ regime, consistent with practical observations.", "questions": "Please refer to the points under Weaknesses.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper is well-written and easy to follow.\n2. The technical contribution is solid. It provides an explanation as to why Adam/AdamW is better in a stochastic setting, i.e., when the batch size is small, and why AdamW allows for a larger regularization parameter.\n3. The author has a good discussion that links the implications to the proofs in the theorems. Specifically, the feature learning vs. noise learning in the two stages.\n\nWeaknesses:\n1. __Regarding the setting.__ I am not very familiar with the data model setting in the paper, and I understand that many practical scenarios are simplified for ease of analysis, common in theorems for modern machine learning. However, how this data model relates to reality and to what extent it can be generalized, e.g., to multiple layers and with wide range of parameters, is unknown to me. Specifically, does the noise model resemble some cases in practice? How realistic is the parameter selection, like $s$, $\\sigma_p$, and $\\\\alpha$, and why are they set to the order in the paper?\n2. __Regarding parameter choices of Adam/AdamW.__ 1) Are the (order of) parameter selections, for example $\\lambda$ and $B$, based on practice or to achieve the best theoretical bounds? 2) For the case of AdamW, the author used a large $\\lambda$; however, they do not have a theoretical negative result on small $\\lambda$. It seems that in practice a small $\\\\lambda$ also works in the experiments? From the theorems, how can we know that AdamW is more robust to parameter selection? For Adam, it seems that it only requires $\\lambda$ to be small enough (what if we set it to be 0?), while AdamW requires the parameter to be tuned to a regime with upper and lower bounds. 3) The batch size choices in Theorem 4.2 and 4.5 are also different, with AdamW's being more restrictive. Then, how should we compare the two results? 4) From an optimization point of view, to achieve the same error, should the time $T$ and step size $\\eta$ be set differently in the mini-batch and large-batch cases? For example, with a large batch, the noise is small, and the algorithm should converge faster. Will early stopping the training of large batches help?\n3. __What happens to other algorithms.__ Under the same data model setting, what would be the performance for SGD, SGD with weight normalization, and Adam without weight normalization? Can we have more insights comparing also with these algorithms or to design better algorithms? This might be future directions.", "was_revised": true, "final_justification": "After the author's detailed explanation during the rebuttal, I believe my original score is reasonable for this paper and I intend to maintain it. The paper makes a solid contribution and employs sound techniques. Therefore, I recommend acceptance.", "rebuttal": "We thank the reviewer for thorough evaluation and constructive feedback, particularly acknowledging the manuscript’s clarity, solid technical contributions, and insightful feature-vs-noise learning discussion. Below, we address each of your questions in turn.\n\n### Weaknesses & Questions\n\n**1. Regarding the setting.** \n\n**a) how this data model relates to reality and to what extent it can be generalized, e.g., to multiple layers and with wide range of parameters, is unknown to me.**\n\nOur feature–noise patch model—widely used in theory (Lines 135-136)—isolates the trade‐off between learning sparse, generalizable features and memorizing noise. Real images can be viewed as a few informative patches amid distractors, and studying a two‐layer CNN reflects the behavior of intermediate layers in deeper nets (Papyan et al., 2017; Yang, 2019). Moreover, experiments validate our theoretical predictions (Figs. 1–6). Extending our proofs to deeper, wider architectures remains a valuable direction.\n\n**b) Specifically, does the noise model resemble some cases in practice?**\n\n Our noise model reflects real images, where sparse label-relevant patches (“signal”) are embedded in large irrelevant regions (“noise”). Empirical evidence shows CNN activations are sparse and treat distant backgrounds as independent noise (Papyan et al., 2017; Yang, 2019), supporting the realism and tractability of our model. Viewing it as an intermediate-layer input further highlights its practical relevance.\n\n**c) How realistic is the parameter selection, like $s$, $\\sigma_p$, and $\\alpha$, and why are they set to the order in the paper?**\n\nWe select $s=\\Theta(\\frac{d^{1/2}}{n^2})$, $\\sigma_p^2=\\Theta(\\frac{1}{s\\cdot\\mathrm{polylog}(n)})$ and $\\alpha = \\Theta(\\sigma_p \\cdot \\mathrm{polylog}(n))$ (App B.1) so that only a tiny fraction of the input (\"signal\") convey class information, while $s$ coordinates serve as iid noise. These settings align with practice.\n\nBy scaling $\\sigma_p$ to match the feature magnitude and letting $\\alpha\\in(0,1)$ model misleading noise, we create a setting in which feature and noise contributions are initially comparable. Lemma B.3 shows that the feature learning $\\langle w_{j,r}, jv\\rangle$ and noise memorization $\\langle w_{y_i,r}, \\xi_i\\rangle$ start at the same order, preventing a bias toward either true features or noise at initialization, highlighting the rigor of our conclusions.\n\n**2. Regarding parameter choices of Adam/AdamW.**\n\n**a) Are the (order of) parameter selections, for example $\\lambda$ and $B$, based on practice or to achieve the best theoretical bounds?**\n\nOur theorems select parameter orders (e.g., $\\frac{n}{B} = \\Theta(1)$ or $\\ge \\Theta(\\log \\epsilon^{-1})$) for Adam, and $\\lambda=\\widetilde\\Omega(\\frac{B^2}{n}\\wedge 1)$ up to $\\widetilde O(1)$ for AdamW) to prove distinct learning regimes rigorously. Crucially, these bounds capture familiar empirical patterns: smaller batch sizes yield better generalization for both Adam and AdamW, and AdamW tolerates much larger $\\lambda$ values. These insights agree with practical observations (Loshchilov & Hutter, 2019) and match our experiments (Figs. 1–2).\n\n**b) For the case of AdamW, the author used a large $\\lambda$; however, they do not have a theoretical negative result on small $\\lambda$. It seems that in practice a small $\\lambda$ also works in the experiments?** \n\nOur bound for AdamW, $\\lambda = \\widetilde\\Omega(\\frac{B^2}{n}\\wedge 1)$, is a sufficient condition to ensure weight decay is strong enough to halt further growth of noise memorization. Concretely, as shown by Eqs. C.5–C.6, noise memorization $\\langle w_{y_i,r}^{(t)}, \\xi_i\\rangle$ grow to $\\widetilde\\Theta(\\frac{Bs\\sigma_p}{n})$ before weight decay takes effect. Lemma D.18 then requires $\\lambda$ of this order to prevent additional noise amplification, and Lemma D.20 shows that only under this regime can mini‑batch AdamW drive test error to near zero.\n\nIn practice, small $\\lambda$ may outperform random guessing, but optimal accuracy typically requires a non‑trivial lower bound (Loshchilov & Hutter, 2019).  We will augment the manuscript to make these explicit.\n\n**c) From the theorems, how can we know that AdamW is more robust to parameter selection? For Adam, it seems that it only requires $\\lambda$ to be small enough (what if we set it to be 0?), while AdamW requires the parameter to be tuned to a regime with upper and lower bounds.**\n\nAdamW's robustness stems from its wide effective $\\lambda$ ranges. Corollary 4.3 requires Adam's $\\lambda$ to be below a tight, initialization‑dependent upper bound (i.e., $\\lambda = \\omega(\\sigma_0^{q-2})$ causes failure). Corollary 4.6 shows AdamW succeeds for any $\\lambda = \\widetilde\\Omega(\\frac{B^2}{n}\\wedge 1)$ up to $\\widetilde O(1)$, a constant‑order range independent of intialization. This explains why AdamW is more robust.\n\nFor Adam w/ $\\lambda=0$: Our conclusions need $\\lambda > 0$; otherwise, the batch size condition for Adam becomes more restrictive. We will clarify this in the next version.\n\nOur analysis shows that weight decay is necessary to prevent noise overfitting by suppressing noise memorization. Without weight decay ($\\lambda=0$), Adam fits both feature and noise, and generalization depends on the relative rates of feature learning ($\\frac{n}{B}$) and noise memorization ($s\\sigma_p$) per epoch. With $\\lambda=0$, the feature learning and noise memorization updates (Eqs. C.3 & C.4) become:\n\n\n$$\n\\begin{align*}\n\\langle w_{j,r}^{(t+1)},jv \\rangle = \\langle w_{j,r}^{(t)},jv \\rangle + j\\eta\\cdot\\mathrm{sgn}( \\sum_{i\\in\\mathcal{I}\\_t}y_i\\ell_{j,i}^{(t)}[\\sigma'(\\langle w_{j,r}^{(t)},y_iv\\rangle)-\\alpha\\sigma'(\\langle w_{j,r}^{(t)},\\xi_i\\rangle)]),  \n\\end{align*}\n$$\n$$\n\\begin{align*}\n\\langle w_{y_i,r}^{(t+1)},\\xi_i\\rangle = \\langle w_{y_i,r}^{(t)},\\xi_i\\rangle +\\eta\\cdot\\sum_{k\\in\\mathcal{B}\\_i}\\langle\\mathrm{sgn}(\\ell_{y_i,i}^{(t)}\\sigma'(\\langle w_{y_i,r}^{(t)},\\xi_i\\rangle)\\xi_i[k]),\\xi_i[k] \\rangle -\\alpha y_i\\eta\\cdot\\mathrm{sgn}(\\sum_{j\\in\\mathcal{I}\\_t}y_j\\ell_{y_i,j}^{(t)}[\\sigma'(\\langle w_{y_i,r}^{(t)},y_jv\\rangle) - \\alpha\\sigma'(\\langle w_{y_i,r}^{(t)},\\xi_j\\rangle)]),  \n\\end{align*}\n$$\nwhere $\\ell_{j,i}^{(t)}:=\\mathbb{1}\\_{y_i=j}-\\mathrm{logit}\\_j(F,x)$ and $\\mathrm{logit}\\_j(F,x) = \\frac{e^{F_{j}(w,x_i)}}{\\sum_{k\\in\\{-1,1\\}} e^{F_k(w,x_i)}}$. Since $\\ell_{y_i,i}>0$ and $\\alpha=o(1)$, both $\\langle w_{j,r}^{(t)},jv \\rangle$ and $\\langle w_{y_i,r}^{(t)},\\xi_i\\rangle$ initially grow at rates $\\Theta(\\frac{n}{B})$ and $\\widetilde{\\Theta}(s \\sigma_p)$, respectively. Later dynamics depend on which grows faster. For mini-batch Adam w/o weight decay to generalize well, the feature learning rate must dominate: $\\frac{n}{B} = \\omega(s \\sigma_p)$, similar to AdamW. Otherwise, noise memorization overtakes feature learning, as in the large-batch case. This condition imposes a stronger batch size constraint on Adam w/o weight decay, so we have $\\lambda > 0$.\n\n**d) The batch size choices in Theorem 4.2 and 4.5 are also different, with AdamW's being more restrictive. Then, how should we compare the two results?**\n\nAs similar in response 2 c), AdamW’s batch size condition is stricter due to proof technicalities. AdamW’s decoupled weight decay acts only in later training (Eq. C.5–C.6). Before that, feature learning and noise memorization both grow (Lemma D.17). Weight decay kicks in once weights are large enough (Lemma D.18), requiring feature learning to outpace noise ($\\frac{n}{B} = \\omega(s\\sigma_p)$) and noise memorization not to harm test error ($\\frac{n}{B} \\ge \\Theta(n^{1/2})$, Lemma D.20, Lines 1615-1616). The main insight is that for both Adam and AdamW, small batch size is key to avoid failure.\n\n**e) From an optimization point of view, to achieve the same error, should the time $T$ and step size $\\eta$ be set differently in the mini-batch and large-batch cases? For example, with a large batch, the noise is small, and the algorithm should converge faster. Will early stopping the training of large batches help?**\n\nGreat point! However, our analysis focuses on the properties of the converged solution from a feature learning perspective, rather than optimization speed.\n\nRegarding early stopping, our analysis shows that large-batch AdamW initially learns features but eventually reverses and overfits to noise (Figure 3a & 4a), and early stopping could potentially halt the training before this reversal (recall the gradient: $\\sigma'(\\langle w_{j,r}^{(t)},y_i v\\rangle)-\\alpha\\sigma'(\\langle w_{j,r}^{(t)},\\xi_i\\rangle)$)—thus preserving some generalization.\n\n**3. What happens to other algorithms.** \n\n**a) Under the same data model setting, what would be the performance for SGD, SGD with weight normalization, and Adam without weight normalization?**\n\nPrior work (Zou et al., 2023) shows that SGD generalizes well in settings where full-batch adaptive methods fail. Regarding SGD with WN, this introduces additional dynamics beyond our current scope, but we agree it is a promising direction for future work. For Adam without weight normalization—if referring to weight decay—its performance depends on the relative rates of feature learning and noise memorization; please see our response to 2 c) for details.\n\n**b)  Can we have more insights comparing also with these algorithms or to design better algorithms? This might be future directions.**\n\nIndeed, our conclusions suggest promising directions for future work, such as designing noise-aware update rules that explicitly suppress sample-specific noise, or adopting stage-wise optimization strategies—e.g., using Adam in early training and switching to SGD later for better generalization."}, {"review_id": "HEHJtDnU7Y", "reviewer": "Reviewer_a6No", "summary": "This paper studies adam and adamw and constructs a setting where the stochastic versions of these algorithms generalize and achieve training loss zero while large batch adam/adamw also get zero training loss but do not generalize. They first present their setting inspired by the classical feature-noise patch concatenation framework, the CNN architecture, the Adam/AdamW algorithms. Then, they present in Theorem 4.1 and Theorem 4.4 the large batch Adam/AdamW results showing that using these algorithms, the test error is large. On the other hand, in Theorems 4.2 and 4.5, they show that when using a mini-batch, the two algorithms achieve perfect test error.  They also show theoretically that AdamW can take larger weight decay than Adam. Finally, they validate their results on Cifar-10 using VGG and Resnet architectures. They swept over batch size in Figure 1 showing that the test error is low for small batch size and swept over weight decay in Figure 2, showing that AdamW can take larger weight decay than Adam.", "questions": "I listed my questions in the weaknesses section.", "limitations": "I think the limitations section is a bit incomplete. I would add the following:\n\n- In the limitations section, I would also mention that important knobs have been dismissed in the analysis: momentum, learning rate schedule, epsilon, gradient clipping. \n\n- Also maybe in the Limitations section, I would maybe mention that current vision non-Transformer architectures actually use large batch sizes? E.g. ConvNext, PeLK, Sparse Large-kernel Network (SLaK). And this seems something that maybe your setting fails to captures. This is totally fine but maybe worth mentioning.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- This paper is very well-written. The related work is appropriately cited, the setting is crystal-clear and the theorems are clear and I appreciated a lot the proof sketch and intuitions that are given. \n- The result is also interesting. To my knowledge,  First formal proof that the mini-batch variant of an adaptive method can strictly outperform its full-batch counterpart in a non-convex neural-network setting. Previous theory had only done this comparison for GD vs SGD or for convex problems.\n- The result on the upper bound for Adam’s weight decay is also very nice. \n\nWeaknesses:\n\n- Here there are two components that are in a way studied: stochastic gradients and the adaptivity. The issue is that in the current paper, the contribution of each is not crystal-clear and I think the authors would gain at clarifying this. For instance, a result that would be interesting would be to show that in your framework, SGD does not get zero test error. I think that adding this result would be very nice, and would allow us to understand the importance of adaptivity in your setting.\n\n- Related to this previous point, I think that it would be nice maybe in the proof sketch section to maybe write the optimization update of Adam so that we clearly understand the contribution of adaptivity at getting rid of the feature noise. \n\n- I think your analysis is interesting regarding batch size. Can you maybe explain what happens in your analysis for intermediate batch size? You cover small batch size and very large batch size.\n\n- Can we derive some optimal batch size B? This optimal batch size would be the middle-ground between speed (large batch size allows for fast updates) and small batch size (small batch size allows for generalization)?\n\n- I think a better job could be made in the experiments section. In particular, large scale vision experiments could be conducted beyond CIFAR-10: maybe training using Resnet-50 on Imagenet. Besides, I think it would be important to do a more fine-grained sweep over beta_1, beta_2 and the learning rate for Figures 1 and 2.\n\n- Lastly, one may wonder how incremental your derivation is from Zou et al. 2023. I think it is important to further highlight the differences of the authors' analysis and Zou et al.'s paper.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for clear and accurate summary of our work, and for highlighting its strengths—namely the manuscript’s clarity, the novel theoretical proof that mini-batch Adam/AdamW can strictly outperform full-batch variants in a non-convex neural-network setting, and the elegant upper bound on Adam’s weight decay. Below, we address each of your questions in turn.\n\n### Weaknesses & Questions\n\n**a) Here there are two components that are in a way studied: stochastic gradients and the adaptivity. The issue is that in the current paper, the contribution of each is not crystal-clear and I think the authors would gain at clarifying this. For instance, a result that would be interesting would be to show that in your framework, SGD does not get zero test error. I think that adding this result would be very nice, and would allow us to understand the importance of adaptivity in your setting.**\n\nThis is a crucial point that we will clarify. Our work primarily studies the interplay between stochastic gradients (batch size) and weight decay within the context of adaptive methods (Adam and AdamW).\n\nPrior work, which we build upon (Zou et al., 2023), has already shown that in a similar setting, full-batch SGD generalizes well, while full-batch Adam fails. This establishes that adaptivity is the source of the large-batch generalization failure. The reason is that adaptivity induces an $\\ell_1$-like normalization—approximating the behavior of SignSGD—which inherently favors sparse gradients. In our theoretical setting, this mechanism inadvertently amplifies sparse noise memorization when the noise level satisfies $s\\sigma_p = \\omega(1)$. As a result, in the full-batch regime (where $s \\gg 1$ and batch stochasticity is absent), the rate of noise memorization outpaces feature learning, ultimately leading to poor generalization. This effect is analytically captured in the update dynamics of AdamW along the noise direction:\n$$\n\\begin{align*}\n\\langle w_{y_i,r}^{(t+1)},\\xi_i\\rangle= (1-\\lambda\\eta)\\langle w_{y_i,r}^{(t)},\\xi_i\\rangle +\\eta\\cdot\\sum_{k\\in\\mathcal{B}\\_i}\\langle\\mathrm{sgn}(\\ell_{y_i,i}^{(t)}\\sigma'(\\langle w_{y_i,r}^{(t)},\\xi_i\\rangle)\\xi_i[k]),\\xi_i[k] \\rangle -\\alpha y_i\\eta\\cdot\\mathrm{sgn}(\\sum_{j\\in\\mathcal{I}\\_t}y_j\\ell_{y_i,j}^{(t)}[\\sigma'(\\langle w_{y_i,r}^{(t)},y_j v\\rangle) - \\alpha\\sigma'(\\langle w_{y_i,r}^{(t)},\\xi_j\\rangle)]). \n\\end{align*}\n$$\nOur key contribution is to show that stochasticity (from mini-batch), when combined with proper weight decay, rescues the adaptive method. It is the synergy between these components that yields good generalization: 1) adaptivity with exponential moving average normalizes gradients coordinate-wise; 2) stochasticity ensures that noise components are sparse and decorrelated across batches; 3) weight decay provides a constant regularization force that suppresses noise components over time. As we explain in the paper (Lines 215-223), this synergy allows Adam to maintain learning on shared features while its normalization and momentum dynamics, aided by weight decay, suppress the learning of sparse noise.\n\n**b) Related to this previous point, I think that it would be nice maybe in the proof sketch section to maybe write the optimization update of Adam so that we clearly understand the contribution of adaptivity at getting rid of the feature noise.**\n\nGood point! While the detailed update rules for Adam and AdamW are in Appendix C (Eqs. C.3–C.6), we agree it would be helpful to highlight them in the main paper. Section 5 builds on SignSGD, which we rigorously justify as an approximation of stochastic Adam (Appendix D, Lemma D.2). Adaptivity manifests through the `sgn` operator, normalizing update magnitudes. In large-batch settings, this amplifies noise; in contrast, in the mini-batch regime, stochasticity slows noise fitting, and weight decay suppresses it—together mitigating feature noise.\n\n**c) I think your analysis is interesting regarding batch size. Can you maybe explain what happens in your analysis for intermediate batch size? You cover small batch size and very large batch size.**\n\nThis is an excellent question. Our theory provides proofs for two distinct regimes to establish a clear, provable separation in behavior. A precise theoretical characterization of the intermediate batch size regime (e.g., for Adam, $\\frac{n}{B} = \\omega(1)$ and $O(\\log\\epsilon^{-1})$) is significantly more complex due to the limitations of our order-wise analysis. One must carefully track how the factors \\($\\beta_1^{\\tau_0}$) and \\($\\beta_2^{\\tau_0/2}$) evolve over \\($\\tau_0\\in[n/B]$) steps to decide whether the sign‐SGD–style gradient term or the decoupled weight-decay term dominates. Concretely, our update in the noise coordinates (Line 1205) reads:\n$$\n\\begin{align*}\n    & \\frac{m_{j,r}^{(t)}[k]}{\\sqrt{v_{j,r}^{(t)}[k]}+\\epsilon} = \\frac{\\beta_1^{\\tau_0}\\left(\\tilde\\Theta\\Big(g_{t-\\tau_0,j,r}^{(t)}[k]\\Big) \\pm \\tilde\\Theta\\bigg(\\frac{\\eta s\\sigma_p|\\ell_{j,i}^{(t)}|}{B}\\bigg)\\pm \\tilde\\Theta(\\lambda\\eta)-\\tilde\\Theta(\\lambda w_{j,r}^{(t)}[k])\\right)+\\tilde\\Theta(\\lambda w_{j,r}^{(t)}[k])\\pm\\tilde\\Theta(\\lambda\\eta) }{\\beta_2^{\\frac{\\tau_0}{2}}\\left(\\Theta\\Big(|g_{t-\\tau_0,j,r}^{(t)}[k]|\\Big) \\pm \\tilde\\Theta\\bigg(\\frac{\\eta s\\sigma_p|\\ell_{j,i}^{(t)}|}{B}\\bigg)\\pm \\tilde\\Theta(\\lambda\\eta)-\\Theta(|\\lambda w_{j,r}^{(t)}[k]|)\\right)+\\tilde\\Theta(|\\lambda w_{j,r}^{(t)}[k]|)\\pm\\tilde\\Theta(\\lambda\\eta)},\n\\end{align*}\n$$\nFor intermediate $B$, our order-wise bounds lose precision: it becomes hard to determine whether momentum gradient or weight decay will dominate, and thus impossible (with current techniques) to prove exactly how dynamics evolves. Nevertheless, our empirical sweep in Figure 1 covers these middle values of \\(B\\) and reveals a decline in test accuracy as batch size grows—suggesting a “critical threshold” beyond which generalization degrades sharply. Fully characterizing this intermediate regime is an important direction for future work.  \n\n**d) Can we derive some optimal batch size B? This optimal batch size would be the middle-ground between speed (large batch size allows for fast updates) and small batch size (small batch size allows for generalization)?**\n\nDeriving an exact optimal batch size is theoretically difficult and remains an open direction. However, our results provide useful guidance: good generalization requires small enough batches to satisfy the conditions in Theorems 4.2 or 4.5. This highlights a trade-off—smaller $B$ supports generalization, while larger $B$ offers faster training. In practice, the optimal batch size lies at this intersection. As shown in Figure 1, there exists a range of small batch sizes that generalize well before a sharp drop in accuracy, allowing practitioners to select a computationally efficient $B$ within this range.\n\n**e) I think a better job could be made in the experiments section. In particular, large scale vision experiments could be conducted beyond CIFAR-10: maybe training using Resnet-50 on Imagenet. Besides, I think it would be important to do a more fine-grained sweep over beta_1, beta_2 and the learning rate for Figures 1 and 2.**\n\nWe appreciate this valuable suggestion. We'll construct an ImageNet-1K subset (randomly sampling 100 training images per class with seed = 0) to ensure a controlled large-batch regime (so that $\\frac{n}{B}$ remains $\\Theta(1)$) while keeping overall compute manageable. We will train ResNet-50 on this subset and, for Figures 1 and 2, also include a systematic sweep over momentum coefficients such like $\\beta_1\\in [0, 0.5, 0.9]$, $\\beta_2\\in [0.5, 0.9, 0.95]$ (learning rates have already been swept; see Lines 274–282). These results will appear in the next version, and if our runs finish in the coming days, we'll report them here.\n\n**f) Lastly, one may wonder how incremental your derivation is from Zou et al. 2023. I think it is important to further highlight the differences of the authors' analysis and Zou et al.'s paper.**\n\nThank you for this insightful question. While we build on Zou et al. (2023), our mini-batch Adam/AdamW analysis introduces two key novelties:\n\n1. Momentum instability control: We use a one-epoch local decomposition (Lemmas D.2, D.11) to rigorously bound how stochastic “noise spikes” in Adam’s moving averages realign with the true gradient after sufficient steps.\n2. Decoupled weight-decay analysis: A refined iteration-drift framework tracks how AdamW’s direct decay on normalized momentum interacts with stochastic gradients, proving that co-tuning batch size and decay strength yields the same noise-suppression phase transition as in mini-batch Adam.\n\nMore details can be found in the response to question a) of the reviewer vcGR.\n\n### Limitations\n\n**I think the limitations section is a bit incomplete. I would add the following:**\n- **In the limitations section, I would also mention that important knobs have been dismissed in the analysis: momentum, learning rate schedule, epsilon, gradient clipping.**\n\nThank you for the insightful suggestion. We acknowledge that other critical hyperparameters—such as momentum ($\\beta$ values), learning rate schedules, and gradient clipping—are not covered in our theory. Understanding their combined effect on generalization remains a complex and largely open question in the community. We will expand on it in the next version.\n\n- **Also maybe in the Limitations section, I would maybe mention that current vision non-Transformer architectures actually use large batch sizes? E.g. ConvNext, PeLK, Sparse Large-kernel Network (SLaK). And this seems something that maybe your setting fails to captures. This is totally fine but maybe worth mentioning.**\n\nThank you for the insight. While our theory predicts challenges with large batches, some modern architectures succeed with them—possibly due to factors like normalization or architectural design that our model simplifies. We will add this as a future research direction."}, {"review_id": "9gjgEebYTm", "reviewer": "Reviewer_vcGR", "summary": "The paper studies the trajectory of Adam and AdamW on a classification problem parametrized by a 2-layer neural network with symmetric weights. The results extend the analysis in previous works to the small batch-size settings. Theorem 4.2 and 4.4 show that, when batch size is sufficiently small compared to sample size, Adam and AdamW can converge into solutions with not only small training error, but also small test error.", "questions": "1) The batch size discussed in Theorem 4.2 is very small $B=O\\left(\\frac{n}{\\log\\frac{1}{\\epsilon}}\\right)$. What happens if $\\epsilon$ is very small and $B$ is larger than the bound? Is the $O\\left(\\frac{1}{\\log\\frac{1}{\\epsilon}}\\right)$ upper bound on $n$ necessary for the $o(1)$ population risk?\n\n2) Could authors summarize a few notes to practitioners about how to choose batchsize?\n\n3) In line 248, why do we need to add feature to the noise vector $\\xi$?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n\n1) The paper tackles an important problem about the convergence of Adam and AdamW. The finding that small batch-size training improves generalization performance aligns with practitioners' findings.\n\n2) Theoretical presentations are clear and understandable.\n\nWeaknesses\n\n1) It seems many technical developments are based on Zou et al. 2023, authors could further elaborate on the technical contribution, especially the novel analytical tools introduced to analyze the mini-batch Adam.\n\n2) According to line 145, the noise $\\xi$ is generated to be sparse. This is usually not the case in practice. Authors should justify such an assumption.", "was_revised": true, "final_justification": "The rebuttal stated the theoretical innovations and the importance to fine-tune weight-decay and batch size together.", "rebuttal": "We thank the reviewer for their clear summary and positive assessment of our work—particularly noting the extension to small batch-size settings and the alignment between our theoretical findings and practitioners’ observations. We’re also glad the presentation of our theorems and proofs resonated as clear and understandable. Below, we address your questions in detail.\n\n### Weaknesses\n\n**a) It seems many technical developments are based on Zou et al. 2023, authors could further elaborate on the technical contribution, especially the novel analytical tools introduced to analyze the mini-batch Adam.**\n\nThank you for this insightful question. While we build on Zou et al. (2023), our analysis of mini-batch Adam/AdamW introduces several novel technical tools to (i) control the stochastic, momentum-based instabilities unique to the mini-batch regime and (ii) handle the distinct, decoupled weight decay mechanism of AdamW:\n\n1. Stochasticity induces momentum instability. Full-batch Adam simply aggregates deterministic gradients, but mini-batch sampling injects both noise and temporal \"inertia\" via Adam’s moving averages. These fluctuations can cause \"noise spikes\" that persist across iterations. We address this with a local, one-epoch decomposition that tracks the dynamics of the momentum buffers over the $\\frac{n}{B}$ steps per epoch. Lemmas D.2 and D.11 then establish a rigorous sign-GD approximation bound, proving that—after sufficient mini-batch steps—the momentum direction realigns with the true gradient rather than lingering noise.\n2. Decoupled weight decay in AdamW. In Adam, weight decay simply augments the raw gradient and is processed through the sign approximation. AdamW instead applies decay directly to the normalized momentum ($m/(\\sqrt{v}+\\epsilon)$), so decay and stochastic gradients jointly determine the update. We handle this via a refined iteration-drift analysis that carefully tracks how the last few momentum updates evolve under decoupled weight decay. This enables us to show that—when batch size and decay strength are co-tuned—mini-batch AdamW exhibits the same noise-suppression \"phase transition\" as mini-batch Adam, with even cleaner control over the regularization path.\n\n**b) According to line 145, the noise $\\xi$ is generated to be sparse. This is usually not the case in practice. Authors should justify such an assumption.**\n\nModeling the noise as sparse with disjoint support is indeed a simplification, but it is a standard abstraction in the theoretical feature-learning literature (citer in Lines 135–136) and lays the groundwork for valuable future extensions—in particular, extending to more general noise models (see below). By isolating a few small, informative “feature patches” within a vast field of irrelevant pixels or sensor noise, we obtain a tractable framework to dissect the competition between feature learning and spurious noise memorization. One can view our two-layer CNN as capturing an intermediate representation in a deeper network whose inputs follow a “sparse + background” distribution (Papyan et al. 2017; Yang 2019).\n\nIn real data, noise is rarely perfectly disjoint. However, our analysis naturally extends to overlapping or structured noise by leveraging concentration bounds on the aggregate noise energy (e.g., modeling each pixel’s perturbation as sub-Gaussian and applying matrix- or vector-Bernstein inequalities). These tools demonstrate that, provided the total noise remains controlled, the same qualitative phase transition—where feature learning dominates below a threshold and noise memorization takes over above it—continues to hold. Extending our model to this general noise setting is a future work of ours.\n\n### Questions\n\n**a) The batch size discussed in Theorem 4.2 is very small $B=O(\\frac{n}{\\log\\frac{1}{\\epsilon}})$. What happens if $\\epsilon$ is very small and $B$ is larger than the bound? Is the $O(\\frac{1}{\\log\\frac{1}{\\epsilon}})$ upper bound $n$ on necessary for the $o(1)$ population risk?**\n\nThank you for this sharp question. Let's clarify the bound and the role of $\\epsilon$. In our paper, we take $\\epsilon=\\Theta(\\lambda\\eta)$. Since $\\eta$ equals to $1/\\mathrm{poly}(n)$, one has that $\\epsilon$ can be polynomially small $1/\\mathrm{poly}(n)$. Therefore, $\\log(1/\\epsilon)=\\Theta(\\log n)$, and the stated bound becomes $B\\le O(n/\\log n)$.\n\n- **If B is larger than this bound:** Our theoretical analysis no longer guarantees good generalization. Indeed, our large-batch theorems (Theorem 4.1 & 4.4) and empirical results (Figure 1) show that performance degrades as $B$ increases towards $n$.\n- **Is the bound necessary?** The bound is a sufficient condition that arises from our proof technique. The core mechanism it ensures is that the number of iterations per epoch ($n/B$) is large enough for Adam's momentum estimates to \"forget\" the gradients from any single, sparse noise patch. If $n/B$ is too small (i.e., $B$ is too large), this forgetting mechanism fails, and noise can accumulate, leading to noise overfitting. So while the bound might not be strictly necessary, our theory clearly demonstrates a phase transition where performance breaks down as the batch size becomes too large.\n\n**b) Could authors summarize a few notes to practitioners about how to choose batchsize?**\n\nCertainly! We offer the following practitioner‐focused guidelines on choosing batch size when using Adam/AdamW:\n\n1. Smaller batches often generalize better. Our theory provides formal support for the common heuristic that smaller batch sizes can improve generalization with Adam and AdamW. Practitioners should be cautious when scaling to very large batch sizes with these optimizers.\n2. Batch size and weight decay must be co‐tuned. These two hyperparameters interact: you cannot optimize one in isolation. Concretely, our analysis shows that mini‐batch Adam/AdamW requires $\\frac{n}{B}\\ge\\log(\\lambda\\eta)^{-1}$ to control noise memorization. Thus, if you increase the batch size $B$ to boost hardware throughput, you should also increase the weight decay $\\lambda$ (albeit only logarithmically) to maintain the same generalization behavior.\n3. Find the “sweet spot” for throughput vs. accuracy.There is a trade‐off between hardware efficiency (which favors large batches) and generalization (which favors small batches). Our experiments (Figure 1) reveal a critical batch‐size threshold beyond which test error rises sharply. We recommend practitioners perform a simple grid search around that threshold to identify the largest batch size that does not compromise test performance.\n\n**c) In line 248, why do we need to add feature to the noise vector $\\xi$?**\n\nThis refers to the term $-\\alpha y v$ added to the noise patch $\\xi$. Its purpose is to make the learning problem more challenging and realistic (avoid complete orthogonality, a trivial case). In real-world data, noise is not always purely random but can contain spurious correlations (e.g., consistent background clutter in images) that are misleading.\n\nTechnically, introducing this negative correlation between the noise patch and the true feature $v$ means that an optimizer that naively fits the noise patch will be actively pushed away from the correct feature direction. This setup allows our theory to crisply differentiate algorithms that successfully isolate the true feature from those that are misled by the noise, which is essential for proving the failure mode of large-batch Adam/AdamW."}, {"review_id": "h8Si484GY9", "reviewer": "Reviewer_qgrL", "summary": "This paper performs a theoretical analysis of the training dynamics of Adam and AdamW on a two-layers CNN trained on a 1 dimensional synthetic dataset. In that particular example, the authors are able to show analytically that mini-batch Adam converges to lower test-error solutions than large-batch Adam. They are also able to show in that setting that AdamW admits a larger beneficial weight decay regularization range than Adam. They show experimentally that their finding hold on CIFAR10 trained with VGG16 and ResNet18.", "questions": "* In Definition 3.1, condition 2 in unclear. From skimming through the literature mentioned, I was able to understand that one wants the noise vector to be perpendicular to the signal vector. Is that what is meant in condition 2?\n\n* I am struggling to understand the significance of Corollary 4.6. Is there a way to rephrase it to make its significance clearer?", "limitations": "Limitations are acknowledged in the conclusion concerning the validity of the analysis to the toy model. However, some wordings sound like overstatements in this context at the same time. For instance in the conclusion we find: \"In this work, we theoretically and empirically analyze the generalization performance of Adam and AdamW under varying batch sizes and weight decay parameters\" although this is only done on a toy example and not in general, as the sentence may let the reader think.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\n* The author were able to carry a fully theoretical analysis of the Adam and AdamW dynamics where the generalization advantage of mini-batch over large-batch becomes apparent\n* The presentation is clear \n\n**Weaknesses:**\n\n* Since the theoretical claims are made in the context of a toy example, expanded experimental verification would give more credibility to the claims, and in particular experiments demonstrating the two learning stages identified (stage 1: pattern learning, and stage 2: regularization)\n* Error bars on the experiment figures are missing; it seems only 1 seed has been used\n* The proofs rely very much on the special structure of the toy example (including the special form of the dataset); this casts some doubt on the general validity of this type of arguments (even in the basic context of a DNN on a generic dataset). A discussion on how obstacles in this direction may be lifted would strengthen the article in my opinion. At present, it is possible to understand the theoretical results presented as being only due to non-generalizable artifacts of the toy model and datasets considered\n* A Discussion/comparison with a recent line of work explaining the implicit bias of Adam toward flat and generalizable solutions in the case of fully general models and datasets is missing; in particular: https://arxiv.org/abs/2309.00079\n* The language at times sounds like overstatements especially in the abstract, intro, and conclusion where the results may sound like they are proven for the general case, although they are valid only in the case of the toy example considered. Although the limitations are acknowledged, I would tone down the claims slightly", "was_revised": true, "final_justification": "Thank you for your detailed answers. I raised my score in consequence.", "rebuttal": "We thank the reviewer for their concise and accurate summary of our work, and for highlighting the fully theoretical analysis revealing the generalization benefit of mini-batch Adam, as well as the broader weight‐decay range of AdamW. We’re pleased the presentation was clear. Below, we address your questions in detail.\n\n### Weaknesses\n\n**a) Since the theoretical claims are made in the context of a toy example, expanded experimental verification would give more credibility to the claims, and in particular experiments demonstrating the two learning stages identified (stage 1: pattern learning, and stage 2: regularization)**\n\nGood point! These experiments are already included in Figures 3 and 4 (Appendix, pp. 21–22), and we will add clearer references to them in the main text.\n\nFigures 3 and 4 (Appendix, pp. 21–22) plot the dynamics of feature learning vs. noise memorization for both mini-batch and large-batch Adam/AdamW, directly visualizing the two-stage process identified by our theory (see Appendix A.2).\n\nIn Figure 4(a), large-batch AdamW initially learns the true feature (measured by $\\max_{r \\in [m]} \\langle w_{j,r}, jv \\rangle$), but is quickly overtaken by noise memorization (measured by $\\min_{i \\in [n]: y_i = j} \\max_{r \\in [m]} \\langle w_{j,r}, xi_i \\rangle$), leading to alignment with the spurious direction $-\\alpha yv$. In Stage 2 (Regularization), noise memorization grows at a logarithmic rate, matching our theoretical predictions.\n\nFigure 4(b) shows that mini-batch AdamW maintains stable feature learning throughout, while noise memorization is suppressed in Stage 2 due to vanishing gradients. These trends align closely with our lemmas and provide direct empirical support for our two-stage analysis.\n\n**b) Error bars on the experiment figures are missing; it seems only 1 seed has been used**\n\nThis is an important point for ensuring experimental statistical significance. We have re-run our CIFAR-10 experiments with five random seeds (0–4), using the best hyperparameters selected based on seed 0. We will include error bars in the next version and provide full result tables (mean ± std) below. These results support and confirm our original claims.\n\n**Corresponding to Figure 1.**\n\nAdam: Test Error (%) vs. Batch Size ($B$)\n\n| Model    | 16         | 32         | 64         | 256        | 1024       | 4096       | 8192       |\n| -------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- |\n| VGG16    | 14.97±0.43 | 15.90±0.85 | 15.29±0.53 | 15.76±0.98 | 19.25±1.40 | 24.42±3.47 | 26.38±3.83 |\n| ResNet18 | 15.82±0.41 | 15.99±0.63 | 15.69±0.40 | 17.23±0.92 | 20.27±1.57 | 27.60±2.17 | 29.96±1.41 |\n\nAdamW: Test Error (%) vs. Batch Size ($B$)\n\n| Model    | 16         | 32         | 64         | 256        | 1024       | 4096       | 8192       |\n| -------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- |\n| VGG16    | 14.18±0.35 | 13.74±0.14 | 13.19±0.36 | 14.20±0.74 | 16.89±1.24 | 21.81±1.07 | 25.04±2.15 |\n| ResNet18 | 15.15±0.44 | 15.16±0.47 | 15.24±0.50 | 15.58±0.49 | 18.51±1.13 | 20.69±0.70 | 25.91±2.54 |\n\n**Corresponding to Figure 2.**\n\nResNet18: Test Error (%) vs. Weight Decay ($\\lambda$)\n\n| Optimizer | 5e-7       | 1e-6       | 5e-6       | 1e-5       | 5e-5       | 1e-4       | 5e-4       | 1e-3       | 5e-3       | 1e-2       | 5e-2       | 1e-1       | 5e-1       |\n| --------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- |\n| Adam      | 15.37±0.47 | 15.66±0.39 | 15.41±0.50 | 15.29±0.53 | 15.56±0.62 | 15.63±0.35 | 15.65±0.57 | 15.89±0.39 | 15.61±0.39 | 15.76±0.67 | 17.35±0.71 | 22.01±1.12 | —          |\n| AdamW     | —          | —          | —          | —          | —          | 15.12±0.26 | 15.10±0.29 | 15.51±0.38 | 15.35±0.48 | 15.15±0.44 | 15.22±0.44 | 15.83±0.37 | 16.46±0.51 |\n\nVGG16: Test Error (%) vs. Weight Decay ($\\lambda$)\n\n| Optimizer | 5e-7       | 1e-6       | 5e-6       | 1e-5       | 5e-5       | 1e-4       | 5e-4       | 1e-3       | 5e-3       | 1e-2       | 5e-2       | 1e-1       | 5e-1       |\n| --------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- |\n| Adam      | 14.46±0.32 | 14.21±0.35 | 14.34±0.32 | 14.51±0.31 | 14.31±0.16 | 14.51±0.51 | 14.97±0.43 | 15.14±0.35 | 15.94±0.58 | 16.51±0.67 | 17.75±0.62 | 23.72±2.43 | —          |\n| AdamW     | —          | —          | —          | —          | —          | 14.10±0.22 | 14.23±0.37 | 14.27±0.36 | 13.97±0.24 | 14.18±0.35 | 14.75±0.66 | 15.36±0.40 | 16.72±0.32 |\n\n**c) The proofs rely very much on the special structure of the toy example (including the special form of the dataset); this casts some doubt on the general validity of this type of arguments (even in the basic context of a DNN on a generic dataset). A discussion on how obstacles in this direction may be lifted would strengthen the article in my opinion. At present, it is possible to understand the theoretical results presented as being only due to non-generalizable artifacts of the toy model and datasets considered**\n\nWe appreciate the concern that our toy model may not fully capture real-world complexity. However, such simplified settings are standard in theoretical deep learning (Lines 135–136). Our feature–noise patch abstraction isolates the tradeoff between learning sparse, task-relevant signal and memorizing irrelevant noise—reflecting empirical patterns in CNNs (Papyan et al., 2017; Yang, 2019). Interpreting our two-layer CNN as modeling an intermediate layer in deeper networks further supports its relevance. Experiments confirm that the mechanism we analyze persists beyond the toy setting.\n\nThe primary obstacles to extending this analysis to general datasets and architectures are: (1) the lack of a simple mathematical structure for real-world data distributions and (2) the highly complex, non-linear dynamics of very deep networks. We will further discuss them in the next version.\n\n**d) A Discussion/comparison with a recent line of work explaining the implicit bias of Adam toward flat and generalizable solutions in the case of fully general models and datasets is missing; in particular: [the arxiv paper]**\n\nThank you for pointing us to the interesting work. We will add a discussion to our Related Works section.\n\nThat work analysis shows that Adam implicitly penalizes the one-norm of (perturbed) gradients, biasing solutions toward flat regions of the loss landscape. Our work offers a complementary and quantitatively precise perspective: within a feature-learning framework, we analyze how Adam/AdamW suppress noise and recover true features, controlled by batch size and weight decay.\n\nWhile their continuous-time view explains the emergence of flat minima, our discrete-time results clarify how this flatness coincides with robustness to noise. Together, these perspectives suggest that mini-batch Adam’s bias toward flat regions is closely tied to its ability to avoid noise overfitting and promote generalization/feature learning.\n\n**e) The language at times sounds like overstatements especially in the abstract, intro, and conclusion where the results may sound like they are proven for the general case, although they are valid only in the case of the toy example considered. Although the limitations are acknowledged, I would tone down the claims slightly**\n\nWe appreciate the feedback and agree that precise language is crucial. We will revise the conclusion to avoid overstatements. We also note that these statements are already mentioned in the abstract (Lines 7–8), introduction (49–51), problem setup (159–164), and limitation (359–362).\n\n### Questions\n\n**a) In Definition 3.1, condition 2 in unclear. From skimming through the literature mentioned, I was able to understand that one wants the noise vector to be perpendicular to the signal vector. Is that what is meant in condition 2?**\n\nThank you for asking for this clarification. The procedure in condition 2 is designed to create a sparse noise vector $\\xi$ where the randomly generated Gaussian noise component has a support that is disjoint from the support of the signal vector $v$. Specifically,$v$ is 1-sparse on coordinate $\\{1\\}$ , while the random Gaussian noise is generated on $s$ coordinates chosen from $\\{2,\\ldots,d\\}$. This disjoint support is a key feature. The final noise vector $\\xi$ is not strictly orthogonal to the signal vector $v$ because of the added feature noise term $−\\alpha yv$, which makes the learning problem more challenging and realistic. \n\n**b) I am struggling to understand the significance of Corollary 4.6. Is there a way to rephrase it to make its significance clearer?**\n\nCertainly. The significance is that it provides a clear theoretical gap between effective weight decay regimes for Adam and AdamW:\n$$\n\\begin{align*}\n\\lambda_{\\mathrm{Adam}} \\sim \\sigma_0^{q-2}\\ll \\frac{B^2}{n}\\wedge 1 \\sim \\lambda_{\\mathrm{AdamW}}\n\\end{align*}\n$$\n\n- For **Adam**, the effective weight decay λ must be extremely small (larger than 0), as its upper bound is tied to an initialization dependent parameter ($\\sigma_0^{q-2}$).\n- For **AdamW**, the effective range is more independent, with any $\\lambda = \\widetilde\\Omega(\\frac{B^2}{n}\\wedge 1)$ up to $\\widetilde O(1)$, i.e., a broad, constant‑order window independent of intialization.\n\nThis result provides the first theoretical explanation for the widely-observed empirical fact that Adam is highly sensitive to weight decay and requires careful tuning, while AdamW's decoupled design makes it significantly more robust (e.g., Loshchilov and Hutter, 2019). Our experiments in Figure 2 validate this finding. We will revise the discussion following the corollary to state this significance more directly.\n\n**limitation please see Weakness e).**"}]

Author rebuttal:

Meta-review: The paper at hand presents a modern, thorough analysis of Adam(W) generalization capabilities on deep neural networks. All reviewers agree that the paper deserves publication, and I think that this work is of great importance towards understanding the fundamental properties of Adam near the solution. Despite this, I find an evident lack of theoretical and practical discussion of weight decay effects on LM training, which was, for instance, explored in [A]. While this setting is different and I consider the paper great, I think the authors should also discuss modern setups in their revised version.

I recommend acceptance, also in light of the novel theoretical tools and techniques presented, which are valuable for the community.

[A] D'Angelo, Francesco, et al. "Why do we need weight decay in modern deep learning?" Advances in Neural Information Processing Systems 37 (2024): 23191-23223.
Decision: Accept (poster)