PaperlayerView exact revision on arXiv

Source-linked AI summary

SGD in Multiclass Logistic Regression: Sequential Learning and Scaling Laws

Konstantinos Christopher Tsiolis, Denny Wu, Christos Thrampoulidis, Murat A. Erdogdu

arXiv:2609.07868v1stat.MLcs.LG

TL;DR

The paper addresses the limited scaling-law theory for cross-entropy classification with many classes. It analyzes gradient-based learning in Gaussian mixtures, finding sequential class learning and compute-optimal scaling under PCA-based capacity control.

  • Problem

    Existing scaling-law theory largely treats squared loss, while cross-entropy next-token prediction involves many classes with highly non-uniform frequencies.

    Source: §1 · pp. 1–2, §1.1 · p. 2

  • Method

    The paper analyzes gradient-based training of high-dimensional K-class logistic regression on Gaussian mixtures with power-law priors and uses PCA projection to study capacity.

    Source: §1.1 · p. 2, §3 · p. 7

  • Results

    Classes are learned in decreasing-frequency order, yielding a three-phase cross-entropy risk law; PCA capacity control decomposes risk into approximation and optimization errors and supports compute-optimal choices.

    Source: §1.1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    The results provide a theoretical scaling-law account for multiclass classification under cross-entropy, complementing existing regression analyses.

    Source: §1 · p. 1, §5 · p. 10

    The analysis assumes orthogonal class means and isotropic within-class covariance, and its SGD analysis concerns one-pass online training.

    Source: §5 · p. 10

Abstract

from arXiv · show

We study the training dynamics of multiclass logistic regression on high-dimensional Gaussian mixture models with a large number of classes and establish precise scaling laws governing the cross-entropy risk under gradient-based optimization. We show that learning proceeds sequentially across classes, from most to least frequent. When the class priors follow a power law distribution, the risk dynamics decompose into three phases: an initial plateau until the first class is learned, a power-law decay regime during which sequential learning occurs, and a final convergence regime. We then analyze how model capacity interacts with optimization under a fixed compute budget. When the effective dimension is restricted via projection onto leading principal components, the risk decomposes into a capacity term (a power law in the retained dimension) and an optimization term (a power law in training time). Optimizing this tradeoff yields a compute-optimal scaling law for logistic regression, with explicit prescriptions for model size and training time as functions of compute. These results extend theoretical scaling laws from linear regression to multiclass classification, while connecting to empirical scaling laws observed in large-scale neural networks.

1 Introduction

The paper addresses the gap between squared-loss scaling theory and cross-entropy prediction by studying multiclass classification with many classes. It analyzes sequential learning, PCA-based capacity control, and online SGD in a high-dimensional Gaussian-mixture setting.

1 Introduction

  • The paper addresses the gap between squared-loss scaling theory and cross-entropy next-token prediction by asking for tractable scaling laws in multiclass problems with diverging class counts.Source: §1 · pp. 1–2

1.3 Preliminaries

  • The setup is K-class logistic regression on a Gaussian mixture with orthonormal means, high-dimensional inputs, and power-law class frequencies.Source: §1.1 · p. 2, §1.3 · p. 3

1.1 Our Contributions

  • Population gradient flow learns classes in decreasing-frequency order at t_i ≍ (log K)/π_i, yielding power-law risk scaling; for noisy GMMs, this persists to T ≍ 1/σ.Source: §1.1 · p. 2
  • PCA projection onto the span of the m most frequent class means controls capacity, and optimizing the dimension–training-time tradeoff under C ≍ mKt yields compute-optimal choices.Source: §1.1 · p. 2
  • With a suitable learning rate, online SGD tracks population gradient flow over the relevant horizon and exhibits similar scaling with effective time t = ηn.Source: §1.1 · p. 2

1.2 Related Work

  • The paper connects sequential-learning scaling to additive-model work [MLGT23] [NFL+24] and distinguishes its diverging-class, power-law setting from majority-first learning [ZT25].Source: §1.2 · p. 3

2 Population Gradient Flow

Population gradient flow tracks multiclass learning through summary statistics, revealing sequential recovery and a three-phase cross-entropy risk trajectory. With nonzero Gaussian noise, the noiseless scaling persists only over a finite horizon, while the theory also characterizes the irreducible risk and the flow’s limit.

