[Paper Review] On the Distributional Properties of Adaptive Gradients

Zhiyi Zhang, Liu Ziyin|arXiv (Cornell University)|May 15, 2021

Generative Adversarial Networks and Image Synthesis38 references4 citations

TL;DR

This paper provides a theoretical analysis of the distributional properties of adaptive gradient methods, particularly Adam, showing that the variance of the parameter update remains finite and bounded over time—even under isotropic Gaussian gradients—contradicting prior claims that divergent variance necessitates warm-up. The study derives the exact update distribution at each time step and validates it experimentally across modern architectures like Transformers.

ABSTRACT

Adaptive gradient methods have achieved remarkable success in training deep neural networks on a wide variety of tasks. However, not much is known about the mathematical and statistical properties of this family of methods. This work aims at providing a series of theoretical analyses of its statistical properties justified by experiments. In particular, we show that when the underlying gradient obeys a normal distribution, the variance of the magnitude of the extit{update} is an increasing and bounded function of time and does not diverge. This work suggests that the divergence of variance is not the cause of the need for warm up of the Adam optimizer, contrary to what is believed in the current literature.

Motivation & Objective

• To investigate the statistical properties of adaptive gradient methods, particularly the variance and distribution of parameter updates during training.
• To challenge the prevailing belief that divergent variance in adaptive methods necessitates learning rate warm-up, especially in models like Transformers.
• To derive the exact analytical distribution of the update in Adam under isotropic Gaussian gradient assumptions, providing a theoretical foundation for observed training dynamics.
• To validate theoretical predictions against empirical data from modern deep learning models, including Transformers and CNNs, at the single-trajectory level.
• To explore conditions under which theoretical predictions deviate from empirical results, such as non-zero mean or heavy-tailed gradient distributions.

Proposed method

• Derives the exact distribution of the parameter update Δθ in Adam under the assumption of i.i.d. isotropic Gaussian gradients and time-averaged preconditioners.
• Proves that the variance of the update magnitude is bounded and monotonically increasing over time, even without distributional assumptions on the gradient.
• Uses analytical derivation and simulation to show that the update distribution transitions from bimodal to unimodal and converges to a Gaussian as training progresses.
• Empirically validates the predicted distribution against real training trajectories on MNIST, ResNet-18, VGG-21, and Transformers, using single-run data.
• Introduces controlled gradient distributions (e.g., time-varying variance, non-zero mean, Cauchy) to study deviations from theory and understand real-world discrepancies.
• Analyzes the impact of gradient correlation and heavy-tailed distributions on update distribution shape, including trimodal structures observed in practice.

Experimental results

Research questions

• RQ1 Does the variance of the adaptive gradient update diverge during training, and if not, what is its time-dependent behavior?
• RQ2 What is the exact analytical form of the update distribution in Adam under isotropic Gaussian gradient assumptions?
• RQ3 Why do real-world training trajectories exhibit bimodal or trimodal update distributions, and how do they relate to gradient distributional properties?
• RQ4 To what extent do theoretical predictions based on i.i.d. gradients hold in practice, especially in deep models like Transformers?
• RQ5 What role do non-zero mean or heavy-tailed gradient distributions play in skewing or distorting the update distribution?

Key findings

• The variance of the adaptive gradient update is always finite and bounded, contradicting the claim in Liu et al. (2019) that divergent variance necessitates warm-up.
• Under isotropic Gaussian gradients, the update distribution is bounded for all finite time steps and transitions from bimodal to unimodal, converging to a Gaussian as t → ∞.
• The variance of the update magnitude increases monotonically with time, indicating that early training steps are more variable, yet the overall update remains stable.
• Theoretical predictions match empirical distributions extremely well across diverse architectures, including Transformers, even at the single-trajectory level.
• Non-zero mean gradients induce left-skewed update distributions that become less skewed and thinner over time, especially when the mean magnitude increases.
• Heavy-tailed gradient distributions, such as Cauchy, lead to trimodal update distributions, suggesting the presence of rare, high-norm gradients or strong noise in training.