← Forthcoming posts

Alpha preview · Work in progress

An early version of a forthcoming blog post. The text, proofs, and examples may still evolve.

The covariance of SGD iterates

An isotropic noise floor, anisotropic relaxation, and the fourth moments hiding in the next term.

With a fixed learning rate, stochastic gradient descent usually does not settle at a single point. It keeps fluctuating. The size and shape of these fluctuations matter: they determine the residual optimization error, and tell us what is lost when a discrete training algorithm is replaced by a continuous approximation.

Least squares offers a particularly transparent example. Under constant conditional observation-noise variance, the stationary covariance of SGD is, to first order,

\[ \Sigma_\eta=\eta\sigma^2 I+O(\eta^2). \]

It is isotropic even when the loss has very different curvatures in different directions. The reason is a cancellation: steeper directions both pull more strongly toward the minimizer and receive more stochastic forcing. At the next order, fourth moments of the features break this balance.

The broader picture is developed beautifully by Aymeric Dieuleveut, Alain Durmus, and Francis Bach [1], who study constant-step SGD through its invariant Markov-chain distribution. Here, a direct least-squares calculation gives exact moment equations and a short proof of the small-step expansion.

A fresh sample at each step

Consider the population loss

\[ f(\theta)=\mathbb E[(X^\top\theta-Y)^2], \qquad \theta\in\mathbb R^d. \]

Assume a well-specified linear model with homoscedastic noise:

\[ Y=X^\top\theta_*+\xi,\qquad \mathbb E[\xi\mid X]=0,\qquad \mathbb E[\xi^2\mid X]=\sigma^2. \]

Let \(G=\mathbb E[XX^\top]\succ0\), assume \(\mathbb E\|X\|^4<\infty\), and take \(\sigma^2>0\). The matrix \(G\) is the second moment of the features, and \(2G\) is the Hessian of \(f\). The features need not be centered or Gaussian, and \(\xi\) need not be independent of \(X\).

At every iteration, draw a new independent pair \((X_{k+1},Y_{k+1})\). With step size \(\eta>0\) and error \(e_k=\theta_k-\theta_*\), SGD reads

\[ e_{k+1}=(I-2\eta X_{k+1}X_{k+1}^\top)e_k +2\eta\xi_{k+1}X_{k+1}. \]

The initialization is independent of these samples and has finite second moment. This is an online, i.i.d. sampling model; sampling without replacement or making deterministic or shuffled passes through a fixed dataset requires a different dependence analysis.

Mean, second moment, and covariance

It is useful to keep two quantities separate:

\[ m_k=\mathbb E[e_k],\qquad M_k=\mathbb E[e_ke_k^\top]. \]

The actual covariance is

\[ \operatorname{Cov}(\theta_k)=M_k-m_km_k^\top. \]

For example, a deterministic initialization away from the optimum has zero covariance, but nonzero \(M_0=e_0e_0^\top\). If \(m_0=0\), the mean stays zero and \(M_k\) is the covariance at every step.

Define the linear fourth-moment operator

\[ \mathcal T(M)=\mathbb E[(X^\top MX)XX^\top]. \]

Despite its fourth-order dependence on \(X\), this is linear in \(M\). Operators of this kind are central to precise least-mean-squares analyses, including the work of Défossez and Bach [2].

Theorem 1 (exact moment recursion). Under the assumptions above,

\[ m_{k+1}=(I-2\eta G)m_k, \] \[ \begin{aligned} M_{k+1}={}&M_k-2\eta(GM_k+M_kG)\\ &+4\eta^2\mathcal T(M_k)+4\eta^2\sigma^2G. \end{aligned} \]
Proof: expand one update, then condition on the fresh features

Write \(A_X=I-2\eta XX^\top\). Independence of the new sample and \(e_k\) gives \(\mathbb E[A_Xe_k]=(I-2\eta G)m_k\). The additive term has zero mean because \(\mathbb E[\xi\mid X]=0\).

Expanding the outer product gives

\[ \begin{aligned} M_{k+1} &=\mathbb E[A_X M_k A_X]+4\eta^2\mathbb E[\xi^2XX^\top]\\ &=M_k-2\eta(GM_k+M_kG) +4\eta^2\mathcal T(M_k)+4\eta^2\sigma^2G. \end{aligned} \]

Both mixed terms vanish after conditioning on \(X,e_k\). No Gaussian integration or continuous-time approximation is involved. \(\square\)

The last term is the noise injected at the optimum. The fourth-moment term is the contribution of the random curvature \(XX^\top\). Replacing this random matrix by its mean \(G\) would change that term and miss part of the finite-step correction.

