A short tutorial · six lessons

How stable is SGD?

We will walk through this slowly. First the SGD update. Then a simple quadratic. Then a real network. Then the tools we need once the linear theory runs out.

1. SGD 2. A quadratic 3. The tails 4. A real network 5. 1-D / 2-D labs 6. Past Lyapunov

Lesson 1. Neural networks are trained with SGD

That is the first fact. We have parameters \(\theta\) and a loss \(L\). At every step we pick a random mini-batch \(B_t\) and take one step downhill on that batch:

\[ \theta_{t+1} \;=\; \theta_t \;-\; \eta\,\nabla L_{B_t}(\theta_t) \]

\(\eta\) is the step size (the learning rate). \(B_t\) is a new random batch each time. That is all SGD is.

Now the natural question: is this stable? Do the parameters stay near a minimum, or do they run away?

A real neural network is too messy to answer that. So we study the simplest loss that still has curvature: a quadratic.

Lesson 2. A quadratic loss, and what “stable” means

The simple model

Let \(x\) be how far we are from the minimum. Each batch sees a slightly different curvature \(h_t\). Think of \(h_t\) as “how steep that batch is.” We sample it from a distribution. The easiest example is a Gaussian:

\[ h_t \;\sim\; \mathcal N(\mu,\, v) \]

\(\mu\) is the average curvature. \(v\) is how much the batches disagree. Smaller batches mean larger \(v\). If we used the whole dataset every time, \(v=0\).

The loss on that batch is a simple quadratic, plus a little noise:

\[ L_{B_t}(x) \;=\; \tfrac12 h_t\, x^2 \;-\; \xi_t\, x \]

Plug this into SGD. The update becomes very short:

\[ x_{t+1} \;=\; (1-\eta h_t)\, x_t \;+\; \eta\sigma\zeta_t \]

Read it in words. Every step we multiply the current position by a random number \(1-\eta h_t\), then add a small Gaussian kick. That is the whole model. A stochastic quadratic. Nothing else.

What do we mean by “stable”?

The iterate \(x_t\) is now random, so we need a metric. The usual one is: the average squared distance from the minimum stays bounded.

\[ \mathbb E[x_t^2]\ \text{ stays finite for all } t \]

People call this mean-square stability.

Why the square? Because it is easy to compute. Square the update. The noise term has mean zero, so it drops out of the cross term, and we get:

\[ \mathbb E[x_{t+1}^2] \;=\; \underbrace{\mathbb E[(1-\eta h)^2]}_{m_2(\eta)}\; \mathbb E[x_t^2] \;+\; \eta^2\sigma^2 \]

Each step multiplies the mean square by one number, \(m_2(\eta)\). If that number is less than 1, the mean square settles down. If it is bigger than 1, the mean square grows forever.

And \(m_2\) has a closed form. It is less than 1 exactly when

\[ \eta \;<\; \eta_{\mathrm{MS}} \;=\; \frac{2\mu}{\mu^2+v} \]
Notice: if the batches all agree (\(v=0\)), this is the classical number \(\eta<2/\mu\). Once batches disagree, the allowed step size gets strictly smaller. That is the first fact.

Try it. Leave the curvature Gaussian. Move the step size across the line \(\eta_{\mathrm{MS}}\). Watch the histogram of \(x\).

Does the mean square stay bounded?

average curvature \(\mu\)1.00
batch disagreement \(v\)0.04
step size \(\eta\)1.80

This is a histogram of \(x\) after many steps. Below the line it stays peaked near zero. Above the line the mass spreads out.

Stability is not one number. It is a spectrum.

The square was only one choice. We could have asked whether \(\mathbb E|x_t|\) stays finite, or \(\mathbb E|x_t|^4\), or \(\mathbb E|x_t|^{1/2}\).

For this simple update, each of those questions has its own answer. Write

\[ m_q(\eta) \;=\; \mathbb E\,|1-\eta h|^{q} \]

