Chapter 15

Optimizers & Schedules

SGD, momentum, Adam, AdamW, warmup, cosine and WSD schedules, clipping, and Muon.

An optimizer turns gradients into parameter updates. In LLM training this small piece of code decides whether a model makes steady progress, diverges, or wastes compute with steps that are too timid. This chapter starts with a quadratic where every update can be solved exactly, then builds the practical tools used in NumPy training loops: momentum, adaptive scaling, AdamW, learning-rate schedules, gradient clipping, and Muon.

15.1 Gradient descent and noise

For $f(\vtheta) = \tfrac12 \vtheta^\T \mH \vtheta$ with symmetric positive definite $\mH$ , gradient descent is

\vtheta_{t+1} = \vtheta_t - \eta \mH\vtheta_t .\tag{15.1}

Along an eigenvector with curvature $\lambda_i$ , the error is multiplied by $1-\eta\lambda_i$ . Stability requires $0 < \eta < 2/\lambda_{\max}$ , but the slow direction shrinks according to $\lambda_{\min}$ . The condition number $\kappa=\lambda_{\max}/\lambda_{\min}$ therefore controls speed; with the best fixed learning rate, the worst-case factor is $(\kappa-1)/(\kappa+1)$ . Bad conditioning means one direction wants tiny steps while another could use large ones.

The quadratic is not a toy distraction. Close to any smooth minimum, the loss looks like a quadratic whose Hessian eigenvalues are local curvatures. A single scalar learning rate must serve all those directions at once, so ill-conditioned problems zigzag in steep directions and crawl in flat ones. Momentum and adaptive methods are different attempts to make that geometry less painful without forming the Hessian.

Listing 15.1 Gradient descent on a quadratic

def quadratic_value(theta, curvature):
    return 0.5 * float(np.sum(curvature * theta * theta))


def quadratic_gradient(theta, curvature):
    return curvature * theta


def gradient_descent(theta, curvature, lr, steps):
    path = [theta.astype(np.float64).copy()]
    for _ in range(steps):
        theta = theta - lr * quadratic_gradient(theta, curvature)
        path.append(theta.copy())
    return np.array(path)

SGD replaces the full gradient with a minibatch estimate: $g_t = \nabla L(\vtheta_t) + \vepsilon_t$ , with $\E[\vepsilon_t$ =0]. The noise can help escape shallow traps, but its variance falls only with batch size. Learning-rate schedules exist because early noisy exploration and late precise convergence want different step sizes.

The key promise is unbiasedness, not accuracy on every step. If a minibatch is sampled uniformly, its gradient points in the right direction on average, but any one draw can be far off. That is why training curves are jagged at the update level and smoother at the epoch level, and why validation should be measured with the model in evaluation mode rather than on one lucky minibatch.

15.2 Momentum, RMSProp, Adam, and AdamW

Heavy-ball momentum keeps a velocity [polyak1964]:

\vv_t = \beta\vv_{t-1} - \eta g_t,\qquad \vtheta_t = \vtheta_{t-1} + \vv_t .\tag{15.2}

Consistent gradients accumulate, while alternating gradient signs cancel. Nesterov momentum evaluates the gradient at a look-ahead point, which changes the correction term but keeps the same idea. RMSProp keeps a moving average of squared gradients and divides by its root, giving coordinates with large recent gradients smaller effective steps.

The state variables are part of the optimizer, not temporary scratch arrays. Resetting momentum or RMSProp’s square average in the middle of training changes the algorithm just as surely as resetting weights would. Conversely, starting a new experiment must create fresh state, or the new run inherits stale velocity from a different loss surface.

Listing 15.2 Momentum and RMSProp

def momentum_state(params):
    return {"velocity": {name: np.zeros_like(value) for name, value in params.items()}}


def heavy_ball_step(params, grads, state, lr=1e-2, beta=0.9):
    for name, value in params.items():
        velocity = state["velocity"][name]
        velocity *= beta
        velocity -= lr * grads[name]
        value += velocity


def rmsprop_state(params):
    return {"square": {name: np.zeros_like(value) for name, value in params.items()}}


def rmsprop_step(params, grads, state, lr=1e-3, decay=0.99, eps=1e-8):
    for name, value in params.items():
        square = state["square"][name]
        square *= decay
        square += (1 - decay) * grads[name] * grads[name]
        value -= lr * grads[name] / (np.sqrt(square) + eps)

Adam combines both moments [kingma2014adam]:

m_t=\beta_1 m_{t-1}+(1-\beta_1)g_t,\quad v_t=\beta_2 v_{t-1}+(1-\beta_2)g_t^2 .\tag{15.3}

Because $m_0=v_0=0$ , a constant gradient would give $m_t=(1-\beta_1^t)g$ and $v_t=(1-\beta_2^t)g^2$ . Dividing by those factors removes the startup bias:

