Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks

AI summary

Authors: Xuan Tang, Han Zhang, Yuan Cao, Difan Zou (School of Computing & Data Science, and Institute of Data Science, The University of Hong Kong) arXiv: 2510.11354v1 [cs.LG], 13 Oct 2025

Overview

Research area: Machine learning theory — specifically the optimization and generalization theory of adaptive gradient methods (Adam and AdamW) trained on neural networks.

Technical level: Advanced. The paper is a theoretical learning-theory paper built on feature-learning analysis, with convergence theorems, corollaries on weight decay bounds, and supporting experiments.

Scope (one sentence): The paper gives the first theoretical characterization of how mini-batch size affects the generalization of stochastic Adam and AdamW, analyzing two-layer over-parameterized CNNs on an image data model and deriving distinct effective weight decay bounds for the two optimizers.

What This Paper Is About

Most existing theory analyzes full-batch Adam, but practitioners train with stochastic (mini-batch) Adam — and unlike SGD, stochastic Adam does not converge to its full-batch counterpart even with infinitesimal learning rates. This leaves a gap: it is unclear from theory how batch size and weight decay jointly govern Adam's generalization. The paper closes that gap by proving when large-batch Adam/AdamW converge to solutions with poor test error (at least 1/2 − o(1)), when mini-batch Adam/AdamW reach near-zero test error, and how much effective weight decay each optimizer can tolerate.

Key Contributions

1. Large-batch failure proved for both optimizers. Theorem 4.1 and Theorem 4.4 rigorously prove that large-batch Adam and large-batch AdamW converge to solutions with poor test error in nonconvex settings, even with proper weight decay. This extends prior full-batch Adam results (Zou et al., 2023b) to AdamW and to the large-batch regime.

2. Mini-batch success proved. Theorem 4.2 and Theorem 4.5 prove that stochastic Adam and AdamW achieve near-zero test error in nonconvex settings with appropriate weight decay. The stated mechanism is twofold: stochastic gradients implicitly regularize the trajectory by slowing noise fitting while preserving feature learning, and weight decay explicitly suppresses residual noise components.

3. A SignSGD/SignSGDW approximation for the stochastic regime. Under constant β₁, β₂ hyperparameters, the paper proves stochastic Adam and AdamW can be rigorously approximated by SignSGD (Bernstein et al., 2018) and SignSGDW (with decoupled decay), respectively — extending the known full-batch Adam-to-SignGD correspondence to stochastic training. Appendix C states this approximation holds precisely when gradient magnitudes dominate optimization noise, e.g. |g_{t,j,r}^{(t)}[k]| ≥ Θ̃(η), where η is the learning rate.

4. Distinct effective weight decay bounds. Corollary 4.3 and Corollary 4.6 derive separate upper bounds: Adam permits a strictly smaller maximum effective λ than AdamW, because Adam's adaptive normalization amplifies the effective impact of weight decay while AdamW's decoupled decay does not. This is offered as a theoretical explanation for why Adam requires more sensitive λ tuning.

Main Findings

• Large-batch Adam overfits. Under η = 1/poly(n), 0 < λ = o(σ₀^{q−2} σ_p / n), T = poly(n)/η epochs, and batch size with n/B = Θ(1), Theorem 4.1 states that with probability at least 1 − n^{−1}, training error is zero (err_S(W^{(T)}) = 0) but test error is high (err_D(W^{(T)}) ≥ 1/2 − o(1)).

• Mini-batch Adam generalizes. Under the same learning rate and λ conditions, but with batch size satisfying n/B ≥ Θ(log ε^{−1}) where ε is the Adam hyperparameter, Theorem 4.2 states that with probability at least 1 − n^{−1}, training error is zero and test error is near-zero (err_D(W^{(T)}) = o(1)).

• The mechanism is asymmetry between shared features and per-sample noise. The feature vector is shared across all data points, so feature learning is robust to batch size; noise vectors are data-specific, so mini-batches expose only a subset of them per update. Adam's coordinate-wise normalization maintains consistent learning rates for shared features while substantially slowing noise memorization.

• Adam has a narrow weight decay window. Corollary 4.3: if λ = ω(σ₀^{q−2}), then with probability at least 1 − n^{−1}, training is stuck at the initialization. The corollary frames σ₀^{q−2} as a theoretical upper bound on effective weight decay for Adam.

• Large-batch AdamW also overfits. Theorem 4.4 (with λ = Ω̃(B²/n ∧ 1), λ = Õ(1), and n/B = Θ(1) or n/B = o(s σ_p)) gives zero training error and test error ≥ 1/2 − o(1). The explanation given is that weight decay has minimal influence in early training, so the model overfits the feature noise −α y v before decoupled decay takes effect; weight decay then guides the model to nearby local minima that preserve the memorized noise.

• Mini-batch AdamW generalizes. Theorem 4.5 (with the same λ conditions plus n/B ≥ Θ(log ε^{−1}) and n/B = ω(s σ_p ∨ n^{1/2})) gives zero training error and near-zero test error. AdamW's decoupled decay penalizes weights independently of gradients, exerting significant regularization mainly in later phases.

• The weight decay gap. Corollary 4.6 states λ_Adam ∼ σ₀^{q−2} ≪ B²/n ∧ 1 ∼ λ_AdamW. Adam's admissible λ is bounded above by an initialization-dependent term; AdamW's effective range runs from Ω̃(B²/n ∧ 1) up to Õ(1), a constant-order window independent of initialization.

• Empirical validation. Training VGG16 and ResNet18 on CIFAR-10 with Adam (λ = 5 × 10^{−4}) and AdamW (λ = 1 × 10^{−2}), selecting the optimal learning rate from {5 × 10^{−4}, 1 × 10^{−4}, 1 × 10^{−5}}, both optimizers show sharp test-error degradation once batch size exceeds a critical threshold (Figure 1). At a fixed batch size of 16 (Figure 2), Adam's error spikes for λ > 0.05 while AdamW remains robust up to λ = 0.5.

