NB2 §1 — Gradient descent as a flow plus corrections

Backward error analysis (§1.2). The paper this notebook is built on.

Is the expansion valid? (§1.3). Why the series diverges, and where training actually sits.

Shuffled-minibatch SGD (§1.4).

Momentum (§1.5). Goh’s first section also does §1.1’s quadratic case.

Also: Rosca, Wu, Qin & Dherin, a continuous-time model of GD and instability (the complex “principal flow”, exact on quadratics; Exercise 1.2.c); Kovachki & Stuart, continuous-time analysis of momentum methods (the damped ODE and slow manifold of §1.5; see the sign-error note there); Cattaneo, Klusowski & Shigida, on the implicit bias of Adam (backward error analysis for Adam and RMSProp; end of §1.5).

On a quadratic loss, gradient descent with step size \eta converges iff \eta S<2, where the sharpness S is the largest eigenvalue of the Hessian H. Near a minimum the same holds for the linearization, so minima sharper than 2/\eta repel GD, while gradient flow (\dot\theta=-\nabla L) converges to them. This threshold recurs in the rest of the notebook and in td2_3_edge_of_stability.

The first section of Goh’s Why momentum really works (in the Readings, under §1.5) does this case with interactive pictures.

Derivations and walkthrough. Notation: GD is \theta_{k+1}=\theta_k-\eta\nabla L(\theta_k); gradient flow uses time t=k\eta, so one step advances time by \eta.

On L(\theta)=\tfrac12\theta^\top H\theta with H symmetric positive semidefinite, GD is linear, \theta_{k+1}=(I-\eta H)\theta_k. In the eigenbasis Hu_i=\lambda_iu_i each coordinate c_i=u_i^\top\theta evolves on its own:

c_i^{(k)}=(1-\eta\lambda_i)^k\,c_i^{(0)}, \qquad\text{while gradient flow gives}\qquad c_i(t)=e^{-\lambda_it}c_i(0),\quad t=k\eta .

One GD step multiplies by 1-\eta\lambda where the flow multiplies by e^{-\eta\lambda}=1-\eta\lambda+\tfrac12(\eta\lambda)^2-\dots. They agree for \eta\lambda\ll1 and part ways as \eta\lambda grows:

• 0<\eta\lambda<1: monotone decay, faster than the flow per unit time;
• \eta\lambda=1: that direction is solved in one step;
• 1<\eta\lambda<2: the coordinate flips sign every step while it decays;
• \eta\lambda=2: a period-2 orbit; \eta\lambda>2: oscillating divergence.

Directions with \lambda_i=0 do not move, so GD converges iff \eta S<2.

Exercise 1.1 — Best constant step; unstable minima.

1. Take a quadratic whose nonzero eigenvalues lie in [\mu,S]. Find the fixed \eta that minimizes the worst per-step contraction factor, and that factor in terms of \kappa=S/\mu. Where does the optimal \eta put \eta S?
2. Let \theta^* be a minimum of a smooth L whose Hessian H^* is positive definite, with sharpness S^*=\lambda_{\max}(H^*). Show that if \eta S^*>2, GD started near \theta^* generically does not converge to it, whereas gradient flow started near it does.

GD does not follow gradient flow exactly, but to first order in the step size \eta it follows gradient flow on a modified loss,

\dot\theta=-\nabla\tilde L,\qquad \tilde L=L+\frac\eta4\|\nabla L\|^2 ,

with time t=k\eta. The extra term penalizes the squared gradient norm: this is implicit gradient regularization (IGR).

Paper: Barrett & Dherin, Implicit gradient regularization, Section 3 (see Readings). The method is standard numerical analysis: Hairer, Lubich and Wanner, Geometric Numerical Integration, ch. IX.

Core idea. Write g=\nabla L and H for the Hessian. Over a time \eta, gradient flow moves by -\eta g+\tfrac{\eta^2}2Hg+O(\eta^3), because the gradient changes along the way. GD uses the gradient from the start of the step and moves by exactly -\eta g, overshooting by -\tfrac{\eta^2}2Hg. A flow that reproduces GD must add that overshoot as a drift, -\tfrac\eta2Hg=-\nabla\big(\tfrac\eta4\|g\|^2\big). Near a minimum \|\nabla L\|^2\approx\delta^\top H^2\delta for a displacement \delta, so the penalty is heaviest in sharp directions. It vanishes at every critical point, so local minima of L stay local minima of \tilde L and IGR does not help GD escape them. It matters when the minima form a family, as in overparametrized models, where it changes which minimum GD reaches (Exercise 1.2.d). That GD ends at flatter minima in general is Barrett and Dherin’s prediction, supported by experiments, not a theorem.

Derivations and walkthrough. Notation: loss L, g=\nabla L, H=\nabla^2L, step \eta; GD is \theta_{k+1}=\theta_k-\eta g(\theta_k) and t=k\eta. For a general one-step numerical method the state is y and the step is h; for GD, y=\theta and h=\eta.

Forward error asks how far the iterates are from the true flow; backward error asks which nearby ODE reproduces them exactly. The difference between that ODE and the original says what the discretization does. (Implicit biases of gradient flow itself, such as max-margin solutions, are a different story.)

The setup. A one-step method with step h maps y\mapsto\Phi_h(y), with

\Phi_h(y)=y+hf(y)+h^2d_2(y)+h^3d_3(y)+\cdots .

GD is \Phi_\eta(\theta)=\theta-\eta\nabla L(\theta): h=\eta, f=-\nabla L, all d_j=0. Look for a modified vector field

\tilde f=f+hf_2+h^2f_3+\cdots

whose exact time-h flow equals \Phi_h as a power series in h. (Indexing as in Hairer–Lubich–Wanner; Barrett and Dherin call the O(h) correction f_1.)

The recipe. The time-h flow of \dot y=F(y) expands as

y(h)=y+hF+\frac{h^2}{2}F'F+\frac{h^3}{6}\big(F''(F,F)+F'F'F\big)+O(h^4),

with F' the Jacobian and F''(u,v) the second derivative applied to two vectors. Substitute F=\tilde f and match powers of h against \Phi_h. At order h^j the unknown f_j enters linearly with coefficient 1 and everything else is known from lower orders, so the modified equation exists and is unique as a formal series. At first order, f_2=d_2-\tfrac12f'f (Exercise 1.2.a); that is all §1.4 and §1.5 use.

