Gradient Regularized Natural Gradients

TL;DR

This work introduces Gradient Regularized Natural Gradients (GRNG), a framework that blends gradient regularization with natural gradient updates to improve convergence speed and generalization in large-scale deep learning. It provides two scalable implementations: a Frequentist variant using Kronecker-factored FIM approximations with Lazy Fisher and Newton updates (RING/RENG), and a Bayesian variant based on Regularized Kalman filtering (R-Kalman). The authors prove convergence guarantees for GRNG and demonstrate through extensive vision and language experiments that GRNG accelerates optimization and enhances generalization relative to SGD, AdamW, and existing second-order baselines. The approach offers a principled, practical path to robust, scalable second-order optimization for transformer-based models and vision architectures across data regimes.

Abstract

Gradient regularization (GR) has been shown to improve the generalizability of trained models. While Natural Gradient Descent has been shown to accelerate optimization in the initial phase of training, little attention has been paid to how the training dynamics of second-order optimizers can benefit from GR. In this work, we propose Gradient-Regularized Natural Gradients (GRNG), a family of scalable second-order optimizers that integrate explicit gradient regularization with natural gradient updates. Our framework provides two complementary algorithms: a frequentist variant that avoids explicit inversion of the Fisher Information Matrix (FIM) via structured approximations, and a Bayesian variant based on a Regularized-Kalman formulation that eliminates the need for FIM inversion entirely. We establish convergence guarantees for GRNG, showing that gradient regularization improves stability and enables convergence to global minima. Empirically, we demonstrate that GRNG consistently enhances both optimization speed and generalization compared to first-order methods (SGD, AdamW) and second-order baselines (K-FAC, Sophia), with strong results on vision and language benchmarks. Our findings highlight gradient regularization as a principled and practical tool to unlock the robustness of natural gradient methods for large-scale deep learning.

Figures (8)

• Figure 1: Average Gradient Norm (left) and Average Pearson Correlation (right) for SAM, AdamW, Sophia, Kalman, R-Kalman, RING, and RENG on the STS-B dataset.
• Figure 2: Validation accuracy of our proposed algorithms (RING, RENG, and R-Kalman) compared to AdamW, Sophia, and NGD on selected language (top row) and vision (bottom row) datasets. The plots illustrate validation accuracy as a function of training iterations. For a comprehensive set of results, please refer to \ref{['sec:Additional Experiments']}.
• Figure 3: Comprehensive results on the image classification benchmarks, including CIFAR-10/100, Oxford-IIIT Pet, Food-101, and ImageNet-100. The plots show the validation accuracy of our proposed algorithms (RING, RENG, and R-Kalman) compared to AdamW, Sophia, and NGD when fine-tuning ViT-B16. Validation accuracy is plotted as a function of training iterations.
• Figure 4: Comprehensive results on the GLUE benchmark, including MNLI-mm, QQP, QNLI, SST-2, CoLA, STS-B, MRPC, and RTE. The plots show the validation accuracy of our proposed algorithms (RING, RENG, and R-Kalman) compared to AdamW and NGD when fine-tuning RoBERTa-base. Validation accuracy is plotted as a function of training iterations.
• Figure 5: Sensitivity analysis of dampening coefficient for RING.

Theorems & Definitions (11)

• Theorem 4.3
• Corollary 4.4
• proof
• proof
• Remark D.1
• Remark D.2
• proof
• proof
• Lemma D.3
• proof