Multi-Token Prediction: Audit the Loss

Multi-token prediction asks one hidden state to predict several future tokens, so a one-column alignment mistake can quietly train the wrong objective. This tutorial builds the shifted targets, masks, and denominator ledger first, then turns leak, off-by-one, wrong-denominator, and empty-head cases into release-blocking tests.

Multi-token prediction starts with alignment

A causal language model normally uses the hidden state at position t to score the token at t+1. Multi-token prediction adds heads for t+2, t+3, and later horizons. The concept is compact; the tensor bookkeeping is not. Every head loses h cells at the right edge, and every padding or document-boundary mask must be shifted into the same coordinate system as its target.

Write the contract before the tensor expression. For a token sequence of length T and horizon h, source positions are 0 through T−h−1 and targets are h through T−1. The nominal cell count is T−h, not T. A target is valid only when both source context and shifted target belong to the declared training example and the target is not ignored.

The multi-token prediction paper reports a particular architecture and empirical program. It motivates predicting several future tokens, but it does not validate an unrelated implementation. Our first deliverable is therefore an alignment receipt, not a throughput promise.

Use a six-token toy sequence and print every source-target pair. If horizon two maps source index three to target index four, the grid is off by one even when the loss decreases. Tiny fixtures make the intended causal relationship visible before batching, compilation, or distributed reduction can obscure it.

Valid geometry for a six-token sequence

Horizon	Nominal cells	Tail masked	Extra rejection
+1	5	1	None
+2	4	2	One boundary
+3	3	3	None

Reading rule: labels, values, patterns, and structure carry the conclusion; color is supplementary.

Build shifted targets without future leakage

Create each target view from immutable token IDs. For horizon h, pair logits[:, :T−h, :] with targets[:, h:T]. Do not roll the array and wrap its first tokens to the tail. A circular shift invents targets across the sequence boundary and gives the last source positions labels drawn from the beginning of the document.

Packed examples need more than a padding mask. A position near the end of document A must never predict a token from document B merely because both occupy one fixed-length training row. Shift a stable segment identifier alongside the target; a cell is valid only when sourceSegment equals targetSegment. This rule also covers end-of-sequence policies explicitly rather than hoping an EOS token absorbs the mistake.

The gradient accumulation token-loss guide handles another coordinate problem: preserving the token-weighted objective across microbatches. Solve target alignment first, then carry each head’s loss sum and valid-cell count through accumulation. Averaging microbatch means repeats the same denominator bug at a larger scale.

A useful leak fixture concatenates two short segments with distinct token ranges. Deliberately allow a cross-segment pair and require the audit to fail. The negative case should name the horizon, source index, target index, and segment mismatch. Multi-token prediction is causal only when the data path can prove what each head was allowed to see.

Make every mask a first-class tensor

Keep three masks distinct until the final conjunction: the geometric tail mask induced by h, the padding or ignore mask, and the segment-continuity mask. Broadcasting one generic attention mask across all future-token heads is attractive, but it hides why a cell disappeared and makes empty-head failures difficult to diagnose.

For each head, emit counts for nominal cells, ignored targets, boundary crossings, and final valid cells. Their ledger must reconcile: nominal minus every mutually exclusive rejection category equals valid. If rejection reasons can overlap, define precedence and also publish an overlap count. A reconciled mask is easier to review than a lone scalar loss.

The same discipline applies when mixture-of-depths token compute routes only selected tokens through expensive blocks. A routing mask and a supervision mask answer different questions; combining them silently changes which tokens define the objective. Name each axis and reason.

This loss normalization contract keeps causal language modeling explicit even when several future-token heads share one hidden state.

Test left padding and right padding separately if both are supported. Test a sequence whose length equals the largest horizon, because it produces zero valid cells for that head. An empty head must be rejected or intentionally skipped under a documented weight-renormalization rule. Returning zero loss without a status flag makes a broken batch look healthy.

Normalize the head-loss ledger correctly

Ask cross-entropy for unreduced per-cell values, mask invalid cells, then keep lossSum and validCount as exact separate outputs. The head mean is lossSum / validCount only when validCount is positive. A batch mean divided by T, a fixed maximum length, or all heads’ combined nominal cells will underweight later horizons because they naturally have fewer targets.

PyTorch’s CrossEntropyLoss reference documents target forms, ignore indexing, weights, and reduction behavior. Verify the installed version and input layout. The reference tells you what the operator computes; your ledger still has to demonstrate that the supplied targets and denominator represent the intended objective.

Once every head has a valid mean, combine heads using explicit weights. Record both raw weight and effective contribution. If weights are normalized to sum to one, say so. If the next-token head remains weight one and auxiliary heads add extra mass, say that instead. These choices change gradient scale and should not be buried in a loop constant.

Figure two’s ledger makes the arithmetic reviewable: horizon, nominal T−h cells, rejected cells, valid denominator, loss sum, mean, and weighted contribution. Multi-token prediction experiments become comparable only when that table travels with learning-rate, batch, and optimizer settings.