\vtheta_t=\vtheta_{t-1} -\eta\,\frac{m_t/(1-\beta_1^t)}{\sqrt{v_t/(1-\beta_2^t)}+\epsilon}.\tag{15.4}

AdamW decouples weight decay from the adaptive gradient step [loshchilov2017decoupled]. Classical L2 regularization adds $\lambda\vtheta$ to the gradient, so Adam’s denominator rescales it coordinate by coordinate. AdamW first shrinks the weights by $1-\eta\lambda$ , then applies the Adam update, making the decay strength independent of the adaptive moments.

The small $\epsilon$ is not a regularizer; it is a numerical guard that prevents division by zero and caps the effective step when the running square is tiny. In float32 code it should be added after the square root, matching the formula above, unless a framework documents a different convention. Seemingly minor differences in this line can make checkpoints diverge after many updates.

Listing 15.3 AdamW with bias correction

def adam_state(params):
    return {
        "t": 0,
        "m": {name: np.zeros_like(value) for name, value in params.items()},
        "v": {name: np.zeros_like(value) for name, value in params.items()},
    }


def adamw_step(params, grads, state, lr=1e-3, beta1=0.9, beta2=0.999,
               eps=1e-8, weight_decay=0.0, decoupled=True):
    state["t"] += 1
    for name, value in params.items():
        grad = grads[name]
        if weight_decay and not decoupled:
            grad = grad + weight_decay * value
        state["m"][name] = beta1 * state["m"][name] + (1 - beta1) * grad
        state["v"][name] = beta2 * state["v"][name] + (1 - beta2) * grad * grad
        m_hat = state["m"][name] / (1 - beta1 ** state["t"])
        v_hat = state["v"][name] / (1 - beta2 ** state["t"])
        if weight_decay and decoupled:
            value *= 1 - lr * weight_decay
        value -= lr * m_hat / (np.sqrt(v_hat) + eps)

15.3 Schedules and clipping

Warmup increases the learning rate from a small value to the target value over the first updates. Cosine decay then smoothly lowers it toward a final value, avoiding a sharp last-step change [loshchilov2016sgdr]. Warmup-stable-decay, or WSD, inserts a plateau between those phases: ramp up, train at the peak, and spend only the tail decaying. The shape is simple, but the schedule must be tied to update count, not epochs, when gradient accumulation changes.

Schedules multiply the optimizer step; they do not change the gradient itself. That distinction matters when resuming training: the optimizer moments, the parameter values, and the current schedule step all need to be restored together. Resuming with the right weights but the wrong schedule can look like an unexplained loss spike.

Global-norm clipping protects the optimizer from rare giant gradients. Compute one norm over all gradient tensors, $\|g\|_2$ , and if it exceeds $c$ , multiply every tensor by $c/(\|g\|_2+\epsilon)$ . The direction is preserved; only the length is capped.

Clip before applying AdamW or momentum unless a recipe says otherwise. Then the optimizer sees the same bounded gradient that the training log reports. Per-tensor clipping is a different operation: it can rotate the combined update because each tensor gets its own scale.

Listing 15.4 Schedules and global-norm clipping

def linear_warmup(step, warmup_steps, peak_lr):
    if warmup_steps <= 0:
        return peak_lr
    return peak_lr * min(1.0, (step + 1) / warmup_steps)


def cosine_decay(step, total_steps, peak_lr, final_lr=0.0):
    if total_steps <= 1:
        return final_lr
    progress = min(1.0, max(0.0, step / (total_steps - 1)))
    weight = 0.5 * (1 + np.cos(np.pi * progress))
    return final_lr + (peak_lr - final_lr) * weight


def wsd_schedule(step, total_steps, warmup_steps, stable_steps, peak_lr,
                 final_lr=0.0):
    if step < warmup_steps:
        return linear_warmup(step, warmup_steps, peak_lr)
    decay_steps = max(1, total_steps - warmup_steps - stable_steps)
    if step < warmup_steps + stable_steps:
        return peak_lr
    return cosine_decay(step - warmup_steps - stable_steps, decay_steps,
                        peak_lr, final_lr)
def global_norm(grads):
    return float(np.sqrt(sum(np.sum(np.asarray(g, dtype=np.float64) ** 2)
                             for g in grads.values())))


def clip_by_global_norm(grads, max_norm, eps=1e-12):
    norm = global_norm(grads)
    scale = min(1.0, max_norm / (norm + eps))
    return {name: grad * scale for name, grad in grads.items()}, norm

15.4 Muon

Muon is a recent optimizer for hidden weight matrices [jordan2024muon]. It keeps momentum, then replaces each 2-D hidden-weight update by its nearest orthogonal direction, the polar factor $\mU\mV^\T$ of the momentum matrix’s SVD. Biases, normalization scales, embeddings, and output heads are usually left to AdamW or SGD; Muon is for interior matrices where update directions can be orthogonalized.

Computing an SVD every step is expensive, so Muon uses a quintic Newton-Schulz iteration. After normalizing a matrix $\mX$ , repeat