2 Population Gradient Flow

  • The population gradient-flow analysis tracks class-wise summary statistics, including θ_ij, the inner product of class weight w_i and class mean μ_j.Source: §2 · p. 4

2.1 Warmup: The Idealized Noiseless Case

  • In the noiseless case, class probabilities rise monotonically, and classes with larger priors are learned faster; this ordering result holds beyond power-law priors.Source: §2.1 · p. 4
  • For power-law priors with α > 1, risk stays Θ(log K) initially, decays as a power law during sequential recovery, then follows K/t after all classes enter the head.The transition between the intermediate and convergence phases includes an uncharacterized window.Source: §2.1 · pp. 4–5

2.2 GMM Case

  • With σ > 0, the noiseless three-phase scaling is established only through T ≍ 1/σ, and depending on σ, trajectories may not reach every phase.The deviation point and subsequent risk behavior are not explained by the current theory.Source: §2.2 · pp. 5–6
  • For 0 < σ = o(1/log K), a unique constrained weight configuration attains the irreducible risk, and zero-initialized gradient flow converges to it.Source: §2.2 · p. 7

3 Compute-Optimal Scaling Laws

The paper uses PCA onto the means of the most frequent classes to impose a capacity bottleneck, decomposing risk into approximation and optimization terms. Balancing these errors under fixed compute yields compute-optimal choices of retained dimension and training time.

Why PCA rather than a Gaussian sketch

  • A Gaussian sketch does not meaningfully reduce logistic-regression expressivity when m ≳ log K, because it largely preserves class-separation geometry.The Johnson–Lindenstrauss result preserves pairwise mean distances, leaving the signal-to-noise ratio effectively unchanged.Source: §3 · p. 7

PCA-based capacity control

  • PCA onto the means of the m most frequent classes collapses the remaining class means, creating a capacity bottleneck that prevents recovery of all means when σ² = 0.In that setting, the convergence phase of the unsketched dynamics disappears.Source: §3 · p. 7

Empirical PCA

  • For empirical PCA with σ² > 0, the scaling result has a finite time horizon and requires enough samples for the PCA step.The theorem states the risk bound with probability at least 1 − δ.Source: §3 · p. 8

Capacity-limited phase

  • After the early phase, risk decomposes into approximation error, a power law in m, plus optimization error, a power law in t.The sketched classes follow the familiar optimization dynamics, while unrepresented classes contribute the capacity bottleneck.Source: §3 · p. 8

Compute-optimal scaling

  • Under compute C ≍ mKt, balancing approximation and optimization errors gives compute-optimal hyperparameters when K log K ≲ C ≲ K^(α+2) log K.The stated range ensures t ≳ log K and m < K.Source: §3 · p. 8

4 Online Stochastic Gradient Descent

The paper shows that online SGD follows the scaling-law behavior of population gradient flow when learning rates and time are appropriately matched. Its compute tradeoff depends on the effective dimension, with smaller sketches allowing faster learning at the cost of approximation error.

Online SGD dynamics

  • Online SGD exhibits scaling behavior analogous to gradient flow when effective optimization time is t = ηn and the learning rate is suitably chosen.The result holds over a fixed time horizon; in the sketched problem, the effective dimension is m rather than d.Source: §4 · p. 9

Online SGD algorithm

  • The SGD iteration draws an independent GMM sample and updates the weights by a gradient step on cross-entropy loss.The algorithm initializes W(0) at zero and outputs the weights after the final iteration.Source: §4 · p. 9

Unsketched SGD

  • For the unsketched model, the theorem gives scaling regimes in effective time t_n = ηn, with high-probability bounds through a finite horizon.The stated horizon is bounded by Tσ = c*/σ.Source: §4 · p. 9

Sketched SGD

  • In the sketched case, SGD has analogous scaling regimes, and its martingale bound depends on m rather than d.The theorem assumes independent samples for empirical PCA and online SGD, and sufficient sample size for constructing the sketch.Source: §4 · pp. 9–10

