To Clip or not to Clip: the Dynamics of SGD with Gradient Clipping in High-Dimensions

TL;DR

This paper provides a theoretical analysis of gradient clipping in high-dimensional stochastic gradient descent, revealing its effects on learning dynamics and practical benefits in noisy settings, supported by experiments.

Contribution

It introduces a deterministic equation modeling clipped SGD in high dimensions and proposes a heuristic for optimal clipping threshold scheduling.

Findings

Clipping does not improve performance with Gaussian noise.

In other noisy environments, clipping can be beneficial with proper tuning.

The proposed heuristic simplifies hyperparameter tuning for clipping thresholds.

Abstract

The success of modern machine learning is due in part to the adaptive optimization methods that have been developed to deal with the difficulties of training large models over complex datasets. One such method is gradient clipping: a practical procedure with limited theoretical underpinnings. In this work, we study clipping in a least squares problem under streaming SGD. We develop a theoretical analysis of the learning dynamics in the limit of large intrinsic dimension-a model and dataset dependent notion of dimensionality. In this limit we find a deterministic equation that describes the evolution of the loss and demonstrate that this equation predicts the path of clipped SGD on synthetic, CIFAR10, and Wikitext2 data. We show that with Gaussian noise clipping cannot improve SGD performance. Yet, in other noisy settings, clipping can provide benefits with tuning of the clipping threshold. We propose a simple heuristic for near optimal scheduling of the clipping threshold which requires the tuning of only one hyperparameter. We conclude with a discussion about the links between high-dimensional clipping and neural network training.