What the truncation buys. Keeping r terms, \tilde f_{[r]}=f+hf_2+\dots+h^{r-1}f_r, gives local error O(h^{r+1}) per step and global error O(h^r) over a fixed time, with constants depending on that time and on derivatives of f. Gradient flow (r=1) is O(\eta) off; the IGR flow (r=2) is O(\eta^2) off. The code after Exercise 1.2.d checks both slopes, and how much of GD’s actual shift the IGR flow accounts for.

Exercise 1.2.a — The recipe.

1. Derive the h^3 term of the flow expansion, \tfrac16\big(F''(F,F)+F'F'F\big).
2. Match orders to show

f_2=d_2-\tfrac12f'f,\qquad f_3=d_3-\tfrac16\big(f''(f,f)+f'f'f\big)-\tfrac12\big(f'f_2+f_2'f\big).

Exercise 1.2.b — Gradient descent.

1. For GD (f=-g, d_2=d_3=0), show f_2=-\tfrac12Hg and that this equals -\nabla\big(\tfrac14\|g\|^2\big). Conclude \dot\theta=-\nabla\big(L+\tfrac\eta4\|\nabla L\|^2\big).
2. (Optional.) Show f_3=-\tfrac13H^2g-\tfrac1{12}\nabla^3L[g,g], where \nabla^3L[g,g]_i=\sum_{jk}\partial_i\partial_j\partial_kL\,g_jg_k.
3. (Optional.) For L(x,y)=\tfrac12x^2y, compute f_3 and its curl \partial_xf_{3,y}-\partial_yf_{3,x}. What does a nonzero curl say about describing GD as “gradient flow on a modified loss”?

Exercise 1.2.c — The quadratic, to all orders.

For L=\tfrac12\lambda x^2 in one dimension:

1. Find the exact modified equation: a linear ODE \dot x=a\,x whose time-\eta flow is exactly one GD step. Expand a in powers of \eta and compare with f_2 and f_3 from 1.2.b.
2. For which \eta\lambda does the series converge? What happens to the exact modified equation when 1<\eta\lambda<2, where GD still converges? Can any real one-dimensional ODE \dot x=F(x) reproduce GD there? In d dimensions, can a real gradient flow?
3. Return to the instability of Exercise 1.1, item 2: a minimum with \eta S^*>2 is locally unstable for GD. Why can no truncation of the modified equation see this?

Code (19 lines)

# Decay rate per unit time of one quadratic direction, in units of lambda, as a function of
# z = eta*lambda: gradient flow (1), the IGR flow (1 + z/2), the next truncation (1 + z/2 + z^2/3),
# and GD's actual rate -log|1 - z| / z (negative = growth). The series for GD's rate converges
# only for z < 1, and GD itself diverges for z > 2.
z = np.linspace(1e-3, 2.4, 1000)
gd_rate = -np.log(np.abs(1 - z)) / z
fig, ax = plt.subplots(figsize=(6.4, 3.6))
ax.plot(z, gd_rate, "k", lw=2, label=r"GD: $-\log|1-\eta\lambda|/(\eta\lambda)$")
ax.plot(z, np.ones_like(z), label="gradient flow")
ax.plot(z, 1 + z / 2, label=r"IGR flow: $1+\eta\lambda/2$")
ax.plot(z, 1 + z / 2 + z**2 / 3, "--", label=r"next order: $+(\eta\lambda)^2/3$")
for x0, txt in [(1, "series radius"), (2, "GD unstable")]:
    ax.axvline(x0, color="gray", ls=":")
    ax.text(x0 + 0.03, -0.8, txt, color="gray", fontsize=8)
ax.set_ylim(-1, 4); ax.set_xlabel(r"$\eta\lambda$"); ax.set_ylabel(r"decay rate / $\lambda$")
ax.legend(fontsize=8, loc="upper right"); plt.tight_layout(); plt.show()
for zz in [0.1, 0.3, 0.6, 0.9]:
    exact = -np.log(1 - zz) / zz
    print(f"eta*lam={zz:.1f}: GD rate {exact:.4f}, IGR {1 + zz/2:.4f} (rel. err {abs(1 + zz/2 - exact)/exact:.1%})")

eta*lam=0.1: GD rate 1.0536, IGR 1.0500 (rel. err 0.3%)
eta*lam=0.3: GD rate 1.1889, IGR 1.1500 (rel. err 3.3%)
eta*lam=0.6: GD rate 1.5272, IGR 1.3000 (rel. err 14.9%)
eta*lam=0.9: GD rate 2.5584, IGR 1.4500 (rel. err 43.3%)

Exercise 1.2.d — IGR picks a flatter minimum.

Use the two-parameter model L(a,b)=\tfrac12(ab-1)^2, a scalar two-layer linear “network” f=ab fitting the target 1. Its minima form the hyperbola ab=1. Write r=ab-1 and D=a^2-b^2 (the imbalance).

1. Show that gradient flow conserves D.
2. Show that one GD step maps D\mapsto(1-\eta^2r^2)\,D exactly.
3. Compute \|\nabla L\|^2 and show that the IGR flow gives \dot D=-\eta r^2D. Check that this matches item 2 at leading order.
4. At a minimum on ab=1, show the Hessian has eigenvalues 0 and a^2+b^2=\sqrt{D^2+4}. Conclude that GD (and the IGR flow) end at a flatter minimum than gradient flow from the same start, and say how the effect scales with \eta.

Code (53 lines)

# GD vs gradient flow vs the IGR flow on L = (ab-1)^2/2, from (a, b) = (2, 0.1), where S ~ 4.1-4.5.
p0, T = np.array([2.0, 0.1]), 2.0
F_gf = lambda p: -grad_ab(p)
F_igr = lambda eta: (lambda p: -(grad_ab(p) + 0.5 * eta * hess_ab(p) @ grad_ab(p)))


def gd(p, eta, K):
    traj = [p]
    for _ in range(K):
        p = p - eta * grad_ab(p); traj.append(p)
    return np.array(traj)


