GradientStabilizer: Fix the Norm, Not the Gradient
Abstract
Training instability in modern deep learning systems is frequently triggered by rare but extreme gradient-norm spikes, which can induce oversized parameter updates, corrupt optimizer state, and lead to slow recovery or divergence. Widely used safeguards such as gradient clipping mitigate these failures but require threshold tuning and indiscriminately truncate large updates. We propose GradientStabilizer, a lightweight, drop-in gradient transform that preserves the instantaneous gradient direction while replacing the update magnitude with a statistically stabilized estimate derived from running gradient-norm statistics. We prove that the resulting stabilized magnitude is uniformly bounded on spike steps, independent of the spike size, and show how this boundedness controls optimizer state evolution in adaptive methods. Across LLM pre-training (FP16), quantization-aware pre-training (FP4), ImageNet classification, reinforcement learning, and time-series forecasting, GradientStabilizer consistently improves training stability, widens stable learning-rate regions, and reduces divergence relative to clipping-based baselines, even substantially reducing Adam’s sensitivity to weight-decay strength.
Lay Summary
Training modern AI models can sometimes be disrupted by sudden, unusually large changes in the learning signal. These rare events may cause the model to take overly large steps, making training slower, unstable, or even causing it to fail. Existing fixes often work by cutting off large updates, but this can require careful manual tuning and may also remove useful learning information. We introduce GradientStabilizer, a simple method that makes training more stable without changing the direction in which the model is learning. Instead of directly using an unusually large update, it replaces only the size of that update with a safer estimate based on historical gradient information. This helps prevent rare gradient spikes from having an outsized effect on the model. We show theoretically that our method keeps these abnormal updates under control, no matter how large the gradient spike is. In experiments across language models, image recognition, decision-making tasks, and forecasting, GradientStabilizer makes training more reliable, reduces failures, and works better across a wider range of training settings.