Multi-token prediction heads

Independent linear-head experiment

Read exact columns by scrolling the table horizontally.

Project one shared state into several horizons

This direct arithmetic inspector uses one editable two-number hidden state and up to six independent linear output heads. Each head has a three-row, two-column weight matrix, so each output is a distribution over token IDs zero, one and two. The head at offset +1 predicts the immediate next token; +2 predicts the token two positions ahead from the same shared state. These heads do not consume one another’s predictions. The probabilities come from actual matrix-vector multiplication followed by softmax, rather than a decorative head-count diagram. The weights are small teaching fixtures, not trained parameters or measured language-model results. Results remain available indefinitely; there is no artificial training or playback clock.

Change each input deliberately

Heads selects the first one to six matrices. Hidden state accepts a JSON array of exactly two finite numbers. Weights accepts all six matrices, each with three rows of two numbers; the advanced disclosure lets you edit and inspect them. Rows correspond to token IDs in ascending order. Entries and hidden coordinates must stay between minus one hundred and one hundred. Targets is a JSON list of one to six future token IDs; null explicitly masks a target, and offsets beyond the list are also masked. A target mask changes supervision, not the predicted distribution. Auxiliary weight sets the contribution of offsets +2 and later between zero and ten. Reset defaults restores every input, including the full matrix fixture.

Reproduce predictions and loss

At defaults, the hidden state is [1,0]. The first head produces logits [1,0,-1], probabilities approximately [0.665241,0.244728,0.090031], and prediction zero. The second head rotates those scores and predicts one; the third predicts two. Default targets match these predictions, so the weighted mean loss is approximately 0.407606 nats at any head count. Set the hidden state to [0,0]: every logit becomes zero, probabilities become one third, and loss becomes log(3), approximately 1.098612. Exact ties choose the lowest token ID for the displayed prediction. This tie convention changes only the argmax label, not the distribution or loss.

Separate masks, weights and objective

For each unmasked head, cross-entropy is minus the natural logarithm of its target probability. The first head has weight one and auxiliary heads have the selected auxiliary weight. This tool divides the weighted sum of valid losses by the sum of their weights. This explicitly chosen weighted mean keeps the objective scale interpretable; other training systems may sum or normalize their losses differently. Try two heads, targets [0,0] and auxiliary weight one: the first loss is about 0.407606, the second about 2.407606, and their mean about 1.407606. Set auxiliary weight to zero and the objective returns to 0.407606. Set targets to [null,null] and no objective is defined. No target is silently replaced with token zero.

Inspect actual gradients and arithmetic cost

Full precision calculation includes logits, probabilities, losses, normalized weights, weight gradients and the accumulated hidden-state gradient. Each logit derivative is its normalized weight times probability minus the one-hot target. Each matrix derivative multiplies that logit derivative by the hidden coordinate; the shared-state derivative sums the transposed matrix contributions across heads. Masked and zero-weight heads contribute zero. These are derivatives at the entered values, not a simulated training history. Log-sum-exp computes finite losses even when an extremely small probability rounds or underflows to zero. The projection count is six scalar multiply-accumulate pairs per head. It excludes the transformer body, softmax, gradient evaluation and memory movement, and therefore cannot predict latency or hardware speedup.

Keep architecture and inference claims bounded

Independent heads on a shared trunk are one multi-token prediction design. This linear toy omits the nonlinear head layers, full sequence training and vocabulary of a production model. DeepSeek-V3 describes a different sequential MTP construction; this page is not an implementation of that architecture. Multiple guesses do not automatically form a correct joint continuation. Inference acceleration requires a specified drafting and verification procedure, with additional computation and potentially rejected tokens. This inspector implements neither speculative decoding nor a throughput benchmark, and does not establish that extra heads improve accuracy. Use the numerical table as the text alternative to the probability meters, scroll its columns on a phone, and record assumptions with the shared notes and baseline controls.

Primary architecture sources

Gloeckle et al.: independent multi-token heads; DeepSeek-V3 technical report: sequential MTP modules.