Research map / Architectures and efficiency / RNNs and normalization

Improving Generalization and Convergence by Enhancing Implicit Regularization

NeurIPS 2024 · Architectures and efficiency › RNNs and normalization · approach: Efficient network compression and training · 1 citations (OpenAlex)

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 Implicit Regularization Enhancement (IRE), which scales up updates only along estimated flat directions (via a diagonal-Fisher mask) to accelerate implicit sharpness reduction, improving generalization for SGD/SAM/AdamW and speeding up Llama pre-training about 2x.

Problem

Implicit sharpness regularization of optimizers like SGD is very slow near minima, and naively raising the learning rate causes instability along sharp directions, limiting generalization and convergence speed.

Approach

Decouple flat and sharp directions: amplify the update only along flat directions (masked by a light-weight diagonal Hessian/Fisher estimate, refreshed every K steps) by a factor (1+kappa), leaving sharp-direction dynamics unchanged, so sharpness reduction speeds up without training instability.

Key results

• AdmIRE achieves about 2.1x speedup over well-tuned AdamW in pre-training Llama (60M, 119M, 229M) on Wikitext-103, Minipile and OpenWebText.
• SAM-IRE improves CIFAR-100 accuracy on WRN-28-10 from 83.05 to 83.50 and ViT-S from 66.93 to 70.47 (CIFAR-100).
• Per-step computational cost of AdmIRE is about 1.12x that of AdamW (K=10); theory shows kappa up to 1/rho gives Theta(1/rho) acceleration for standard SAM.

Measured effects reported in the paper

Surprising findings

• IRE accelerates LLM pre-training convergence, although it was motivated only by sharpness reduction; the authors note the acceleration may come from non-zero small-curvature directions and leave the explanation to future work.
• The final loss of AdmIRE after 50k steps matches AdamW's loss after 100k steps (2.47), while reaching flatter solutions (trace 88.86 vs 120.41).

Limitations

• Theoretical analysis relies on local near-manifold assumptions and a two-phase algorithm; the reason for the observed LLM speedup is not explained.
• Hyperparameters gamma and kappa are tuned via grid search, and the LLM experiments use small models and limited scale.
• Downstream performance of IRE-trained LLMs is not evaluated.

Experimental setup

Datasets: CIFAR-10 CIFAR-100 ImageNet Wikitext-2 Wikitext-103 Minipile OpenWebText

Benchmarks: CIFAR-10 CIFAR-100 ImageNet Wikitext-2 Wikitext-103 Minipile OpenWebText

Metrics: top-1 accuracy top-5 accuracy training loss trace of Hessian wall-clock time per step

Baselines: SGD SAM AdamW

Builds on: sharpness-aware minimization AdamW SGD implicit regularization theory empirical diagonal Fisher

Closest papers by meaning

• Implicit Gradient Regularization (ICLR 2021 · similarity 0.89)
• Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems (NeurIPS 2024 · similarity 0.89)
• Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training (ICLR 2025 · similarity 0.88)
• Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization (NeurIPS 2024 · similarity 0.87)
• Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning (ICLR 2024 · similarity 0.87)
• Avoiding spurious sharpness minimization broadens applicability of SAM (ICML 2025 · similarity 0.87)