• Additional experiments. The paper reports feature-learning and noise-memorization dynamics (Figures 3, 4), weight decay sensitivity (Figures 5, 6), error bars across random seeds and momentum parameters (Figures 7–10), and large-scale vision experiments with ResNet-50 on ImageNet-1K (Figures 11, 12), all in Appendix D. Synthetic-data experiments in Appendix D are said to confirm that the large-batch failure stems from noise-dominated solutions.

Methodology in Plain English

The authors study a deliberately simplified but analyzable setting. Data points are vectors in R^{2d} made of two patches: one is a signal patch y · v where v is 1-sparse in R^d (a single meaningful coordinate), and the other is a noise patch ξ drawn from a masked Gaussian with feature noise α y v subtracted. The sparsity and scale parameters are set as s = Θ(d^{1/2}/n²), σ_p² = Θ(1/(s · polylog(n))), and α = Θ(σ_p · polylog(n)). A randomly chosen one of the two patches holds the signal; the other holds the noise. This isolates the competition between learning real features and memorizing noise.

The model is a two-layer CNN with width m = polylog(n), activation σ(x) = [x]_+^q with q ≥ 3, weights initialized from N(0, σ₀² I_d) with σ₀ = Θ(d^{−1/4}), and a logit index j ∈ {±1} matched to the label. Training minimizes cross-entropy; Adam's objective includes a (λ/2)‖W‖_F² term, while AdamW's objective has no explicit penalty, since its decay is applied directly to the weights in the update step.

The theoretical work then tracks two quantities through training: feature learning ⟨w_{j,r}, j v⟩ (alignment with the true feature direction) and noise memorization ⟨w_{y_i,r}, ξ_i⟩ (correlation with an individual sample's noise patch). The proof follows a two-stage framework: Stage I, pattern learning, where regularization is negligible and the model rapidly learns dominant patterns; Stage II, regularization, where gradients shrink and regularization dominates, driving the model to a local minimum that retains stage-I patterns. Because Adam resembles SignGD when the learning rate is small or β₁, β₂ approach zero, the authors first analyze SignSGD and then extend to stochastic Adam in Appendix C.

Note: the provided paper content is truncated partway through Section 5, so the full proof outlines, Appendix A gradient calculations, Appendix B proofs of the remaining theorems, Appendix C Adam analysis, and Appendix D experimental details are not included here; the paper states proofs appear in Appendices A–C.

Why This Matters

Impact on research. Prior theory largely analyzed full-batch Adam, and the practical stochastic variant can behave fundamentally differently even at vanishingly small learning rates. This work supplies matching large-batch and mini-batch guarantees on a concrete nonconvex architecture, extends the Adam-to-SignGD approximation story from the full-batch regime to the stochastic one, and separates the effective weight decay profiles of Adam and AdamW with explicit bounds rather than empirical folklore.

Real-world applications:

• Image classification under memory-constrained batch sizes. The paper trains VGG16 and ResNet18 on CIFAR-10 and ResNet-50 on ImageNet-1K, giving a direct read on how batch size choices affect test error for vision workflows.
• Large language model training. The paper notes that Adam and its variants are used to train LLMs such as GPT, LLaMA, and Deepseek; its analysis speaks to batch-size and weight-decay choices in those pipelines.
• Optimizer selection between Adam and AdamW. The corollaries indicate Adam tolerates only a small effective λ while AdamW tolerates a window up to Õ(1), which informs which optimizer needs tight tuning.
• Distributed or large-batch training. When batch size is scaled up for hardware throughput, the theorems identify the regime (n/B = Θ(1)) where test error degrades toward 1/2 − o(1).

Industry relevance. The results argue that "just use a bigger batch" is not generalization-neutral for Adam-family optimizers, and that mini-batch training is not only a memory concession but an implicit regularizer. They also quantify why weight decay tuning for Adam is costly and fragile (spiking for λ > 0.05 in the reported experiments) while AdamW remains robust up to λ = 0.5, which matters for automated hyperparameter search budgets.

Future Directions

• Extending beyond the single-patch data model. The paper states that results can be extended to multi-patch or denser features and noises by assuming sub-Gaussian noise and using concentration inequalities such as Bernstein bounds, with similar qualitative behavior expected as long as total noise remains controlled — but this is not carried out here.
• Architectures beyond two-layer CNNs. Related work cited by the authors analyzes sign-gradient descent in two-layer Transformers; whether the batch-size and weight decay bounds carry over to Transformers, deeper networks, or diffusion models remains open.
• Refining the momentum theory. The paper proves the SignSGD/SignSGDW approximation under constant β₁, β₂ and reports momentum-parameter experiments (Figures 9, 10), leaving the precise dependence on momentum in stochastic Adam an open direction.
• Bridging to LLM-scale training. The paper motivates itself with the use of Adam in GPT, LLaMA, and Deepseek style models, but its theory is developed for a two-layer CNN on a synthetic image model; how the effective weight decay bounds scale in that setting is not reported.

Target Audience

The primary audience is optimization and learning-theory researchers working on implicit bias, feature learning, and adaptive gradient methods. It is also relevant to machine learning practitioners and engineers who tune batch size and weight decay for Adam versus AdamW in vision or language model training and want a theoretical account of why the two optimizers behave differently. Readers need comfort with convergence analysis, sparsity-based data models, and asymptotic notation; the paper is not introductory.

Authors’ 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 $λ$ 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 $λ$ tuning. Extensive experiments validate our findings, demonstrating the critical role of batch size and weight decay in Adam's generalization performance.