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

Thank you very much for the constructive feedback. Below we would like to address your concerns and elaborate on some of our key points.

Weaknesses

1. Indeed it is important to understand how clipping will affect training in other settings perhaps over other convex losses. We believe that this is possible in a manner similar to [1] which provides results for unclipped SGD in terms of the ambient dimension (rather than the intrinsic dimension) over a class of non-quadratic losses. However, the analysis required for these settings is rather long and so we felt it important that a full understanding of the quadratic setting was important.

Additionally, understanding the behavior over quadratic losses can tell a lot about even non-quadratic cases (see for example [2]). Even with a description for the behavior of clipping for non-quadratic problems the most natural question to ask would be if it differs significantly from the quadratics.

2. We have attempted to give an intuition for these factors after introducing them (see lines 155-161). We also give upper and lower bounds for these factors such that up to constants our bounds are equivalent (Equation 10a, 10b). We’ll elaborate on the intuition now and improve the clarity in revision:

• The mean reduction factor measures how much, in expectation, clipping reduces the length of the gradient.
• The variance reduction factor measures how much, in expectation, clipping reduces the variance of the norm-squared of the gradient.

In our original submission, we showed that these factors can indeed be estimated in an online manner in practical neural networks, using statistics of the stochastic gradients already computed during training (Appendix C.2). We demonstrated their computation on ResNet18 and ViT. We think detailed exploration of $\mu$, $\nu$, and their usefulness for setting practical clipping schedules is an exciting avenue for future work.

Questions:

1. We think that this sensitivity question is interesting and have made some progress to explore it. Given your suggestion, we’ve provided some results in Appendix C.4 showing the relative insensitivity of the heuristic to changes in the optimal hyper-parameter.

2. We hope we have address the question of extending our result in the text above responding to listed weaknesses.

Reference: [1] Collins-Woodfin et al., "Hitting the high-dimensional notes: An ode for SGD learning dynamics on GLMs and multi-index models," 2023. [2] Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross entropy in classification tasks. ICLR2021.

Thank you for the constructive feedback. Below we would like to address your concerns and elaborate on some of our key points.

Weaknesses

1. Indeed it is important to understand how clipping will affect training in other settings perhaps over other convex losses. We believe that this is possible in a manner similar to [1] which provides results for unclipped SGD in terms of the ambient dimension (rather than the intrinsic dimension) over a class of non-quadratic losses. However, the analysis required for these settings is rather long and so we felt it important that a full understanding of the quadratic setting was important.

Additionally, understanding the behavior over quadratic losses can tell a lot about even non-quadratic cases (see for example [2]). Even with a description for the behavior of clipping for non-quadratic problems the most natural question to ask would be if it differs significantly from the quadratics.

2. Certainly it is desirable to precisely understand the form of the $\mu$ reduction factor for non-Gaussian data. However, it is likely that it is not meaningfully very different from the Gaussian case. This might be seen in two ways:

• We provide evidence that the theory is effective in reproducing real-world results across two distinct settings: vision and NLP with the CIFAR10 and Wikitext2 datasets.

• It seems that a universality principle is likely to apply. Universality is a common phenomenon seen in high-dimensionals where large-scale behaviors do not depend on the precise distributions of the individual components. As a simple example, the gradient coefficient $\ell_{\boldsymbol{\theta}}$ should see CLT effects as (ambient) dimension grows.

3. Our results are proven to hold for subgaussian data (Theorem 7). Subgaussians need not be “close” to Gaussians save for their tail behavior. For example, any distribution with finite support is subgaussian. This includes any vision dataset (consider RGB images whose values range from 0 to 255) or any tokenized text dataset with a finite vocabulary that does not scale with dimension.

4. There is a lot of notation to handle in this paper. We’ll work to add a section in the Appendix to help ease the struggle.

Questions:

1. Roughly the coefficients of C-HSGD are chosen to match the first and second moments of the risk (or distance to optimality) under C-SGD. Importantly though they are not exactly the same. The high-intrinsic-dimension limit allows for a significant simplification of these terms as any errors introduced by the simplifications vanish in the limit. That said, it is correct, to a first approximation, to think of matching first and second moments.

2. We believe that rigorously finding the optimal clipping schedule in a simpler setting is the first step to developing better heuristics in the practical setting. For example, neural networks are known to linearize as their width approaches infinity [3] implying that we could describe the behaviour of a wide neural network trained with C-SGD.

We found that the reduction factors $\mu$ and $\nu$ are sufficient to compute optimal clipping schedules in our setting. We hypothesize that $\mu$ and $\nu$ play an important role in the more general non-linear setting as well, and in our original submission we described an efficient way to estimate them online for practical networks. We also found simpler heuristics for the optimal clipping schedule based on the risk and noise estimation. We hope to investigate these quantities in practical models in followup work.

3. Stability is considered but only implicitly. If the CCC is satisfied the CSC must be satisfied as well (since $0 \leq \mu \leq 1$). Therefore in all cases in Section 5 where the CCC is met the CSC is implicitly satisfied as well.

References:

[1] Collins-Woodfin et al., "Hitting the high-dimensional notes: An ode for SGD learning dynamics on GLMs and multi-index models," 2023.

[2] Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross entropy in classification tasks. ICLR2021.

[3] Jacot et al., “Neural Tangent Kernel: Convergence and Generalization in Neural Networks“ NeurIPS 2018.

Thank you for your feedback, and for your careful reading of our paper. You’re correct about the missing floor and we are happy to address your major concerns and clarify some points.

We see three main points made in your weaknesses and questions that we would like to address separately:

Restriction to linear regression with MSE loss

