M7.4 - Multi-Token Prediction Foundations: Objectives, Gradients, and Prediction Heads

M7.4 - Multi-Token Prediction Foundations: Objectives, Gradients, and Prediction Heads

0x00 Overview And Background

Next-token prediction asks one hidden state to explain one future label. Multi-token prediction (MTP) asks the same causal prefix to explain several future labels.

The idea has two histories that should not be collapsed. One line uses future-token heads to reduce serial decoding work: blockwise parallel decoding trained auxiliary predictors and accepted the longest continuation validated by the base model. A second line uses future targets as a representation-learning objective: ProphetNet predicted future n-grams in sequence-to-sequence pretraining, while the 2024 Better & Faster Large Language Models via Multi-token Prediction study moved independent future heads onto a shared decoder-only trunk and tested the objective at billion-parameter scale.

Those histories meet at an interface, not an identity:

training claim:
    one causal prefix supervises several future-token marginals

inference claim:
    auxiliary predictions may become proposals for a verifier

non-claim:
    several trained heads directly emit a verified autoregressive path

The distinction matters because the training graph, the probability model, and the serving state machine answer different questions. MTP can change what the trunk learns even if every auxiliary head is deleted after pretraining. The same heads can also be retained, but then M7.1’s target-governed verification contract becomes mandatory before multiple tokens are committed.

That sounds like a small loss-function change. It is not. It changes:

training labels:
    one shifted target -> several horizon-specific targets

gradient flow:
    one loss path into the trunk -> several future-horizon paths

head architecture:
    one output path -> independent heads, recurrent heads, or modules

training memory:
    one vocabulary-logit tensor -> potentially K tensors

inference interface:
    one next-token distribution -> candidate futures that still need a contract

The owner question is therefore:

What exactly changes when a causal model predicts several future tokens during training, and what additional mechanism is required before those predictions can accelerate inference?

M7.4 derives that foundation. M7.5 owns EAGLE-style feature drafting. M7.6 owns DeepSeek’s sequential MTP modules.

The route keeps its historical deepseek-mtp suffix so existing links do not break. The article owner changed when the material was split: DeepSeek-specific module, checkpoint, teacher-forcing, and cache mechanics now belong to M7.6.

Background knowledge: auxiliary future losses change representation learning before they change decoding

Standard causal language modeling supervises each hidden row with the immediately following token. Multi-token prediction adds losses for farther future offsets. The shared trunk is then trained to retain information useful for several horizons instead of only one next-token boundary.

one-step objective:
    h_t -> label x_(t+1)

multi-token objective:
    h_t -> labels x_(t+1), x_(t+2), ..., x_(t+K)

inference question:
    are the extra heads discarded, retained as candidates, or connected causally?

Training several horizons does not by itself define a lossless decoding algorithm. Independent future heads may be mutually inconsistent because farther heads do not condition on sampled nearer tokens. Verification, a causal head chain, or another runtime contract is still required for acceleration.

0x01 Claim Status And Scope

This article does not use “predict K tokens” to mean “emit K verified output tokens.” Prediction, proposal, verification, and commitment remain separate operations.

0x01A Immutable Training Coordinate

Before deriving a loss, freeze the object being trained. Otherwise the same formula can conceal different label supports, normalization denominators, or packed-document leakage.

The second batch row is not one seven-token document:

batch 0:
    tokens   [11, 12, 13, 14, 15, PAD, PAD]
    segment  [ 0,  0,  0,  0,  0,  -1,  -1]

batch 1:
    tokens   [21, 22, 23, 31, 32, 33, 34]
    segment  [ 0,  0,  0,  1,  1,  1,  1]

For source row $(b,t)$ and horizon $k$, label ownership is valid only if:

$$m_{b,t}^{(k)} = M_{b,t} M_{b,t+k} \mathbf{1}[G_{b,t}=G_{b,t+k}].$$

Renderer-safe form:

valid[b,t,k] =
    active[b,t]
  & active[b,t+k]
  & (segment[b,t] == segment[b,t+k])

This coordinate contains 9 valid horizon-1 pairs, 6 horizon-2 pairs, and 3 horizon-3 pairs. A shape-only slice would instead expose 12, 10, and 8 rows. The difference is padding and document ownership, not a minor masking detail.

The immutable checkpoint/training manifest therefore records:

tokenizer and vocabulary identity
padding id and ignored-label id
packing and segment-id convention
K and semantic horizon ids
head architecture and parameter keys
shared-unembedding alias identity
lambda_k and loss-normalization policy
activation, logit, reduction, and gradient-accumulator dtypes
distributed vocabulary and parameter ownership

Changing any item creates a different training contract even if every tensor still has a loadable shape.

0x02 Derive MTP From The Causal Language-Model Objective

Let a sequence be:

$$x_{1:T} = (x_1, x_2, \ldots, x_T)$$

A causal trunk $f_\theta$ maps each prefix position to a hidden state:

$$h_t = f_\theta(x_{\le t}), \qquad h_t \in \mathbb{R}^{d}$$

The ordinary next-token head defines:

$$p^{(1)}(x_{t+1} \mid x_{\le t}) = \operatorname{softmax}(W_U g_1(h_t))$$

where $g_1$ may be an identity mapping, a normalization, or a richer prediction head, and $W_U \in \mathbb{R}^{V \times d}$ is the vocabulary projection.

The next-token loss over valid positions is:

$$\mathcal{L}_{NTP} = -\sum_{t=1}^{T-1} \log p^{(1)}(x_{t+1} \mid x_{\le t})$$