# (1) global error at fixed time T against eta: expect slope 1 (GF) and slope 2 (IGR flow)
etas = T / np.array([20, 40, 80, 160, 320, 640])
err_gf, err_igr = [], []
for eta in etas:
    pK = gd(p0, eta, int(round(T / eta)))[-1]
    err_gf.append(np.linalg.norm(pK - rk4(F_gf, p0, T)))
    err_igr.append(np.linalg.norm(pK - rk4(F_igr(eta), p0, T)))
print(f"log-log slope of |GD - flow| at t = {T}: gradient flow {slope(etas, err_gf):.2f}, "
      f"IGR flow {slope(etas, err_igr):.2f}")

# (2) where the long-time runs end up, and how much of GD's shift IGR accounts for
T_long = 40.0
p_gf = rk4(F_gf, p0, T_long, 8000)
fig, axs = plt.subplots(1, 3, figsize=(13, 3.8))
S_of = lambda p: np.linalg.eigvalsh(hess_ab(p))[-1]
t_gf = np.linspace(0, 8, 80)
axs[0].plot(t_gf, [S_of(rk4(F_gf, p0, t, 400)) for t in t_gf], "k", label="gradient flow")
for eta, col in [(0.05, "C0"), (0.2, "C1"), (0.35, "C3")]:
    tr = gd(p0, eta, int(T_long / eta))
    p_igr = rk4(F_igr(eta), p0, T_long, 8000)
    shift = np.linalg.norm(tr[-1] - p_gf)
    captured = np.dot(p_igr - p_gf, tr[-1] - p_gf) / shift**2
    S_end = np.linalg.eigvalsh(hess_ab(tr[-1]))[-1]
    print(f"eta={eta:<5} eta*S_init={eta * np.linalg.eigvalsh(hess_ab(p0))[-1]:.2f}  "
          f"final S: GF {np.linalg.eigvalsh(hess_ab(p_gf))[-1]:.3f}, GD {S_end:.3f};  "
          f"IGR flow captures {captured:.0%} of GD's shift away from the GF endpoint")
    ts = np.arange(len(tr)) * eta
    keep = ts <= 8
    axs[0].plot(ts[keep], [S_of(p) for p in tr[keep]], color=col, label=f"GD, eta={eta}")
    axs[2].plot(ts, tr[:, 0]**2 - tr[:, 1]**2, color=col, label=f"GD, eta={eta}")
axs[0].set_xlabel("t = k eta"); axs[0].set_ylabel("sharpness S")
axs[0].legend(fontsize=7); axs[0].set_title("sharpness along the run")
axs[1].loglog(etas, err_gf, "o-", label="GD vs gradient flow")
axs[1].loglog(etas, err_igr, "s-", label="GD vs IGR flow")
axs[1].set_xlabel(r"$\eta$"); axs[1].set_ylabel(f"error at t={T}"); axs[1].legend(fontsize=8)
axs[1].set_title("global error vs step size")
axs[2].axhline(p0[0]**2 - p0[1]**2, color="k", label="gradient flow (conserved)")
axs[2].set_xlabel("t = k eta"); axs[2].set_ylabel(r"imbalance $a^2-b^2$"); axs[2].legend(fontsize=7)
axs[2].set_title("GD shrinks the imbalance")
plt.tight_layout(); plt.show()

log-log slope of |GD - flow| at t = 2.0: gradient flow 1.06, IGR flow 2.01
eta=0.05  eta*S_init=0.20  final S: GF 4.463, GD 4.448;  IGR flow captures 81% of GD's shift away from the GF endpoint
eta=0.2   eta*S_init=0.82  final S: GF 4.463, GD 4.369;  IGR flow captures 42% of GD's shift away from the GF endpoint
eta=0.35  eta*S_init=1.43  final S: GF 4.463, GD 4.125;  IGR flow captures 16% of GD's shift away from the GF endpoint

The slopes come out at 1 and 2: the IGR flow tracks GD one order better than gradient flow. The second printout is more sobering. The IGR flow accounts for most of GD’s displacement from the gradient-flow endpoint at \eta S\approx0.2, less than half at \eta S\approx0.8, and a small fraction at \eta S\approx1.4. The first-order correction gets the direction of the effect right (toward balance and flatness) well before it gets the size right. §1.3 asks why.

Only as an asymptotic series. The modified equation of §1.2 is a power series in the step size \eta that usually diverges; truncations come with rigorous error bounds, controlled roughly by \eta S (step size times sharpness, the largest Hessian eigenvalue). The kept first correction degrades as \eta S approaches 1, and no truncation sees the GD instability at \eta S=2, which is where full-batch training of neural networks typically sits.

Papers: Hairer & Lubich (2000) on asymptotic expansions; Cohen et al. (2021) on the edge of stability (see Readings).

Derivations and walkthrough. Notation: GD with step \eta on L, gradient g, Hessian H, sharpness S. A one-step method y\mapsto y+hf(y)+O(h^2) has modified field \tilde f=f+hf_2+h^2f_3+\cdots (§1.2) and truncation \tilde f_{[r]}=f+\dots+h^{r-1}f_r. For GD, h=\eta, f=-g, f_2=-\tfrac12Hg.

It usually diverges. Even for analytic f the series typically has zero radius of convergence; the quadratic of Exercise 1.2.c is the exception. For GD on L=y^3/3 (Exercise 1.3.a and the code below) the coefficients grow like (k-1)!/(2\pi)^k over 70 computed orders.

What a truncation guarantees.

1. Fixed order: local error O(h^{r+1}), global error O(h^r) over a fixed time, with constants involving derivatives of f up to order about r+1. The global constant can grow like e^{\ell T}, with \ell the rate at which nearby solutions separate: errors do not pile up near a minimum with positive definite Hessian, but do near saddles.
2. Optimal truncation: for analytic f, truncating at order N\sim c/h makes the one-step error O(e^{-\gamma/h}) for some \gamma>0 (Benettin–Giorgilli 1994, Hairer–Lubich 1997, Reich 1999; HLW §IX.7). Roughly, \gamma is the radius of a complex neighbourhood of the trajectory on which f is bounded, divided by that bound.

What sets the scale. On a quadratic the parameter is exactly \eta\lambda per eigendirection, so \eta S near a minimum. In general higher derivatives enter too, and small \eta S is not sufficient: L(x)=x+K^{-2}\sin(Kx) has |L''|\le1, yet at x=0 the formulas of Exercise 1.2.b give f_2=0 and f_3=\tfrac{K}{12}(1+\tfrac1K)^2. The real comparison is the step length \eta\|g\| against the distance over which the Hessian changes (here 1/K).