Compute-optimal tradeoff

  • At fixed compute C ≍ mKn, balancing approximation and optimization errors determines the compute-optimal sketch dimension and iteration count.The stated range also requires ηn* ≤ T.Source: §4 · p. 10

5 Conclusion

The paper derives scaling laws for multiclass logistic regression with power-law class frequencies, including sequential class learning, capacity-optimization tradeoffs, and online SGD dynamics. Its analysis leaves open extensions to nonzero-noise dynamics beyond a perturbative horizon, less restrictive data assumptions, other optimizers, and data reuse.

Conclusion

  • Multiclass logistic regression with many power-law-frequency classes exhibits sequential learning, a three-phase risk scaling law, and a compute-optimal tradeoff between capacity and optimization.The paper also studies time-discretized dynamics and scaling laws for online SGD.Source: §5 · p. 10

Open directions

  • The analysis assumes orthogonal class means and isotropic within-class covariance, and controls nonzero-noise dynamics perturbatively only up to a noise-limited horizon.Other open directions include different optimizers and data reuse across multiple epochs.Source: §5 · p. 10

A Proofs for Unsketched Population Gradient Flow

The appendix proves the unsketched population gradient-flow results by tracking class-level probabilities and weights, first in an idealized zero-noise setting. It then decomposes cross-entropy risk into head and tail contributions to bound their evolution.

A Proofs for Unsketched Population Gradient Flow

  • The proof tracks θ_ij, p_i|j, and P_i|j because the gradient-flow expectation has no closed form.The appendix derives summary-statistic dynamics using these quantities and the orthonormality assumption.Source: A Proofs for Unsketched Population… · p. 16

A.1.1 Class-wise Learning Dynamics

  • For σ² = 0, class-wise dynamics decouple, and the conservation law relates each off-class statistic θ_ij to the corresponding on-class statistic θ_jj.The relation is θ_ij(t) = −θ_jj(t)/(K−1) for i ≠ j.Source: §A.1 · p. 17, §A.1.1 · p. 17
  • The idealized class-wise dynamics yield a hitting-time expression for when a class probability reaches a threshold.The proof obtains the expression by solving the class probability dynamics; it later uses t_i without an argument for the threshold 1/2.Source: §A.1.1 · pp. 17–18

A.1.2 Risk Decomposition

  • Cross-entropy risk is split into head classes with p_i|i(t) > 1/2 and the remaining tail classes.The split is indexed by k*(t), with head classes i ≤ k*(t).Source: §A.1.2 · p. 18

A.1.3 Tail Risk

  • The tail-risk proof bounds the remaining classes using monotonicity of −log p_i|i(t) and an integral approximation.The trivial upper bound is log K; a sharper upper bound applies once k*(t) ≥ 1.Source: §A.1.3 · pp. 18–19

A.1.4 Head Risk

  • When t − t_K ≳ t, the head-risk upper bound collapses to K/t and matches the lower bound.The proof controls head classes using their recovery times and bounds on 1 − p_i|i(t).Source: §A.1.4 · pp. 19–20

A.1.5 The Case α > 1

For power-law class priors, the risk has an early plateau, an intermediate regime, and a convergence regime. The intermediate behavior differs across α > 1 and α < 1, while both cases converge at order K/t after the recovery horizon.

A.1.5 The Case α > 1

  • When α > 1, the normalization factor is constant order and k*(t) ≍ (t/log K)^(1/α).This scaling describes the number of recovered classes in that case.Source: §A.1.5 · p. 20
  • For t < t1, k*(t) = 0 and Rhead(t) = 0; during this early phase, the tail risk satisfies Rtail(t) ≲ log K and Rtail(t) ≳ log K.Source: §A.1.5 · p. 20
  • In the intermediate phase, R(t) ≍ (log K)^(2−1/α) · t^(−(1−1/α)); the head contribution is lower order than the tail risk.Source: §A.1.5 · p. 21
  • When t > (1 + δ)tK for δ > 0, the tail risk is zero, the head risk is Θ(K/t), and total risk satisfies R(t) ≍ K/t.Source: §A.1.5 · p. 21

