Optimization: steps, momentum, AdamW, and honest experiments

Course map · DL006 · Optimization and numerical stability

The story

A derivative describes the loss near the current parameters. An optimizer chooses a step; a schedule changes its size over time; a validation protocol asks whether the resulting model is useful. These are different claims. Neither “SGD always generalizes better” nor “Adam removes learning-rate tuning” is a reliable decision rule.

We will derive exact quadratic behaviour, implement SGD, momentum, Adam and AdamW, then reproduce numerical failure cases safely. Modern optimizer comparisons remain tied to particular algorithms and studies, not declared winners on every architecture.

Prerequisites

DL004's gradient and DL005's stable logit loss. Write the objective's reduction—sum, example mean or valid-token mean—beside every learning rate. A learning rate copied from another reduction can change the update by a large factor.

1 · Gradient descent — the whole idea in one line

First-order expansion predicts a decrease about $\eta\|g\|^2$ for small positive η. On a smooth function this ignores curvature and higher-order terms; at a stationary point it predicts no progress at all. The gradient is steepest in the Euclidean parameter metric, not necessarily in a rescaled coordinate system.

1.1 A tiny 1D example

For $L(w)=(w-3)^2$, define error $e=w-3$. Then $e_{t+1}=(1-2\eta)e_t$. From w=0 with η=.1, the first weights are .6,1.08,1.464,1.7712. With η=1 they alternate0 and6. With η=1.1 they begin6.6,−1.32,8.184: the original−.72 and7.464 arithmetic was wrong. At η=.5 the minimum is reached in one step, not by bouncing. For .5 < η < 1, errors alternate sign while shrinking.

For the general positive-curvature quadratic $L=\tfrac12aw^2$, the multiplier is1−ηa, so convergence from a nonzero start requires $0<\eta<2/a$. The following exact retained program tests a=4, including a divergent control.

import numpy as np
 
def run(lr):
    w = 1.
    for _ in range(20):
        w -= lr * 4*w
    return w
 
assert abs(run(.1)) < .001
assert abs(run(.4)) < .001
assert abs(run(.6)) > 100
print('quadratic endpoints', [round(run(lr), 6) for lr in (.1, .4, .6)])

The interactive GradientLab contract

The existing page-mounted experiment uses $L(w)=(w-2)^2$, initial w=−1.5 and η=.2. Its first gradient is−7, next weight−.1 and next loss4.41. Error contracts by.6 each step; η=.5 reaches2 in one step, η=1 oscillates, and η>1 grows away from2 except at the exact minimum. This is the same curvature2 as the earlier shifted example, not the a=4 program.

The weight slider selects−3 through5; the rate slider .05 through1.1. The component stops at100 steps or a nonfinite/over100-magnitude next weight, labels off-chart points, and resets weight, rate and history. These are teaching/UI guards, not optimizer convergence tests. Tab and arrow keys operate the existing sliders. No MDX import replaces this component; its slug-based page attachment and implementation are preserved. The reviewed source contract is distinct from parent-owned browser/hydration acceptance.

2 · Stochastic vs batch vs mini-batch

A full-batch gradient differentiates the empirical training objective, not an unknowable population “true gradient.” A one-example or minibatch estimate can be unbiased under suitable random sampling and weighting. Correlated or reweighted sampling changes that statement. The batch controls variance, memory and hardware use; there is no universal32–512 optimum or95% noise benefit.

A full-batch pass can be accumulated over chunks without keeping every activation simultaneously. Compare examples/tokens processed and wall time, not steps alone: a batch4096 step consumes128 times the examples of a batch32 step. For microbatches of sizes3 and1 with mean gradients2 and6, the joint mean is3, not4. Weight by valid counts and flush the final incomplete accumulation group.

2.1 Why noise can help and hurt

Sampling noise can move a trajectory out of some narrow regions, but does not guarantee escaping saddles, finding a global optimum or better generalization. Small batches can also give noisy, unstable updates. Flatness depends on parameterization, and one two-dimensional landscape cannot classify the obstacles in an arbitrary neural network. Compare validation performance across seeds and budgets rather than attributing every difference to “noise regularizes.”