MTP adds prediction horizons $k=1,\ldots,K$. Head $k$ receives a representation of the same causal prefix and predicts the token $k$ positions ahead:

$$p^{(k)}(x_{t+k} \mid x_{\le t}) = \operatorname{softmax} \left( W_U g_k(h_{\le t}) \right)$$

The general weighted objective is:

$$\mathcal{L}_{MTP} = -\sum_{k=1}^{K} \lambda_k \sum_{t=1}^{T-k} \log p^{(k)}(x_{t+k} \mid x_{\le t})$$

Renderer-safe form:

L_MTP =
    - sum_{k=1..K} lambda_k
      sum_{t=1..T-k} log p^(k)(x_{t+k} | x_<=t)

0x02A Why The Upper Bound Is (T-k)

At position $t$, horizon $k$ needs label $x_{t+k}$. The label exists only when:

$$t+k \le T$$

Therefore:

$$t \le T-k$$

This is not bookkeeping trivia. If every head is averaged over the same padded sequence length without masking its invalid tail, later horizons receive a different effective loss scale and may train against padding or an incorrect label.

There are two defensible normalizations:

per-head normalization:
    average head k over its T-k valid positions,
    then combine heads with lambda_k

global-token normalization:
    sum all valid horizon losses,
    divide by the total number of valid (position, horizon) pairs

They are not numerically identical. The selected reduction belongs in the training configuration.

0x02B MTP Predicts Horizon Marginals, Not A Joint Continuation

An independent horizon-$k$ head models:

$$p^{(k)}(x_{t+k} \mid x_{\le t})$$

It does not generally model:

$$p(x_{t+k} \mid x_{\le t}, x_{t+1:t+k-1})$$

The missing intermediate tokens matter. A proper autoregressive joint continuation factors as:

$$p(x_{t+1:t+K} \mid x_{\le t}) = \prod_{k=1}^{K} p(x_{t+k} \mid x_{\le t+k-1})$$

Multiplying independent head probabilities instead gives:

$$\prod_{k=1}^{K} p^{(k)}(x_{t+k} \mid x_{\le t})$$

These products condition on different information. This is why independent MTP heads can propose mutually inconsistent tokens even when every individual head has reasonable top-token accuracy.

The distinction motivates three later design choices:

1. build a tree of candidate combinations and let the target verify paths;
2. feed sampled intermediate tokens into a recurrent drafter, as EAGLE does;
3. chain causal prediction modules, as DeepSeek MTP does.

0x02C Support Counts And Reduction Policies

For one unpadded segment of length $L$, horizon $k$ owns:

$$N_k(L)=\max(L-k,0)$$

label pairs. For a packed batch with segment lengths $L_{b,j}$:

$$N_k = \sum_b\sum_j \max(L_{b,j}-k,0).$$

The total valid support is therefore triangular even before padding:

$$N_{all} = \sum_{k=1}^{K}N_k.$$

The frozen two-row coordinate has segment lengths [5] and [3,4]:

k=1: (5-1) + (3-1) + (4-1) = 9
k=2: (5-2) + (3-2) + (4-2) = 6
k=3: (5-3) + (3-3) + (4-3) = 3
total valid pairs = 18

This count explains the upper bound $T-k$: the final $k$ source positions do not own a label $k$ steps ahead. Padding and packing can only remove more pairs; they cannot create legal support beyond that bound.

Now distinguish four reductions that are often all called “weighted MTP”:

$$\mathcal L_{sum} = \sum_{k=1}^{K}\lambda_k \sum_{b,t}m_{b,t}^{(k)}\ell_{b,t}^{(k)},$$ $$\mathcal L_{global} = \frac{\sum_k\lambda_k\sum_{b,t}m_{b,t}^{(k)}\ell_{b,t}^{(k)}} {\sum_k\sum_{b,t}m_{b,t}^{(k)}},$$ $$\mathcal L_{weighted\ pair} = \frac{\sum_k\lambda_k\sum_{b,t}m_{b,t}^{(k)}\ell_{b,t}^{(k)}} {\sum_k\lambda_k\sum_{b,t}m_{b,t}^{(k)}},$$

and:

$$\mathcal L_{head\ mean} = \sum_{k=1}^{K}\lambda_k \frac{\sum_{b,t}m_{b,t}^{(k)}\ell_{b,t}^{(k)}}{N_k}.$$

They produce different effective horizon weights. In head_mean, a distant horizon with only three rows can influence the aggregate as much as a near horizon with nine rows. In global, the near horizon contributes more simply because it owns more valid labels. Neither policy is universally correct; the checkpoint and optimizer trajectory are inseparable from the chosen denominator.

The 2024 paper writes a summed objective and describes independent heads. A reimplementation that averages tokens or batches must record its reduction rather than claiming byte-for-byte equivalence from the symbolic sum alone.

0x02D A Concrete Marginal Counterexample

Let horizon 1 predict {A,B} and horizon 2 predict {X,Y}. Consider two true continuation distributions:

Both have identical horizon marginals:

P(x[t+1]=A) = P(x[t+1]=B) = 0.5
P(x[t+2]=X) = P(x[t+2]=Y) = 0.5

Independent heads therefore cannot tell which joint generated the data. Their product assigns 0.25 to all four paths, including two impossible paths under either true joint. This is an identifiability problem, not merely low head accuracy.

