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

Mar 2, 2026

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

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.

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

Ask PDF

Feb 6, 2026

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

We study optimal learning rate (LR) schedules under the functional scaling law (FSL) framework (Li et al., 2025), which decomposes training dynamics into signal learning and noise forgetting. In power-law kernel regression, these two components are governed by a source exponent

s>0

and a capacity exponent

q>1

, respectively, with smaller

s

corresponding to harder tasks. For a fixed training horizon

N

, we characterize the schedules that minimize the final-step loss under a stability constraint and reveal a sharp phase transition. In the easy-task regime

s>1-1/q

, the optimal schedule follows power decay from the beginning of training; in the hard-task regime

s<1-1/q

, it becomes warmup-stable-decay (WSD)-like (Hu et al., 2024), staying at the largest admissible LR for most of training before a final decay. In both regimes, the decay exponent is

2q-1

: task difficulty determines when to decay, while model capacity determines how to decay. Beyond the exact optimum, we study fractional schedules, whose shape is defined over relative training progress. We show that precise tuning of the decay shape is often unnecessary: a broad class of profiles attains the optimal convergence rate, while overly slow terminal decay leads to schedule-induced capacity saturation. Finally, for one-pass SGD in kernel regression, FSL-motivated power-decay schedules achieve optimal last-iterate rates. Experiments support the theoretical predictions and the task-dependent transition between early and delayed decay.

Binghui Li, Zilin Wang, Fengling Chen +3

Center for Machine Learning Research, Peking University · School of Mathematical Sciences, Peking University · AI for Science Institute, Beijing

Learning Rate Scheduling Kernel Regression Stochastic Gradient Descent

Ask PDF

Jul 12, 2026

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

Standard learning rate schedules such as cosine annealing are tied to a fixed training horizon, limiting their ability to accommodate post hoc horizon extension. Warmup-stable-decay (WSD) partially addresses this issue by maintaining a long constant-rate phase before a short linear cooldown, allowing training to resume from a pre-decay checkpoint. However, its peak learning rate is still tuned based on the original training horizon and can become suboptimal when training is extended. Motivated by stochastic convex optimization, we propose WSqD (Warmup with Square-root base and linear Decay), a learning rate schedule that replaces WSD's constant stable phase with a shifted inverse-square-root base while retaining the final linear cooldown. In the stochastic convex setting, WSqD provably attains the minimax-optimal

O(1/\sqrt{T})

last-iterate convergence rate. Importantly, its base learning rate schedule is horizon-independent, and the training horizon is needed only to determine when to begin the final cooldown. Empirically, on language-model pretraining using the SlimPajama corpus, WSqD matches or outperforms carefully tuned WSD and other baselines across multiple training horizons while reusing a single peak learning rate.

Jianhao Ma, Yuxin Chen

Department of Statistics and Data Science, University of Pennsylvania

Learning Rate Scheduling

Ask PDF