3 · Momentum — the ball that remembers

One convention is $v_t=\beta v_{t-1}+g_t$, then $\theta\leftarrow\theta-\eta v_t$. Another uses $(1-\beta)g_t$. Their steady-state scales differ by1/(1−β); learning-rate conventions must move with the equation. β=.9 gives exponentially decaying weights with a rough10-step scale, not a literal fixed10-step memory.

For alternating gradients, history can partially cancel oscillations; for consistent gradients it builds up movement. It can also overshoot. The ball analogy describes inertia, not a proof that momentum always fixes a narrow valley or that β=.999 is always bad.

3.1 Code

The shared optimizer implementation below uses the unscaled momentum convention, with one velocity array per parameter. Gradients are computed from the current parameter set, validated, and then applied; mismatched lists must not be silently truncated by zip.

4 · Adam — adaptive first and second moments

4.1 What each piece does

m is an exponential gradient average; v is an uncentred second moment, not the variance $E[g^2]-E[g]^2$. Under a stationary mean, zero initialization introduces a factor1−β^t; dividing removes that startup bias. It does not make nonstationary estimates unbiased in all senses.

If a gradient history scales by c>0 and epsilon is negligible, m scales by c and sqrt(v) by c, so their ratio is approximately invariant. That is why “100 times the gradient gives100 times smaller updates” was wrong: the numerator changes too. Epsilon's placement and size matter near zero. Adam has convergence results under stated assumptions and known counterexamples outside them; neither “provably converges everywhere” nor “has no theory” is accurate.

4.2 Code

For a first gradient2 with β1=.9 and β2=.999, m=.2,v=.004; corrected values2 and4 give an approximately η-sized first step. The exact retained program also includes decoupled decay.

import numpy as np
 
w = np.array([1., -2.])
g = np.array([2., -4.])
m = np.zeros_like(w); v = np.zeros_like(w)
lr, decay, beta1, beta2 = .1, .01, .9, .999
m = beta1*m + (1-beta1)*g
v = beta2*v + (1-beta2)*g*g
mhat = m / (1-beta1)
vhat = v / (1-beta2)
new = w - lr*mhat/(np.sqrt(vhat)+1e-8) - lr*decay*w
np.testing.assert_allclose(new, [.899, -1.898], atol=1e-8)
print('first AdamW step', new.tolist())

4.3 AdamW — decouple the regularizer

Adding λθ to the gradient before Adam moments changes the moment histories and coordinate-normalized update. AdamW instead applies a separate factor $(1-\eta\lambda)\theta$ in the learning-rate-scaled convention used here. No general statement that low-v coordinates receive less coupled decay follows; the interaction depends on both moments and history. Check which parameters receive decay—biases and normalization scales are often handled separately by particular recipes.

import numpy as np
 
class Optimizer:
    def __init__(self, params, kind, lr, beta=.9, betas=(.9,.999), eps=1e-8, decay=0.):
        assert kind in ('sgd','momentum','adam','adamw') and lr > 0
        assert 0 <= beta < 1 and all(0 <= b < 1 for b in betas) and eps > 0
        self.params=list(params); self.kind=kind; self.lr=lr; self.beta=beta
        self.b1,self.b2=betas; self.eps=eps; self.decay=decay; self.t=0
        self.m=[np.zeros_like(p) for p in self.params]; self.v=[np.zeros_like(p) for p in self.params]
    def step(self, grads):
        assert len(grads)==len(self.params)
        for p,g in zip(self.params,grads):
            assert p.shape==g.shape and np.isfinite(g).all()
        self.t+=1
        for i,(p,g) in enumerate(zip(self.params,grads)):
            if self.kind=='sgd': update=g
            elif self.kind=='momentum':
                self.m[i]=self.beta*self.m[i]+g; update=self.m[i]
            else:
                self.m[i]=self.b1*self.m[i]+(1-self.b1)*g
                self.v[i]=self.b2*self.v[i]+(1-self.b2)*g*g
                mh=self.m[i]/(1-self.b1**self.t); vh=self.v[i]/(1-self.b2**self.t)
                update=mh/(np.sqrt(vh)+self.eps)
            if self.kind=='adamw': p *= 1-self.lr*self.decay
            p -= self.lr*update
 