A.1.6 The Case α < 1

  • For α < 1, the early and intermediate phases both have R(t) ≍ log K, with the tail risk dominating in the intermediate phase.The intermediate regime runs from t_1 to a constant fraction of t_K.Source: §A.1.6 · pp. 21–22
  • For α < 1 and t ≥ (1 + δ)t_K, R(t) ≍ K/t.The convergence bound holds for any positive constant δ.Source: §A.1.6 · p. 22

A.1.7 The Case α = 1

The appendix treats α = 1 and then extends the analysis to noisy Gaussian mixtures, where the zero-noise class decoupling no longer applies. It handles the noise as a perturbation and bounds the time horizon over which noisy dynamics stay close to the idealized case.

A.1.7 The Case α = 1

  • The α = 1 class-recovery scaling is summarized through t_i and k*(t) in terms of log K.The appendix gives the corresponding relations in (A.57).Source: §A.1.7 · p. 22
  • For α = 1, the risk is R(t) ≍ log K before the first recovery, and the tail risk dominates the head risk in the intermediate phase.The convergence regime has R(t) ≍ K/t for t ≥ (1 + δ)t_K.Source: §A.1.7 · p. 22

A.2 GMM Case

  • With nonzero noise, the summary-statistic dynamics do not decouple as in the σ² = 0 case, so the proof treats noise as an additive perturbation.The perturbation analysis compares noisy dynamics with the idealized dynamics.Source: §A.2 · p. 22

A.2.1 Perturbation Lemmas

  • The noisy dynamics remain close to the idealized σ² = 0 case over a time horizon T = c*/σ for 0 < σ ≤ 1.The proof relates this horizon to the growth of the largest weight norm.Source: §A.2.1 · p. 23
  • The perturbation lemmas compare probabilities at class means with their conditional expectations under Gaussian noise.The proof uses a Gaussian moment-generating-function formula and a Taylor expansion to bound the softmax perturbation.Source: §A.2.1 · pp. 24–25

A.2.2 Class-wise Learning Dynamics

For the noisy GMM, the appendix establishes class-wise recovery and post-recovery decay within a bounded time horizon, then uses these results to analyze aggregate risk across phases. It also opens the proof for the sketched setting, focusing on α > 1.

A.2.2 Class-wise Learning Dynamics

  • For class i recovered within the horizon, t_i is bounded above and below by constant multiples of (log K)/π_i, and recovery is guaranteed when the upper bound fits within T.The result is stated for constants 0 < c < C.Source: §A.2.2 · p. 27
  • After recovery, p_i|i(t) stays at least 1/2 and its error decays as 1 − p_i|i(t) ≍ 1/[1 + π_i(t − t_i)].The result holds uniformly for t_i ≤ t ≤ T.Source: §A.2.2 · p. 28

A.2.3 Tail Risk

  • The noisy-GMM analysis separates head and tail risks to assess their contributions to aggregate risk.The proof then combines these bounds in its decomposition into risk phases.Source: §A.2.2 · p. 28, §A.2.5 · p. 30

A.2.5 Decomposition into Risk Phases

  • A matching upper bound gives R(t) ≍ log K in that phase; separately, tail risk dominates in the cited regime, where other head classes contribute at most 1/log K.Source: §A.2.5 · pp. 30–31
  • For α > 1, the noisy-GMM intermediate-phase risk is R(t) ≍ (log K)^(2−1/α) · t^(−(1−1/α)).The proof finds the tail-risk terms dominate the head-risk terms in this regime.Source: §A.2.5 · p. 30

B Proofs for Sketched Population Gradient Flow

  • The appendix introduces the sketched-gradient-flow proof for α > 1 using summary statistics formed from Sμ_j and probabilities at Sμ_j.The sketching matrix is S ∈ R^(m×d).Source: B Proofs for Sketched Population… · p. 31

B.1 Idealized Case

In the idealized projection, retained classes keep their unsketched dynamics while collapsed classes remain indistinguishable. With noise, the analysis bounds the risk of collapsed classes and preserves phased risk decay.

B.1 Idealized Case

  • For σ²=0, classes i≤m retain their unsketched recovery times, while i>m remain uniform and contribute time-constant approximation error.The retained classes follow the same class-wise dynamics as in the unsketched case.Source: §B.1 · p. 32
  • The idealized risk is log K before t1, decays as (log K)^(2−1/α)t^−(1−1/α), then scales as m/t after (1+c2)tm.The power-law regime extends through c1tm.Source: §B.1 · p. 32