How it fails, most important first.

1. The kept correction stops being enough. Its relative error is O((\eta S)^2): the IGR flow accounts for 81%, 42% and 16% of GD’s shift at \eta S\approx0.2, 0.8, 1.4 (code after Exercise 1.2.d).
2. No gradient flow reproduces the step once a quadratic multiplier 1-\eta\lambda is \le0 (Exercise 1.2.c). For y^3/3 the GD map stops being invertible where \Phi_h'(y)=1-2hy=1-\eta S=0, and the truncation errors in the code below deteriorate as \eta S\to1.
3. Divergence. Harmless in practice: optimally truncated, the series is extremely accurate when \eta S is small (least term e^{-4\pi/(\eta S)} for y^3/3, Exercise 1.3.b), and nobody uses high orders anyway.

Where training sits. Beyond \eta S=2 a minimum is unstable for GD (Exercise 1.1), a threshold invisible to every truncation (Exercise 1.2.c, item 3). Full-batch training typically sits at \eta S\approx2 (Cohen et al.), where the expansion parameter is about 2 in the sharpest direction but still small for a much flatter bulk; td2_3_edge_of_stability treats the oscillating directions separately. Under minibatch SGD the full-batch sharpness settles below 2/\eta, while a batch-level curvature, \mathbb E_B\big[g_B^\top H_Bg_B/\|g_B\|^2\big] with g_B, H_B the gradient and Hessian of a random minibatch’s loss, hovers near 2/\eta (Andreyev and Beneventano, v5).

Exercise 1.3.a — A divergent modified equation.

Take GD on L(y)=y^3/3 for y>0, i.e. explicit Euler for \dot y=-y^2: \Phi_h(y)=y-hy^2.

1. Show that the exact modified field F (if it exists) satisfies F(\Phi_h(y))=\Phi_h'(y)\,F(y).
2. Write F(y)=\sum_{k\ge2}a_kh^{k-2}y^k with a_2=-1, and set u=hy, so F=h^{-2}\varphi(u) with \varphi(u)=\sum_ka_ku^k. Turn item 1 into an equation for \varphi, then into a recursion for the a_k. Compute a_3,a_4,a_5, and check a_3 against f_2=-\tfrac12f'f.

Exercise 1.3.b — The least term.

F(y)=\sum_{k\ge2}a_kh^{k-2}y^k is the modified field of GD with step h on L(y)=y^3/3, with a_k from Exercise 1.3.a. Suppose |a_k| grows like (k-1)!/(2\pi)^k up to factors varying slower than exponentially in k (the next code cell supports this over k\le70; a fit, not a theorem). At a point y the k-th term of F has size about y^2|a_k|u^{k-2}, with u=hy.

1. Find the order k^* of the smallest term and its size, to exponential accuracy.
2. Translate into the sharpness S=L''(y)=2y. How small is the least term at \eta S=0.1, 0.5, 0.8? Why should the estimate not be trusted as \eta S\to1?

Code (58 lines)

# Modified-equation coefficients for GD on L = y^3/3 (Euler for ydot = -y^2), from the
# recursion of Exercise 1.3.a in exact rational arithmetic.
N = 70
a = {2: Fraction(-1)}
for n in range(4, N + 2):
    s = sum(a[k] * comb(k, n - k) * (-1) ** (n - k) for k in a if k <= n - 2 and n - k <= k)
    a[n - 1] = s / (n - 3)
print("a_2..a_7 =", [str(a[k]) for k in range(2, 8)])
growth = {k: (np.log(abs(float(a[k]))) - np.sum(np.log(np.arange(1, k)))) / k for k in (30, 50, 70)}
print("(log|a_k| - log (k-1)!)/k at k = 30, 50, 70:", [round(float(v), 3) for v in growth.values()],
      f"-> drifting toward -log(2 pi) = {-np.log(2*np.pi):.3f} (a numerical fit)")

# One-step error of the truncated modified flow vs the GD step, as a function of the order n.
# A 1D flow moving down from y0 reaches y after time T(y) = int_y^y0 dz / (-F_n(z)); we find the
# endpoint T(y) = h by bisection (Gauss-Legendre for T). Orders whose field points the wrong way
# at y0 never reach the GD point at all; they are left blank (nan). The flow cannot cross an
# equilibrium of F_n, so the search stops above the first one.
af = np.array([float(a[k]) for k in range(2, N + 1)])
orders, pw = np.arange(1, 50), np.arange(N + 1)
xg, wg = np.polynomial.legendre.leggauss(120)


def truncation_errors(h, y0=1.0):
    y1 = y0 - h * y0**2                              # one GD step
    grid = np.linspace(y0, y1 - 0.5 * h, 4001)       # search window, descending from y0
    out = np.full(len(orders), np.nan)
    for i, n in enumerate(orders):
        c = np.zeros(N + 1)
        c[2:n + 2] = af[:n] * h ** np.arange(n)      # F_n(y) = sum_k a_k h^(k-2) y^k
        F = lambda y: (np.asarray(y)[..., None] ** pw) @ c
        bad = np.nonzero(F(grid) >= 0)[0]
        if len(bad) and bad[0] == 0:
            continue                                 # wrong direction at y0
        lo = grid[bad[0] - 1] if len(bad) else grid[-1]
        T = lambda y: 0.5 * (y0 - y) * np.sum(wg / -F(0.5 * (y0 + y) + 0.5 * (y0 - y) * xg))
        if T(lo) < h:
            continue
        top, bot = y0, lo
        for _ in range(60):
            mid = 0.5 * (top + bot)
            top, bot = (top, mid) if T(mid) > h else (mid, bot)
        out[i] = abs(0.5 * (top + bot) - y1)
    return out


fig, axs = plt.subplots(1, 2, figsize=(11, 3.6))
axs[0].semilogy(np.arange(2, N + 1), np.abs(af), label=r"$|a_k|$")
axs[0].semilogy(np.arange(2, N + 1), [np.exp(np.sum(np.log(np.arange(1, k))) - k * np.log(2 * np.pi))
                                       for k in range(2, N + 1)], "--", label=r"$(k-1)!/(2\pi)^k$")
