Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

A training run has a shape in time as well as in the loss curve: data preparation that nothing else can start without, a warmup whose length is set by the step size, a long stable phase, and a cooldown whose purpose is to end at a point the optimiser has settled into.

This note uses a Gantt chart, which is the diagram type for anything where the order and duration of phases is the content. The bars below are illustrative rather than measured, and the durations are the ones a small decoder run would typically use.

gantt
    title A training schedule for a small decoder
    dateFormat YYYY-MM-DD
    axisFormat %b %d
    section Data
    Tokenise and shard          :done,    data,   2026-09-01, 3d
    section Training
    Warmup                      :done,    warm,   after data, 2d
    Stable phase                :active,  stable, after warm, 9d
    Cooldown                    :         cool,   after stable, 3d
    section Evaluation
    Checkpoint sweeps           :         eval,   after cool, 2d

The dependency arrows matter more than the bar lengths: the schedule encodes the constraint that evaluation cannot begin before a checkpoint exists, and the warmup cannot be skipped because the step size at initialisation is larger relative to the curvature than it is later in the run.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

The number a language model reports while training is the mean negative log-likelihood of the token that actually came next: softmax cross-entropy over the vocabulary at every position. Its backward pass has an unusually clean closed form — the softmax Jacobian and the logarithm cancel, leaving the probabilities minus the one-hot target, divided by the number of predictions.

The code below is the full training-side loop for that quantity: a forward pass that keeps the probabilities its backward will need, the matching backward pass, and a finite-difference check that compares the analytic gradient against central differences one parameter at a time.

import numpy as np


def softmax_cross_entropy(logits, targets):
    """Forward pass for next-token prediction.

    logits:  (B, T, V)  score for every vocabulary entry at each position
    targets: (B, T)     the token that actually came next
    Returns the mean negative log-likelihood and a cache for the backward.
    """
    shifted = logits - logits.max(axis=-1, keepdims=True)   # stays finite
    p = np.exp(shifted)
    p /= p.sum(axis=-1, keepdims=True)  # rows sum to 1

    B, T, V = logits.shape
    flat = p.reshape(B * T, V)
    picked = flat[np.arange(B * T), targets.reshape(-1)]
    return -np.log(picked).mean(), {"p": p, "targets": targets}


def softmax_cross_entropy_backward(cache):
    """dL/dlogits = (p - one_hot(target)) / (B * T).

    The softmax Jacobian and the logarithm cancel, so the gradient of
    the loss needs no extra matrix product.
    """
    p, targets = cache["p"], cache["targets"]
    B, T, V = p.shape
    grad = p.copy()
    flat = grad.reshape(B * T, V)  # a view: grad changes with it
    flat[np.arange(B * T), targets.reshape(-1)] -= 1.0
    return grad / (B * T)


def finite_difference(f, x, eps=1e-5):
    """Central-difference gradient, one parameter at a time."""
    grad = np.zeros_like(x)
    for index in np.ndindex(x.shape):
        original = x[index]
        x[index] = original + eps
        high = f(x)
        x[index] = original - eps
        low = f(x)
        x[index] = original
        grad[index] = (high - low) / (2 * eps)
    return grad


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    logits = rng.normal(size=(2, 4, 7))
    targets = rng.integers(0, 7, size=(2, 4))

    loss, cache = softmax_cross_entropy(logits, targets)
    analytic = softmax_cross_entropy_backward(cache)
    numeric = finite_difference(
        lambda x: softmax_cross_entropy(x, targets)[0], logits.copy()
    )

    gap = np.abs(analytic - numeric).max()
    print(f"loss {loss:.4f}")
    print(f"max |analytic - numeric| = {gap:.2e}")

Run as printed, the loss is 2.0409 on random scores — just above \(\ln 7 \approx 1.95\), the value of a uniform guess over seven tokens — and the two gradients agree to 2.43e-11, so the closed form is correct to the precision of the step size. A check like this is the cheapest way to catch a wrong axis or a missing division before a training run.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

Step sizes in deep learning are chosen as if the loss were smooth in its parameters, and that is an assumption, not a theorem. It is the smoothness postulate: the gradient does not change arbitrarily fast, so a small step in parameter space produces a predictable change in loss.

The assumption is what turns a learning rate from a guess into a bound. Where it fails — across a sharp loss barrier, or at a spike — the same learning rate produces the visible training instabilities that motivate warmup and clipping.

Postulate (L-smooth loss). Let \(L(\theta)\) be the training loss as a function of the parameters. The smoothness postulate assumes a constant \(L \gt 0\) with

for all parameter vectors, which is equivalent to the loss lying below its tangent parabolas:

Reading the bound as a function of the step shows the largest safe step for descent is \(1/L\).

Nothing about a transformer guarantees such a constant exists globally; the postulate is used locally, inside a trust region where it is a fair description of the loss, which is the same logic that motivates clipping and warmup in the first place.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

A proof is the part of a note that cannot be skipped, and this one is short enough to show the whole pattern: write the definition, factor out the constant, use that it cancels between numerator and denominator.