Target verification can resolve proposed paths because it evaluates $p(x_{t+2}\mid x_{\le t},x_{t+1})$. Independent heads instead estimate $p^{(2)}(x_{t+2}\mid x_{\le t})$. The missing sampled token is the precise conditionality gap handed to M7.2, M7.5, or M7.6.

0x02E Conditional Information Decomposition

The MTP paper offers an interpretation for two future tokens. Restore the conditioning context $C$ that its compact derivation omits. The entropy chain rule gives:

$$H(X\mid C) = H(X\mid Y,C)+I(X;Y\mid C),$$

and symmetrically:

$$H(Y\mid C) = H(Y\mid X,C)+I(X;Y\mid C).$$

Adding them yields:

$$H(X\mid C)+H(Y\mid C) = H(X\mid Y,C) +2I(X;Y\mid C) +H(Y\mid X,C).$$

The coefficient two is algebraic. The interpretation is more limited: adding the second-horizon target can reward a representation that retains information shared by adjacent future tokens. It does not prove that every model learns a better planner, nor that an empirical cross-entropy equals ground-truth entropy before the model distribution matches the data. The paper appendix introduces a relative cross-entropy/mutual-information decomposition to make that model mismatch explicit.

0x02F Why Choice Points Receive More Terms

Source: Gloeckle et al., Figure 9, arXiv v1. The figure counts how a difficult transition and its consequences appear in several future-horizon losses. It is an explanatory model, not a measured token-importance oracle.

Under the paper’s simplified trajectory argument, a bad choice influences one future target from the nearest preceding source row, two from the next, and so on through $K$ horizons. The count is:

$$W_{choice} = 1+2+\cdots+K = \frac{K(K+1)}{2}.$$

An inconsequential transition appears once per horizon:

$$W_{local}=K.$$

The relative implicit weight is therefore:

$$\frac{W_{choice}}{W_{local}} = \frac{K+1}{2}.$$

For $K=4$, the teaching model counts 10 terms around a consequential choice and 4 around an inconsequential transition, a ratio of 2.5. This conclusion depends on the simplified assumptions: one error changes all modeled consequences, other transitions are easy, and every valid term receives equal explicit weight. Real token dependencies, masks, and $\lambda_k$ can change the count.

0x03 Tensor-Level Label Construction

Assume:

tokens: [B, S]
hidden: [B, S, d]
K:      number of prediction horizons
V:      vocabulary size

For horizon $k$, valid hidden rows and labels are:

$$H^{(k)} = H[:, :S-k, :]$$ $$Y^{(k)} = X[:, k:S]$$

Both have sequence extent (S-k).

For tokens = [11, 22, 33, 44, 55, 66] and K=3:

Equivalent slices:

head 1:
    hidden[:, 0:5] -> labels tokens[:, 1:6]

head 2:
    hidden[:, 0:4] -> labels tokens[:, 2:6]

head 3:
    hidden[:, 0:3] -> labels tokens[:, 3:6]

0x03A Padding And Packed Sequences

With padded batches, the validity mask for horizon $k$ must shift with the label:

valid_k = attention_mask[:, :S-k] & attention_mask[:, k:S]

For packed sequences, a simple slice is insufficient because a horizon may cross a document boundary. The training pipeline needs segment identifiers or cumulative sequence lengths:

same_segment = segment_id[:, :S-k] == segment_id[:, k:S]
valid_k = base_valid[:, :S-k] & base_valid[:, k:S] & same_segment

Without same_segment, a token near the end of one packed document can be trained to predict the beginning of the next document as if the two were one continuation.

0x03B Shape-Checked Objective

This early sketch is for padded, single-document (unpacked) sequences. Its active-token mask alone does not protect packed document boundaries. For packed batches, include the same_segment predicate above; the later integrated packed trace supplies that contract explicitly.

def mtp_loss(tokens, attention_mask, trunk, heads, unembed, lambdas):
    # tokens: [B, S]
    hidden = trunk(tokens, attention_mask)       # [B, S, d]
    total = hidden.new_zeros(())
    valid_pairs = 0

    for k, (head, weight) in enumerate(zip(heads, lambdas), start=1):
        source = hidden[:, :-k, :]               # [B, S-k, d]
        labels = tokens[:, k:]                    # [B, S-k]
        valid = attention_mask[:, :-k] & attention_mask[:, k:]

        head_hidden = head(source)                # [B, S-k, d]
        logits = unembed(head_hidden)             # [B, S-k, V]

        loss_sum = cross_entropy(
            logits.flatten(0, 1),
            labels.flatten(),
            reduction="none",
        ).view_as(labels)

        total = total + weight * (loss_sum * valid).sum()
        valid_pairs += valid.sum()

    return total / valid_pairs.clamp_min(1)

Implementation assertions:

assert source.shape[:2] == labels.shape
assert logits.shape[:2] == labels.shape
assert logits.shape[-1] == vocab_size
assert not horizon_crosses_packed_document_boundary(valid)

0x04 What The Shared Trunk Actually Learns

Let the loss from horizon $k$ be $\mathcal{L}_k$. The shared-trunk gradient is:

$$\nabla_\theta \mathcal{L}_{MTP} = \sum_{k=1}^{K} \lambda_k \nabla_\theta \mathcal{L}_k$$

Applying the chain rule through hidden state (H):

$$\nabla_\theta \mathcal{L}_{MTP} = \sum_{k=1}^{K} \lambda_k J_\theta(H)^\top \nabla_H \mathcal{L}_k$$

The trunk therefore receives gradient information about several future horizons at each valid position.