The \(q\)-th moment of \(x_t\) stays finite when \(m_q(\eta)<1\), and blows up when \(m_q(\eta)>1\). So each \(q\) has its own step-size threshold.

Two named cases:

\[ \begin{aligned} \text{mean-square stability:}&\quad \mathbb E|1-\eta h|^{2} \;<&\; 1\\[4pt] \text{harmonic stability:}&\quad \mathbb E|1-\eta h|^{-1} \;<&\; 1 \end{aligned} \]

Mean-square asks about the second moment. Harmonic asks about the inverse moment: how often the random multiplier is small, i.e. how often a step pulls \(x\) back toward zero.

If the step size is below the mean-square line, every moment of order \(2\) or less is stable. Moments of higher order can already be unstable. That is what “spectrum” means: we do not get a yes/no. We get a list of which averages survive and which ones blow up.

The picture below is the same Gaussian as above. Move \(\eta\). The blue curve is \(m_q\). Moments under the red line \(m_q=1\) are the stable ones.

Which moments are still finite?

step size \(\eta\)1.96

The green mark is \(\kappa\): the last stable moment. Everything to the right of it diverges.

Lesson 3. What does the distribution of \(x\) look like?

So far we only asked whether averages stay finite. We can also ask: after many steps, what does a histogram of \(x\) look like?

When a typical run still contracts, the iterates settle into a fixed shape. That shape is not a Gaussian. Sometimes a few unlucky batches multiply \(x\) up several times in a row. Those rare runs make a heavy tail.

The 1-D Kesten theorem says exactly how heavy. Find the number \(\kappa>0\) that solves

\[ \mathbb E\,|1-\eta h|^{\kappa} \;=\; 1 \]

Then the tail of the stationary law is a power law with that same \(\kappa\):

\[ \mathbb P(|x|>t) \;\approx\; C\, t^{-\kappa} \qquad\text{for large }t. \]
This \(\kappa\) is the same number as in Lesson 2. A tail \(t^{-\kappa}\) has a finite \(q\)-th moment exactly when \(q<\kappa\). So “which moments survive” and “how heavy is the tail” are one fact.

Change the step size, or change how the batch curvature is sampled. The histogram changes. That is all this lesson is.

Change the curvature and the step size. Watch the histogram.

step size \(\eta\)1.96

The peak stays near zero. The tails get heavier as \(\eta\) grows. For more laws and a longer run, open the 1-D lab.

Lesson 4. Is this all? What about a neural network?

Everything so far was a one-dimensional quadratic. A real network has thousands of parameters and a loss that is not a bowl. So a fair question is:

Can the stability of SGD tell us what happens in a neural network? Or did we just solve a toy problem?

Here is one experiment. Take a trained network and freeze a checkpoint \(\theta_\star\). Keep running plain SGD from there. After each step, look only along the sharpest direction: the top eigenvector \(u_1\) of the Hessian. Call that one number

\[ z_t \;=\; \langle \theta_t-\theta_\star,\; u_1\rangle \]

If the quadratic story were the whole story, then past the Lyapunov line this coordinate should run away. It does not. It sits in a stationary cloud. Raise the step size further and that cloud becomes two lobes.

The pictures below are not a simulation of a toy map. They are a real 2-block transformer, continued with full-batch SGD, then projected onto \(u_1\). The number \(\eta\lambda_1\) is the step size times the largest curvature. The classical edge is \(2\). The Lyapunov number \(\chi_{\mathrm{lin}}=\mathbb E\log|1-\eta\, u_1^\top H_B u_1|\) turns positive at the same place.

A real network, probed along its sharpest direction

step size \(\eta\lambda_1\)2.50
loading the recorded experiment…

Histogram of \(z=\langle\theta-\theta_\star,u_1\rangle\). Same checkpoint, only the step size changes.

Raising \(\eta\) walks through the whole range: a spike at zero, then a spread cloud past Lyapunov, then two lobes. Open the checkpoint lab for other nets and batch sizes.