The proof is worth writing out because it is the template for a family of invariance arguments. The same move — factor a common term out of a normalised expression — proves invariance to scale in other normalisations and shows where such invariances stop.

Proof. Let \(z \in \mathbb{R}^{n}\) and \(c \in \mathbb{R}\). By definition,

where the third equality uses \(e^{z_j + c} = e^{c} e^{z_j}\) and the factor \(e^{c}\), which is positive and finite for every real \(c\), cancels. Since this holds for every index \(i\), the two vectors are equal.

Two remarks on the hypotheses: the argument needs \(e^{c}\) to be a common factor, which is why a shift by a constant vector works and a shift by different amounts per coordinate does not; and it needs \(0 \lt e^{c} \lt \infty\), which fails in floating point once \(c\) is large enough that the exponentials overflow.

That last remark is where the lemma stops being a statement about code: mathematically the invariance is exact, while numerically it holds only for the range of \(c\) the representation can express, which is the reason the shift is chosen as the negative of the maximum rather than any other constant.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

A counterexample is the cheapest way to kill a plausible statement, and the one worth having ready concerns the order of matrix multiplication. The tempting assertion that \(AB = BA\) fails for almost every pair of matrices, and it fails on the smallest possible example.

The consequence for models is concrete: two linear layers composed in one order compute a different function from the same two layers in the other order, so the order in a transformer block is a modelling decision rather than a coding detail.

Counterexample. Take

Then

which are different matrices, so \(AB \neq BA\). Both products are invertible with determinant one, so this is not a degenerate case in any respect.

What is true, and sometimes confused with commutation, is that ranks and traces of products are order-independent when the shapes allow both orders: \(\operatorname{tr}(AB) = \operatorname{tr}(BA)\) holds even here, with both equal to three.

The example also shows why reordering two projections with different roles — queries and keys, say — changes the model rather than merely the notation.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

The log-partition function appears wherever a normalising constant is needed, and it is trapped between two elementary bounds that are enough for most arguments: the largest term below, the largest term plus the log of the number of terms above.

The bounds are used for more than estimates. They are the reason the log-partition is Lipschitz with constant one in the maximum norm, which is what makes it safe to compute in log space in the first place.

Lemma (bounds on the log-sum-exp). For \(z \in \mathbb{R}^{n}\),

The lower bound is the largest term alone; the upper bound follows because every term is at most \(e^{\max_i z_i}\). Both bounds are attained: the lower one when the other terms vanish, the upper one when all terms are equal.

A gap of at most \(\log n\) is what makes the log-partition usable as a normaliser even when the individual terms differ by hundreds of nats, since the bounds do not depend on that spread.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

Architectures are usually chosen for what they can compute, but the sharper criterion is what they cannot distinguish: matching the symmetries of the problem removes the need to learn them from data. This principle is why convolution appears in vision and why attention appears in sequences.

It is a principle rather than a theorem because the symmetry of the problem is a modelling judgement. Once it is identified, however, the consequence is exact: a layer that commutes with the symmetry group reduces the function space being searched rather than merely biasing the search.

Principle (equivariance). Let \(G\) be a group acting on inputs and outputs, and let \(T_g\) denote the action of \(g \in G\). A layer \(f\) is equivariant when

and invariant when \(f(T_g x) = f(x)\). The principle is to choose \(f\) from the equivariant class when the task itself satisfies \(f^{\star}(T_g x) = T_g f^{\star}(x)\). Standard instances: translation for a convolution, permutation for a set encoder, rotation for a point-cloud network — and for a transformer with no position embedding, permutation of the tokens.

The principle also predicts where an architecture will struggle: a model that is equivariant to a symmetry the task does not have must spend capacity learning to break it, which is why position embeddings are needed at all.

Demo note. This note demonstrates what the SciDraft Hugo theme can do for academic and scientific publishing — for research institutions, researchers and students. Its content is demonstration material only: readers should ignore whether the demo content is correct.

A criterion is a condition that decides whether an argument or a method applies, and this one is the usual substitute for convexity in the analysis of gradient descent: it bounds the loss by the size of its gradient, so a small gradient implies a small excess loss.

The criterion is attractive because it is weaker than strong convexity and still sufficient for linear convergence, and because it is stated in a form a training run can be checked against by measuring both quantities.

Criterion (Polyak-Lojasiewicz). Let \(L\) be differentiable with minimum \(L^{\star} = \min_\theta L(\theta)\). The PL criterion holds with constant \(\mu \gt 0\) when

If \(L\) is also \(L_s\)-smooth, gradient descent with step \(\eta = 1 / L_s\) satisfies

so the excess loss decays geometrically and the ratio \(\mu / L_s\) sets the rate.

The criterion is a condition to be checked rather than assumed: both \(L(\theta) - L^{\star}\) and \(\lVert \nabla L(\theta) \rVert\) are measurable during training, so the inequality can be tested on a run and its constant estimated from the data.