Scaling Laws of SignSGD in Linear Regression: When Does It Outperform SGD?

AI summary

Overview

Research area: Machine learning theory — specifically neural scaling laws and stochastic optimization theory, analyzed through the power-law random features (PLRF) model of linear regression.

Technical level: Advanced. The paper is a theoretical analysis built on asymptotic risk decompositions, self-consistent integral equations, and phase diagrams in a two-parameter (feature decay, target decay) plane. The introduction, framing, and the qualitative conclusions are readable without the derivations, but the core results are stated as exponent formulas.

Scope in one sentence: The paper derives the scaling law of one-pass signSGD in the PLRF model, computes its compute-optimal exponents, and identifies exactly which regions of the (α, β) parameter plane let signSGD beat the SGD scaling law of Paquette et al. (2024).

What This Paper Is About

Theory explaining neural scaling laws has mostly analyzed stochastic gradient descent (SGD), yet state-of-the-art large language models are trained with Adam and its variants, which are often approximated in theory by the simpler signSGD (Bernstein et al., 2018). This paper asks how the scaling law changes when SGD is replaced by signSGD. It derives a closed-form four-term risk formula for signSGD under a power-law random features model, then asks under what optimal learning rate and model-size allocation signSGD achieves a steeper compute-optimal slope than SGD.

Key Contributions

1. The paper derives a scaling law for one-pass signSGD with constant learning rates, expressed as a function of model size M, training steps N, learning rate γ₀, feature decay α, and target decay β (Equation 12). Comparing against the SGD scaling law of Paquette et al. (2024) and Lin et al. (2024) isolates two effects unique to signSGD: a drift-normalization effect and a noise-reshaping effect.

2. Under a fixed compute budget 𝔣 = MN, the authors balance model size against training steps and optimize over the learning rate, parameterized as γ₀ = M^{−e}. This yields the compute-optimal loss decay exponent η(α, β), the optimal model size M*, and the optimal learning rate exponent e* across six subphases (Phase A split into Aa, Ab, Ac, Ad; Phase B split into Ba, Bb), summarized in Table 1 and Table 2.

3. The paper analyzes the warmup-stable-decay (WSD) schedule of Hu et al. (2024) for signSGD. By keeping drift velocity during the stable interval and cutting stochastic noise during the polynomial decay interval, the schedule yields a strictly larger compute-optimal slope in a region the authors call Area Aa* (large α, small β).

4. The authors report empirical validation of the theory in Figure 1 and Appendix C, with additional exponent checks in Figures 8 to 12.

Main Findings

• SignSGD risk formula. For −α + 0.5 < β < α + 0.5, the risk decomposes into four terms: the approximation error 𝒜(M) = M^{−2α+max(0,1−2β)}; an aligned feature loss 𝒟_al^sign = (M^{min(α,0.5)} N γ₀)^{−2(2α+2β−1)/(2α−2β+1)}; a distorted feature loss 𝒟_dis^sign = M^{−(6α−1)/(2α+1)} (N γ₀)^{−2(2α−1)/(2α+1)}; and a noise term 𝒩^sign = γ₀² M^{2−min(1,2α)}.

• Drift-normalization effect. The signSGD drift is 4γ_k / (π √L(k)) · λ_i(K̄), versus 2γ_k λ_i(K) for SGD. The diagonal preconditioning in K̄ = diag(K)^{−1/2} K contributes an extra factor M^{min(α, 1/2)}, and the 1/√L(k) normalization replaces the effective flow time N γ₀ with γ₀ ∫₀^N L(u)^{−1/2} du, accelerating progress whenever L(u) ≲ 1. This produces a self-consistent equation that does not arise for SGD.

• Noise-reshaping effect. The signSGD quadratic noise is (2γ_k²/π) w_iᵀ K_σ K u_i, whereas SGD noise carries a multiplicative L(k) factor. Removing that factor eliminates the Volterra structure in Paquette et al. (2024) and yields a noise term that does not decay with N — the (Nγ₀)^{−(4α−1)/(2α)} factor present in SGD noise is gone. Consequently, the signSGD noise term increases for all (α, β) as the learning rate grows, while SGD noise decreases for α > 0.5 and increases for α < 0.5.

• Faster N-decay in the drift terms. When the absolute value of the N-exponent is x for SGD, it becomes 2x/(2−x) in signSGD, which is strictly larger than x. So the aligned and distorted feature losses decrease faster in N under signSGD.