A stationary law, not a limiting point

Choose a finite constant \(R\) such that

\[ \mathbb E[\|X\|^2XX^\top]\preceq RG. \]

Such an \(R\) exists because \(G\succ0\) and the fourth moment is finite. If \(\|X\|^2\le R\) almost surely, the same constant works. The order \(U\preceq V\) means that \(V-U\) is positive semidefinite.

Theorem 2 (stationarity at a safe step size). If \(0<\eta<1/R\), SGD has a unique stationary law with finite second moment. From any admissible initialization, its law converges to this stationary law in \(W_2\). Its mean is exactly \(\theta_*\), and its covariance \(\Sigma_\eta\) satisfies

\[ G\Sigma_\eta+\Sigma_\eta G =2\eta\sigma^2G+2\eta\mathcal T(\Sigma_\eta). \]

This is an exact linear matrix equation for \(\Sigma_\eta\). The step-size condition is sufficient, not asserted to be sharp for every feature distribution.

Proof: couple two runs with the same samples

Drive two copies of the algorithm with the same fresh \((X,\xi)\). Their difference \(\Delta\) evolves as \(\Delta^+=(I-2\eta XX^\top)\Delta\). Hence, with \(\mu=\lambda_{\min}(G)\),

\[ \begin{aligned} \mathbb E[\|\Delta^+\|^2\mid\Delta] &\le\|\Delta\|^2-4\eta(1-\eta R)\Delta^\top G\Delta\\ &\le q\|\Delta\|^2, \end{aligned} \qquad q=1-4\eta(1-\eta R)\mu<1. \]

Here \(q\ge0\), as is also clear from the nonnegative left-hand side. Thus one SGD step contracts the quadratic Wasserstein distance by at most \(\sqrt q\). The quadratic Wasserstein distance \(W_2\) is the smallest root-mean-square separation over couplings of two laws.

Why the contraction gives an invariant probability law

The one-step update sends finite-second-moment laws to finite-second-moment laws: the multiplicative term uses \(\mathbb E\|X\|^4<\infty\), and the additive term uses \(\mathbb E[\xi^2\|X\|^2]=\sigma^2\operatorname{tr}(G)<\infty\). The space \((\mathcal P_2(\mathbb R^d),W_2)\) is complete. Applying the contraction mapping theorem to this one-step map gives a unique fixed law in \(\mathcal P_2\) and convergence to it. In particular, first and second moments converge.

At stationarity the mean equation gives \(Gm=0\), hence \(m=0\). Setting \(M_{k+1}=M_k=\Sigma_\eta\) in Theorem 1 and dividing by \(2\eta\) gives the covariance equation. \(\square\)

Why the leading covariance is isotropic

The stationary covariance is of order \(\eta\), so \(2\eta\mathcal T(\Sigma_\eta)\) only enters at order \(\eta^2\). At first order, the equation becomes

\[ G\Sigma_\eta+\Sigma_\eta G\simeq2\eta\sigma^2G. \]

The solution is \(\eta\sigma^2I\). Both the restoring force and the injected-noise covariance contain the same \(G\), and they cancel in the stationary balance. This is a feature of this noise model, not a universal property of SGD.

Theorem 3 (small-step covariance expansion). As \(\eta\to0\),

\[ \Sigma_\eta=\eta\sigma^2I+\eta^2D+O(\eta^3), \]

where the remainder is in matrix norm and the positive-semidefinite matrix \(D\) is the unique solution of

\[ GD+DG=2\sigma^2\mathbb E[\|X\|^2XX^\top]. \]

Only fourth moments of \(X\) are needed. In coordinates where \(G=\operatorname{diag}(\lambda_1,\ldots,\lambda_d)\), the correction is explicit:

\[ D_{ij}=\frac{2\sigma^2\mathbb E[\|X\|^2X_iX_j]}{\lambda_i+\lambda_j}. \]
Proof: expand an invertible linear operator

Let \(\mathcal L(M)=GM+MG\). The stationary equation is

\[ (\mathcal L-2\eta\mathcal T)\Sigma_\eta=2\eta\sigma^2G. \]

In an eigenbasis of \(G\), \(\mathcal L\) multiplies entry \((i,j)\) by \(\lambda_i+\lambda_j>0\). It is therefore invertible. Since \(\mathcal T\) is a fixed finite-dimensional linear operator, the inverse of \(\mathcal L-2\eta\mathcal T\) has a convergent Taylor expansion near zero. In particular \(\Sigma_\eta=O(\eta)\), and substitution gives successively

\[ \mathcal L(\sigma^2I)=2\sigma^2G, \qquad \mathcal L(D)=2\sigma^2\mathcal T(I). \]