for kind in ('sgd','momentum','adam','adamw'):
    p=np.array([1.,-2.]); opt=Optimizer([p],kind,.1,decay=.01)
    opt.step([np.array([2.,-4.])])
    expected=[.8,-1.6] if kind in ('sgd','momentum') else ([.899,-1.898] if kind=='adamw' else [.9,-1.9])
    np.testing.assert_allclose(p,expected,atol=1e-8)
p=np.array([1.]);opt=Optimizer([p],'momentum',.1); opt.step([np.array([2.])]);opt.step([np.array([-2.])])
np.testing.assert_allclose(p,[.82])  # velocity .9*2-2 = -.2
for rate in (.1,.5,.75,1.,1.1):
    w=0.
    for _ in range(20): w-=rate*2*(w-3)
    assert np.isclose(w-3,-3*(1-2*rate)**20)
assert np.isclose(-1.5-.2*2*(-1.5-2),-.1)
assert np.isclose((-.1-2)**2,4.41)
# Axis-aligned anisotropic quadratic: curvature, not a universal net bound.
curvature=np.array([1.,100.]); p=np.array([1.,1.])
for _ in range(100): p-=.01*curvature*p
assert p[0] > .3 and p[1]==0  # flat direction is still slow
print('optimizer first/two-step oracles, recurrence and interactive arithmetic passed')

5 · SGD vs Adam — when to pick which?

SGD without momentum has no moment tensor; momentum commonly has one; Adam has two. This is an optimizer-state comparison, not a fixed fraction of total training memory, which also includes weights, gradients, activations, master copies and temporary buffers. Factored or quantized states trade approximation and implementation constraints for memory.

Use a validated nearby recipe as a baseline, then tune learning rates and schedules for each candidate under matched data budgets. CNNs do not universally have equal gradient scales; transformers are not mathematically untrainable without AdamW. A convex toy fit says nothing about comparative generalization on ImageNet or language modeling. Check both validation quality and memory/time-to-target; a faster step can require more steps.

6 · The stability tricks — log-sum-exp, log-softmax, and friends

6.1 The problem — softmax overflows

Float64 exp(1000) overflows even though softmax(1000,999,998) is perfectly well-defined. Invalid intermediates can destroy an otherwise finite result. Assert finiteness at boundaries and find the first nonfinite operation.

6.2 The shift trick

Subtract the row maximum before exponentiation. For finite logits the largest exponent is0, so its exponential is1; normalization then sums values at most1. This does not promise universal precision or handle an all-masked row of−∞, since −∞−(−∞) is undefined.

6.3 log-sum-exp — the same idea for log(sum(exp(x)))

Compute $m+\log\sum_i e^{z_i-m}$. For log probabilities, keep the shifted expression $(z_i-m)-\log\sum_j e^{z_j-m}$ to avoid cancellation from huge common offsets. keepdims=True preserves the reduction axis for broadcasting back.

6.4 log-softmax instead of log(softmax(x))

A normalized probability may underflow to0 even when its log is a moderate representable negative number. Log-softmax avoids that materialization. All-zero logits produce a uniform valid distribution, not a masked-row failure. Detect all-disallowed attention rows explicitly; adding epsilon to an undefined subtraction is not a general fix.

7 · fp16 and finite range

IEEE binary16 has max finite65504, minimum positive normal about6.10e−5 and minimum positive subnormal about5.96e−8. These are different limits. Binary32 has a much larger exponent range; bfloat16 uses8 exponent bits like binary32 but fewer significand bits. BF16 still has overflow, underflow and rounding; it is not “immune.” Hardware may flush subnormals, and operations may accumulate in a wider dtype.

7.1 Loss scaling — the fp16 workaround

Scaling the loss by S scales its gradients by S. After all microbatch contributions, unscale once before checking/clipping and updating. A dynamic scaler detects overflow and adjusts scale/skips an update according to its API. It cannot recover a value that already overflowed in the forward pass. A scaled fp16 square of1000 is still invalid; cast to a wider dtype before squaring/reducing.