• Simpler phase structure. For signSGD, all four terms are dominant only when α > 0.5 and β > 0.5 (Phase B); otherwise the dominant terms are 𝒜, 𝒟_al^sign, and 𝒩^sign (Phase A). Subphase boundaries occur at α = 0.5, β = 0.5, and β = α + 0.5, giving six subphases. Paquette et al. (2024) partition the same plane into four phases with six subphases for the optimal learning rate and seven for the maximal learning rate.

• Compute-optimal exponents (Table 1). Phase Aa: γ₀* = M^{−(α+β)}, M* ≂ 𝔣^{1/(2α+1)}, R ≂ 𝔣^{−(2α+2β−1)/(2α+1)}. Phase Ab: γ₀* = M^{−(2β+1)/2}, M* ≂ 𝔣^{1/2}, R ≂ 𝔣^{−(2α+2β−1)/2}. Phase Ac: γ₀* = M^{−1}, M* ≂ 𝔣^{(2α+2β−1)/(2(2β−α(2β−3)−1))}, R ≂ 𝔣^{−α(2α+2β−1)/(2β−α(2β−3)−1)}. Phase Ad: γ₀* = M^{−1}, M* ≂ 𝔣^{1/(2−α)}, R ≂ 𝔣^{−2α/(2−α)}. Phase Ba: γ₀* = M^{−(2α+4β−1)/(4β)}, M* ≂ 𝔣^{β/(α+β)}, R ≂ 𝔣^{−(2α+2β−1)/(2α+2β)}. Phase Bb: γ₀* = M^{−(6α+1)/(4α+2)}, M* ≂ 𝔣^{(2α+1)/(4α+1)}, R ≂ 𝔣^{−4α/(4α+1)}.

• Where signSGD wins with a constant learning rate. For the intersection of Phases Aa, Ab, Ac, Ba with Phases I, II of the SGD analysis, the compute-optimal slope and optimal model size are the same for signSGD and SGD. In Phases III and IV — the SGD noise-bottleneck regimes — excluding the region 0.25 < α < 1/3, β > (1−α)(1−2α)/(2(1−3α)), signSGD achieves a steeper compute-optimal slope and a larger optimal model size. The authors call this region Area III-IV_sub.

• Learning rate is always smaller. The optimal learning-rate exponent e* is always larger for signSGD than for SGD, meaning signSGD uses a smaller optimal learning rate.

• Warmup-stable-decay improves the slope further. Under the schedule γ_k = γ₀ f(k) with linear warmup k/(wN), constant 1 on wN ≤ k ≤ pN, and polynomial decay (1 + τ(k − pN))^{−c} for k > pN (with w, p, c ∈ (0,1), τ > 0, w < p/2), the scheduled risk bound is R_f ≲ 𝔣^{−2(4α−1)(2α+2β−1)/(16α²+8αβ+2α−2β−1)}. This exponent exceeds the constant-learning-rate compute-optimal slope when α > 0.5 and 0.5 − α < β < (2α−1)/(2(4α−1)) — the region called Area Aa*, marked dark blue in Figure 2. Scheduling does not improve the SGD compute-optimal exponent in Phases I–II, so in Area Aa* scheduled signSGD beats SGD.

• Why the gains happen. Phases III–IV are the SGD noise-bottleneck regimes; noise-reshaping in signSGD alleviates that bottleneck with a suitable learning-rate choice. In Phase Aa*, WSD reduces the noise term enough to sharpen the slope.

• Constant learning rates do not work for signSGD. With γ₀ ≂ 1, 𝒩^sign = γ₀² M^{2−min(1,2α)} is much larger than 𝒩^SGD = γ₀(Nγ₀)^{−(4α−1)/(2α)}, making the compute-optimal slope asymptotically zero. This is why the learning rate is scaled as γ₀ = M^{−e}. The authors also report in Appendix F.1 that the maximal stable learning rate leads to a zero compute-optimal slope for signSGD.

Methodology in Plain English

The authors work in an idealized model of linear regression that is designed to reproduce power-law scaling. Features have covariance eigenvalues decaying as 1^{−2α}, …, d^{−2α}, and the target coefficients along the corresponding eigenvectors decay as i^{−β}; these two parameters α and β control how "hard" the problem is. Instead of d raw features, the model uses M sketched features produced by a fixed random Gaussian matrix S, where M ≤ d is the model size. This is the same power-law random features setup used by Paquette et al. (2024) for SGD, so signSGD and SGD can be compared under identical assumptions.

Training is one-pass: at each step a fresh sample is drawn and a signSGD step is taken, which updates each coordinate only by ±γ_k. To analyze this, the authors take a second-order Taylor expansion of the per-mode loss and use sign–Gaussian identities to get a one-step expected change for each mode, decomposed into a drift term and a quadratic noise term. They then convert this to a continuous-time ODE and solve it with a variation-of-constants formula, which splits the loss into a drift part that depends on initialization, a noise part that depends on accumulated variance injection, and an optimizer-independent approximation error.

