Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks

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.

Weaknesses

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)

Good 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.

Figures 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).

In 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.

Figure 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.

b) Error bars on the experiment figures are missing; it seems only 1 seed has been used

This 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.

Corresponding to Figure 1.

Adam: Test Error (%) vs. Batch Size ($B$)

AdamW: Test Error (%) vs. Batch Size ($B$)

Corresponding to Figure 2.

ResNet18: Test Error (%) vs. Weight Decay ($\lambda$)

VGG16: Test Error (%) vs. Weight Decay ($\lambda$)

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

We 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.

The 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.

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]

Thank you for pointing us to the interesting work. We will add a discussion to our Related Works section.

That 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.

While 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.

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

We 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).

Questions

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?

Thank 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.

b) I am struggling to understand the significance of Corollary 4.6. Is there a way to rephrase it to make its significance clearer?

Certainly. The significance is that it provides a clear theoretical gap between effective weight decay regimes for Adam and AdamW:

• 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}$).
• 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.

This 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.

limitation please see Weakness e).

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.

Weaknesses

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.

Thank 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:

1. 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.
2. 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.

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.

Modeling 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).

In 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.

Questions

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?

Thank 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)$.

• 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$.
• 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.

b) Could authors summarize a few notes to practitioners about how to choose batchsize?

Certainly! We offer the following practitioner‐focused guidelines on choosing batch size when using Adam/AdamW:

1. 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.
2. 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.
3. 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.

c) In line 248, why do we need to add feature to the noise vector $\xi$?

This 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.

Technically, 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.

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.

Weaknesses & Questions

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.

This 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).

Prior 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:

Our 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.

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.

Good 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.

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.

This 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:

For 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.

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)?

Deriving 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.

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.

We 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.

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.

Thank you for this insightful question. While we build on Zou et al. (2023), our mini-batch Adam/AdamW analysis introduces two key novelties:

1. 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.
2. 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.

More details can be found in the response to question a) of the reviewer vcGR.

Limitations

I think the limitations section is a bit incomplete. I would add the following:

• In the limitations section, I would also mention that important knobs have been dismissed in the analysis: momentum, learning rate schedule, epsilon, gradient clipping.

Thank 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.

• 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.

Thank 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.

We would like to add more experimental results based on your previous questions. It takes some time and we just finished the experiments.

e) I think a better job ... large-scale vision experiments could be conducted beyond CIFAR-10 ..., more fine-grained sweep over beta_1, beta_2 and the learning rate for Figures 1 and 2.

These experiments further confirm our conclusions. The details are provided below:

1. Large-Scale Vision Experiments with ResNet-50 on ImageNet-1K Subset

We constructed an ImageNet-1K subset (100 training images per class, seed=0) to maintain a controlled large-batch regime ($\frac{n}{B} = \Theta(1)$) while managing compute. We trained ResNet-50 on this subset for 90 training epochs, with standard ImageNet preprocessing. The key results as follows:

• Top5 Validation Error vs. Batch Size: We tested learning rates (1e-4) with $\beta_1=0.9$, $\beta_2=0.99$, and weight decay settings (Adam: 1e-4; AdamW: 1e-2). Results are shown as follows, confirming our conclusions on batch size sensitivity for Adam and AdamW.

ResNet50: Top5 Val Err (%) vs. Batch Size

• Top5 Validation Error vs. Weight Decay: With batch size=64, lr=1e-3, and $\beta_1=0.9$, $\beta_2=0.99$, we analyzed weight decay ($\lambda$) effects. The following table show the results, confirming AdamW’s stability across broader decay ranges.

ResNet50: Top5 Val Err (%) vs. Weight Decay

2. Fine-Grained Sweep over $\beta_1$ and $\beta_2$ for Figures 1 and 2

As suggested, we performed a systematic sweep over momentum coefficients, with $\beta_1\in[0, 0.5, 0.9]$, $\beta_2\in[0.5, 0.9, 0.95]$, plus the original $\beta_1=0.9$, $\beta_2=0.99$ (total 9 combinations, satisfying $\beta_1^2 < \beta_2$). Learning rates were already swept (Lines 274–282), and key results for Figures 1 and 2 are as follows:

• Figure 1 (Test Error vs. Batch Size): Results (mean±std) for Adam and AdamW on VGG16/ResNet18 confirm that our conclusions on batch size behavior hold across $\beta$ combinations.

Adam: Test Err (%) vs. Batch Size

AdamW: Test Err (%) vs. Batch Size

• Figure 2 (Test Error vs. Weight Decay): Results (mean±std) for ResNet18/VGG16 show AdamW’s robustness across $\lambda$ ranges, consistent with our original findings.

ResNet18: Test Err (%) vs. Weight Decay

VGG16: Test Err (%) vs. Weight Decay

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.

Weaknesses & Questions

1. Regarding the setting.

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.

Our 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.

b) Specifically, does the noise model resemble some cases in practice?

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.

c) How realistic is the parameter selection, like $s$, $\sigma_p$, and $\alpha$, and why are they set to the order in the paper?

We 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.

By 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.

2. Regarding parameter choices of Adam/AdamW.

a) Are the (order of) parameter selections, for example $\lambda$ and $B$, based on practice or to achieve the best theoretical bounds?

Our 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).

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?

Our 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.

In 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.

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.

AdamW'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.

For 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.

Our 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:

where $\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$.

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?

As 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.

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?

Great point! However, our analysis focuses on the properties of the converged solution from a feature learning perspective, rather than optimization speed.

Regarding 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.

3. What happens to other algorithms.

a) Under the same data model setting, what would be the performance for SGD, SGD with weight normalization, and Adam without weight normalization?

Prior 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.

b) Can we have more insights comparing also with these algorithms or to design better algorithms? This might be future directions.

Indeed, 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.