7.2 bf16 — wider exponent, lower precision

BF16 often avoids the need for fp16-style loss scaling in specific recipes, but still needs finite checks and a precision budget. It cannot represent every small difference at a large magnitude. Storage, compute and accumulation dtypes can differ. No GPU throughput multiplier follows solely from choosing a dtype.

8 · Diagnosing "loss went to NaN at step X"

1. Check finite inputs, targets, masks and denominators; isnan alone misses infinities.
2. Save the failing batch identity and last finite model/optimizer state. Inspect logits, loss numerator, valid count and gradients at each stage.
3. Re-run the identical small batch in float32/float64; isolate exp/log/divide/sqrt and invalid indices before changing the architecture.
4. Probe a smaller rate from the same pre-failure state. Improvement is evidence, not proof that rate was the only cause.
5. Distinguish forward overflow, backward overflow and fp16 underflow using dtype ranges, scaler logs and gradient norms before/after unscale.
6. Correct the cause. Silently skipping every bad batch can bias training; record exclusions and investigate repeated failures. In distributed training all ranks must take consistent control flow.

Global clipping rescales g by min(1,c/||g||). For(3,4),c=2 it gives(1.2,1.6), not coordinatewise(2,2). Accumulate and unscale first. Clipping bounds the supplied gradient, not the optimizer's later momentum-normalized update or the next loss; it cannot repair NaNs or vanished gradients. Even finite coordinates can overflow a naive sum of squares: (1e308,1e308) has a finite float64 norm, but squaring either entry overflows. The helper below first scales by the largest absolute coordinate and forms the clipped direction without ever constructing those squares at the original scale.

import numpy as np
 
def logsoftmax(z):
    z=np.asarray(z,float); assert np.isfinite(z).all()
    s=z-z.max(-1,keepdims=True)
    return s-np.log(np.exp(s).sum(-1,keepdims=True))
def clip(g,c):
    g=np.asarray(g,dtype=np.float64)
    assert g.size and np.isfinite(c) and c>0 and np.isfinite(g).all()
    scale=np.max(np.abs(g))
    if scale==0: return g.copy()
    direction=g/scale
    unit_norm=np.sqrt(np.sum(direction*direction))
    if scale <= c/unit_norm: return g.copy()
    return direction*(c/unit_norm)
np.testing.assert_allclose(np.exp(logsoftmax(np.zeros((2,3)))),np.full((2,3),1/3))
assert np.isfinite(logsoftmax([1000.,-1000.])).all()
np.testing.assert_allclose(clip(np.array([3.,4.]),2),[1.2,1.6])
huge=np.array([1e308,1e308])
np.testing.assert_allclose(clip(huge,2),np.full(2,np.sqrt(2)))
with np.errstate(over="ignore"):
    naive=huge*min(1.,2/np.linalg.norm(huge))
assert np.all(naive==0) and not np.allclose(naive,clip(huge,2))
np.testing.assert_array_equal(clip(np.zeros(2),2),np.zeros(2))
np.testing.assert_array_equal(clip(np.array([.3,.4]),2),[.3,.4])
g=np.array([3.,4.]);S=128
np.testing.assert_allclose(clip((S*g)/S,2),clip(g,2))
assert not np.allclose(clip(S*g,2)/S,clip(g,2))
assert np.finfo(np.float16).max==65504
assert np.nextafter(np.float16(0),np.float16(1))==2**-24
with np.errstate(over='ignore',under='ignore'):
    assert np.isinf(np.float16(1000)*np.float16(1000))
    assert np.float16(1e-8)==0
assert np.float32(1000)**2==1e6
try: clip(np.array([np.nan,1.]),1)
except AssertionError: pass
else: raise AssertionError('nonfinite gradient was accepted')
print('stable logits, mask premise, unscale order and finite-range controls passed')

8b · Modern optimizer mechanisms, not a winner list

8b.1 · AdamW — a separate decay path

