SGD's stationary cloud vs. the posterior

WB2 §8–6.4. A 2-D quadratic basin L(θ) = ½ θ^THθ, H = diag(1, κ), gradient noise C = R_φ H^γ R_φ^T. Is SGD's stationary covariance shaped like the posterior's?

SGD, SDE (OU) law Σ SGD, discrete-time exact Σ_η posterior (T/n)H⁻¹ simulated SGD cloud, empirical cov.

All ellipses are 2σ (contain 86% of a 2-D Gaussian). Gray: loss contours. Axes equal-scaled; θ₁ is the flat direction (curvature 1), θ₂ the sharp one (curvature κ).

noise exponent γ (C ∝ H^γ)

−10123

jump to:

condition number κ

learning rate η

batch size B

noise rotation φ (C's eigenbasis vs H's)

posterior temperature T (n = 1000)

keep T = T* (match trace)

excess risk ½Tr(HΣ)
SGD, SDE exact for OU: ηTr C / 4B
SGD, discrete exact discrete Lyapunov
simulated cloud Monte Carlo, 1500 chains
posterior T·p / 2n, p = 2
shape
KL* (SDE Σ ‖ best-T posterior) nats; 0 iff same shape
KL at current T
T* = n·Tr(HΣ)/p
anisotropy of HΣ eig ratio; posterior has 1
‖Σ_η − Σ‖/‖Σ‖ discrete vs SDE

Shape mismatch as a function of γ

KL* = min over T of KL(N(0,Σ) ‖ N(0,(T/n)H⁻¹)) = ½[p log(Tr(HΣ)/p) − log det(HΣ)], at the current κ, φ (and η for the discrete curve). Click to set γ.

SDE Σ discrete-time Σ_η 1 ≤ γ ≤ 2: diagonal range reported by Zhang–Yang–Tu (2026)

What to look at. Slide γ. At γ = 0 (isotropic noise, C ∝ I) the orange SGD ellipse has exactly the posterior's shape H⁻¹: long along the flat direction. At γ = 1 (C ∝ H, the Fisher identity of WB2 §7) the SGD law is a circle, Σ = (η/2B)·I. At γ = 2 it is ∝ H: more spread along the sharp direction than the flat one, the opposite of the posterior. With φ = 0, Σ ∝ H^γ−1; KL* is symmetric about γ = 0 in log-eigenvalue and vanishes only there. Rotating the noise (φ ≠ 0) tilts Σ, and then no temperature makes it match. The excess risk ½Tr(HΣ) = ηTr(C)/4B holds for any C, and choosing T = T* makes the posterior's excess risk T p/2n equal SGD's; this T* is also the KL-optimal temperature.

Exact vs. simulated. The solid orange ellipse solves the continuous Lyapunov equation HΣ + ΣH = (η/B)C exactly. The dashed one solves the exact stationary equation of discrete SGD θ ← θ − η(Hθ + ξ), Cov ξ = C/B, namely Σ_η = (I − ηH)Σ_η(I − ηH) + η²C/B. They differ at O(ηκ): the discrete law is wider in the sharp direction and diverges at ηκ ≥ 2. The green cloud is a direct simulation of that discrete recursion (Gaussian ξ, 1500 independent chains started at θ = 0) and should converge onto the dashed ellipse.

Claim boundary. This is a local, linear Ornstein–Uhlenbeck model near one minimum with a fixed, state-independent C and Gaussian noise. It does not show loss-dependent noise (e.g. C ∝ L·H, Mori et al.), which gives heavy-tailed stationary laws, nor non-quadratic basins, interpolation (C → 0, WB2 §8), momentum or adaptive optimizers. The 1 ≤ γ ≤ 2 band is Zhang–Yang–Tu's (arXiv 2602.05600) empirical finding for diagonal entries C_ii ∝ H_ii^γ in their setups; it is not a claim that C is exactly a power of H. "SGD is approximately Bayesian" in this model is the claim γ ≈ 0, or that the noise has been preconditioned by C⁻¹ (Mandt et al.).