B.2 GMM Case

  • For σ²>0, optimal risk has an initial log K plateau, power-law decay for log K≲t≤c1m^α log K, and an m/t tail once t≥c2m^α log K.Source: §B.2 · p. 33
  • In noisy GMM, classes i>m stay unlearned with p_i|i=1/K, and approximation risk is at least the order of optimal risk by c1m^α log K.Source: §B.2 · pp. 33–34

B.3 Empirical PCA

This section studies when empirical PCA preserves the paper’s scaling-law behavior and how gradient-flow guarantees transfer to discrete and stochastic optimization, analyzing empirical-sketch, discretization, and stochastic-gradient stability. It also notes that population gradient-flow risk is non-increasing and that a constant in a condition on η can be chosen sufficiently small.

B.3 Empirical PCA

  • Sequential recovery and post-recovery decay carry over to empirical PCA when εPCA ≤ c0πm, with the recovered class probability remaining at least 1/2 over the time horizon.The post-recovery risk bound includes an additive sketch-error contribution proportional to εPCA/πi.Source: §B.3 · pp. 35–36
  • Under the theorem’s assumptions, empirical PCA retains the risk’s plateau, power-law decay, and final convergence regimes, with an additional m^(1−α) term in the head-risk bound.The head-risk bound applies through time c1 m^α log K; for later times, the stated bound is Ropt(t) ≲ m/t + m^(1−α).Source: §B.3 · p. 37

C Proofs for Online SGD

  • The proof framework treats sketched and unsketched optimization together, using effective optimization time and second-moment scales to control discretization and stochastic-gradient error.The framework assumes a fixed sketch with operator norm at most one; L controls smoothness and discretization error, while M controls stochastic-gradient variance.Source: C Proofs for Online SGD · p. 38

C.1 Risk and Gradient Descent Properties

  • The population risk is convex and L-smooth, and population gradient descent is non-expansive for step sizes 0 < η ≤ 2/L.The gradient norm is also non-increasing along population gradient flow and population gradient descent.Source: §C.1 · pp. 38–39

C.2 From Gradient Flow to Gradient Descent

  • For 0 < η ≤ 1/L, the gradient-descent trajectory stays close to gradient flow over the specified finite horizon, with error bounded by a universal constant times ηLT(1 + LT).The discretization result assumes ηN ≤ T.Source: §C.2 · p. 40

C.3 From Population to Stochastic Gradient Descent

  • The cited lemmas concern conditional stochastic-gradient variance and uniform online SGD stability; the stability result assumes 0 < η < 1/L, T > 0, δ ∈ (0, 1), and ηN ≤ T.Source: §C.3 · p. 41

C.4 From Population Gradient Flow to Online SGD

  • Combining discretization and SGD stability transfers gradient-flow guarantees uniformly to online SGD with probability at least 1 − δ.The theorem assumes 0 < η ≤ 1/L and ηN ≤ T; a separate corollary gives a learning-rate condition for the transfer.Source: §C.4 · p. 42
  • The constant c in the condition on η can be chosen sufficiently small.The excerpt does not provide the condition itself.Source: §C.4 · p. 42
  • Risk is non-increasing along population gradient flow.The excerpt then says this gives the result, but does not specify the result further.Source: §C.4 · pp. 42–43

C.5 Unsketched GMM Case

Online SGD is shown to inherit the unsketched population risk phases over a stated horizon. With empirical PCA, the corresponding result holds under assumptions on sketch accuracy and sample independence.

C.5 Unsketched GMM Case

  • With probability at least 1−δ, online SGD follows the unsketched population phases for every n with t_n≤T, subject to the theorem’s stated condition.Source: §C.5 · p. 43
  • The unsketched SGD risk has a log K plateau, then follows f_K,α(t_n), and later scales as K/t_n across the stated time ranges.Source: §C.5 · p. 43

C.6 Empirical Sketch Case

  • With independent PCA and SGD samples, α>1, m≤K/2, and ε_PCA≤c0 min{π_m,σ log K}, empirical-PCA SGD has the stated risk bound uniformly for t_n≤T.Source: §C.6 · p. 44