That does not mean “K times more independent data.”

The horizon losses:

• reuse the same input tokens;
• share the same trunk activations;
• predict correlated labels from one sequence;
• may produce aligned, redundant, or conflicting gradients.

A useful diagnostic is gradient cosine similarity:

$$\rho_{i,j} = \frac{ \langle g_i, g_j \rangle }{ \lVert g_i \rVert_2 \lVert g_j \rVert_2 }$$

where $g_k = \nabla_\theta \mathcal{L}_k$ for a selected layer or parameter group.

rho near +1:
    horizons reinforce similar updates

rho near 0:
    horizons provide different local signals

rho below 0:
    horizons conflict for that parameter group

This makes $\lambda_k$ a gradient-allocation policy, not decorative notation.

0x04A Why Distant Horizons Can Emphasize Choice Points

Suppose one token determines several later tokens. An error at that decision point affects losses at multiple horizons, while a locally interchangeable token may affect fewer future predictions. MTP can therefore weight consequential choice points more heavily through the number of future losses that depend on them.

This is an intuition, not a universal proof that every future-token loss improves reasoning. The Better & Faster MTP paper reports stronger gains at larger model sizes and on generative coding tasks, plus synthetic evidence related to induction heads. Those are source-reported experimental findings.

0x04B Gradient Instrumentation

per_horizon = []
for k in range(K):
    gk = grad(losses[k], monitored_parameters, retain_graph=True)
    per_horizon.append(flatten(gk))

for i in range(K):
    for j in range(i + 1, K):
        log_cosine(i, j, cosine(per_horizon[i], per_horizon[j]))

log_scalar("mtp/loss_by_horizon", [x.item() for x in losses])
log_scalar("mtp/token_count_by_horizon", valid_counts)

The loss should be logged by horizon. A single aggregate can hide a collapsed distant head or an incorrectly shifted label tensor.

0x04C Parameter-Local Gradient Ownership

Let $\theta$ denote shared-trunk parameters, $\phi_k$ the parameters unique to head $k$, and $U$ the shared unembedding. For:

$$\mathcal L = \sum_k\lambda_k\mathcal L_k(\theta,\phi_k,U),$$

linearity of differentiation gives:

$$\nabla_\theta\mathcal L = \sum_k\lambda_k\nabla_\theta\mathcal L_k,$$ $$\nabla_U\mathcal L = \sum_k\lambda_k\nabla_U\mathcal L_k,$$

while one private head receives only its path:

$$\nabla_{\phi_j}\mathcal L = \lambda_j\nabla_{\phi_j}\mathcal L_j.$$

This is the exact ownership reason that detach() changes the training method. If every head reads H.detach(), then:

head gradients:               still present
shared-unembedding gradient:  still present if U is trainable
trunk gradient from MTP:       exactly zero

Such a run may train competent proposal heads, but it does not train the trunk with a multi-token objective.

For a scalar reconciliation model:

$$h=\theta x, \qquad \hat y_k=U a_k h, \qquad \mathcal L_k=\frac{1}{2}(\hat y_k-y_k)^2,$$

the horizon contribution to the trunk gradient is:

$$\frac{\partial(\lambda_k\mathcal L_k)}{\partial\theta} = \lambda_k(\hat y_k-y_k)Ua_kx.$$

Summing those terms exactly matches a finite-difference derivative of the aggregate loss in the article’s executable oracle. That equality is simple but important: sequential head backward passes are valid only if they accumulate the same trunk and shared-parameter gradients.

Gradient cosine telemetry diagnoses interaction; it does not change the objective. Reweighting, clipping, projecting conflicting gradients, or detaching paths are alternative optimization policies and must be recorded as method changes, not harmless instrumentation.

0x05 Prediction-Head Design Choices

“K output heads” does not identify one architecture.

Source: Better & Faster Large Language Models via Multi-token Prediction, top panel of Figure 1. A shared trunk feeds horizon-specific heads during training. Retaining, discarding, or using those heads for self-speculation is a separate inference decision.

0x05A Linear Horizon Heads

The smallest design uses independent projections:

$$z_t^{(k)} = W_k h_t + b_k$$

This is simple, but a full $V \times d$ matrix per horizon is expensive:

$$P_{heads} = K V d$$

For $V=128{,}000$, $d=4096$, and $K=4$, four untied vocabulary projections contain about 2.10 billion parameters before biases.

0x05B Horizon Transformer Plus Shared Unembedding

The Better & Faster MTP architecture gives each horizon an independent transformer head and shares the vocabulary projection:

$$u_t^{(k)} = g_k(h_{\le t})$$ $$z_t^{(k)} = W_U u_t^{(k)}$$

Parameter ledger:

$$P_{MTP} = P_{trunk} + K P_{head\ block} + P_{W_U}$$

rather than:

$$P_{trunk} + K(P_{head\ block}+P_{W_U})$$

Sharing $W_U$ matters because the vocabulary projection is often one of the largest dense matrices in the model.

0x05C Parallel Heads Versus A Causal Head Chain

Independent heads all read the shared trunk and can run in parallel:

             -> head 1 -> x[t+1] distribution
shared trunk -> head 2 -> x[t+2] distribution
             -> head 3 -> x[t+3] distribution

A chained design instead feeds one prediction depth into the next:

trunk -> module 1 -> module 2 -> module 3

The chain can condition deeper horizons on an intermediate state and shifted tokens, but loses head-level parallelism. M7.6 derives this DeepSeek-style alternative.