The primary AdamW paper distinguishes decay from L2 under adaptive scaling. It does not establish a universal1–3 perplexity improvement or that all modern LLMs use it. Preserve optimizer moments, step counters and decay conventions when resuming; weights alone are not the same trajectory.

8b.2 · Lion — sign of a momentum blend

The official implementation first computes $u=\operatorname{sign}(\beta_1m+(1-\beta_1)g)$ from the old m, decays/updates parameters, then sets $m\leftarrow\beta_2m+(1-\beta_2)g$. The old lesson swapped these roles and updated m too early. Lion stores one moment rather than Adam's two: about half the moment-state storage, not half all memory. Sign-based step magnitude changes the appropriate rate/decay tuning; paper results are not universal wall-time guarantees.

8b.3 · Sophia — curvature estimators differ

Sophia uses a momentum-like gradient estimate and a diagonal curvature estimate with clipped coordinate updates. Sophia-H uses a Hutchinson-style Hessian estimator; Sophia-G uses a Gauss–Newton–Bartlett estimator. Curvature is updated periodically in the study. A claimed training efficiency belongs to that estimator, model/data budget and accounting, not merely the name “second order.”

8b.4 · Shampoo and SOAP

Shampoo accumulates matrix gradient statistics associated with tensor dimensions and uses matrix-root preconditioning; it does not store the full parameter-by-parameter Hessian. SOAP performs Adam-like optimization in a Shampoo-derived eigenbasis. Matrix factorizations add compute/state cost; update frequency and block sizes matter. Neither is simply an autograd gradient-of-gradient calculation, and no unverified lab-adoption rumour is needed to teach them.

8b.5 · Muon — matrix update geometry

Muon takes momentum updates for appropriate matrix parameters and uses Newton–Schulz iterations to approximate a polar/orthogonalized update. Rectangular matrices cannot have both all rows and all columns orthonormal; interpretation depends on the smaller dimension. Embeddings, output heads and scalar/vector parameters may use a different optimizer in a recipe. This layer selection is why comparing “Muon versus Adam” without parameter groups is underspecified. We describe the author algorithm, not proprietary training internals or a reproduced speedrun.

8b.6 · Learning-rate schedules — a second clock

Warmup gradually increases the rate; cosine decay smoothly reduces it; warmup-stable-decay retains a plateau before a chosen cooldown. Specify whether time means optimizer updates, epochs or processed tokens. Skipped updates and accumulation change these clocks. The MiniCPM study supports WSD as a studied training strategy, not a universal national-lab default.

Training a smaller model beyond a training-compute-optimal token budget can be rational when lifetime inference cost matters. Muennighoff et al.'s data-constrained study examines repeated/limited data; it is not interchangeable with every inference-optimal scaling claim. Match the study's resource constraint before importing a conclusion.

8b.7 · Clipping and loss scaling in the loop

The complete order is forward loss with the intended denominator, scaled backward if needed, accumulation, one unscale, finite check, global clipping, optimizer step, scaler update and correctly clocked schedule. A threshold1 or scale2^15 is a recipe setting, not a law. Log skipped steps and resume state. None of this proves model quality.

10 · Failure stories replaced by testable hypotheses

On a narrow positive-definite quadratic with Hessian eigenvalues λmin and λmax, each eigencoordinate contracts by1−ηλ. Stability is limited by λmax while a safe rate may make progress slow along λmin; the condition number explains that mismatch. This local model does not establish that every NaN is excessive curvature, every plateau is conditioning, or every high-dimensional critical point is a saddle. Random symmetric matrices are not a sample from trained-network Hessians.

An enormous rate, an invalid batch, a schedule discontinuity and a forward fp16 square can all produce alarming curves. Conversely a large rate can remain finite on a small separable logistic problem. Inspect intermediates and state; do not “verify” a preconceived explosion. Norm clipping may help a large finite gradient but cannot undo the original invalid computation.

11 · Retention scaffold