axs[0].set_xlabel("k"); axs[0].set_title("modified-equation coefficients"); axs[0].legend()
for h in [0.25, 0.3, 0.35, 0.4, 0.45]:
    e = truncation_errors(h)
    axs[1].semilogy(orders, e, label=f"eta*S = {2*h:.1f}")
    print(f"eta*S={2*h:.1f}: best one-step error {np.nanmin(e):.1e} at n={orders[np.nanargmin(e)]};"
          f"  exp(-2 pi/h) = {np.exp(-2*np.pi/h):.1e};  {np.isnan(e).sum()} orders never reach the GD point")
axs[1].set_xlabel("truncation order n"); axs[1].set_ylabel("one-step error")
axs[1].set_title("truncated modified flow vs one GD step"); axs[1].legend(fontsize=7)
plt.tight_layout(); plt.show()

a_2..a_7 = ['-1', '-1', '-3/2', '-8/3', '-31/6', '-157/15']
(log|a_k| - log (k-1)!)/k at k = 30, 50, 70: [-1.684, -1.767, -1.798] -> drifting toward -log(2 pi) = -1.838 (a numerical fit)
eta*S=0.5: best one-step error 1.4e-11 at n=31;  exp(-2 pi/h) = 1.2e-11;  0 orders never reach the GD point
eta*S=0.6: best one-step error 3.3e-09 at n=27;  exp(-2 pi/h) = 8.0e-10;  0 orders never reach the GD point
eta*S=0.7: best one-step error 1.1e-08 at n=27;  exp(-2 pi/h) = 1.6e-08;  2 orders never reach the GD point
eta*S=0.8: best one-step error 6.5e-07 at n=28;  exp(-2 pi/h) = 1.5e-07;  7 orders never reach the GD point
eta*S=0.9: best one-step error 1.0e-04 at n=28;  exp(-2 pi/h) = 8.6e-07;  9 orders never reach the GD point

The error first falls with the truncation order, then stops improving and grows; gaps are orders whose truncated field points the wrong way. The step starts at y_0=1, where S=2, so the least-term estimate of Exercise 1.3.b is e^{-2\pi/h}. Up to \eta S=0.8 the best error lands within a factor of a few of it: consistent with the exponent, not a confirmation of the prefactor (the left panel shows how far the coefficients sit from the fitted curve). At \eta S=0.9 the best error is far above the estimate, as the step approaches the point where it is no longer invertible.

SGD with randomly shuffled minibatches, averaged over the order of the batches, follows to second order in the step size \eta a gradient flow on

\mathbb E\,\tilde L_{\rm SGD}=L+\frac\eta4\|\nabla L\|^2+\frac{N-B}{N-1}\,\frac{\eta}{4B}\,\mathrm{tr}\,\Sigma ,

with N examples, batch size B, and \Sigma the covariance of per-example gradients. The first correction is GD’s implicit gradient regularization (§1.2) and depends on \eta alone; the second penalizes the spread of per-example gradients and scales like \eta/B.

Paper: Smith, Dherin, Barrett & De, On the origin of implicit regularization in SGD (see Readings).

Core idea. To first order in \eta, one epoch of SGD equals one GD step of size m\eta (m=N/B batches): with each example used once per epoch, the gradient noise cancels at O(\eta) and first shows up at O(\eta^2), as the bias above. The result is about the mean iterate over batch orders, and it needs the epoch length m\eta times the minibatch curvatures to be small, a much stronger condition than for GD.

Derivations and walkthrough. Setup: N examples with losses \ell_i, L=\frac1N\sum_i\ell_i, split into m=N/B minibatches of size B with losses \hat L_k=\frac1B\sum_{i\in\text{batch }k}\ell_i, so L=\frac1m\sum_k\hat L_k. One epoch of SGD with random shuffling visits each batch once, in a random order \pi: \theta_{j+1}=\theta_j-\eta\nabla\hat L_{\pi(j)}(\theta_j). Per-example gradient covariance: \Sigma=\frac1N\sum_i(\nabla\ell_i-\nabla L)(\nabla\ell_i-\nabla L)^\top.

Smith et al. apply backward error analysis to a whole epoch. Averaged over batch orders (batches held fixed), the iterate after one epoch follows, up to O(m^3\eta^3), gradient flow for time m\eta on

\tilde L_{\rm SGD}=L+\frac{\eta}{4m}\sum_{k=0}^{m-1}\|\nabla\hat L_k\|^2 =L+\frac\eta4\|\nabla L\|^2+\frac{\eta}{4m}\sum_k\|\nabla\hat L_k-\nabla L\|^2 .

Averaging also over the random split into batches gives \mathbb E\,\tilde L_{\rm SGD} above (Exercise 1.4.b). Individual epochs differ from the mean at order m^2\eta^2, through the order-dependent term in Exercise 1.4.a.

Exercise 1.4.a — One epoch.

Write \hat g_k=\nabla\hat L_k(\theta_0) and \hat H_k=\nabla^2\hat L_k(\theta_0) for the gradient and Hessian of batch k at the start of the epoch, so \sum_k\hat g_k=mg and \sum_k\hat H_k=mH, where g and H are the full-batch gradient and Hessian at \theta_0. Here k labels batches and j labels steps within the epoch.

1. For a fixed order (batches 0,1,\dots,m-1), show \theta_m=\theta_0-\eta\sum_k\hat g_k+\eta^2\sum_{j}\sum_{k<j}\hat H_j\hat g_k+O(m^3\eta^3).
2. Average over all m! orders. Show \mathbb E\,\theta_m=\theta_0-m\eta g+\tfrac{\eta^2}2\big(m^2Hg-\sum_k\hat H_k\hat g_k\big)+O(m^3\eta^3).
3. Match against the flow \dot\theta=-g+\eta f_2 run for time m\eta (the §1.2 notation: f_2 is the O(\eta) correction), and conclude \tilde L_{\rm SGD}=L+\frac{\eta}{4m}\sum_k\|\nabla\hat L_k\|^2. Then derive the second form of \tilde L_{\rm SGD}.

Exercise 1.4.b — The \eta/B term.

1. A batch is a uniformly random subset of size B, drawn without replacement. Show \mathbb E\|\nabla\hat L-\nabla L\|^2=\frac{N-B}{N-1}\frac{\mathrm{tr}\,\Sigma}{B}, and conclude the formula for \mathbb E\,\tilde L_{\rm SGD}.
2. At fixed \eta, what does halving B do to each correction term? Which term does the linear scaling rule (\eta\propto B) hold fixed, and which does it change?