This proves the expansion and the entrywise expression. Repeated compositions of \(\mathcal T\) do not require higher moments of \(X\): they are compositions of an already defined linear operator, not higher powers of the same sampled feature.

Why the correction is positive semidefinite

For any symmetric matrix \(H\), the solution of \(GU+UG=H\) is

\[ U=\int_0^\infty e^{-tG}H e^{-tG}\,dt. \]

Differentiate the integrand to verify the equation. If \(H\succeq0\), every integrand is positive semidefinite. Taking \(H=2\sigma^2\mathcal T(I)\) gives \(D\succeq0\). \(\square\)

There is even an exact inequality throughout the safe step-size regime:

\[ \Sigma_\eta\succeq\eta\sigma^2I. \]

Indeed, the stationary equation gives \(\mathcal L(\Sigma_\eta-\eta\sigma^2I)=2\eta\mathcal T(\Sigma_\eta)\succeq0\), and the integral formula in the proof shows that \(\mathcal L^{-1}\) preserves positive semidefiniteness. Thus the leading isotropic approximation underestimates the variance in every direction.

Two explicit examples

One dimension: kurtosis controls the correction

Write \(g=\mathbb E[X^2]>0\), \(h=\mathbb E[X^4]\). The stationary equation is just \(2gv=2\eta\sigma^2g+2\eta hv\), so

\[ v=\frac{\eta\sigma^2}{1-\eta h/g}, \qquad 0<\eta<g/h. \]

Two feature distributions with the same second moment have the same leading covariance, but different fourth moments can produce very different finite-step noise floors. For example, \(X=\pm\sqrt g\) with equal probabilities gives \(h=g^2\), whereas \(X\sim\mathcal N(0,g)\) gives \(h=3g^2\).

Gaussian features in any dimension

If \(X\sim\mathcal N(0,G)\), Gaussian fourth moments give

\[ \mathcal T(M)=2GMG+\operatorname{tr}(GM)G. \]

In particular, \(\mathcal T(I)=2G^2+\operatorname{tr}(G)G\), and therefore

\[ D=\sigma^2[\operatorname{tr}(G)I+2G]. \]

At second order, directions with larger curvature have larger variance. A sufficient stability constant is \(R=\operatorname{tr}(G)+2\lambda_{\max}(G)\).

For isotropic features, \(G=\lambda I_d\), the exact answer is

\[ \Sigma_\eta=\frac{\eta\sigma^2}{1-\eta\lambda(d+2)}I_d, \qquad 0<\eta<\frac1{\lambda(d+2)}. \]
Proof: fourth-moment pairing and a scalar fixed point

For a centered Gaussian vector,

\[ \mathbb E[X_iX_jX_rX_s] =G_{ij}G_{rs}+G_{ir}G_{js}+G_{is}G_{jr}. \]

Contracting with the entries of a symmetric \(M\) yields the formula for \(\mathcal T\). The expression for \(D\) follows by substitution into its Lyapunov equation.

When \(G=\lambda I_d\), start from \(M_0=0\). The exact moment recursion preserves the form \(M_k=v_kI_d\), so convergence from Theorem 2 implies \(\Sigma_\eta=vI_d\). Substitution gives \(v=\eta\sigma^2+\eta\lambda(d+2)v\). Solving proves the formula. No rotational symmetry of the full conditional noise law is needed. \(\square\)

The denominator reveals an instability invisible in the first-order approximation. As \(\eta\lambda(d+2)\) approaches one, the second moment blows up, even though the mean recursion may still be stable.

The transient: flat directions take longer

Start at \(\theta_0=\theta_*\). The covariance initially vanishes and gradually builds up. On the time scale \(t=k\eta\), the leading behavior is especially simple.

Proposition 4 (small-step transient). For each fixed \(T<\infty\), uniformly over integers \(k\) with \(0\le k\eta\le T\),

\[ \operatorname{Cov}(\theta_k) =\eta\sigma^2(I-e^{-4k\eta G})+O_T(\eta^2). \]

In eigendirection \(i\), the variance is approximately \(\eta\sigma^2(1-e^{-4\lambda_i t})\). The limiting variance is the same in every direction at first order, but a direction with small \(\lambda_i\) takes longer to reach it.

Proof: a scaled moment equation and Euler’s method

Since \(m_0=0\), write \(N_k=M_k/\eta\). The exact recursion becomes

\[ N_{k+1}=N_k-2\eta\mathcal L(N_k) +4\eta^2\mathcal T(N_k)+4\eta\sigma^2G. \]

