Exact risk curves of signSGD in High-Dimensions: quantifying preconditioning and noise-compression effects

Thank you for the review, we’re glad that you appreciate our contributions. Our expectation that $K_\sigma$ inherits a power-law spectrum from $K$ is mostly speculative. We should qualify that this claim is basis dependent. If $K$ is diagonal then $K_\sigma$ is the identity, which collapses any power-law structure. So this is more about `very non-diagonal' matrices.

This is probably best defined as having eigenvectors which are very spread out over the space, such as in CIFAR-10 or in an i.i.d. random matrix. In the case of CIFAR-10, the map appears to preserve the power-law structure but change the power. We've included an additional plot to illustrate this effect: https://anonymous.4open.science/r/revision-2025-stuff-F275/cifar_eigenvalues.pdf

We don't know of an existing theorem demonstrating this effect.

We’ll revise this section to make ourselves clearer.

Thanks for your comments, we address some key questions here.

NTK Regime

If we write consider the mean square loss for a general neural network $f(\theta, x)$ in the NTK regime such that $f(\theta, x) = f(\theta_0, x) + \nabla_{\theta} f(\theta_0, x) ^T (\theta_k - \theta_*)$ then the risk, in a student-teacher setup, is given by

$$\mathcal{R}(\theta_k) = E [ (f(\theta_k, x) - f(\theta_*, x)^2 ] = E[( \langle \nabla_{\theta} f(\theta_0, x), \theta_k - \theta_*\rangle^2] = (\theta_k - \theta_*)^T E[\Theta(x,x)] (\theta_k - \theta_*)$$

where $\Theta$ is the NTK. Therefore in order to describe the NTK regime, we need to understand the expected signSGD updates where the “data” is given by $\nabla_{\theta} f(\theta_0, x)$. Right now, we must assume this is Gaussian. Then, if the model is such that $\nabla_{\theta} f(\theta_0, x)$ is Gaussian then our results will hold as written. In the more likely case that $\nabla_{\theta} f(\theta_0, x)$ isn’t Gaussian but is sufficiently regular (such as subgaussian) then if we can formulate our results for this larger class of data distribution then we can capture the behaviour of this NTK regime.

Quite possibly our equations, if not our method of proof, already hold for subgaussian data given the good agreement with the CIFAR-10 and IMDB datasets which are not themselves Gaussian.

Exponential blowup in T

This $e^T$ term appears because we allow for very flexible learning rate scheduling, including schedules where the risks diverge. With an additional restriction on the smallness of the learning rate, one could improve the exponential explosion in $T$. It will not go away entirely, however -- optimistically one could put a $polylog(T)$. On very long time-frames there will always be a chance for the stochastic optimizer to explode.

We don't have a full picture of how a proof like this would go, but roughly you would want to show that the projection of the errors in each eigendirection of $\bar{K}$ are $O(1/\sqrt(d))$ and crucially remain bounded that way in time due to the contractive nature of the dynamics in that eigendirection. This is also the reason the small stepsize is needed -- to ensure the dynamics tends to contract in each eigendirection, and hence the approximation errors are shrunk as well. (Our current proof uses a resolvent argument, which is an average over all eigendirections; this is technically simpler, but would not be good as the contractive-ness of the optimizer dynamics is lost).

We thank the author for their detailed review and address weaknesses and questions below.

Weaknesses

Noise reshaping

For the $K \to K_\sigma$ mapping, there is strong dependence on the eigenbasis of $K$. For example if $K$ is diagonal then $K_\sigma = Id$. Otherwise, if the eigenbasis is relatively random (such as what is seen in CIFAR and in Figure 4), the operation can preserve some spectral properties of $K$. For example, in CIFAR-10, $K_\sigma$ appears to preserve the power-law structure of $K$, albeit with a different power (plots here).

One major point (discussed below in response to the learning-rate) is that the the trace of $K_\sigma$ can be much larger than the trace of $K$, which will cause sign-SGD gradient noise to be a major slow-down Rigorously saying more is probably difficult. But, we'll rewrite this section in order to be clearer.

Assumption 5

Thank you for catching this. We like your suggestion to instead assume $v^T(\theta_0 - \theta_*)$ is subgaussian with constant $O(d^{-1/2})$. You're right that general deterministic initialization and targets do not fit the assumption. For the record, the assumption is for all $z \in B(||K||)$.

Questions

Batch sizes

This is a very good question -- thanks for bringing it up. We have worked out some of the mathematics for batch sizes small compared to $d$, and plan to add the following results to an appendix.

So first, just to be clear, we are following usual optimization (and standard package) conventions, and we are applying the minibatch average prior to applying the sign-function.

We will take a fixed batch size b independent of the dimension of the problem. One could also look at other scaling regimes, and we expect (based on the fixed batch size case) that the behavior continues to hold for batch sizes that satisfy (b/d) to 0, but we’re only 100% confident for fixed $b,$ and $d \to \infty$. The good news is that even in this limited setting, the story is already interesting.

The bias-term (or descent-term) is the only term affected by minibatching. The gradient noise term (the Brownian term in homogenized sgd) is unaffected.

The homogenized-SGD equation is changed to

$$\overline{\mathbf{K}} \left(\Theta_t-\theta^{*}\right) \mathrm{d}t$$

$+ \eta_t \sqrt{\frac{2 \mathbf{K}_\sigma}{\pi d}} \mathrm{~d} \mathbf{B}_t$

The $N_b$ is a numerical constant (the mean of a chi-variable with $b$ degrees of freedom, to the mean of a chi variable with $1$ degree of freedom). It is asymptotic to $\sqrt{b}$, consistent with what is shown in Malladi et al.

The $\varphi^{(b)}$ is the $\varphi$ of a new convolved noise. It is the $\varphi$ that results from looking at $\sum_{i=1}^b \epsilon^{(i)} \omega_i$ where $\epsilon^{(i)}$ are iid copies of the label noise and where $\omega$ is an independent random vector, drawn uniformly from the sphere of dimension $b$.

We have run some simulations implementing this (code is available in the anonymous code repository for the paper). Simulations at fixed learning rate, varying batch size are available here and with $\sqrt{b}$ scaling here.

Learning rate scaling

Indeed, the choice of $\eta’ = O(1/d)$ is needed to prevent (gradient) noise domination. See the Hessian term of equation 94 where the Hilbert-Schmidt inner product is $O(Tr(K_\sigma))$. In this setting, if $\eta’$ decays slower than $1/d$ this noise term would dominate. Had $\eta’$ decayed faster, the noise term would vanish and we would be performing gradient flow in the high-dimensional limit.

For SGD it is possible to instead rescale by the “intrinsic dimension” of the data, defined as the trace of covariance normalized by its operator norm ( $Tr(K) / ||K||$). We think a similar change could be made for signSGD but some care has to be taken since signSGD reshapes the gradient covariance to $K_\sigma$. This connects to our discuss about the reshaping of gradient noise. Given a diagonal $K$ with intrinsic dimension $O(\sqrt{d})$, signSGD still has intrinsic dimension $O(d)$ since $K_\sigma = Id$. This leads to signSGD having learning rates $O(d^{1/2})$ smaller than vanilla SGD to guarantee convergence, leading to slow dynamics.

Let us know if we can clarify anything else; if we have addressed your concerns satisfactorily, we hope that you will consider revisiting your review score.

We thank the reviewer for their comments, and address some key points below.

Minbatching:

We have added more details regarding batch sizes $b$ in our response to Reviewer ugE8. One particular consequence of minibatching is that our condition on the behaviour of the noise near $0$ relaxes significantly. Averaging multiple sources of noise effectively conditions the noise even for $b = 2$. Numerical results can be found here for fixed learning rate, varying batch size.

Questions:

First question

(27) is comparing the risk obtained by the optimally scheduled signSGD to optimally scheduled SGD when data is isotropic. We don’t know the optimal learning rate schedule in non-isotropic settings but what we expect is that whether or not signSGD outperforms SGD depends generally on the magnitude of $\psi$.

The argument why can be sketched like: for any signSGD learning rate $\eta^{S}(t)$ run SGD with a learning rate of $\eta(t) =\eta^{S}(t)\phi(R_t) / \sqrt{R_t}$. Now the descent terms are the same up to the $K$ vs $\overline{K}$ distinction. With this learning rate we see the function $\psi$ in the variance term on the SGD risk equation. If the various matrices involved ($K, \overline{K}, K_\sigma$) are similar enough (as they are when $K=Id$) then whether or not signSGD is beaten by SGD with this learning rate is again determined by the magnitude of $\psi$. This “similar enough” statement is what we mean by well-conditioned settings.

Second question

For SGD in the same setting, the limiting risk can be found in Equation (254) and is

$$\frac{\eta v^2 Tr(K) }{2(2d-Tr(K)\eta)}.$$

For sufficiently small $\eta$, the limiting SGD risk is $O(\eta v^2 Tr(K) / d)$. This is consistent with 'neighborhood' convergence for fixed step-size SGD when $v > 0$ (e.g. Theorem 5.5 of Garrigos-Gower '23, Handbook of Convergence Theorems) and convergence with fixed step-size in the case $v=0$.

In contrast we can use Theorem 3 to approximate the limiting signSGD risk as up to constants:

$$\frac{\eta Tr(\overline{K})}{d} \max \left (\frac{\eta Tr(\overline{K})}{d}, v \right)$$

Notably, even if $v = 0$ this limiting risk is not $0$. The simple explanation is that with SGD, the average magnitude of the gradient naturally decreases with the risk and so we effectively take smaller steps; however for signSGD the magnitude of the gradient is always $\sqrt{d}$ and does not change with the risk. To achieve similar results with signSGD the stepsize needs to be decreased. This aligns with what is seen in practice for algorithms like Adam.

Other:

"In a nonconvex setting, identifying ...." paragraph. We have attempted to improve the wording for clarity, but we’d be happy for additional guidance here. Our intention is to say that dividing by the norm of the gradient could be reasonable in vanishing gradient situations, such as near saddles.