Source-linked AI summary
Penalizing Gradient Norm for Efficiently Improving Generalization in Deep Learning
Yang Zhao, Hao Zhang, Xiuyuan Hu
TL;DR
Severely overparameterized networks make it difficult to guide optimization toward minima that generalize well. The paper penalizes the loss-gradient norm using a Hessian-free first-order approximation, and experiments report generally improved performance, with gains up to 70% greater than SAM’s improvements over standard training. The method requires careful hyperparameter selection to avoid losing approximation precision.
Problem
Severely overparameterized DNNs contain many minima with diverse generalization ability, so guiding optimizers toward minima with better generalization remains a central challenge.
Method
The method adds a loss-gradient-norm penalty and approximates its gradient using first-order derivatives without computing the Hessian.
Results
The scheme generally improves performance across Cifar and ImageNet models, with improvements up to 70% greater than SAM’s improvements over standard training.
Takeaways & Limitations
SAM is a special case of the proposed gradient-norm penalty, while tuning its separate parameters can yield stronger reported performance.
Takeaways & Limitations
The penalty parameters require careful setting because inappropriate choices can reduce approximation precision, and some training settings can fail or require repeated trials.
Abstract
from arXiv · showhide
How to train deep neural networks (DNNs) to generalize well is a central concern in deep learning, especially for severely overparameterized networks nowadays. In this paper, we propose an effective method to improve the model generalization by additionally penalizing the gradient norm of loss function during optimization. We demonstrate that confining the gradient norm of loss function could help lead the optimizers towards finding flat minima. We leverage the first-order approximation to efficiently implement the corresponding gradient to fit well in the gradient descent framework. In our experiments, we confirm that when using our methods, generalization performance of various models could be improved on different datasets. Also, we show that the recent sharpness-aware minimization method (Foret et al., 2021) is a special, but not the best, case of our method, where the best case of our method could give new state-of-art performance on these tasks. Code is available at {https://github.com/zhaoyang-0204/gnp}.
1. Introduction
Severely overparameterized DNNs contain many minima with differing generalization ability, making optimizer guidance important. The paper penalizes loss-gradient norms and uses a first-order approximation to efficiently encourage flatter minima, reporting improved performance across models and datasets.
- Motivation: Severely overparameterized DNNs contain many minima with diverse generalization ability, making it difficult to train models toward optimal minima.The paper frames optimizer guidance as increasingly important because training loss minimization alone does not ensure satisfactory minima.
- Motivation: Penalizing the loss-gradient norm is proposed to encourage optimization toward relatively flat minima associated with better generalization.The paper motivates this penalty through an intuitive toy example and a later Lipschitz-continuity analysis.
- Method: A first-order approximation computes the gradient-norm penalty without Hessian calculation, enabling practical implementation with basic first-order derivatives.The approximation is designed to fit the gradient-descent framework efficiently.
- Empirical evaluation: Across Cifar-10, Cifar-100, and ImageNet experiments, the optimization scheme generally improves performance across convolutional and vision-transformer architectures.The experiments include both simple and complex convolutional networks as well as recent vision transformers.
2. Related Works
The related work connects the proposed approach to regularization and flat-minima research. Existing methods constrain model properties or implicitly and explicitly bias optimization toward smoother surfaces and flatter minima.
- Regularization: Regularization methods include penalty functions, data augmentation, dropout, and normalization, with penalty methods imposing constraints on specific model properties.Examples include weight L2-norm penalties for model capacity and spectral-norm penalties for reducing input sensitivity.
- Flat minima: Prior work associates flat minima with better generalization from both empirical and theoretical perspectives.The paper situates its approach within research beginning with the observation that well-generalized models may have flat minima.
- Flat minima: SGD and momentum can implicitly favor flat minima, while researchers also seek explicit optimization biases toward smoother surfaces and flatter minima.This motivates methods that directly steer optimization rather than relying only on implicit regularization.
3. Method
The method adds a gradient-norm penalty to the loss to encourage flatter minima, connecting local gradient behavior with Lipschitz continuity and generalization. A first-order approximation avoids direct Hessian computation and supports practical optimization, while exposing tunable parameters and SAM as a special case.
- Objective: The objective adds a penalty λ · ||∇θLS(θ)||p to the empirical loss, with λ controlling regularization strength and the L2 norm used in the main demonstration.The penalty is formulated alongside the standard loss; λ is investigated as positive and, experimentally, negative.
- Gradient norm and flatness: Reducing the loss gradient norm is expected to reduce the local Lipschitz constant, encouraging convergence toward flatter minima.The paper relates smaller local output changes to flatter loss landscapes and connects flat minima with generalization.
- Scope: The gradient-norm penalty operates in weight space, unlike related Lipschitz regularization methods that target stability to perturbations in input space.The distinction concerns where perturbations and regularization act, not merely the use of a Lipschitz-related quantity.
- Approximation: A first-order approximation replaces direct Hessian evaluation with algebraic operations involving first-order derivatives, making gradient-norm optimization practical for DNNs.The approximation uses a small perturbation r, which must balance truncation error against excessively similar perturbed and reference gradients.
- Optimization scheme: The scheme perturbs parameters by r times the normalized loss gradient, recomputes the loss gradient, and combines the two gradients using balance coefficient α.For SGD, the combined gradient updates the model parameters; the procedure can be adapted to other optimizers before the update step.
- Relation to SAM: SAM is recovered when α = 1, but tying its penalty coefficient λ to r restricts tuning because the parameters serve distinct roles.The paper states that SAM may not be the best implementation and separately identifies λ as regularization strength and r as the Hessian-approximation perturbation.
4. Experiments
Across CNN and ViT experiments on Cifar10, Cifar100, and ImageNet, the proposed gradient-norm penalty generally improves testing performance over standard training and SAM. Results also show sensitivity to hyperparameter choices and occasional training instability in competing schemes.
- Training stability: SAM can be substantially less stable under some settings, including cutout augmentation, where selected learning rates and searched r values failed to train from the start.The proposed method could instead use the common optimal training hyperparameters with an appropriate balance coefficient.
- CNN results: For WideResNet, the proposed scheme’s improvements are 38% on average and up to 70% greater than SAM’s improvements over standard training.The reported maximum occurs on Cifar10 with cutout augmentation, where the proposed method achieves 0.65 versus SAM’s 0.38 improvement comparison.
- CNN results: Our method improves CNN performance across the evaluated Cifar10 and Cifar100 architectures, including WideResNet, PyramidNet, and batch-normalized VGG16.For WideResNet, gains reach 1% on Cifar10 and nearly 3% on Cifar100 versus standard training; batch-normalized VGG16 reaches as low as 4% testing error on Cifar10.
- ViT results: The proposed scheme improves ViT performance on both Cifar10 and Cifar100, although ViTs remain worse than CNNs under heavy augmentation.The authors interpret these results as evidence of broad applicability across architectures and datasets.
- Parameter study: Hyperparameter selection is consequential: α = 0.8 with r ∈{0.05, 0.1} is recommended, while large or nonpositive penalty settings can destabilize or prevent convergence.For α = 2.0, the gradient norm drops near zero immediately; nonpositive λ caused complete convergence failure in the reported trials.
5. Conclusion
The paper introduces an efficient gradient-norm penalty scheme and reports improved performance over standard training and SAM, with gains up to 70% greater than SAM’s improvements.
- The scheme penalizes the loss gradient norm during optimization without Hessian computation, enabling practical implementation.
- Across Cifar and ImageNet image-classification experiments, the scheme outperforms standard training and SAM, achieving several new state-of-the-art results.
- 70% greater improvement than SAM is the reported maximum gain from the proposed scheme.
A. Simplification Process of Equation 6
The derivation simplifies the gradient-norm penalty using the parameter vector’s 2-norm and the loss gradient, supporting an efficient first-order implementation.
- The derivation focuses on simplifying the gradient of the gradient-norm penalty term.
- For θ represented as a vector of parameters, the 2-norm provides the norm used in the simplification.
- Defining h(θ) as the loss gradient enables simplification of ∇θ(||∇θL(θ)||2).
- The resulting expression is substituted back into Equation 6 to simplify the optimization calculation.