The paper appendix distinguishes a narrower causal head composition:

$$p_i = \operatorname{softmax}\circ f_U\circ f_{h_i}\circ\cdots\circ f_{h_1}\circ f_s.$$

Here later head transforms consume earlier hidden transforms, not sampled token ids. It is architecturally dependent, but it still does not equal the autoregressive factor:

$$p(x_{t+i}\mid x_{\le t+i-1}).$$

That distinction prevents the word “causal” from silently claiming joint-path conditioning. DeepSeek’s modules in M7.6 add shifted ground-truth token embeddings during training and sampled-token state at inference; they are a different chain contract.

The paper also tests an anticausal composition:

$$p_i = \operatorname{softmax}\circ f_U\circ f_{h_i}\circ\cdots\circ f_{h_K}\circ f_s,$$

which begins from the farthest-head transformation and refines toward nearer horizons. It changes hidden dependency order without changing target-label ownership.

Source: Gloeckle et al., Appendix Figure S11, arXiv v1. The arrows describe activation and gradient scheduling through composed prediction heads. They do not describe autoregressive token commitment.

0x05E Replicated Versus Shared Unembedding

Replicating a vocabulary matrix for every horizon is mechanically simple:

$$P_{replicated\ U}=K V d.$$

With one shared matrix:

$$P_{shared\ U}=Vd.$$

The saved parameter count is:

$$\Delta P_U=(K-1)Vd.$$

For $K=4$, $V=128{,}000$, and $d=4096$:

one unembedding:          524,288,000 parameters
four replicated copies: 2,097,152,000 parameters
saved by sharing:        1,572,864,000 parameters

Sharing parameters does not share output values: every horizon still applies $U$ to a different head activation and materializes a different logit tensor.

0x05F Architecture Contract Table

The paper reports that linear, causal, and anticausal alternatives improved on the next-token baseline in some experiments but were not consistently better than its parallel transformer heads. Those are source-reported results, not a proof that the parallel form is universally optimal.

For compute-matched model-scaling experiments, the paper removes $K-1$ layers from the shared trunk when adding $K-1$ auxiliary transformer heads. A model that simply appends heads to an unchanged trunk has a different parameter and training-compute budget.

0x05G Medusa And EAGLE Are Different Design Points

The shared phrase “multi-token” does not make their state transitions equivalent.

0x06 Training Memory And Compute

For batch $B$, sequence length $S$, vocabulary $V$, dtype bytes $b$, and $K$ heads, materializing every full logit tensor at once costs approximately:

$$M_{logits} \approx b B V \sum_{k=1}^{K}(S-k)$$

For $B=8$, $S=4096$, $V=128{,}000$, $K=4$, and BF16 logits:

sum_k(S-k) = 4095 + 4094 + 4093 + 4092 = 16,374

M_logits ~= 2 * 8 * 128,000 * 16,374 bytes
         ~= 31.2 GiB

That is only the raw logit payload. It excludes softmax workspaces, head activations, gradients, and distributed communication.

0x06A Sequential Head Evaluation

For independent heads, finish each head’s forward and backward before starting the next. Sequential forward calls alone do not establish this memory bound:

Source: Better & Faster Large Language Models via Multi-token Prediction, Figure 2. Sequential head forward/backward evaluation avoids materializing every vocabulary-sized logit gradient at once. This is a training-memory schedule, not sequential token generation.

# Compact dense-batch sketch: 1 <= K < S, with valid rows in every horizon.
# Head-specific causal masks and the loss reduction are part of the contract.
optimizer.zero_grad(set_to_none=True)
trunk_hidden = trunk(tokens)
head_leaf = trunk_hidden.detach().requires_grad_(True)

for k, head in enumerate(heads, start=1):
    head_hidden = head(head_leaf[:, :-k])
    loss_k = vocab_parallel_cross_entropy(
        shared_unembedding,
        head_hidden,
        tokens[:, k:],
    )
    (lambda_k[k] * loss_k).backward()
    del head_hidden, loss_k

# Propagate the sum of all head contributions through the trunk exactly once.
trunk_hidden.backward(head_leaf.grad)
optimizer.step()

This ordering allows each horizon’s vocabulary-loss saved tensors to be released before the next horizon, provided backward does not retain that graph and no logging/container reference keeps its tensors alive. Optimizer gradients and the detached leaf gradient accumulate across horizons; neither is zeroed inside the loop. The final trunk vector-Jacobian product is essential: detach is a scheduling boundary here, not permission to discard the trunk gradient.

By contrast, summing graph-connected loss_k tensors and calling backward only after the loop can retain saved vocabulary-loss state from every horizon. Checkpointing the head alone does not establish release of the projection/loss state outside that checkpoint. A fused or recomputing loss needs its own explicit saved-tensor contract before claiming a smaller peak.

The bytes below describe logical live logit payload, not measured device allocation, allocator-reserved memory, or an entire training step. They do not make the heads free:

• each head still performs its transformation;
• the shared unembedding is applied once per horizon;
• the trunk graph must remain available or be recomputed;
• vocabulary-parallel loss may communicate once per horizon;
• optimizer state exists for all trainable head parameters.

The Better & Faster MTP paper reports no training-time overhead for its measured setup through a memory-efficient implementation. That result is source-reported, not a hardware-independent consequence of MTP.

0x06B Distributed Vocabulary Projection

When $W_U$ is sharded across tensor-parallel ranks:

local logits:
    [B, S-k, V / TP]

cross-entropy:
    needs distributed max and normalization reductions