D The Irreducible Risk

This section characterizes irreducible risk by establishing the population-risk minimizer and bounding its risk, and reports related results for noiseless multi-output squared-loss regression. It describes the regression dynamics across early, intermediate, and late training phases and gives α-dependent irreducible-risk scaling under the stated condition.

D The Irreducible Risk

  • Zero-initialized population gradient flow converges to the unique minimizer of the risk on the centered weight space, which is also globally minimizing by softmax translation invariance.The risk is strictly convex on that space; the infimum is nonzero when σ² > 0.Source: §2.2 · p. 7, D The Irreducible Risk · pp. 44–46
  • The Bayes posterior entropy lower-bounds logistic-regression cross-entropy risk, and no closed-form expression for the Bayes risk is available, motivating tight bounds.The risk decomposes into posterior entropy plus a nonnegative KL divergence.Source: D The Irreducible Risk · pp. 44, 46
  • Because the model has no biases, it cannot attain the Bayes-optimal classifier for unequal priors; adjusted weights approximate the Bayes decision boundary with additive error o(σ²).The proposed construction chooses class weights proportional to class means, with adjustments based on the log priors.Source: D The Irreducible Risk · p. 49
  • The section states upper and lower bounds on the best achievable logistic-regression risk when 0 < σ ≤ c and σBπ ≤ c, for positive constants c and C.The supplied passages do not expose the bound formulas.Source: D The Irreducible Risk · p. 52
  • The irreducible risk scales as K for 0 < α < 1, K/log K for α = 1, K^(2−α) for 1 < α < 2, (log K)^2 for α = 2, and 1 for α > 2.Source: D The Irreducible Risk · p. 53
  • Under σ = o(1/log K), the irreducible risk is strictly lower order than the population gradient-flow risk for t ≤ T.The comparison passage identifies log K during Phase 1 and Kσ during Phase 3 as larger-order risks.Source: D The Irreducible Risk · p. 53

E Squared Loss in the Idealized Case

E.1 Multi-Output Linear Regression

  • Multi-output regression replaces the initial log K plateau and terminal K/t phase with a constant-order plateau and exponential convergence, while retaining the intermediate exponent 1−1/α.Each regression coordinate converges exponentially and independently; summing over power-law-distributed relaxation rates produces the intermediate power law.Source: §E.1 · p. 56

E.2 Squared Loss with Softmax Activation

With squared-loss training on a softmax activation, evaluated using cross-entropy risk, the idealized scaling law has early, intermediate, and convergence phases. Compared with cross-entropy training, squared loss has slower learning and convergence in the stated regimes.

E.2 Squared Loss with Softmax Activation

  • The comparison uses squared loss on the softmax activation, evaluated by cross-entropy risk; squared loss on logits does not permit direct cross-entropy comparison.Source: §E.2 · pp. 56, 58
  • For population gradient flow on squared loss with softmax activation in the GMM setting, assuming σ²=0 and α>1, risk has three phases: log K, intermediate decay, then t^(-1/2)gα(K).The phases apply for 0≤t≤t1; t1≤t≤c1tK, with risk K^(1−1/α)log K·t^(−(1−1/α)) + t^(−1/2)gα((t/K)^(1/α)); and (1+c2)tK≤t, respectively.Source: §E.2 · p. 58
  • For α>1, the scaling law’s gα factor depends on α: it scales as m^(1−α/2) for 1<α<2, log(1+m) for α=2, and 1 for α>2.Source: §E.2 · p. 58
  • Ignoring log factors, the early and intermediate phases last longer by a factor of K than under cross-entropy training.The squared-loss dynamics include an extra factor K/(K−1)·pi|i(1−pi|i), which the text identifies as slowing learning.Source: §E.2 · pp. 57, 60
  • When α>2, squared-loss convergence has rate Θ(t^(−1/2)) for t ≳ K^(1+α), versus Θ(K/t) under cross-entropy training.The text attributes the slower rate to cubic rather than quadratic dependence on (1−pi|i) in the key ODE.Source: §E.2 · p. 60

Related papers

Loading 2609.07868v1…