Exercise 1.4.c — Check it numerically (code).

The cell below uses N=8 examples with loss \ell_i=\log\cosh(x_i^\top w-y_i), w\in\mathbb R^2, and B=2 (m=4 batches, 24 orders). Complete mean_epoch (the end-of-epoch iterate averaged exactly over all orders) and the "sgd" branch of flow_end (the field -\nabla\tilde L_{\rm SGD}). The cell compares the mean iterate with gradient flow, the IGR flow and the flow on \tilde L_{\rm SGD}, each run for time m\eta. This one-epoch time shrinks with \eta, so each error is one power of \eta higher than in the fixed-time §1.2 cell. Why does the IGR flow not improve the order of the error, while the flow on \tilde L_{\rm SGD} does?

rng = np.random.default_rng(0)
Nex, B = 8, 2
m = Nex // B
X, yv = rng.normal(size=(Nex, 2)), rng.normal(size=Nex)
batches = [np.arange(k * B, (k + 1) * B) for k in range(m)]
everyone = np.arange(Nex)


def g_set(w, idx):        # gradient of the mean of log cosh(x_i.w - y_i) over examples idx
    r = X[idx] @ w - yv[idx]
    return (np.tanh(r)[:, None] * X[idx]).mean(0)


def H_set(w, idx):        # Hessian of the same
    r = X[idx] @ w - yv[idx]
    s = 1 / np.cosh(r) ** 2
    return (s[:, None, None] * X[idx][:, :, None] * X[idx][:, None, :]).mean(0)


def mean_epoch(w0, eta):
    # TODO: average over all m! orders (itertools.permutations) of one epoch of SGD:
    # for each order, take m steps w <- w - eta * g_set(w, batches[k]); return the mean final w.
    raise NotImplementedError


def flow_end(w0, eta, kind):
    def F(w):
        g = g_set(w, everyone)
        if kind == "gf":
            return -g
        if kind == "igr":
            return -(g + 0.5 * eta * H_set(w, everyone) @ g)
        # TODO (kind == "sgd"): return -grad of L_SGD = L + (eta / 4m) sum_k |grad L_k|^2,
        # using g_set(w, b) and H_set(w, b) for each batch b.
        raise NotImplementedError
    return rk4(F, w0, m * eta, 400)


w0 = np.array([1.5, -1.0])
etas_sgd = np.array([0.1, 0.05, 0.025, 0.0125])
# errs = {kind: [np.linalg.norm(mean_epoch(w0, eta) - flow_end(w0, eta, kind)) for eta in etas_sgd]
#         for kind in ["gf", "igr", "sgd"]}
# for kind, e in errs.items(): print(kind, f"slope {slope(etas_sgd, e):.2f}")

rng = np.random.default_rng(0)
Nex, B = 8, 2
m = Nex // B
X, yv = rng.normal(size=(Nex, 2)), rng.normal(size=Nex)
batches = [np.arange(k * B, (k + 1) * B) for k in range(m)]
everyone = np.arange(Nex)


def g_set(w, idx):        # gradient of the mean of log cosh(x_i.w - y_i) over examples idx
    r = X[idx] @ w - yv[idx]
    return (np.tanh(r)[:, None] * X[idx]).mean(0)


def H_set(w, idx):        # Hessian of the same
    r = X[idx] @ w - yv[idx]
    s = 1 / np.cosh(r) ** 2
    return (s[:, None, None] * X[idx][:, :, None] * X[idx][:, None, :]).mean(0)


def mean_epoch(w0, eta):
    finals = []
    for order in permutations(range(m)):
        w = w0.copy()
        for k in order:
            w = w - eta * g_set(w, batches[k])
        finals.append(w)
    return np.mean(finals, axis=0)


def flow_end(w0, eta, kind):
    def F(w):
        g = g_set(w, everyone)
        if kind == "gf":
            return -g
        if kind == "igr":
            return -(g + 0.5 * eta * H_set(w, everyone) @ g)
        # grad of (eta/4m) sum_k |g_k|^2 is (eta/2m) sum_k H_k g_k
        return -(g + eta / (2 * m) * sum(H_set(w, b) @ g_set(w, b) for b in batches))
    return rk4(F, w0, m * eta, 400)


w0 = np.array([1.5, -1.0])
etas_sgd = np.array([0.1, 0.05, 0.025, 0.0125])
errs = {kind: [np.linalg.norm(mean_epoch(w0, eta) - flow_end(w0, eta, kind)) for eta in etas_sgd]
        for kind in ["gf", "igr", "sgd"]}
names = {"gf": "gradient flow on L", "igr": "IGR flow (GD's modified loss)", "sgd": "flow on L_SGD"}
for kind, e in errs.items():
    print(f"{names[kind]:<32} slope {slope(etas_sgd, e):.2f}   errors {np.array2string(np.array(e), precision=1)}")