On any bounded time interval, the matrices \(N_k\) are uniformly bounded by a discrete Gronwall estimate. Thus this is Euler’s method, with an additional local \(O(\eta^2)\) term, for

\[ \dot N=-2(GN+NG)+4\sigma^2G,\qquad N(0)=0. \]

Its solution is \(N(t)=\sigma^2(I-e^{-4tG})\). Euler’s global error is \(O_T(\eta)\); multiplying by \(\eta\) proves the stated \(O_T(\eta^2)\) covariance error. This argument uses only the finite fourth-moment operator. \(\square\)

The same leading covariance equation is realized by the Ornstein–Uhlenbeck surrogate

\[ de_t=-2Ge_t\,dt+2\sqrt\eta\,\sigma G^{1/2}dB_t, \]

where \(B_t\) is standard Brownian motion. Its invariant covariance is exactly \(\eta\sigma^2I\). This gives useful intuition, but the discrete fourth-moment correction is absent from this surrogate. The OU perspective is prominent in Mandt, Hoffman, and Blei [3]; stochastic modified equations make finite-step corrections systematic [4]. A finite-time diffusion approximation does not, by itself, justify exchanging the small-step and infinite-time limits.

Numerical experimentation

The figure uses two-dimensional Gaussian features with eigenvalues \(1\) and \(1/\kappa\), observation-noise variance \(\sigma^2=1\), and initialization \(\theta_*\). Vary the curvature ratio \(\kappa\), the safe step fraction \(\eta R\), and the number of iterations. The exact covariance is computed from the moment recursion, not estimated from Monte Carlo runs.

All variances are divided by \(\eta\sigma^2\), so the first-order stationary approximation is the unit circle. Solid curves show exact moments; dashed curves show the first-order transient. The stationary comparison includes the second-order expansion, with relative errors measured in Frobenius norm. The safe bound \(\eta R<1\) is sufficient, not necessarily the exact stability threshold for anisotropic features.

What changes beyond this model?

Nonconstant conditional noise. Keep \(\mathbb E[\xi\mid X]=0\), assume \(\mathbb E[\xi^2]<\infty\) so that the loss remains finite, and suppose \(C=\mathbb E[\xi^2XX^\top]\) is finite. Then replace \(\sigma^2G\) by \(C\) in the exact equations. The expansion becomes

\[ \Sigma_\eta=\eta S+\eta^2D+O(\eta^3), \] \[ GS+SG=2C,\qquad GD+DG=2\mathcal T(S). \]

The first term need not be isotropic. It is the alignment of \(C\) with \(G\), not quadraticity alone, that produced our cancellation.

Nonquadratic objectives. Dieuleveut, Durmus, and Bach [1, Proposition 2 and Theorem 4] establish invariant-law convergence and small-step moment expansions under suitable strong-convexity, regularity, and noise assumptions. Their leading second moment solves a Lyapunov equation involving the Hessian and the gradient-noise covariance at the minimizer. Unlike least squares, nonlinear objectives generally also have a step-size-dependent stationary mean bias. Their theory is broader; the direct argument here needs only the stated least-squares moment assumptions.

Averaging iterates is a different question. We have studied the covariance of a single late iterate. The covariance of \(\bar\theta_n=n^{-1}\sum_{k=1}^n\theta_k\) also depends on correlations between times. In general, it is not \(\Sigma_\eta/n\). Precise least-squares averaging results appear in [2], and the general Markov-chain approach also analyzes this distinction [1].

The common lesson, also motivating our work on diffusion discretization [5], is to retain enough of the discrete dynamics to see which parameters the leading continuous picture hides. Here, the first term explains the isotropic noise floor; the next term reveals the role of feature geometry and the limits of that approximation.

References

  1. Aymeric Dieuleveut, Alain Durmus, and Francis Bach. Bridging the gap between constant step size stochastic gradient descent and Markov chains. Annals of Statistics 48(3), 1348–1382, 2020. Open-access preprint.
  2. Alexandre Défossez and Francis Bach. Averaged Least-Mean-Squares: Bias-Variance Trade-offs and Optimal Sampling Distributions. AISTATS, Proceedings of Machine Learning Research 38, 205–213, 2015.
  3. Stephan Mandt, Matthew D. Hoffman, and David M. Blei. Stochastic Gradient Descent as Approximate Bayesian Inference. Journal of Machine Learning Research 18(134), 1–35, 2017.
  4. Qianxiao Li, Cheng Tai, and Weinan E. Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations. Journal of Machine Learning Research 20(40), 1–47, 2019.
  5. Samuel Hurault, Thomas Moreau, and Gabriel Peyré. Geometry-Aware Discretization Error of Diffusion Models. arXiv:2605.08392, 2026.