repeated horizons:
    repeat or fuse that distributed work K times

An implementation claiming “only an auxiliary loss” must still account for these kernels and collectives.

0x06C Optimizer And Checkpoint State

The checkpoint must record:

number of horizons K
horizon-to-label offset convention
head architecture per horizon
whether W_U is shared or untied
lambda_k values and loss reduction
packed-sequence boundary rule
whether heads are retained for inference

Changing K or the head ordering while loading a checkpoint can silently attach the wrong semantic horizon to a weight tensor.

0x06D Full Training Tensor Trace

The immutable teaching batch compiles ragged legal rows before any head runs:

The same detach/accumulate/propagate ordering extends to this ragged-row sketch. Here head(source) is row-local; a context-bearing Transformer head must preserve full sequence/segment visibility before selecting loss rows. The sketch assumes every listed horizon has valid rows. Empty-horizon policy and distributed synchronization require an explicit implementation contract:

# One optimizer zeroing operation for the whole MTP step.
optimizer.zero_grad(set_to_none=True)

trunk_hidden = trunk(tokens, attention_mask)       # [B,S,d]

# The head graph terminates at this leaf. Its gradient will be sent through the
# trunk once, after all private head backward passes have accumulated into it.
head_leaf = trunk_hidden.detach().requires_grad_(True)

for k, head in semantic_horizon_order():
    source_index, labels = compile_valid_pairs(
        tokens=tokens,
        active=attention_mask,
        segment_id=segment_id,
        horizon=k,
    )
    source = gather_rows(head_leaf, source_index)  # [N_k,d]
    u_k = head(source)                             # [N_k,d]
    loss_k = vocab_parallel_cross_entropy(
        shared_unembedding,
        u_k,
        labels,
        reduction_manifest,
    )

    # Accumulates into head_leaf.grad, shared W_U.grad, and this head's params.
    (lambda_k[k] * loss_k).backward()

    del u_k, loss_k                                # release horizon-local state

# One vector-Jacobian product propagates the accumulated leaf gradient through
# the shared trunk graph.
trunk_hidden.backward(head_leaf.grad)
optimizer.step()

If each horizon gathers a different source-row subset, autograd scatter-adds those gathered gradients into the common [B,S,d] leaf. The trunk then receives the exact sum from every legal (batch, source position, horizon) pair.

The schedule has four non-negotiable invariants:

1. optimizer gradients are zeroed once, not between horizons;
2. head_leaf.grad accumulates rather than being overwritten;
3. shared-unembedding gradients accumulate across all heads;
4. the trunk vector-Jacobian product runs once with the complete leaf gradient.

The primary-paper figure shown earlier uses this detach/accumulate/propagate ordering. Detach is a graph-scheduling technique there; omitting the final trunk_hidden.backward(head_leaf.grad) would turn it into detached-head training and change the objective.

0x06E Liveness, Not Only Tensor Shape

At the teaching coordinate with FP32 logits:

horizon-1 logits: 9 * 128 * 4 bytes = 4,608 bytes
horizon-2 logits: 6 * 128 * 4 bytes = 3,072 bytes
horizon-3 logits: 3 * 128 * 4 bytes = 1,536 bytes

all horizons live together: 9,216 bytes
sequential peak:            4,608 bytes

The toy sizes are small; the liveness relation is the point. Sequential backward with horizon-local graph release changes this raw-logit payload estimate from a sum over horizons to a max; it does not measure allocator peak:

$$M_{parallel\ logits} = b_{logit}V\sum_k N_k,$$ $$M_{sequential\ logits} = b_{logit}V\max_k N_k.$$

It does not remove trunk activations, private head parameters, shared unembedding weights, gradients, optimizer state, or collective work.

0x06F Distributed Vocabulary Cross-Entropy Per Horizon

For vocabulary-parallel degree $P_V$, rank $r$ owns a shard $U_r\in\mathbb R^{(V/P_V)\times d}$ and computes:

$$Z_r^{(k)}=H_kU_r^\top, \qquad Z_r^{(k)}\in\mathbb R^{N_k\times(V/P_V)}.$$

An exact softmax cross-entropy needs global row statistics:

1. all-reduce MAX of local row maxima
2. all-reduce SUM of exp(local_logits - global_max)
3. owner rank contributes each target logit
4. all-reduce SUM of target-logit contributions
5. form -target_logit + global_max + log(global_exp_sum)

The sequence repeats for every horizon unless an implementation fuses multiple heads. Sharing the unembedding parameter does not collapse these distinct logit and collective computations.

0x06G Realistic Memory And Compute Ledger

For $B=8$, $S=4096$, $V=128{,}000$, $d=4096$, $K=4$, and BF16 logit storage:

valid dense rows, ignoring padding:
    8 * (4095 + 4094 + 4093 + 4092) = 130,992

all logits live:
    130,992 * 128,000 * 2 bytes = 31.23 GiB

largest horizon live:
    (8 * 4095) * 128,000 * 2 bytes = 7.81 GiB

Under a convention counting one multiply and one add as two FLOPs, the shared unembedding work alone is:

$$F_U = 2dV\sum_k N_k.$$

At this coordinate that is about 137.36 TFLOPs, or 0.137 PFLOPs, per training batch across all four horizon projections, before head-transformer, softmax, backward, and communication work. This does not imply 0.137 PFLOPs of extra wall-clock-visible work relative to a baseline: the baseline already owns one next-token head, model parallelism distributes the matrices, and the paper uses compute-matched architectures. It is the logical projection ledger.