1. SGD and its knobs? Subtract ηg. Batch sampling and size determine g's statistics; reduction determines scale. For curvature10,η=.2 gives multiplier−1: oscillation, not convergence.
2. Adam from memory? Two EMAs, two startup corrections, then divide first by square-root second plus epsilon. For first gradient2, corrected values2 and4 give a step nearη. v is not centred variance.
3. NaN versus oscillation? Check data/domain/dtype for NaN; rate/schedule, minibatch variability and momentum convention for oscillation. These are hypotheses, not unique diagnoses.
4. Which optimizer converges fastest? There is no architecture-free answer. Match data/time budget, tune each, state parameter groups and compare time-to-target; Muon's matrix-only role changes the comparison.
5. When choose momentum SGD? A validated recipe, competitive held-out results, implementation simplicity or a measured state-memory bottleneck can justify it. “Always deeper convergence” cannot.
6. Four-step instability experiment? Capture one failing batch and pre-failure state; inspect finite values/dtypes; rerun it in wider precision; then probe rate and scaling separately while logging unscaled norms and actual updates. Keep data fixed so the comparison can distinguish causes.

Recall — from memory

Explain the difference between clipping g and bounding the full Adam step; between equal minibatch means and a count-weighted joint mean; and between L2-through-moments and independent decay. Reproduce(3,4)→(1.2,1.6) without running code.

Stretch — for one extra hour

Train a tiny two-layer tanh classifier with the same initialization/data order under all three optimizers. Use a small rate grid per method and a held-out split; report observed curves rather than requiring Adam to win first and momentum to win last. Here is a bounded CPU experiment with no downloads. It uses one illustrative rate per method, not the rate grid requested in the exercise; extending the comparison fairly requires tuning each method on validation data while keeping the final test set untouched. Its printed losses are run results, not a benchmark conclusion.

import numpy as np
rng=np.random.default_rng(61)
X=rng.normal(size=(80,2)); y=(X[:,0]*X[:,1]>0).astype(float)
train=np.arange(60); valid=np.arange(60,80)
initial=[rng.normal(scale=.3,size=(2,8)),np.zeros(8),rng.normal(scale=.3,size=(8,)),np.array(0.)]
def objective(params,xx,yy):
    W,b,v,c=params; h=np.tanh(xx@W+b);z=h@v+c
    return float(np.mean(np.maximum(z,0)-yy*z+np.log1p(np.exp(-np.abs(z)))))
for kind,lr in [('sgd',.3),('momentum',.03),('adam',.01)]:
    ps=[p.copy() for p in initial];opt=Optimizer(ps,kind,lr)
    first=objective(ps,X[train],y[train])
    for _ in range(250):
        W,b,v,c=ps;xx=X[train];yy=y[train];h=np.tanh(xx@W+b);z=h@v+c
        e=np.exp(-np.abs(z));prob=np.where(z>=0,1/(1+e),e/(1+e)); dz=(prob-yy)/len(yy)
        dh=dz[:,None]*v*(1-h*h)
        grads=[xx.T@dh,dh.sum(0),h.T@dz,np.array(dz.sum())]
        opt.step(grads)
    last=objective(ps,X[train],y[train]); held=objective(ps,X[valid],y[valid])
    assert np.isfinite(last) and last < first
    print(kind,'train',round(first,5),'->',round(last,5),'validation',round(held,5))

In your own words

Explain why adaptive normalization can help with coordinate scale without removing rate, objective and evaluation choices.

Spaced-review pointer

Tomorrow derive the1−ηa multiplier and trace two momentum steps; next week reproduce the nonfinite-input and wrong-clipping-order controls.

Next-session teaser

DL007 connects a stable probability loss, its exact gradient and an optimizer to a single neuron, then proves the representation fails on XOR.

What to bring back tomorrow — sticky note

Check the objective before the optimizer. Record dtype, reduction, rate convention, valid counts, state and schedule clock. Stable arithmetic and clipped gradients do not certify convergence.

Primary references and boundary

Adam, AdamW, PyTorch2.9 AMP, Lion implementation, Sophia, Shampoo, SOAP, Muon, MiniCPM and data-constrained scaling support the scoped mechanisms. Only the exact small CPU programs are executed. No AMP/GPU run, proprietary optimizer attribution or universal speed/generalization ranking is claimed.