Why Training Can Keep Improving at the Edge of Stability
Training a machine-learning model means adjusting numbers called parameters so its predictions make fewer errors. Yet a run can improve over many updates while its measured error rises and falls from one update to the next. The puzzle is how the same update rule can produce both oscillation and progress.
To guide those changes, we need a numerical measure of error, called a loss. Suppose a model uses weather measurements to predict temperature. An example's observed temperature is 10 degrees and the model predicts 8. One possible loss is squared error: (8 − 10)² = 4. A prediction of 9 reduces that loss to 1. Averaging losses across the training examples gives a training loss: the score the training process tries to reduce. Different tasks use different loss formulas, but the score serves the same role in choosing parameter updates.
The next problem is deciding which parameters to change, and by how much. The gradient describes how the loss changes under small changes to the parameters. With one parameter, it is a slope. With many, it collects a slope for each parameter. Gradient descent updates the parameters in the opposite direction, scaled by a number called the learning rate. This update rule is one example of an optimizer.
The gradient describes local behavior, while an actual update moves a finite distance. A larger learning rate can carry the parameters past a nearby minimum into a region with higher loss. A smaller one may make progress slowly. Training therefore produces a sequence of parameter settings and losses, which we can plot to see how the process develops.
In full-batch gradient descent, every update uses the entire training dataset. In some neural-network training runs, loss fluctuates over short intervals while declining over longer ones. At the same time, the largest local curvature—the fastest increase in slope along any parameter direction—stays near a threshold associated with instability in a fixed-curvature model. This regime is called the edge of stability.
Explaining continued progress requires identifying both what makes an oscillation grow and what could push back against that growth.
A threshold we can calculate
Begin with one adjustable parameter, x, and a loss that measures its distance from the desired value of zero:
L(x) = ½Sx².
The positive number S sets the curvature. A larger S means the slope changes more quickly as we move away from zero. Gradient descent measures that slope, Sx, and takes a step in the opposite direction. The learning rate, η, controls the size of that step:
xₜ₊₁ = xₜ − ηSxₜ = (1 − ηS)xₜ.
The entire behavior fits inside the multiplier, 1 − ηS. If its magnitude is smaller than one, each update shrinks the distance from zero. If the multiplier is negative, the parameter also changes sides. Crossing the minimum is therefore compatible with convergence: the crossing can get smaller every time.
Take S = 2 and start at x = 1. With η = 0.75, the multiplier is −0.5. The first few positions are 1, −0.5, 0.25, and −0.125. Here the loss equals x², so it falls from 1 to 0.25, then 0.0625. The parameter oscillates while the loss decreases at every step.
Increase η to 1. Now the multiplier is −1. The parameter alternates between 1 and −1, and the loss stays at 1. This is the boundary: motion continues without progress.
At η = 1.1, the multiplier becomes −1.2. The positions begin 1, −1.2, 1.44, and −1.728. Each crossing carries the parameter farther from zero. There is nothing in this particular objective that can interrupt the growth.
Figure 1. Compare shrinking oscillations, the neutral boundary, and growth on identical axes. The warm arrow shows the subtraction along the parameter axis; the green dot marks an exact computed state. Step advances one update.
For this positive quadratic, a positive learning rate, and a nonzero starting point, convergence to zero requires 0 < ηS < 2. Equality at two gives the repeating orbit above; exceeding two gives growing oscillations. The threshold comes from the update rule.
This quadratic separates crossing the minimum from making progress. Its parameter can alternate sides while its loss falls, stays fixed, or grows. It cannot yet explain why a growing excursion might later subside: its curvature never changes.
Curvature changes with position
With many parameters, a position means a particular setting of all those numbers. A direction means changing some combination of them together. The loss can bend gently in one direction and sharply in another. Its largest local curvature is called sharpness, S, here. Formally, this is the largest eigenvalue of the Hessian, the matrix describing how the gradient changes with position. This is the measurement observed near 2/η in the edge-of-stability regime.
The quadratic calculation then raises a concrete concern: a sufficiently sharp direction can produce expanding oscillations at the chosen learning rate. But a quadratic approximation describes the neighborhood of a position. Extending its prediction indefinitely also assumes that the curvature governing the updates will remain unchanged.
When curvature varies with position, excursions along a sharp direction can contribute an update toward lower sharpness. A third-order expansion of the loss exposes a contribution proportional to the squared displacement. It therefore survives a change in the displacement's sign. This is the local feedback behind self-stabilization.
The objective itself stays fixed. The optimizer moves through parameter space, where different locations have different curvature.
How an oscillation can reduce curvature
Consider a local example with two coordinates:
L(x,y) = ½(2 + y)x² − 0.1y.
The first term is a quadratic in x whose curvature depends on y. The second term supplies a tendency to increase y: at x = 0, its gradient with respect to y is −0.1, so gradient descent steps toward larger y. Moving that way makes the x direction sharper because its curvature is 2 + y.
Figure 2. One fixed surface, three locations. Scan y from −1 to 0 to 1: the x-curvature rises from 1 to 2 to 3. The highlighted slice appears below on identical axes, measured above its own minimum. Moving the slice changes which part of the surface we inspect.
At y = 0, write the gradient, ∇L, as an ordered pair: the x component first, then the y component. It has three useful pieces:
∇L(x,0) = (0, −0.1) + (2x, 0) + (0, ½x²).
The baseline term encourages increasing y. The middle term changes sign with x and drives the crossing motion. The last term is always nonnegative in the y component. Because gradient descent subtracts the gradient, this term contributes motion toward smaller y, opposing the baseline tendency to sharpen.
At y = 0, this decomposition is exact.
Choose two symmetric sample locations, x = +0.2 and x = −0.2. Their x-gradient components are +0.4 and −0.4. Average them and they cancel. Their squared-displacement contributions do not: each contributes 0.02 to the y gradient. Adding the baseline gives −0.1 + 0.02 = −0.08.
At η = 1.1, the average y update is therefore +0.088. These small excursions supply a correction, but it is too weak to reverse the sharpening tendency.
Now sample x = +0.6 and x = −0.6. The x-gradient components are +1.2 and −1.2, which again cancel. Each squared-displacement contribution is now 0.18. Adding the baseline gives a positive y gradient: −0.1 + 0.18 = +0.08. Subtracting it at the same learning rate gives an average y update of −0.088.
Increasing the excursion has changed the direction of the average update. The factor that caused the change is visible in the formula: the displacement tripled, so its squared contribution increased ninefold.
Figure 3. Compare the amplitudes, then step through the positive sample, negative sample, paired mean, and average update. The arrows show gradient contributions. Averaging cancels the alternating x term and retains the squared y term. Subtracting the mean determines the update's direction. The final stage compares this with the gradient at the average position.
These are two chosen sample positions, not consecutive optimizer states. Their average update demonstrates a local correction without determining a training trajectory. The polynomial is only a local illustration: along x = 0 it decreases without bound as y increases.
Nor does one such update restore stability. At the center after the larger-amplitude average update, the x curvature would be 2 − 0.088 = 1.912. The quadratic threshold at η = 1.1 is approximately 1.818. The correction points in the relevant direction without yet crossing that threshold.
There is also a distinction between coordinate curvature and sharpness. At the center x = 0, near y = 0, the largest Hessian eigenvalue is 2 + y. Away from that center, the Hessian also has off-diagonal terms. Reading only the x curvature would omit their contribution to the largest eigenvalue.
The feedback can now be stated directly. Sharpening can enlarge an oscillation; larger excursions strengthen the correction that opposes sharpening. If these effects balance, training can remain near the stability boundary while its longer-term path continues moving. The calculation exposes the opposing effects; whether they balance requires more than this pair of samples.
What the average has to retain
The example presents a problem for a tempting simplification. Suppose we replace the two sampled positions by their average position. The average of (+0.6,0) and (−0.6,0) is (0,0). At that center, the gradient is (0,−0.1), which encourages increasing y.
But the average of the two gradients is (0,+0.08), which encourages decreasing y. Evaluating the gradient at the average position has given us the opposite y direction from averaging the gradients at the sampled positions.
The difference comes from the square. Averaging x first produces zero, and squaring that zero removes the contribution. Squaring each displacement first produces a positive value on both sides, which survives averaging. There is no contradiction: these are different operations on a nonlinear expression.
Central flow models a locally averaged path through parameter space while retaining the influence of the oscillations around it. Ordinary gradient flow idealizes gradient descent as continuous motion in the downhill direction. That idealization can miss the curvature-dependent correction produced by finite steps. Central flow has matched training trajectories in studied settings; it remains an approximation, not a general convergence guarantee.
For the mathematical connection, write a nearby parameter vector as w̄ + xu. Here w̄ is the center, u is a vector of length one along the Hessian's sharpest direction, and x is displacement along it. Assume a smooth loss and a largest Hessian eigenvalue separated from the others. Neglecting higher-order terms gives:
∇L(w̄ + xu) ≈ ∇L(w̄) + xS(w̄)u + ½x²∇S(w̄).
Treat the center and local eigendirection as approximately fixed over a short averaging interval. If the signed displacement averages to zero, the linear term cancels while the term involving the average of x² remains. Subtracting that surviving term contributes motion against the gradient of sharpness.
Retaining this mean square is only part of the model. At the stability boundary, when the baseline motion would keep sharpening, the single-direction central flow determines an oscillation variance—the mean squared displacement about the center—whose curvature-reducing effect balances that tendency. This supplies a rule for the averaged path beyond merely smoothing an already recorded trajectory.
A smoothed loss chart is a different object again. It averages recorded scalar losses. It neither constructs the average parameter vector nor calculates the loss of that vector. Our example makes this distinction numerical: both samples at ±0.6 have loss 0.36, whereas their average position has loss zero. The average loss is 0.36; the loss at the average position is zero.
Smoothing can make a trend readable. It cannot, by itself, supply the dynamics that produced the trend.
Reading a training run more carefully
The first question is what the run measures. If each point reports a different minibatch—a subset of the training examples—a change in the recorded loss includes a change in the evaluated examples. In the full-batch setting that motivates the basic argument here, the objective being evaluated stays the same between updates. Those observations require different interpretations.
Next comes the relation between the short interval and the longer one. Retain the unsmoothed measurements alongside the trend, with the smoothing window stated. A display that hides every excursion cannot help distinguish a bounded pattern from excursions whose magnitude keeps expanding.
Where curvature measurements are practical, loss and sharpness provide complementary evidence. For plain gradient descent, the relevant comparison is between sharpness and 2/η. A claim about self-stabilization is stronger when it connects an excursion to a subsequent curvature reduction and continued progress. A jagged line alone provides none of those connections.
The local correction may be insufficient. The approximation may stop describing the region reached by a large step. Evidence of this mechanism must come from the run being evaluated.
The optimizer matters too. Momentum carries information from earlier updates; minibatch sampling changes the gradient from step to step; AdamW also rescales coordinates and applies weight decay. These mechanisms change the stability calculation. The threshold 2/η cannot simply be carried over. The local expansion also becomes less reliable when large steps reach regions where higher-order effects matter, or where the assumed smoothness fails.
Finally, training progress concerns the objective being optimized. Generalization concerns performance beyond the training data. Different parameterizations can represent the same predictions while changing parameter-space curvature. Sharpness alone therefore cannot establish predictive performance. Validation evidence remains necessary.
On a fixed objective, a loss spike tells us that an update made the score worse. It does not tell us whether later updates will recover. The useful evidence is what follows: whether curvature falls, whether excursions remain contained, and whether the same objective continues improving over a longer interval. The edge of stability gives us a mechanism for interpreting that evidence. Recovery still has to appear in the run.