The drift part produces a self-consistent integral equation — the loss appears on both sides because the drift is normalized by 1/√L(k). Solving it yields the aligned and distorted feature loss terms. The noise part is evaluated by taking the limit loss L_∞, giving the noise term. The four pieces are combined and verified to satisfy the implicit integral equation. Finally, the authors substitute M = 𝔣^x, N = 𝔣^{1−x}, and γ₀ = M^{−e}, and minimize the risk over e and x; since each of the four terms scales as 𝔣^{−ℓ_i(e,x)}, minimizing the risk reduces to maximizing the minimum of the four exponents, and the optimum is found by balancing three active exponents. The same machinery is repeated for the warmup-stable-decay learning-rate schedule, where noise bounds are refined for the decaying phase.

Why This Matters

Impact on research. The paper extends the SGD scaling-law theory of Paquette et al. (2024) to a second, qualitatively different optimizer and shows that scaling exponents are not optimizer-agnostic. It identifies a concrete regime (SGD noise bottlenecks) where the choice of optimizer changes not just the constant, but the exponent of the compute-optimal scaling law — and shows this depends on where the problem sits in the (α, β) plane. It also gives a theoretical account of why the warmup-stable-decay schedule is useful that is distinct from, and narrower than, "WSD is always better": the benefit is restricted to a specific region.

Real-world applications (as implications of the theory):

• Budgeting training compute for models trained with Adam-like adaptive optimizers, where the paper's exponents suggest a different split between model size and training steps than SGD-based scaling laws would imply.
• Choosing between adaptive and non-adaptive optimizers for a given data/problem regime, based on where that regime falls in the (α, β) plane.
• Designing learning-rate schedules — specifically, deciding when a warmup-stable-decay schedule is theoretically justified in terms of compute-optimal exponent.
• Setting learning-rate magnitude relative to model size, since the paper finds the optimal signSGD learning rate is always smaller than the SGD one under the same allocation.

Industry relevance. Practitioners who fit empirical scaling laws from training runs and then extrapolate to decide model sizes and budgets are implicitly assuming an optimizer. This work shows that assumption matters for the exponent in noise-dominated regimes, which is where large-scale training frequently sits.

Future Directions

• Tightening the scheduled bounds. The authors state explicitly that whether scheduling benefits other signSGD regions or other SGD phases remains open, because for both methods the scheduled noise upper and lower bounds do not match tightly, even up to constant factors.
• Bridging to Adam. The paper motivates signSGD as a tractable stand-in for Adam's coordinate-wise adaptivity, so extending the analysis to Adam or other adaptive methods is a natural next step.
• Multi-pass and data-reuse regimes. Lin et al. (2025) cover multi-pass SGD scaling laws and the effect of data reuse; the analogous signSGD analysis is not treated here.
• Other models and larger regions. The authors note that Kunstner and Bach (2025) compare gradient descent and sign descent in a linear bigram model, and that model/training-protocol variations are discussed in the appendix — leaving open how these conclusions transfer beyond the PLRF model and beyond the regime −α + 0.5 < β < α + 0.5 that the four-term formula covers.

Target Audience

This paper is for machine learning theorists working on scaling laws, stochastic optimization, or learning theory for adaptive optimizers, and for mathematically comfortable practitioners who fit and extrapolate scaling laws. Readers need comfort with asymptotic notation (≂, ≲), risk decompositions, and phase-diagram reasoning to get full value from Sections 3 and 4. Those mainly interested in the practical takeaway — when signSGD beats SGD, and when WSD sharpens the slope — can read the introduction, Section 1.1, Section 4.3, Section 5, and Figure 2 without the derivations.

Authors’ abstract

We study scaling laws of signSGD under a power-law random features (PLRF) model that accounts for both feature and target decay. We analyze the population risk of a linear model trained with one-pass signSGD on Gaussian-sketched features. We express the risk as a function of model size, training steps, learning rate, and the feature and target decay parameters. Comparing against the SGD risk analyzed by Paquette et al. (2024), we identify a drift-normalization effect and a noise-reshaping effect unique to signSGD. We then obtain compute-optimal scaling laws under the optimal choice of learning rate. Our analysis shows that the noise-reshaping effect can make the compute-optimal slope of signSGD steeper than that of SGD in regimes where noise is dominant. Finally, we observe that the widely used warmup-stable-decay (WSD) schedule further reduces the noise term and sharpens the compute-optimal slope, when feature decay is fast but target decay is slow.