The paper’s training table reports relative times between 1.02 and 1.22 for most multi-head configurations in its measured systems, with larger relative overhead on smaller models. Its caption attributes part of the overhead to lost FSDP communication/computation overlap under separate backward passes. The often repeated phrase “no training overhead” must therefore remain tied to the paper’s rounded system-level conclusion, not treated as an arithmetic identity.

0x06H Optimizer, Alias, And Resume Contract

For a trainable parameter count $P$ with BF16 weights/gradients and FP32 Adam moments, a simplified persistent state ledger is:

BF16 parameter:       2P bytes
BF16 gradient:        2P bytes
FP32 first moment:    4P bytes
FP32 second moment:   4P bytes
optional FP32 master: 4P bytes

Actual frameworks shard, offload, fuse, or omit some copies. The checkpoint must nevertheless preserve semantic identity:

head key -> horizon offset
head architecture and depth
shared W_U alias target
K and lambda vector
reduction denominator
packing/segment convention
optimizer slots by parameter key
scheduler and global update count
dtype and distributed-sharding metadata

Loading head-3 weights into head 2 is shape-correct and semantically wrong. Duplicating the shared unembedding on load can also create two parameter objects that start equal but receive separate optimizer states. Resume parity requires alias restoration before the optimizer binds parameter identities.

0x07 Training Details That Change The Result

0x07A Loss Weights

Uniform weights are only one option:

$$\lambda_k = 1$$

Alternatives include distance decay:

$$\lambda_k = \gamma^{k-1}, \qquad 0 < \gamma \le 1$$

or a curriculum that introduces longer horizons later. Distance decay recognizes that distant marginal predictions are harder and noisier. A curriculum changes the optimization path. Neither should be presented as equivalent to the paper’s exact setup unless reproduced.

0x07B Head Initialization

Possible initializations include:

copy the final trunk block:
    starts each head near a language-model transformation

independent random initialization:
    preserves horizon specialization but may create large early loss

copy one head into all horizons:
    identical start, then horizon-specific gradients separate them

The choice affects early gradient magnitude into the trunk. It belongs in the experiment record.

0x07C Shared Or Detached Trunk

If future losses are detached from the trunk:

head_input = hidden.detach()

the heads learn future prediction, but the main model receives no multi-horizon gradient. That is a draft-head training scheme, not integrated MTP pretraining.

The distinction is easy to miss because both versions produce future-token logits.

0x07D Teacher Forcing Boundary

Classical independent MTP heads do not consume their own sampled intermediate tokens during training. They predict horizon marginals directly from the true prefix. A sequential MTP module may consume ground-truth shifted tokens during training and sampled tokens at inference, creating an exposure gap. That belongs to M7.6 and must not be projected backward onto every MTP architecture.

0x08 From Training Objective To Inference Mechanism

There are three deployment choices.

The choice is part of the release artifact, not an automatic consequence of the checkpoint. One MTP-trained checkpoint can legitimately produce three different serving packages.

0x08A Discard Auxiliary Heads

served model:
    trunk + ordinary next-token head

effect:
    possible training-quality benefit
    no direct decode-step reduction

This is valid. An MTP-trained model does not need to serve its auxiliary heads.

“Discard” means remove horizon heads 2 through $K$ while retaining the complete next-token path: trunk, head 1, normalization, shared unembedding, tokenizer, and generation configuration. In the paper’s compute-matched architecture, head 1 is part of the trained next-token model; deleting every head would leave no language-model output path.

The stripped release should record the parent training checkpoint and the exact removed keys. This permits evaluation against the unstripped checkpoint and prevents a later loader from silently creating randomly initialized auxiliary heads.

0x08B Use Heads As Candidate Generators

At decode position $t$, each head proposes candidates for its horizon. Those candidates can be combined into paths or a tree. The target model then verifies the paths under M7.1’s acceptance contract.

The runtime state is:

committed prefix state:
    target-model KV for accepted tokens

candidate state:
    horizon logits, candidate token ids, tree metadata

verification state:
    target probabilities and temporary branch KV

commit:
    accepted target-verified prefix only

Independent head predictions are proposals, not committed state.

A typed proposal buffer needs at least:

prefix generation/version
semantic horizon id for every row
candidate token ids and proposal scores
head/checkpoint identity
candidate topology and parent indices
target tokenizer/vocabulary identity

For top-$b_k$ candidates per horizon, the number of full Cartesian paths is:

$$N_{paths}=\prod_{k=1}^{K}b_k.$$

The corresponding full prefix tree has:

$$N_{nodes} = \sum_{j=1}^{K}\prod_{k=1}^{j}b_k.$$

At widths [4,3,2], that means 24 complete paths and 40 nodes before pruning. M7.2 and M7.7 own how such candidates become an ancestry-correct sparse tree.

Per-horizon top-1 accuracy does not determine path accuracy. If the correctness events have probabilities $a_1,\ldots,a_K$, then without a dependence model:

$$\max\left(0,\sum_ka_k-K+1\right) \le P(\text{all K correct}) \le \min_k a_k.$$

For [0.9,0.8,0.7], full-path correctness can lie anywhere from 0.4 to 0.7. Multiplying the marginals gives 0.504 only under an independence assumption that the shared prefix, trunk, labels, and training objective do not justify.

0x08C Retain A Sequential Draft Module

A causal module chain or feature drafter can autoregress over candidates before target verification. This may produce a more coherent proposal path than independent horizon marginals, but introduces serial draft work and its own KV state. M7.5 and M7.6 own those designs.