gradient flow on L               slope 2.01   errors [2.4e-03 5.7e-04 1.4e-04 3.6e-05]
IGR flow (GD's modified loss)    slope 1.85   errors [2.3e-03 7.2e-04 1.9e-04 5.0e-05]
flow on L_SGD                    slope 2.98   errors [3.5e-04 4.5e-05 5.7e-06 7.2e-07]

Heavy-ball momentum, \theta_{k+1}=\theta_k-\eta\nabla L(\theta_k)+\beta(\theta_k-\theta_{k-1}) with \beta\in[0,1), does three things relative to GD with step \eta. On a quadratic it stays stable up to \eta S<2(1+\beta) (S the sharpness). For small steps it follows gradient flow sped up by 1/(1-\beta): an effective learning rate \eta/(1-\beta). And at matched effective learning rate its implicit gradient penalty is (1+\beta)/(1-\beta) times GD’s (19 times at \beta=0.9), though only for curvatures far smaller than those where GD’s applies.

Papers: Goh (Distill, 2017); Ghosh, Lyu, Zhang & Wang (2023) (see Readings).

Derivations and walkthrough. Notation: loss L, gradient g=\nabla L, Hessian H, sharpness S=\lambda_{\max}(H), step \eta, momentum \beta\in[0,1), time t=k\eta. Recall from §1.1–1.2: for GD (\beta=0) a quadratic direction of curvature \lambda is stable iff 0<\eta\lambda<2, and to first order GD follows gradient flow on L+\tfrac\eta4\|\nabla L\|^2 (IGR).

Heavy ball (Polyak 1964) keeps a velocity v:

v_{k+1}=\beta v_k-\eta\,g(\theta_k),\qquad \theta_{k+1}=\theta_k+v_{k+1}.

PyTorch’s SGD with momentum (b\leftarrow\beta b+g, \theta\leftarrow\theta-\eta b) and Goh’s z\leftarrow\beta z+\nabla f, w\leftarrow w-\alpha z are the same method (v=-\eta b, \alpha=\eta); Adam-style averaging, m\leftarrow\beta m+(1-\beta)g, is the same with \eta\to\eta(1-\beta).

On a quadratic (Exercise 1.5.a and the code after it). The characteristic roots of \theta_{k+1}=(1+\beta-\eta\lambda)\theta_k-\beta\theta_{k-1} give the stability region 0<\eta\lambda<2(1+\beta), a spiral-in band with rate \sqrt\beta, and at critical damping a rate 1-\sqrt{\eta\lambda} instead of GD’s 1-\eta\lambda: that square root is the acceleration (Goh derives the optimal tuning).

The flow limit (Exercise 1.5.b). A slowly varying velocity gives v\approx-\eta g/(1-\beta), i.e. gradient flow sped up by 1/(1-\beta); at matched effective learning rate momentum tolerates (1+\beta)/(1-\beta) times sharper directions. The next Taylor term gives a damped particle with mass O(\eta) (Kovachki and Stuart), a good picture of the start-up transient but globally only O(\eta) accurate.

The slow manifold (Exercise 1.5.c). Heavy ball maps pairs (\theta,v), so §1.2’s recipe does not apply directly. But a mismatch in v decays like \beta^k, and after a few times 1/(1-\beta) steps the velocity is slaved to the position. On this attracting invariant “slow manifold” (Kovachki and Stuart prove it exists) the step is a function of the current point, \theta_{k+1}-\theta_k=D(\theta_k), and §1.2 applies: a map \theta\mapsto\theta+hf+h^2d_2+O(h^3) is followed, to O(h^2) on fixed time intervals, by \dot\theta=f+hf_2 with f_2=d_2-\tfrac12f'f (f' the Jacobian). The result (Ghosh et al.) is

\dot\theta=-\frac1{1-\beta}\nabla\Big(L+\frac{\eta(1+\beta)}{4(1-\beta)^2}\|\nabla L\|^2\Big),

i.e. gradient flow on L+\frac{\eta_{\rm eff}}4\cdot\frac{1+\beta}{1-\beta}\|\nabla L\|^2 in units of the effective step \eta_{\rm eff}=\eta/(1-\beta).

Range of validity. The velocity must relax much faster than the position drifts. On a direction of curvature \lambda the slow root z\approx1-\eta\lambda/(1-\beta) and the fast root near \beta meet at \eta\lambda=(1-\sqrt\beta)^2, so the expansion needs \eta\lambda\ll(1-\sqrt\beta)^2, i.e. \eta_{\rm eff}\lambda\ll(1-\sqrt\beta)/(1+\sqrt\beta)\approx0.026 at \beta=0.9, against \eta\lambda\ll1 for GD. The 19-fold penalty applies to curvatures about 40 times smaller than GD’s IGR does. Starting from rest (v_0=0) the velocity’s initial deficit, \eta g/(1-\beta), decays like \beta^k but leaves an O(\eta) displacement of about \eta\beta g/(1-\beta)^2, so the formula is then only O(\eta) accurate; Ghosh et al.’s theorem, which starts from rest, recovers O(\eta^2) by keeping explicit \beta^k terms. The last code cell below shows both cases.

Exercise 1.5.a — Roots and regimes.

1. Derive the characteristic polynomial and show the stability region is 0<\eta\lambda<2(1+\beta) for 0\le\beta<1.
2. Show that the roots are complex exactly for (1-\sqrt\beta)^2<\eta\lambda<(1+\sqrt\beta)^2, with |z|=\sqrt\beta there.
3. For small \eta\lambda, find the root closest to 1 to first order in \eta\lambda, and read off the effective learning rate. For 0<\eta\lambda\le1, at critical damping \beta=(1-\sqrt{\eta\lambda})^2, find the double root.

Code (30 lines)

# Spectral radius max|z| of heavy ball on a quadratic over the (eta*lambda, beta) plane.
el = np.linspace(0.0, 4.4, 441)            # eta * lambda
be = np.linspace(0.0, 0.99, 199)           # beta
EL, BE = np.meshgrid(el, be)
tau = 1 + BE - EL
disc = np.sqrt((tau**2 - 4 * BE).astype(complex))
rho = np.maximum(np.abs((tau + disc) / 2), np.abs((tau - disc) / 2))

fig, ax = plt.subplots(figsize=(7, 4.2))
im = ax.pcolormesh(EL, BE, np.ma.masked_where(rho >= 1, rho), shading="auto", cmap="viridis", vmin=0, vmax=1)
fig.colorbar(im, ax=ax, label=r"spectral radius $\max|z|$ (blank: unstable)")
ax.plot(2 * (1 + be), be, "r--", lw=1.2, label=r"$\eta\lambda=2(1+\beta)$")
ax.plot((1 - np.sqrt(be))**2, be, "w:", lw=1.2, label=r"$(1\mp\sqrt{\beta})^2$: complex band")
ax.plot((1 + np.sqrt(be))**2, be, "w:", lw=1.2)
ax.set_xlim(0, 4.4); ax.set_xlabel(r"$\eta\lambda$"); ax.set_ylabel(r"$\beta$")
ax.set_title("heavy ball on a quadratic: rate where stable")
ax.legend(fontsize=8, loc="upper left"); plt.tight_layout(); plt.show()