Lesson 5. The same picture in 1-D and 2-D

The network probe is high-dimensional. The same family of histograms appears in the small models you already have as labs. Those models add the next term in the loss — a cubic and a quintic — so the iterate is pulled back when it gets large. That is why a stationary cloud can exist past the Lyapunov line.

\[ x_{t+1} \;=\; (1-\eta h_t)\,x_t \;+\; \alpha x_t^{3} \;-\; \beta x_t^{5} \;+\; \sigma\zeta_t \]

The first term is the quadratic we already know. The extra powers are the self-stabilization. Move \(\eta\) from the mean-square line, through Lyapunov, into the two-lobe region.

1-D experiment · the scalar map from the lab

step size \(\eta\)2.55

Full controls — curvature law, \(\alpha\), \(\beta\) — live in the 1-D lab.

In two dimensions there is a sharp coordinate \(z\) and a tangent coordinate \(y\). The same walk: a blob at the origin, then a cloud past Lyapunov, then two lobes along the sharp axis.

2-D experiment · sharp axis \(z\), tangent axis \(y\)

step size \(\eta\)2.50

Left: the \((z,y)\) cloud. Right: the histogram along the sharp axis. Full coupling knobs are in the 2-D lab. The rotating-axes version is the matrix lab.

Start from the quadratic. SGD on a quadratic, and what breaks after builds the classical picture, shows a network that violates it, then walks the three labs.
1-D

Scalar cocycle

One curvature, cubic and quintic return, the Fig. 1 strip.
2-D

Sharp–tangent

Two coupled coordinates. Lobes live on the sharp axis.
a real net

Checkpoint probe

SGD from \(\theta_\star\), occupancy in the top Hessian plane.

Lesson 6. Do we have tools past the Lyapunov line?

At the Lyapunov line the typical one-step gain stops contracting:

\[ \chi(\eta) \;=\; \mathbb E\log|1-\eta h| \;\ge\; 0 \]

Kesten’s theorem needs \(\chi<0\). So the linear / quadratic theory stops here. It cannot tell us the shape of the cloud we just saw in the network, or in the 1-D and 2-D labs.

Do we have any tool left for SGD once we cross Lyapunov? Yes — but we have to look at a different kind of moment. Not “is \(x\) large on average?”, but “how often is the gain small?”

That is a negative moment. For \(\alpha>0\),

\[ m_{-\alpha}(\eta) \;=\; \mathbb E\,|1-\eta h|^{-\alpha} \]

If \(m_{-\alpha}<1\), the gain is small often enough that a typical excursion is pulled back toward zero. Two named cases:

\[ \begin{aligned} \text{negative-moment stability:}&\quad \mathbb E\,|1-\eta h|^{-\alpha} \;<&\; 1 \quad\text{for some }\alpha>0\\[6pt] \text{harmonic stability:}&\quad \mathbb E\,|1-\eta h|^{-1} \;<&\; 1 \end{aligned} \]

Harmonic is just the case \(\alpha=1\). It is the last of these inverse-moment edges. In the labs, the two-lobe histograms appear around there — not at the Lyapunov line, and not at the mean-square line.

Self-stabilization is what keeps the iterate from escaping after Lyapunov. Negative-moment / harmonic stability is the linear diagnostic for when the return toward zero becomes rare. Together they are the language for the bimodal region.

The inverse-moment spectrum

step size \(\eta\)2.12

Now \(q\) runs through negative values too. The crossing of \(m_q=1\) on the left is the inverse-moment edge. For a Gaussian the curve blows up as \(q\to-1\), which is why the papers use a fractional \(\alpha<1\).

That is the end of this first pass. The quadratic gave us a spectrum of ordinary moments and a Kesten tail. A real network, and the 1-D / 2-D maps, keep a stationary shape past Lyapunov. The new tools there are self-stabilization and the inverse moments — mean-square, then Lyapunov, then negative-moment, then harmonic.