From text to one gradient update

A composition of functions carries one row to its loss: h ↦ z = Wh + b ↦ p = softmax(z) ↦ ℓ = −log p_y. Mathematically ℓ = logsumexp(z) − z_y. It is computed in shifted coordinates s = z − max z: ℓ = log Σ_c exp(s_c) − s_y and p = exp(s) / Σ_c exp(s_c). No large positive number is exponentiated, and a large common offset in z is never added back before z_y is subtracted, which would cancel the digits the loss needs. A result that double precision cannot hold, such as an overflowing logit spread, sum or updated parameter, is rejected rather than rounded to infinity.

The objective. M = Σ_i,t m_it counts the targets after the shift. If M > 0, L = (1/M) Σ_i,t m_it ℓ_it. If M = 0 there is no mean: the helper reports that there are no supervised targets, returns no loss, and leaves W and b unchanged. It does not call the loss zero.

The chain rule, one entry at a time. For one row, ∂ℓ/∂z_c = p_c − [c = y], the softmax cross-entropy derivative. z_c depends on W only through row c: ∂z_c/∂W_cj = h_j and ∂z_c/∂b_c = 1. Multiplying along the chain and summing over the rows that count,

∂L/∂W_cj = Σ_i,t (m_it/M)(p_itc − [c = y_it]) h_itj = Σ_i,t δ_itc h_itj.

For each row that is the outer product δhᵀ, a V × 1 column times a 1 × d row, which has the shape of W; and ∇b = Σ δ.

Microbatches. With microbatches k = 1, …, K holding summed losses S_k over M_k targets, the token mean is (Σ_k S_k) / (Σ_k M_k). The mean of microbatch means is (1/K) Σ_k S_k/M_k, which gives each target the weight 1/(KM_k). They agree when every M_k is equal. An empty microbatch adds nothing to either total of the token mean and has no mean of its own.

What is left out. In a trainable network the same δ continues as Wᵀδ into h, and from there into attention and embeddings, including through prompt positions that do not count. Here H is frozen, so that path is not computed.

Related lessons: Cross-Entropy for the p − e_y derivative, Backpropagation for a received signal times a local derivative and the weight-row outer product, and Gradient Descent for the step.

Array code over each sequence’s rows rather than one row at a time. The browser runner supplies your current mask and grouping as the dictionary case. Run without a case in Python 3 and NumPy, it prints the completion-only, one-batch reference and checks with assert that two microbatches divided once give the one-batch gradient. The final JSON line keeps full precision for comparison.

import numpy as np

# Output classes 0 red, 1 blue, 2 green (V = 3). PAD marks padding; it is not a class.
RED, BLUE, GREEN, PAD = 0, 1, 2, "PAD"
V, d, eta = 3, 2, 0.3

# Two right-padded sequences, B = 2 and T = 4. Roles mark the completion; they are not tokens.
tokens = [[GREEN, GREEN, RED, RED],
          [GREEN, BLUE, PAD, PAD]]
roles = [["prompt", "prompt", "completion", "completion"],
         ["prompt", "completion", None, None]]