band = (EL > (1 - np.sqrt(BE))**2 + 1e-9) & (EL < (1 + np.sqrt(BE))**2 - 1e-9)
print("max |rho - sqrt(beta)| in the complex band:", np.abs(rho - np.sqrt(BE))[band].max())
stable = rho < 1
pred = (EL > 0) & (EL < 2 * (1 + BE))
print("grid points where the stability region disagrees with 0 < eta*lam < 2(1+beta):",
      int((stable != pred)[(np.abs(EL - 2 * (1 + BE)) > 0.02) & (EL > 0.02)].sum()),
      "(excluding a thin strip around the boundary, where rounding decides)")
# Critical damping: at beta = (1 - sqrt(eta*lam))^2 the rate is 1 - sqrt(eta*lam), vs GD's 1 - eta*lam.
for e in [0.01, 0.1]:
    b = (1 - np.sqrt(e))**2
    zs = np.roots([1, -(1 + b - e), b])
    print(f"eta*lam={e}: critical beta={b:.3f}, rate {np.abs(zs).max():.4f}; GD rate {1 - e:.4f}")

max |rho - sqrt(beta)| in the complex band: 2.220446049250313e-16
grid points where the stability region disagrees with 0 < eta*lam < 2(1+beta): 0 (excluding a thin strip around the boundary, where rounding decides)
eta*lam=0.01: critical beta=0.810, rate 0.9000; GD rate 0.9900
eta*lam=0.1: critical beta=0.468, rate 0.6838; GD rate 0.9000

Exercise 1.5.b — The flow limit.

1. Show that if the iterates follow a smooth curve \theta(t), t=k\eta, then to first order \dot\theta=-\nabla L/(1-\beta).
2. Keep the next order of the Taylor expansion and derive \tfrac{\eta(1+\beta)}2\ddot\theta+(1-\beta)\dot\theta=-\nabla L+O(\eta^2). What are the mass, the friction, and the velocity relaxation time in steps?
3. (Optional; Exercise 1.5.c is the cleaner route.) After the transient, the inertia term is small and can be treated perturbatively. Solve the ODE of item 2 for \dot\theta to first order in \eta (substitute the leading-order \ddot\theta) and compare with the modified flow quoted above.

Exercise 1.5.c — The heavy-ball modified flow.

After transients, assume the step is a function of the current point: \theta_{k+1}-\theta_k=D(\theta_k) with D=\eta D_1+\eta^2D_2+O(\eta^3) (an invariant “slow manifold” in (\theta,v) space).

1. Use the recursion to find D_1=-g/(1-\beta) and D_2=-\beta Hg/(1-\beta)^3.
2. Apply the recipe of §1.2 to the one-step map \theta\mapsto\theta+D(\theta) to obtain the modified flow above.
3. Compare with GD at the same effective learning rate.

Code (40 lines)

# Heavy ball vs its modified flow on L = (ab-1)^2/2 at beta = 0.5, error at a fixed time t = k*eta.
# Starting on the slow manifold, the rescaled gradient flow is O(eta) off and Ghosh et al.'s modified
# flow O(eta^2). Starting from rest (v0 = 0), the start-up transient shifts the trajectory by O(eta),
# so the modified flow is again only O(eta) accurate at this short time.
beta, T_hb = 0.5, 0.5
p0 = np.array([2.0, 0.1])


def heavy_ball(p, eta, K, v):
    for _ in range(K):
        v = beta * v - eta * grad_ab(p)
        p = p + v
    return p


def hb_flow(eta):
    c = eta * (1 + beta) / (2 * (1 - beta)**2)   # grad of eta(1+b)/(4(1-b)^2)|g|^2 is c * H g
    return lambda p: -(grad_ab(p) + c * hess_ab(p) @ grad_ab(p)) / (1 - beta)


etas_hb = T_hb / np.array([25, 50, 100, 200, 400, 800])
res = {"rescaled GF, slow-manifold start": [], "modified flow, slow-manifold start": [],
       "modified flow, start from rest": []}
for eta in etas_hb:
    K = int(round(T_hb / eta))
    g0, H0 = grad_ab(p0), hess_ab(p0)
    D1 = -g0 / (1 - beta)
    D2 = -beta * H0 @ g0 / (1 - beta)**3
    v_slow = eta * D1 + eta**2 * (D2 - (-H0 / (1 - beta)) @ D1)   # previous step D(theta_{-1}) to O(eta^2)
    p_slow = heavy_ball(p0, eta, K, v_slow)
    p_rest = heavy_ball(p0, eta, K, np.zeros(2))
    res["rescaled GF, slow-manifold start"].append(np.linalg.norm(p_slow - rk4(hb_flow(0.0), p0, T_hb, 2000)))
    res["modified flow, slow-manifold start"].append(np.linalg.norm(p_slow - rk4(hb_flow(eta), p0, T_hb, 2000)))
    res["modified flow, start from rest"].append(np.linalg.norm(p_rest - rk4(hb_flow(eta), p0, T_hb, 2000)))
fig, ax = plt.subplots(figsize=(5.6, 3.6))
for name, e in res.items():
    ax.loglog(etas_hb, e, "o-", label=name)
    print(f"{name:<38} slope (smallest 4 eta) {slope(etas_hb[2:], e[2:]):.2f}")
ax.set_xlabel(r"$\eta$"); ax.set_ylabel(f"|heavy ball - flow| at t={T_hb}"); ax.legend(fontsize=7)
plt.tight_layout(); plt.show()

rescaled GF, slow-manifold start       slope (smallest 4 eta) 1.05
modified flow, slow-manifold start     slope (smallest 4 eta) 2.00
modified flow, start from rest         slope (smallest 4 eta) 0.86

Read the slopes, not the heights. From a slow-manifold start the rescaled gradient flow is off by O(\eta) and the modified flow by O(\eta^2). From rest, the modified flow carries the O(\eta) start-up offset derived above. Its constant happens to be small here, so the from-rest curve sits below the slow-start curve at larger \eta and only overtakes it at the smallest steps, where its slope of about 1 shows.

Adam (quoted, not derived). Cattaneo, Klusowski and Shigida apply the same machinery to Adam. For full-batch Adam after the bias-correction transient, with gradient components large compared with \sqrt\epsilon, they find a sign-like flow on L+\frac\eta2\big(\frac{1+\beta}{1-\beta}-\frac{1+\rho}{1-\rho}\big)\|\nabla L\|_1 (their eq. (5); \beta, \rho the first- and second-moment parameters). With the usual \rho>\beta the coefficient is negative: Adam anti-penalizes the \ell_1 norm of the gradient.