Research map / Optimization / Optimizers
Enhancing Optimizer Stability: Momentum Adaptation of The NGN Step-size
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes NGN-M (NGN step-size with heavy-ball momentum) and diagonal variants NGN-MD, giving O(1/sqrt(K)) convex convergence without interpolation or bounded-gradient assumptions and improved step-size robustness on vision and language tasks.
Problem
Adaptive optimizers such as Adam are highly sensitive to the step-size hyperparameter, and existing Polyak-type step-size methods with momentum need strong assumptions (interpolation, bounded gradients) or lack diagonal (coordinate-wise) adaptivity.
Approach
Combine the NGN Gauss-Newton-derived step-size with heavy-ball momentum via an iterative-moving-average view (NGN-M), and add RMSprop-style diagonal preconditioning (NGN-MD) so the step-size adapts per coordinate while staying robust to the hyperparameter c.
Key results
- NGN-M and NGN-MDv1 match best accuracy of baselines on CIFAR10/CIFAR100 and ImageNet while tolerating step-sizes 1-2 orders of magnitude larger than the optimum without significant performance drop.
- NGN-MDv1 matches Adam perplexity across 70M to 1B SlimPajama language models and remains competitive at c=1e-2 where Adam's performance drops significantly.
- NGN-M attains O(1/sqrt(K)) convergence in the convex setting without interpolation, bounded gradient, or bounded iterate assumptions.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use NGN-M version 1 instead of version 2 (ablation) | NGN with momentum · Diabetes quadratic | convergence | Helped |
| use NGN-M version 1 instead of version 2 (ablation) | NGN with momentum · Diabetes quadratic | stability: c ∈ {10^3} | Helped |
| add momentum using NGN-M | NGN · Rosenbrock function | stability across step-size hyperparameters: 10^{-2}; c=10^2 | Helped |
| add momentum using NGN-M | NGN · tested standard benchmark settings | stability to step-size hyperparameter | Helped |
| use NGN-M instead of NGN | NGN · ResNet18 on ImageNet1k | validation accuracy | Helped |
| use NGN-M instead of NGN or Momo | optimizer training ResNet18 · ImageNet1k | robustness to larger step-sizes | Helped |
| use NGN-MDv1 instead of NGN-MDv2 (ablation) | NGN-MD · empirical evaluations | performance | Helped |
| use NGN-MDv1 instead of Adam or Momo-Adam | adaptive optimizer · vision transformers | validation accuracy | Helped |
| use NGN-MDv1 instead of Adam | adaptive optimizer · SlimPajama-627B language modeling across 70M, 160M, 410M, and 1B models | validation perplexity | No clear effect |
| use NGN-MDv1 with c=10^{-2} rather than relying on Adam's tuned learning rate | language model optimizer · SlimPajama-627B | validation perplexity: c=10^{-2} | Helped |
| use NGN-MDv1 instead of Adam at larger learning rates | adaptive optimizer · vision transformers | training stability | Helped |
Surprising findings
- Theoretical analysis recommends small momentum beta, yet the default beta=0.9 works well empirically; the authors attribute this to proof techniques rather than the algorithm.
Limitations
- Convergence theory limited to convex setting; non-convex analysis for NGN-M deferred to future work.
- Theory requires small momentum beta while practice uses beta=0.9; large-beta behavior not analyzed.
- Some residual sensitivity to step-size remains for NGN-MDv1 with weight decay; NGN step-size may add runtime overhead at scale.
Experimental setup
Datasets: CIFAR10 CIFAR100 ImageNet1k ImageNet32 SlimPajama-627B Diabetes
Benchmarks: CIFAR10 CIFAR100 ImageNet1k SlimPajama-627B
Metrics: validation accuracy validation perplexity step-size robustness across c convergence rate
Baselines: SGDM Adam Momo Momo-Adam MomSPS ALR-SMAG NGN AdamW
Builds on: NGN step-size Stochastic Polyak Step-size SGD with momentum Adam Momo
Closest papers by meaning
- Lean and Mean Adaptive Optimization via Subset-Norm and Subspace-Momentum with Convergence Guarantees (ICML 2025 · similarity 0.89)
- Flatland: The Adventures of Gradient Descent with Large Step Sizes (ICML 2026 · similarity 0.89)
- Step Size Matters in Deep Learning (NeurIPS 2018 · similarity 0.88)
- A geometric framework for momentum-based optimizers for low-rank training (NeurIPS 2025 · similarity 0.88)
- On the Parameterization of Second-Order Optimization Effective towards the Infinite Width (ICLR 2024 · similarity 0.87)
- The Marginal Value of Momentum for Small Learning Rate SGD (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.