Though our initial work was restricted to MSE loss, we believe it’s possible to extend our work to other losses (including cross-entropy loss), and some classes of non-linear models. Recent work established a high-dimensional theory for streaming SGD with Gaussian data on a large class of loss functions, and a broader class of models [1]. This approach required large ambient dimension, and the assumption that intrinsic dimension $\approx$ ambient dimension.

We believe that a similar approach could be used to understand clipped SGD as well. However this requires very detailed technical work ([1] is 98 pages!) that benefits from an understanding of simpler cases. We hope to pursue this avenue in the near future.

We also believe analyzing $\mu$ and $\nu$ in more realistic models is tractable and promising (see Appendix C.2 of our original submission); this is another area of future work.

Applicability of the heuristics to real data

We thank the reviewer for bringing up this point. We have revised our submission to test the heuristics on real-world datasets (Appendix C.4, on wikitext2). We find that the heuristics work well even for non-Gaussian data. Indeed we find that if the mean reduction factor $\mu$ is well-characterized by the probability of clipping, the heuristic proof (Theorem 8 and Corollary 1) extends to any data distribution that is sufficiently regular around 0 (formally, any distribution with an integrable characteristic function and positive density at 0).

It may be possible to derive a universality result which extends the heuristic to a large class of subgaussian data distributions; the clipping functions behavior is determined by the relationship between the gradient norm and clipping threshold. In high dimensions the gradient norm becomes more and more independent of distributional details, increasing the utility of our heuristic. We leave exploration of this universal behavior for future work.

Calculating the mean reduction factors $\mu$ and $\nu$ in general settings

In our original submission, we described an efficient online estimator of $\mu$ and $\nu$ which uses gradient statistics already collected during training (Appendix C.2). We implemented and ran the estimator during training runs of ResNet18 and ViT, to prove that the estimator works.

We neglected to mention that the estimators can be run in parallel for many values of the clipping threshold c, with negligible overhead. This lets us cheaply optimize the CCC (and therefore the clipping threshold) at every step; we simply maximize using the estimators over our fixed grid of c.

For MSE loss we can also use our heuristics from Equation 10 as estimators. We believe other estimators can be derived for common losses like cross-entropy.

Reference: [1] Collins-Woodfin et al., "Hitting the high-dimensional notes: An ode for SGD learning dynamics on GLMs and multi-index models," 2023.

Thank you to the authors for their detailed responses. I agree that the universality of the heuristics seems reasonable under mild conditions. Based on the clarifications during the rebuttal, I’m raising my score from 5 to 6. Good luck!

Thank you for your thoughtful review. We’re glad you appreciate our work and the value of our results to understanding and improving the usage of gradient clipping. Below, we address your specific comments and concerns in detail.

Weaknesses

1. Our approach does leverage foundational techniques from Paquette et al. (2022) and builds on their homogenization framework. However, we provide a novel technical contribution by considering learning rate scalings by the intrinsic dimension, better reflecting the truth and providing a non-asymptotic statement. This is a non-trivial contribution even for unclipped SGD. When the data lie in a low dimensional subspace SGD does not concentrate. In comparison to previous work our paper captures this behavior.

Additionally, our theory does apply beyond the Gaussian setting (see Theorem 7) and it is only the identification of $\mu$ with the probability of not clipping that requires Gaussianity.

Together, these contributions provide a more complete picture of how clipping affects SGD, particularly in high-dimensional, noisy environments, than what was previously known.

2. We agree that estimating whether CCC $> 1$ can be challenging a priori. The aim of this result (Theorem 4) is to provide a theoretical tool, along with intuitive guidance for when clipped SGD is effective. Interpreting this CCC result leads to the guiding principle: if the variance of the gradients can be drastically reduced relative to the probability of clipping at any time during training (for any risk value one is likely to achieve) then clipping will be effective. Our result formalizes this idea.

Additionally, this criteria is not so difficult to measure in practice. The definitions of $\mu$ and $\nu$ in Equation (7) requires an estimation of stochastic gradients and their norms; both of these are already computed during training. Therefore a sort of online estimation of $\mu$ and $\nu$ are readily available during training at negligible cost. We have elaborated on this point and performed preliminary experiments in Appendix C.2. A comprehensive study implementing these results is reserved for future work.

3. Your observation about the limitations of applying our results to the NTK (Neural Tangent Kernel) regime is well-taken. We chose this linear approximation framework for tractability and to lay a foundation for understanding how clipping interacts with learning dynamics in a simplified setting. However, it is clear that linear models cannot fully capture the entire dynamics of SGD.

We believe it is important to first check that this simpler setting provides meaningful results. Having laid the groundwork, future results regarding the effects of non-linearity can be compared in relation to this null model of linear models.

Questions

1. It should be noted that if the CCC > 1 then the CSC > 1 as well (since $0 \leq \mu \leq 1$). For this reason we chose to aim our discussion at the former. If the CSC > 1 but the CCC is not (at every point along the training path), then clipped SGD won’t converge faster than unclipped SGD. In this case, clipped SGD would have the (still noteable) effect of increasing robustness to a sub-optimal choice of learning rate, increasing the rate of values under which the optimizer will converge.

2. We did not include explicit comparisons between unclipped and clipped SGD under heavy-tailed noise, such as Cauchy noise, because unclipped SGD diverges. This can be seen given the explicit dependence of the risk on the variance of the noise term. A proof of this fact has been given by [1] in Remark 1.

Thank you again for your feedback and for highlighting areas where we can improve our submission.

[1] Zhang, J., Karimireddy, S., Veit, A., Kim, S., Reddi, S., Kumar, S., & Sra, S. (2020). Why are Adaptive Methods Good for Attention Models?. In Advances in Neural Information Processing Systems (pp. 15383–15393)