Turn common mistakes into negative fixtures

Positive fixtures show that one expected path works. Negative fixtures prove that the audit notices the failures you fear. Freeze at least four. The leak fixture crosses a packed-document boundary. The off-by-one fixture shifts horizon two by one token. The wrong-denominator fixture divides its sum by T instead of validCount. The empty-head fixture requests a horizon with no valid target.

Make each mutation surgical: start from the same known-good sequence, change one rule, and require a specific error code. A test that merely produces a different floating-point value is weak because unrelated implementation changes can also move that value. A strong receipt says expected BOUNDARY_LEAK and observed BOUNDARY_LEAK at a named cell.

For numeric checks, construct logits whose per-target negative log likelihood can be recomputed independently. Compare the implementation’s unreduced values, loss sum, and mean with a small oracle. Do not use the same helper for production and test arithmetic. That would only prove that one bug is internally consistent.

The lab includes all four mutations and exports their expected and observed outcomes. It deliberately uses bounded synthetic logits, not claims about language-model quality. Multi-token prediction correctness begins with these invariants, while usefulness still requires a real training and evaluation program.

Separate the training objective from decoding

Extra future-token heads change supervision during training. They do not automatically change the deployed decoding loop. A system may discard auxiliary heads and decode one token at a time, or it may connect them to a separately designed proposal-and-verification method. Those are different product and correctness contracts.

The DeepSeek-V3 technical report describes multi-token prediction within a larger model recipe and discusses a potential role in speculative decoding. Read that evidence within its architecture, data, training, and serving context. It is not proof that adding auxiliary heads to an arbitrary checkpoint yields the same result.

Likewise, model distillation contracts can add auxiliary supervision without granting an inference-speed claim; name which heads and losses survive deployment.

Figure three draws a hard boundary. On the training side sit shifted targets, per-head masks, loss sums, weights, and gradients. On the decoding side sit proposal generation, acceptance rules, cache updates, latency, and output parity. Crossing the line requires another specification and another receipt.

This boundary also clarifies evaluation. First verify that the implemented objective matches its declaration. Then compare matched training runs for convergence and downstream quality. Only after a decoding design exists should you measure accepted tokens, speed, and exact output behavior. Multi-token prediction can be valuable without every head appearing in production inference.

Instrument the real training run

Promote the toy ledger into production telemetry without emitting raw training text. Per step or sampled interval, log each head’s loss sum, valid count, mean, configured weight, weighted contribution, empty status, and rejection counts. Aggregate sums and counts across workers before dividing so distributed means preserve the token-weighted objective.

Watch ratios, not just losses. A sudden fall in validCount may indicate a packing, padding, or sequence-length regression. Unexpected boundary rejection can reveal a data-collator change. A later head whose count exactly equals the next-token head probably failed to remove its tail. These invariants detect plumbing failures earlier than a downstream benchmark.

When targets depend on tokenizer IDs, pair the run with the tokenizer drift release checklist. A changed tokenizer, EOS policy, or document packer can alter segment and ignore semantics even if the model code is unchanged. Pin identities in the receipt.

Protect privacy by logging dimensions, counts, hashes, and aggregate arithmetic rather than token content. Retain a synthetic canary batch whose full alignment grid is safe to inspect. Run it before and after compiler, kernel, framework, and data-loader upgrades. A reproducible canary lets multi-token prediction keep a stable semantic anchor while the training stack evolves.

Ship the loss receipt before the experiment

The local lab accepts a bounded token sequence, segment IDs, ignored positions, up to four horizons, head weights, and deterministic logits. It constructs shifted targets, lists every accepted and rejected cell, computes per-head cross-entropy sums and valid denominators, and exports JSON and CSV. The same input always produces the same ordered receipt.

The lab is an alignment and arithmetic instrument. It does not reproduce published scaling results, train a network, select auxiliary weights, model distributed performance, or promise faster decoding. Those boundaries are printed beside the export so a screenshot cannot turn a unit test into a benchmark claim.

Before approving a run, require a receipt for the synthetic canary and one aggregate receipt from the actual data pipeline. Review that every horizon contains T−h nominal cells per unpadded segment, all boundary crossings are rejected, empty heads are explicit, and the combined objective can be reconstructed from head-level values.

Then attach the objective declaration to the experiment: tokenizer and packing identity, horizon set, head architecture, weights, reduction order, distributed aggregation, optimizer scale, and evaluation plan. Multi-token prediction is not difficult because the formula is long. It is difficult because several plausible tensors produce a scalar. The receipt makes only one of them the declared objective.

Runnable local artifact — The lab audits target alignment and reduction arithmetic on bounded synthetic logits; it does not reproduce a paper result, train a model, choose auxiliary weights, or predict downstream quality or decoding speed.

Enter bounded tokens, horizons, ignore positions, and logits; derive T−h targets, apply the declared mask, sum unreduced cross-entropy, divide by valid cells, run leak/off-by-one/denominator/empty-head fixtures, and export the receipt.