\mX \leftarrow a\mX + (b\mX\mX^\T + c(\mX\mX^\T)^2)\mX .\tag{15.5}

The constants below are a stable quintic Newton-Schulz polar iteration; the tests compare the result with the SVD polar factor $\mU\mV^\T$ on small matrices.

The orthogonalized update keeps the update’s row or column directions balanced. It is still an optimizer step, not a constraint on the weights themselves: the parameter matrix is updated by subtracting the orthogonalized momentum direction, and the next gradient is computed normally. That is why Muon can share a training loop with AdamW parameter groups.

Listing 15.5 Muon orthogonalization

def orthogonalize_newton_schulz(matrix, steps=20, eps=1e-7):
    if matrix.ndim != 2:
        raise ValueError("Muon orthogonalization expects a 2-D matrix")
    transposed = matrix.shape[0] > matrix.shape[1]
    x = matrix.T.copy() if transposed else matrix.copy()
    x = x.astype(np.float64, copy=False)
    x /= np.linalg.norm(x) + eps
    a, b, c = 15 / 8, -10 / 8, 3 / 8
    for _ in range(steps):
        gram = x @ x.T
        x = a * x + (b * gram + c * gram @ gram) @ x
    return x.T if transposed else x


def muon_update(gradient, momentum, beta=0.95, steps=20):
    if gradient.ndim != 2:
        raise ValueError("Muon applies to 2-D hidden weight matrices")
    momentum *= beta
    momentum += (1 - beta) * gradient
    return orthogonalize_newton_schulz(momentum, steps=steps)

In practice

Most 2024-2026 LLM training recipes still center on AdamW, warmup, decay, and clipping because they are robust across scales. Momentum and RMSProp remain useful baselines and building blocks. Muon is newer and promising for hidden matrices, with follow-up work studying its LLM scaling behavior [liu2025muon], but it is not a drop-in replacement for every parameter group.

Key equations

\vtheta_{t+1}=\vtheta_t-\eta g_t

\vv_t=\beta\vv_{t-1}-\eta g_t

\hat{m}_t=m_t/(1-\beta_1^t),\quad \hat{v}_t=v_t/(1-\beta_2^t)

\vtheta \leftarrow (1-\eta\lambda)\vtheta -\eta\,\hat{m}/(\sqrt{\hat{v}}+\epsilon)

g \leftarrow g\min(1,c/(\|g\|_2+\epsilon))

15.5 Teach it

The one-sentence version: an optimizer is a filtered, scaled, and scheduled gradient. Analogy: plain SGD is walking downhill, momentum adds a flywheel, Adam adds per-coordinate shock absorbers, clipping adds a speed limit, and the schedule changes the throttle over the trip. Board steps: (1) solve one quadratic eigen-direction; (2) add momentum’s velocity; (3) show Adam’s two moving averages and bias correction; (4) separate AdamW decay from the gradient. Misconceptions: Adam does not remove the need for a learning rate; clipping fixes step length, not a wrong gradient; weight decay and L2 are identical for SGD but not for Adam. Check for understanding: why does $m_t$ need division by $1-\beta_1^t$ early in training?

15.6 Exercises

Exercise 15.1 ★ Condition number

For a quadratic with eigenvalues $1$ and $25$ , what learning-rate constraint keeps plain gradient descent stable? Why does the small-curvature direction still move slowly?

Solution in Appendix D

Exercise 15.2 ★★ Adam bias correction

Assume the scalar gradient is the same value $g$ at every step. Derive $m_t=(1-\beta_1^t)g$ and $v_t=(1-\beta_2^t)g^2$ , then explain the correction factors.

Solution in Appendix D

Exercise 15.3 ★★ AdamW and clipping

Show why Adam with L2 regularization does not decay weights the same way as AdamW. Then write the global-norm clipping scale for gradient tensors whose combined norm is $13$ and cap is $5$ .

Solution in Appendix D

Exercise 15.4 ★★★ Muon implementation

Implement the quintic Newton-Schulz orthogonalizer for a 2-D matrix, compare it with the SVD polar factor $\mU\mV^\T$ , and explain why the chapter applies Muon only to hidden weight matrices.

Solution in Appendix D

References

[polyak1964] B. T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics 4(5), 1–17, 1964.
[jordan2024muon] K. Jordan et al. Muon: An optimizer for hidden layers in neural networks. Blog post, 2024. https://kellerjordan.github.io/posts/muon/
[kingma2014adam] D. P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. 2014. arXiv:1412.6980
[liu2025muon] J. Liu et al. Muon is Scalable for LLM Training. 2025. arXiv:2502.16982
[loshchilov2016sgdr] I. Loshchilov and F. Hutter. SGDR: Stochastic Gradient Descent with Warm Restarts. 2016. arXiv:1608.03983
[loshchilov2017decoupled] I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. 2017. arXiv:1711.05101