# Authored, frozen hidden rows, one per input position: shape (B, T, d).
H = np.array([[[0.0, 1.0], [1.0, 0.0], [1.0, 0.0], [0.0, 0.0]],
              [[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]])
W0, b0 = np.zeros((V, d)), np.zeros(V)


def align(policy):
    """Shift first: row (i, t) is scored on token t + 1 of the same sequence.
    Returns targets y (a class, or None) and eligibility m (0 or 1), each B x T."""
    y, m = [], []
    for seq, role in zip(tokens, roles):
        nxt, nxt_role = seq[1:] + [None], role[1:] + [None]  # the last row has no successor
        y.append([None if tok in (None, PAD) else tok for tok in nxt])
        m.append([int(tok not in (None, PAD) and (policy == "all" or r == "completion"))
                  for tok, r in zip(nxt, nxt_role)])
    return y, m


def finite(*values):
    """Reject any result float64 cannot hold, instead of carrying inf or nan onward."""
    for value in values:
        if not np.isfinite(value).all():
            raise FloatingPointError("a result is not representable in float64")
    return values


def sums(seqs, y, m, W, b):
    """Summed loss, eligible count and summed gradients over the rows of the given sequences.
    Only rows with m = 1 are read; every other row has delta = 0."""
    S, M, GW, Gb = 0.0, 0, np.zeros(W.shape), np.zeros(b.shape)
    for i in seqs:
        rows = [t for t in range(len(y[i])) if m[i][t] == 1]
        if not rows:
            continue
        Hr = np.array([H[i, t].tolist() for t in rows])   # (R, d) eligible hidden rows
        Z = Hr @ W.T + b                                   # (R, d) @ (d, V) + (V,) -> (R, V) logits
        Zs = Z - Z.max(axis=1, keepdims=True)              # shifted: each row's largest is 0
        finite(Z, Zs)
        E = np.exp(Zs)
        total = E.sum(axis=1, keepdims=True)               # (R, 1), between 1 and V
        P = E / total                                      # (R, V) probabilities
        Y = np.array([[float(y[i][t] == c) for c in range(len(b))] for t in rows])  # one-hot targets
        # Loss in shifted coordinates: log sum exp(z - max) - (z_y - max). The offset never returns.
        S += float((np.log(total) - (Zs * Y).sum(axis=1, keepdims=True)).sum())
        M += len(rows)
        R = P - Y                                          # (R, V): p - e_y
        GW = GW + R.T @ Hr                                 # (V, R) @ (R, d) = sum of (p - e_y) h^T
        Gb = Gb + R.sum(axis=0)
        finite(S, GW, Gb)                                  # summed, not yet divided
    return S, M, GW, Gb


def token_mean(parts, y, m, W, b):
    """The declared objective: add every microbatch's sums and counts, then divide once."""
    S, M, GW, Gb = 0.0, 0, np.zeros(W.shape), np.zeros(b.shape)
    for part in parts:
        s, k, gw, gb = sums(part, y, m, W, b)
        S, M, GW, Gb = S + s, M + k, GW + gw, Gb + gb
    finite(S, GW, Gb)
    if M == 0:
        return None                                        # no supervised target: no mean, no step
    return S / M, GW / M, Gb / M


def mean_of_means(parts, y, m, W, b):
    """A different objective: each microbatch's own mean, averaged with equal weight."""
    pieces = [sums(part, y, m, W, b) for part in parts]
    if any(k == 0 for _, k, _, _ in pieces):
        return None                                        # an empty microbatch has no mean
    finite(sum(s for s, _, _, _ in pieces))                # the summed raw loss, as the ledger reports it
    K = len(pieces)                                        # each target weighs 1 / (K k)
    return finite(sum(s / k / K for s, k, _, _ in pieces),
                  sum(gw / k / K for _, k, gw, _ in pieces),
                  sum(gb / k / K for _, k, _, gb in pieces))


def sgd_step(W, b, gW, gb, eta):
    """One plain step. A parameter float64 cannot hold is rejected, never clamped."""
    return finite(W - eta * gW, b - eta * gb)


def readout(W, b, h):
    """Shifted logits, probabilities and log sum exp(z - max) for one hidden vector."""
    z = W @ h + b                                          # (V,)
    zs = z - z.max()
    finite(z, zs)
    e = np.exp(zs)
    return zs, e / e.sum(), float(np.log(e.sum()))


def loss_at(W, b, h, target):
    zs, _, log_total = readout(W, b, h)
    return log_total - float(zs[target])


def show(a):
    return (np.round(a, 6) + 0.0).tolist()                 # + 0.0 prints -0.0 as 0.0


REFERENCE_CASE = {"example": "masked-readout-one-step-v1", "policy": "completion",
                  "grouping": "one-batch", "eta": 0.3}


def reproduce_case(case):
    """This version fixes tokens, H, zero W/b and one step. Export only its declared controls."""
    if (not isinstance(case, dict) or set(case) != set(REFERENCE_CASE)
            or case["example"] != REFERENCE_CASE["example"]
            or case["policy"] not in ("completion", "all")
            or case["grouping"] not in ("one-batch", "split-token-mean", "split-mean-of-means")
            or type(case["eta"]) not in (int, float) or case["eta"] != eta):
        raise ValueError("This masked-readout case is not supported by this code")
    y, m = align(case["policy"])
    parts = [[0, 1]] if case["grouping"] == "one-batch" else [[0], [1]]
    reduction = mean_of_means if case["grouping"] == "split-mean-of-means" else token_mean
    L, gW, gb = reduction(parts, y, m, W0, b0)
    W1, b1 = sgd_step(W0, b0, gW, gb, case["eta"])
    before, _, _ = token_mean([[0, 1]], y, m, W0, b0)
    after, _, _ = token_mean([[0, 1]], y, m, W1, b1)
    h = np.array([1.0, 0.0])
    _, p1, _ = readout(W1, b1, h)
    return {**case, "y": y, "m": m, "targets": sum(map(sum, m)),
            "objective": float(L), "numerator": float(sums([0, 1], y, m, W0, b0)[0]),
            "training_before": float(before), "training_after": float(after),
            "grad_W": gW.tolist(), "grad_b": gb.tolist(), "W_after": W1.tolist(), "b_after": b1.tolist(),
            "row_loss_after": [float(loss_at(W1, b1, H[i, t], y[i][t]))
                               for part in parts for i in part for t in range(len(y[i])) if m[i][t]],
            "held_out_before": float(loss_at(W0, b0, h, BLUE)),
            "held_out_after": float(loss_at(W1, b1, h, BLUE)), "probabilities_after": p1.tolist()}


# Standalone reference demonstration; a supplied case is reproduced below instead.
if "case" not in globals():
    y, m = align("completion")
    L, gW, gb = token_mean([[0, 1]], y, m, W0, b0)
    W1, b1 = sgd_step(W0, b0, gW, gb, eta)                     # one SGD step
    L1, _, _ = token_mean([[0, 1]], y, m, W1, b1)
    h = np.array([1.0, 0.0])
    _, p1, _ = readout(W1, b1, h)

    print("targets M =", sum(map(sum, m)), "| loss", round(L, 6), "->", round(L1, 6))
    print("grad_W", show(gW), "grad_b", show(gb))
    print("W'", show(W1), "b'", show(b1))
    print("p' at h = (1, 0):", show(p1))

    # The same rows in two microbatches, then one step: add sums, divide once.
    _, gW_split, gb_split = token_mean([[0], [1]], y, m, W0, b0)
    assert np.allclose(gW_split, gW) and np.allclose(gb_split, gb)
    # Averaging the two microbatch means starts at the same loss but moves differently.
    L_mm, gW_mm, gb_mm = mean_of_means([[0], [1]], y, m, W0, b0)
    print("mean of means: loss", round(L_mm, 6), "grad_b", show(gb_mm))

    # A held-out case declared in advance: red prompt -> blue completion, also h = (1, 0).
    print("held-out loss", round(loss_at(W0, b0, h, BLUE), 6), "->", round(loss_at(W1, b1, h, BLUE), 6))

    # The other declared target set, and a batch with nothing to learn from.
    y_all, m_all = align("all")
    L_all, gW_all, gb_all = token_mean([[0, 1]], y_all, m_all, W0, b0)
    print("all targets: M =", sum(map(sum, m_all)), "grad_b", show(gb_all))
    assert token_mean([[0, 1]], [[None] * 4] * 2, [[0] * 4] * 2, W0, b0) is None

# The browser sets case from JSON. Locally, use runpy.run_path(..., init_globals={"case": data}).
# Full-precision JSON is the comparison contract; the preceding reference printout is rounded.
import json
result = reproduce_case(globals().get("case", REFERENCE_CASE))
if "case" in globals():
    print("Case:", result["policy"], "|", result["grouping"], "| one step, eta =", result["eta"])
    print("targets M =", result["targets"], "| training mean", result["training_before"], "->", result["training_after"])
    print("grad_W", result["grad_W"], "grad_b", result["grad_b"])
    print("W'", result["W_after"], "b'", result["b_after"])
    print("p' at h = (1, 0):", result["probabilities_after"])
print("MASKED_READOUT_CASE=" + json.dumps(result, allow_nan=False, separators=(",", ":")))

Two short checks on the same example. A hint comes before the answer. What you open is noted beside each item on this page only; nothing is saved, and reloading clears it.

1. Add a third sequence: green (prompt), red (prompt), blue (completion), PAD. With completion tokens as the targets, what is M for all three sequences together?

   M =CheckHintShow the answer

2. Now put each of the three sequences in its own microbatch and average the three microbatch means. What weight does the third sequence’s single target get? The token mean gives every target 1/4.

   weight =CheckHintShow the answer