When Will Gradient Regularization Be Harmful?

TL;DR

Gradient regularization (GR) aims to bias optimization toward flat minima by penalizing gradient norms, but this work shows that GR can destabilize training when used with adaptive optimizers and learning-rate warmup. Through empirical studies on Vision Transformers across CIFAR-10/100 and larger datasets, alongside theoretical arguments, the authors identify that GR can induce instability in early gradient statistics and increase adaptive learning-rate variance. To address this, they introduce three GR warmup strategies—r-warmup, lambda-warmup, and zero-warmup—that temper GR during the warmup phase and then resume full GR, achieving improved performance, particularly for larger models (e.g., up to about 3% on CIFAR-10). The findings suggest practical guidance for deploying GR with adaptive methods and pave the way for more stable regularization in large-scale transformer training.

Abstract

Gradient regularization (GR), which aims to penalize the gradient norm atop the loss function, has shown promising results in training modern over-parameterized deep neural networks. However, can we trust this powerful technique? This paper reveals that GR can cause performance degeneration in adaptive optimization scenarios, particularly with learning rate warmup. Our empirical and theoretical analyses suggest this is due to GR inducing instability and divergence in gradient statistics of adaptive optimizers at the initial training stage. Inspired by the warmup heuristic, we propose three GR warmup strategies, each relaxing the regularization effect to a certain extent during the warmup course to ensure the accurate and stable accumulation of gradients. With experiments on Vision Transformer family, we confirm the three GR warmup strategies can effectively circumvent these issues, thereby largely improving the model performance. Meanwhile, we note that scalable models tend to rely more on the GR warmup, where the performance can be improved by up to 3\% on Cifar10 compared to baseline GR. Code is available at \href{https://github.com/zhaoyang-0204/gnp}{https://github.com/zhaoyang-0204/gnp}.

Figures (4)

• Figure 1: Comparison of test error rates (lower values are preferable) of the ViT-B model on the Cifar10 dataset under Base training, GR, and our three proposed GR warmup strategies. All the training instances have also applied the LR warmup. Notably, the performance with LR warmup and normal GR (red line) can be worse compared to training with only LR warmup (blue line).
• Figure 2: Evolution of gradient norm during training for ViT-Ti, ViT-S and ViT-B models on Cifar-{10, 100}, comparing training with GR (red line) and without GR (blue line).
• Figure 3: Plot of the Equation \ref{['eqn : variance computation']}
• Figure 4: Evolution of gradient norm during training for ViT-Ti, ViT-S and ViT-B models on Cifar-{10, 100}, comparing training with $r$-warmup GR strategy (dark brown line), $\lambda$-Warmup GR strategy (green line) and zero-warmup GR strategy (purple line).

Theorems & Definitions (6)

• Lemma 4.1
• Theorem 4.2
• Theorem 4.3
• Lemma 1.1
• Theorem 1.2
• Theorem 1.3