Asymptotic Risk of Learning-Rate Schedules for Least-Squares Regression with Discretely Distributed Features

Noé Colomban; Francis Bach; Fabian Schaipp

cs.LGOct 12, 2026

Asymptotic Risk of Learning-Rate Schedules for Least-Squares Regression with Discretely Distributed Features

Authors: Noé Colomban, Francis Bach, Fabian Schaipp

Organizations: Inria - École Normale Supérieure PSL Research University

Abstract

We study the last iterate of stochastic gradient descent (SGD) on streaming least squares with discretely distributed features (Slock's model), which decouples the dynamics across eigendirections. Under power-law source and capacity conditions, we derive the exact finite-time risk and asymptotic equivalents of constant and Warmup-Stable-Decay (WSD) schedules. With a tuned base learning-rate, the resulting scaling law decays as $Θ(T^{-q})$ , where $T$ is the horizon, and $q$ is independent of the cooldown length $c \in (0,1]$ but is strictly smaller for the constant schedule: any cooldown of fixed relative length $c > 0$ , however short, improves the risk by a polynomial factor, while its length only affects the constant, which is minimized by linear decay ( $c = 1$ ). Finally, we show that the optimal risk depends on the batch size $b$ and horizon $T$ only through $bT$ , up to a critical batch size.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison of the empirical SGD excess risk against the theoretical final risk from Corollary F.4 . Both experiments are averaged over $n=50$ independent runs with parameters $\gamma=0.1,L=\Delta=1$ , $\alpha=1.5$ , $\beta=2$ , $\sigma=0.1$ and $d=100$ . For the linear schedule (b), the plot displays the entire risk trajectories over time, where each trajectory is computed with a schedule decaying to zero at its specific horizon $T$ .

Sched.	Bias $B_{T}$	Variance $V_{T}$
Cst.	$L\Delta\,(2L\gamma T)^{-\omega}$	$\sigma^{2}L\gamma$
WSD α<2	$L\Delta\,\big((2-c)L\gamma T\big)^{-\omega}$	$\mathcal{W}(\alpha,c)\sigma^{2}(L\gamma)^{1/\alpha}T^{1/\alpha-1}$
WSD α>2	$L\Delta\,\big((2-c)L\gamma T\big)^{-\omega}$	$\sigma^{2}(L\gamma/c)^{1/2}\,T^{-1/2}$

Table 1: Equivalents of the bias and the variance for the studied schedules (constant: both regimes

1<\alpha<2

and

\alpha>2

), with

\omega=(\alpha+\beta-1)/\alpha

. The cooldown length is

c\in(0,1]

. Constants depending only on

(\alpha,\beta)

are omitted.

\mathcal{W}(\alpha,c)

is defined in Eq. 16 .

Schedule	Regime	$\boldsymbol{\widetilde{\gamma}^{*}_{T}}$
Constant	both	$\big(\frac{C_{\mathrm{cst}}\,\Delta}{\sigma^{2}L^{\omega}}\big)^{\frac{1}{1+\omega}}T^{-\frac{\omega}{1+\omega}}$
WSD	$1<\alpha<2$	$\big(\frac{C_{\mathrm{wsd}<2}\,\Delta}{\sigma^{2}L^{\beta/\alpha}}\big)^{\frac{\alpha}{1+\alpha\omega}}T^{-\frac{\beta}{1+\alpha\omega}}$
WSD	$\alpha>2$	$\big(\frac{C_{\mathrm{wsd}>2}\,\Delta}{\sigma^{2}L^{\omega-1/2}}\big)^{\frac{2}{1+2\omega}}T^{-\frac{2\omega-1}{1+2\omega}}$

Table 2: Optimal learning-rate

\widetilde{\gamma}^{*}_{T}

for each schedule and regime of

\alpha

. All exponents of

T

are negative; in particular, for

\alpha>2

\omega>1-1/\alpha>1/2

Panel of Figure 2 — Figure 2: Excess risk as a function of the learning-rate $\gamma$ for different values of $\alpha\in\mathcal{S}_{\alpha}$ , with $T=10^{5}$ , $\beta=2$ , $d=100$ , $\sigma=0.1$ , and $L=\Delta=1$ . The vertical dashed lines correspond to the theoretical asymptotic optima $\gamma=\widetilde{\gamma}^{*}_{T}$ . While these approximations appear to coincide with the minima of the risk curves under a constant schedule, a slight discrepancy is visible for the linear decay. This error arises because the asymptotic derivation requires $T\to\infty$ , an effect that becomes pronounced near the regime change as $\alpha\to 2$ .

Sched.	Optimal risk
Cst.	$(L\Delta)^{\frac{1}{1+\omega}}\,(\sigma^{2}/T)^{\frac{\omega}{1+\omega}}$
WSD α<2	$(L\Delta)^{\frac{1}{1+\alpha\omega}}\big(\mathcal{W}(\alpha,c)\,(2-c)^{-1/\alpha}\,\sigma^{2}/T\big)^{\frac{\alpha\omega}{1+\alpha\omega}}$
WSD α>2	$(L\Delta)^{\frac{1}{1+2\omega}}\big(\big(c(2-c)\big)^{-1/2}\,\sigma^{2}/T\big)^{\frac{2\omega}{1+2\omega}}$

Table 3: Equivalents obtained using the rate

\widetilde{\gamma}^{*}

optimized for each schedule. The exponent of

\sigma^{2}/T

is greater for WSD schedules than for a constant schedule, in both regimes of

\alpha

. Constants depending only on

(\alpha,\beta)

are omitted except where they depend on

c

Panel of Figure 3 — Figure 3: Log-ratio $\log(\gamma^{*}(1))-\log(\gamma^{*}(c))$ of the optimal base learning-rates as a function of the cooldown length $c$ . (a) Asymptotically optimal $\widetilde{\gamma}^{*}_{T}$ from Table 2 ; the value of $T$ is inconsequential since the ratio is independent of $T$ by Eq. 19 . (b) Empirical $\gamma^{*}$ and $\gamma^{*}$ obtained from the convex bound of Schaipp et al. (2025) (their $\sqrt{\mathcal{T}_{1}/\mathcal{T}_{2}}$ , with constants $G=D=1$ ).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Effect of the batch size $b$ under a linear decay schedule with the asymptotically optimal base learning-rate $\widetilde{\gamma}^{*}_{T,b}$ , parameters $L=\Delta=\sigma=1$ , $\beta=2$ , $d=100$ . 4(a) : Optimal excess risk as a function of $b$ at fixed total sample budget $N:=bT=200000$ , normalized by its value at $b=1$ . 4(b) : Number of steps $T$ needed to reach a fixed risk $\mathcal{R}^{(b)}_{T}=r:=10^{-5}$ , as a function of $b$ . The dashed lines show the horizon $T$ at which $b=\widetilde{b}_{\mathrm{max}}(T)$ . At a fixed target risk $r$ ( Fig. 4(b) ), the required number of steps initially decreases inversely proportionally to $b$ . This scaling stops for a batch size $b$ a few times the theoretical boundary $b=\widetilde{b}_{\mathrm{max}}(T)$ (dashed lines), beyond which the trend reverses and $T$ grows with $b$ : the learning-rate prescribed by the $b^{p}$ rule is then no longer admissible ( Eq. 24 ).

$\alpha$	$\widetilde{b}_{\mathrm{crit}}(N)$
$1.4$	$2653$
$1.8$	$2245$
$2.2$	$1633$
$2.6$	$1021$

Appendix

Table 4: Critical batch size

\widetilde{b}_{\mathrm{crit}}

for a sample budget of

N=2\times 10^{5}

and various exponents

\alpha

, computed by testing all values until the limit

\widetilde{b}_{\mathrm{max}}(N/b)

is exceeded.

Panel of Figure 5 — Figure 5: Ratio of the bias ( Fig. 5(a) ) and variance ( Fig. 5(b) ) evaluated under Slock’s model relative to the standard Normal distribution, with a linear decay schedule. Computed with parameters $d=100,\gamma=10^{-2},\beta=2,L=\Delta=1,\sigma=1$ . As the accumulation of coupled terms increases with $T$ , the bias ratio drifts from 1.

Panel of Figure 6 — Figure 6: Ratio of the bias ( Fig. 6(a) ) and variance ( Fig. 6(b) ) evaluated under Slock’s model relative to the standard Normal distribution, with a linear decay schedule. Computed with parameters $d=100,\gamma=\widetilde{\gamma}^{*}_{T},\beta=2,L=\Delta=1,\sigma=1$ .

Panel of Figure 7 — Figure 7: For the constant schedule, ratio of the asymptotic equivalents to the exact values for the bias (a) and variance (b) components of the risk, for $\gamma=10^{-3},L=\Delta=1,\beta=2,\sigma=0.1$ , $d=1000$ and $\mathcal{S}_{\alpha}=\{1.4,1.8,2.2,2.6\}$ . As the total number of steps $T$ increases, the ratios systematically converge to 1 across different spectral exponents $\alpha$ , empirically validating the theoretical approximations.

Panel of Figure 8 — Figure 8: For the WSD schedule ( $c=0.2$ ), ratio of the asymptotic equivalents to the exact values for the bias (a) and variance (b) components of the risk, for $\gamma=10^{-3},L=\Delta=1,\beta=2,\sigma=0.1$ , $d=1000$ and $\mathcal{S}_{\alpha}=\{1.4,1.8,2.2,2.6\}$ .

Panel of Figure 9 — Figure 9: Comparison of risk trajectories for constant and WSD schedules ( $c=0.2$ ) over a total horizon $T=10^{5}$ . Hyperparameters are fixed to $\alpha\in\mathcal{S}_{\alpha}=\{1.4,1.8,2.2,2.6\}$ , $\beta=2$ , $\sigma=0.1,d=100$ and $L=\Delta=1$ . The parameter $\gamma=\widetilde{\gamma}^{*}_{T}$ is optimized for the constant schedule, and this exact same value is applied to both methods. As expected, the risks remain identical up to $t=0.8T=80,000$ , after which the WSD cooldown phase effectively reduces the final risk.

Explore similar work

Scaling Laws of SignSGD in Linear Regression: When Does It Outperform SGD?
Mar 2, 2026Jihwan Kim, Dogyoon Song, Chulhee YunSign-Based Optimization Linear Regression
Seoul National University & KAIST InnoCORE LLM · University of California, Davis · KAIST
Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay
Feb 6, 2026Binghui Li, Zilin Wang, Fengling Chen +3Learning Rate Scheduling Kernel Regression
Center for Machine Learning Research, Peking University · School of Mathematical Sciences, Peking University · AI for Science Institute, Beijing
WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training
Jul 12, 2026Jianhao Ma, Yuxin ChenLearning Rate Scheduling
Department of Statistics and Data Science, University of Pennsylvania

Mar 2, 2026J/K move · Enter open · S save

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

Jihwan Kim, Dogyoon Song, Chulhee Yun

Seoul National University & KAIST InnoCORE LLM · University of California, Davis · KAIST

Sign-Based Optimization Linear Regression Scaling Laws Stochastic Gradient Descent

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay

WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training