0x08D Release Decision Matrix

0x09 Inference Cost Ledger

Let:

C_T:       one ordinary target decode step at the comparison coordinate
C_heads:   future-head candidate cost
C_pack:    candidate tree packing and masks
C_verify:  full target verification cost for the round, not extra overhead
C_state:   separate proposal/cache/publication work not counted above
G:         newly emitted output tokens, including any correction/bonus after stopping

For an additive, non-overlapping cost ledger at one workload coordinate, the round beats ordinary decoding when:

$$C_{heads}+C_{pack}+C_{verify}+C_{state} < \mathbb E[G]C_T.$$

Equivalently, a conceptual speedup is:

$$S = \frac{\mathbb E[G]C_T} {C_{heads}+C_{pack}+C_{verify}+C_{state}}.$$

Here $C_{verify}$ is the complete target-forward verification charge, including the baseline-equivalent work; do not add $C_T$ to that denominator again. If a profile instead reports only an incremental verification charge, first form C_verify = C_T + incremental_verify at the same comparison coordinate. Neither the increment alone nor the full charge plus another baseline is valid.

$G$ counts output actually published, not just accepted draft candidates: a correction or bonus counts only if emitted, and EOS/length stopping may suppress later tokens. $C_{state}$ excludes target-forward work already charged to verification. The terms must use one latency or service-demand denominator; mixing FLOPs with milliseconds produces no valid speedup. If stages overlap, use a measured critical-path round cost instead of this additive sum.

This ledger explains why training improvements and serving improvements can move independently:

better trunk, heads discarded:
    quality may improve; decode cost is unchanged

accurate heads, inefficient verification:
    acceptance improves; wall time may not

cheap heads, incoherent paths:
    drafting is cheap; accepted progress remains low

large batch:
    ordinary decode may already use hardware efficiently;
    candidate expansion can reduce throughput even if batch-1 latency improves

For a synthetic coordinate with C_T=10 ms, heads 2 ms, packing 1 ms, full target verification 12 ms, state publication 1 ms, and expected emitted progress 2.5 tokens:

ordinary work for 2.5 tokens: 25 ms
proposal/verification round: 2 + 1 + 12 + 1 = 16 ms
conceptual speedup:           25 / 16 = 1.5625x

same verification charge as baseline plus increment:
    10 + 2 = 12 ms
same round: 2 + 1 + (10 + 2) + 1 = 16 ms

Using only the 2 ms increment would undercount the round as 6 ms. Adding the 10 ms baseline on top of the full 12 ms verification charge would double-count it and produce 26 ms. Both disagree with the declared 16 ms round.

If expected emitted progress falls to 1.2 tokens, the same round gives 12/16 = 0.75x. This is a synthetic sensitivity example, not a benchmark. M7.8 owns measured batch, queue, SLO, release, and rollback evaluation.

0x09A Integrated Training-To-Release Oracle

The article-local oracle composes four independent checks:

label compiler:
    validates active source, active target, and matching segment

loss oracle:
    computes stable cross-entropy over exactly those rows

gradient oracle:
    reconciles summed horizon trunk gradients with finite differences

checkpoint oracle:
    hashes horizon keys, offsets, aliases, normalization, and optimizer step

It also injects same-shape failures:

The oracle is not a model-quality test. It proves that labels, losses, gradients, and persisted identities describe the same training program.

0x0A Variant Comparison

The comparison is architectural. Benchmark results must remain tied to model, batch size, sampling settings, hardware, and runtime implementation.

0x0B Failure Lab

0x0C Observability And Reproducibility

Training metrics:

loss and accuracy by horizon
valid token count by horizon
gradient norm by horizon
gradient cosine similarity between horizons
head activation memory
unembedding time and communication by horizon
trunk recomputation time

Inference metrics when heads are retained:

candidate count by horizon
candidate-tree nodes
accepted length
head time
packing/mask time
target verification time
temporary branch-KV bytes
end-to-end latency and throughput

Without both sets, an experiment cannot tell whether MTP helped representation learning, candidate accuracy, or actual serving.

0x0D Boundary To M7.5 And M7.6

M7.4 established four contracts:

1. every horizon needs a correctly shifted, bounded label tensor;
2. future losses combine through the shared-trunk gradient;
3. independent horizon distributions are not a joint autoregressive continuation;
4. faster inference requires candidate generation plus target-governed verification.

M7.5 asks how EAGLE predicts second-to-top-layer features, conditions on an advanced token sequence, trains against target features, and evolves through EAGLE-2, HASS, and EAGLE-3.

M7.6 asks how DeepSeek constructs a sequential MTP module from normalized hidden state, shifted token embedding, projection, a Transformer block, and a shared output head.

M7.7 compares those trained proposal mechanisms with retrieval, self-drafting, dependent heads, multi-candidate selectors, and hardware-aware trees. M7.8 owns continuous-batch integration, dynamic speculative depth, evaluation, release, and rollback.

0x0E Publish Gate

0xFF References

• Transformer Series A33 - DeepSeek MTP
• M7.1 - Speculative Decoding
• M7.2 - Medusa
• M7.7 - Draft Sources and Hardware-Aware Trees
• M7.8 - Speculative Serving
• Better & Faster Large Language Models via Multi-token Prediction
• Better & Faster MTP - pinned arXiv v1
• ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training
• Blockwise Parallel Decoding for Deep Autoregressive Models
• Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads