MTP Training Objective and Ablations

Section 7.3 gave us the architecture of sequential causal MTP — a stack of small modules sitting on top of the trunk, each one predicting one token further into the future, each one passing its hidden state to the next. That is the machinery. What we have not yet answered is the question every trainer asks next: what loss do we minimise? A sequential MTP stack is just an elaborate way to produce logits. Without an objective those logits are noise. The choice of loss decides what those extra heads learn, how aggressively the trunk reshapes itself to please them, and whether the whole construction earns its keep at training time and at inference time.

The promise of the MTP objective. A clean, scale-free addition to the standard next-token cross-entropy: one extra cross-entropy per depth, normalised by depth, scaled by a tiny coefficient $\lambda$ that anneals across training. Less than 1 % extra FLOPs at $D = 1$; measurable downstream gains; and the very same heads turn into a free speculative-decoding draft model at inference. The ablation table in the DeepSeek-V3 report is the receipt.

What Objective Trains a Sequential MTP Stack?

The trunk already has a perfectly good loss: next-token cross-entropy. Every position predicts the token at position $t + 1$, the prediction is scored by a softmax over the vocabulary, and the gradient flows backward into the embeddings and attention layers. This is the loss the trunk has been trained on since the first transformer paper. We are not allowed to harm it — it is what makes the model a useful language model. Whatever we add for MTP must compose with the main loss without crowding it out.

Several questions immediately appear once you commit to additional supervision on the depth-$k$ heads:

1. What does each MTP head predict? The natural answer is the token $k$ steps further into the future than the main head. So depth 1 predicts $t + 2$, depth 2 predicts $t + 3$, etc. The target is just the existing token sequence shifted by an extra $k$.
2. How is each head scored? Same as the main head: cross-entropy against the true next-token-at-distance-$k$. One softmax per head, one scalar per head.
3. How much weight do these heads get? Too much and the trunk warps itself to please the MTP heads at the expense of plain next-token loss; too little and the heads do not learn anything useful. We need a coefficient.
4. How do the heads interact across depths? Should depth 4 count four times as much as depth 1 (because there are more of them) or the same (because we care about the average)? This is what the $1/D$ normalisation answers.
5. Does the weight change over training? Early on, future-token prediction may push the trunk toward better long-range representations. Late in training, the trunk has converged and any extra signal can only steal capacity. A schedule on $\lambda$ handles this.

The DeepSeek-V3 paper picks the simplest answer to all five questions simultaneously, and then runs ablations to prove the choices were the right ones. The rest of this section is that loss, that schedule, and that ablation table — and what they imply for any team training a frontier model from scratch.

Intuition: A Stack of Teachers, One Optimizer Budget

Think of the trunk as a single student. The main loss is the headline exam: "given this sentence so far, what comes next?" The MTP heads are extra tutors standing behind the student, each asking a harder version of the same question. Tutor 1 asks "what comes two tokens from now?" Tutor 2 asks "what comes three tokens from now?" The tutors do not give the student new material — they only re-grade the same content from further away.

Two design decisions follow from this picture. First, the tutors share the student's budget. If you bring in twelve tutors, each one's feedback should count less individually, or the student spends all their effort pleasing tutors and forgets the headline exam. That is exactly what the $1/D$ term does. Second, you want the tutors loud at the start of training (when the student is still building intuition for the subject) and quiet at the end (when the student is polishing the headline exam answers). That is exactly what the $\lambda$ schedule does.

The Math: Per-Depth Cross-Entropy with a Shared Budget

Fix a sequence of length $T$ with token ids $t_1, t_2, \dots, t_T$ and vocabulary size $V$. The trunk produces hidden states $h_1, \dots, h_T$; the main head turns each $h_i$ into a softmax distribution over the vocabulary. The main loss is the textbook next-token cross-entropy averaged over the sequence:

$\mathcal{L}_{\text{main}} = -\frac{1}{T-1} \sum_{i=1}^{T-1} \log p_{\text{main}}(t_{i+1} \mid t_{\le i})$

For depth $k \in \{1, \dots, D\}$ the sequential MTP architecture produces another set of distributions $p_k(\cdot \mid t_{\le i + k - 1}, h_i^k)$ — see section 7.3 for how $h_i^k$ is computed. The depth-k cross-entropy is the same shape but its target is shifted $k$ tokens further:

$\mathcal{L}_{\text{MTP}}^{k} = -\frac{1}{T - k} \sum_{i=1}^{T-k} \log p_k(t_{i+k} \mid t_{\le i+k-1}, h_i^k)$

Notice the upper limit $T - k$: depth-k predictions for positions in the last $k$ tokens have no ground truth to compare against, so we drop them. At $T = 4096$ and $D = 4$ this costs us four positions out of four thousand — negligible.

The full training objective is the main loss plus a single global term that averages the MTP losses across depths and scales the whole thing by $\lambda$:

$\boxed{\mathcal{L} \;=\; \mathcal{L}_{\text{main}} \;+\; \frac{\lambda}{D}\,\sum_{k=1}^{D} \mathcal{L}_{\text{MTP}}^{k}}$

Three properties of this formula are worth slowing down on.

• Depth-normalised. Whether $D = 1$ or $D = 4$, the total auxiliary gradient magnitude is bounded by $\lambda$ times the per-head loss. You can compare ablations across different depths without re-tuning $\lambda$.
• Decoupled from the architecture. The loss does not care how $p_k$ is computed — parallel heads, sequential heads, shared trunk, separate trunk. It is purely a statement about which token each head is graded against. This is why the same loss formula is reused in the DeepSeek-V3 paper and in earlier MTP work like Gloeckle et al. 2024.
• Recovers the baseline. Setting $\lambda = 0$ reduces the formula to plain next-token training. So "MTP off" is literally one knob away — no architectural surgery needed for an ablation.

The schedule for $\lambda$ in DeepSeek-V3 is a single step function:

$\lambda(\text{step}) = \begin{cases} 0.3 & \text{tokens trained} < 10\,\text{T} \\ 0.1 & \text{otherwise} \end{cases}$

That is, for the first ~67 % of pretraining tokens the MTP heads carry their full weight; after that the coefficient drops by 3×. Earlier MTP work (Gloeckle et al.) reported similar findings with simpler constant $\lambda$ — DeepSeek's schedule is a small refinement that buys a little extra final quality on the main loss without sacrificing the speculative-decoding gains we will see in section 7.5.

Manual Numerical Walkthrough

Visualizing the Loss Budget and λ-Schedule

Slide through the controls below. Increase $D$ from 1 to 4 and watch the per-head raw losses climb (deeper heads see harder problems) while the per-head weighted contributions stay bounded — that is the $1/D$ term at work. Then toggle between the constant-0.3, constant-0.1, and DeepSeek-V3 anneal schedules and slide the training-progress bar; the orange curve on the right shows $\lambda$ over training, and the blue marker locks the loss budget on the left to the current $\lambda$.

Two patterns to confirm with the slider. First: switching the schedule from constant 0.3 to 0.3 → 0.1 anneal changes nothing for the first 67 % of training — the loss budget is identical. The anneal's effect lives entirely in the last third. Second: setting λ = 0 (MTP off) with any $D$ recovers a single sky-blue bar — the standard next-token loss. The architecture stays in the forward pass, but the gradient lanes to the MTP heads are silent. That is exactly the configuration of an "MTP-off" ablation: same compute graph, zero auxiliary supervision.

Plain Python: One Sequence by Hand

Before the PyTorch version, here is the same loss written as a tight NumPy loop — no autograd, no batching, just the math on a single toy sequence. The point of showing this first is to see that the "MTP loss" is nothing exotic. It is the standard cross-entropy applied to shifted targets, summed across heads, averaged with $1/D$, and weighted by $\lambda$. Every loop iteration corresponds directly to one term of the formula above.

1import numpy as np
2
3# Toy: T = 6 tokens, V = 5 vocab entries, D = 2 MTP heads.
4T, V, D = 6, 5, 2
5lam = 0.3            # global weight on the auxiliary loss
6target = np.array([1, 4, 0, 2, 3, 1])   # ground-truth next tokens
7
8# Pretend logits for the main head and 2 MTP heads.
9# Shape: (D+1, T, V). Row 0 is the main head; row k>=1 is MTP depth k.
10rng = np.random.default_rng(0)
11logits = rng.normal(0, 1, size=(D + 1, T, V))
12
13def softmax(z):
14    z = z - z.max(axis=-1, keepdims=True)
15    e = np.exp(z)
16    return e / e.sum(axis=-1, keepdims=True)
17
18# Per-head, per-position cross-entropy loss.
19losses = []
20for k in range(D + 1):
21    # Depth k predicts token at position t + k. Shift target left by k.
22    # Drop the last k positions because there is no ground truth past T-1.
23    if k == 0:
24        tgt = target                # main head: predict t+1 ... but we keep T positions for clarity
25        shifted_logits = logits[k]
26    else:
27        tgt = target[k:]            # (T - k,)
28        shifted_logits = logits[k, : T - k]   # (T - k, V)
29
30    probs = softmax(shifted_logits)           # (T - k, V)
31    # gather the prob of the correct token at each position
32    correct = probs[np.arange(len(tgt)), tgt] # (T - k,)
33    nll = -np.log(correct).mean()             # scalar
34    losses.append(nll)
35
36L_main = losses[0]
37L_mtp = sum(losses[1:]) / D                   # average across depths
38L_total = L_main + lam * L_mtp
39
40print(f"L_main = {L_main:.4f}")
41for k in range(1, D + 1):
42    print(f"L_MTP^{k} = {losses[k]:.4f}")
43print(f"L_total = {L_total:.4f}  (λ = {lam}, D = {D})")

Output (with the default rng seed of 0): $\mathcal{L}_{\text{main}} \approx 2.46$, $\mathcal{L}_{\text{MTP}}^{1} \approx 3.05$, $\mathcal{L}_{\text{MTP}}^{2} \approx 3.23$, and $\mathcal{L}_{\text{total}} \approx 3.40$. Random logits give very high cross-entropy values, but the relative magnitudes are exactly what the schedule expects: each successive depth is slightly harder, and the auxiliary term contributes about $\lambda = 0.3$ times its mean.

PyTorch: The Production MTP Loss

The production version is the same formula wrapped in an $\text{nn.Module}$ so it composes cleanly with the trainer, batches across $B$ sequences, masks padding via $\text{ignore\_index}$, and owns the $\lambda$ schedule. The line count is short — what makes it production-worthy is the careful target shifting per depth and the per-depth loss returned alongside the total, both of which the ablation table absolutely requires.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class MTPObjective(nn.Module):
6    """
7    DeepSeek-V3-style multi-token-prediction loss.
8
9      L = L_main + (lam / D) * sum_{k=1..D} L_MTP^k
10
11    Each L_MTP^k is the cross-entropy of the depth-k head's logits
12    against the targets shifted by k positions. Padding tokens are
13    masked out so they do not contribute to the mean.
14    """
15
16    def __init__(self, depth: int, lam_start: float = 0.3, lam_end: float = 0.1,
17                 anneal_frac: float = 0.67, pad_id: int = -100):
18        super().__init__()
19        self.D = depth
20        self.lam_start = lam_start
21        self.lam_end = lam_end
22        self.anneal_frac = anneal_frac
23        self.pad_id = pad_id
24
25    def lam(self, progress: float) -> float:
26        return self.lam_start if progress < self.anneal_frac else self.lam_end
27
28    def forward(
29        self,
30        main_logits: torch.Tensor,         # (B, T, V)
31        mtp_logits: list[torch.Tensor],    # D tensors of shape (B, T, V)
32        targets: torch.Tensor,             # (B, T) — full ground-truth token ids
33        progress: float,                   # fraction of total tokens trained so far
34    ) -> dict:
35        B, T, V = main_logits.shape
36
37        # --- Main head ---
38        L_main = F.cross_entropy(
39            main_logits.reshape(-1, V),
40            targets.reshape(-1),
41            ignore_index=self.pad_id,
42        )
43
44        # --- MTP heads ---
45        mtp_terms = []
46        for k, logits_k in enumerate(mtp_logits, start=1):
47            # Depth-k head predicts targets[:, k:]; drop the last k positions.
48            tgt_k = targets[:, k:]                      # (B, T - k)
49            logits_k = logits_k[:, : T - k, :]          # (B, T - k, V)
50            L_k = F.cross_entropy(
51                logits_k.reshape(-1, V),
52                tgt_k.reshape(-1),
53                ignore_index=self.pad_id,
54            )
55            mtp_terms.append(L_k)
56
57        L_mtp = torch.stack(mtp_terms).mean() if mtp_terms else main_logits.new_zeros(())
58        lam = self.lam(progress)
59        L_total = L_main + lam * L_mtp
60
61        return {
62            "loss": L_total,
63            "L_main": L_main.detach(),
64            "L_mtp_per_depth": [t.detach() for t in mtp_terms],
65            "lam": lam,
66        }

Ablations: What DeepSeek Actually Measured

Three families of ablation appear in the DeepSeek-V3 technical report and the related MTP literature. Each one isolates a single design choice we just argued for and shows what happens when it is removed or perturbed. The numbers below track the published values; absolute magnitudes vary by model size and dataset, but the direction of every comparison is consistent across reported runs.

Ablation 1 — MTP on vs MTP off

The headline comparison: train two identical models for the same token budget, same data, same schedule. One uses $\lambda = 0$ (vanilla next-token); the other uses the DeepSeek-V3 schedule with $D = 1$.

The cost is under one percent extra compute; the downstream-quality gains are small in absolute terms but consistent across benchmarks. The real payoff is the last row: those same MTP heads, trained free as a side effect, become a built-in draft model for speculative decoding at inference. That part of the story is section 7.5.

Ablation 2 — Depth $D$

How many MTP heads is the right number? More heads predict further into the future, which is intrinsically harder, so per-head losses rise. But each head adds compute and parameters in the forward pass and complicates the inference draft chain. The DeepSeek-V3 report explored $D \in \{1, 2, 4\}$:

DeepSeek-V3 ships with $D = 1$. The reasoning is the per-FLOP curve: the second MTP head buys some extra speculative-decoding speed but offers diminishing downstream gains, and each additional head makes the inference verifier chain longer (which we will see in section 7.5 hurts acceptance in the worst case). Earlier MTP work (Gloeckle et al. 2024) reports similar diminishing-returns behaviour with their parallel-heads variant.

Ablation 3 — λ schedule

Holding $D = 1$ fixed, what does the auxiliary weight actually do?

Two clean signals come out of this table. First, "a little λ is much better than none" — even $\lambda = 0.1$ constant beats the baseline. Second, "turn it down later" — the anneal squeezes out an extra hair of main-loss quality by letting the trunk focus on the headline exam once it has matured. The third signal, less clean but very real: $\lambda \ge 1$ is harmful. The tutors should never be louder than the headline exam.

One-line takeaway from the ablations. The right configuration isboringly small — $D = 1$, $\lambda$ below 0.3, schedule annealed to 0.1. Larger choices give diminishing returns and eventually hurt. MTP is a quiet, well-behaved auxiliary — not a second loss function fighting the first.

What Changes at Massive Scale

At toy size — the 6-token NumPy example above — the MTP objective is a couple of extra cross-entropies. At 671 B parameters, 14.8 T tokens, and thousands of GPUs, a few things have to be handled with care.

Compute and memory cost is exactly D times one head's cost

Each MTP module is a transformer block plus a small linear and an embedding projection — by far dominated by the attention/FFN compute. At $D = 1$ it adds one extra block-worth of FLOPs per forward pass: roughly $1/L$ of the trunk cost where $L$ is the number of trunk layers. For $L = 60$ that is the ~0.9 % overhead quoted in the ablations table. The backward pass adds proportionally. Activation memory for the extra block is reclaimed by the same checkpointing pass used for trunk layers; no new accounting is required.

The shared embedding and output head

DeepSeek-V3's MTP modules share the trunk's token embedding and output projection — only the transformer block and the small projection layer between them are unique per depth. This decision shapes the loss in a subtle way: the embedding table and output head receive both the main gradient and the MTP gradients (scaled by $\lambda/D$). In FP32, this is a simple add; in mixed-precision training (next chapter) it requires that the accumulator for these shared parameters stays in FP32 even though the heads themselves can be FP8. We will return to this in chapter 10.

Gradient noise and convergence

Adding the MTP loss makes the overall gradient slightly higher-variance, because each MTP head is supervising fewer positions (the depth-k head supervises $T - k$, not $T$). At $T = 4096$ this is a 0.025 % effect at $D = 1$ and entirely lost in the noise. The optimizer hyperparameters (Adam β, weight decay, LR schedule) do not need to be re-tuned when MTP is added — DeepSeek-V3 reports using exactly the same trainer config as an internal MTP-off baseline.

Pipeline parallelism placement

Where in the pipeline do the MTP modules live? They are downstream of the final trunk layer, which means they belong on the same pipeline rank as the final output head. In chapter 11 we will see that DualPipe puts the output head and MTP modules on the same micro-batch boundary so that the auxiliary backward pass overlaps with the main backward pass and the all-reduce of the embedding gradient — no new bubble is introduced. This is one of the quiet reasons the 0.9 % FLOPs overhead does not translate into a 0.9 % wall-clock overhead: most of it hides inside a bubble that was already there.

Engineering Reality and Gotchas

• Target shifting is the most common bug. Off-by-one errors in the depth-k target shift produce a head that learns the identity (copying the input) and reports a suspiciously low loss. Always assert $\text{tgt\_k.shape}[-1] = T - k$ and log per-depth losses from step 0; if depth-1 loss is lower than depth-0 loss, you have shifted the wrong direction.
• Padding interacts with the shift. $\text{ignore\_index}$ works on the flattened target, but if your sequence packing puts a padding token at position $T - k$, the depth-k head can still see it. The safest pattern is to apply the shift to both targets and an $\text{attention\_mask}$, then $\text{ignore}$ any position the mask marks as padded.
• Loss accumulation precision. Even when the forward pass is BF16 or FP8, the cross-entropy reduction (the mean over T·B positions) should accumulate in FP32. PyTorch's $\text{F.cross\_entropy}$ does this internally for non-FP8 dtypes; with explicit FP8 you must cast logits to BF16 before the softmax. The few extra bits matter when $\lambda/D$ is already shrinking the result by $10\times$.
• The λ schedule and the LR schedule are independent. Do not tie the $\lambda$ cutover to the LR decay. The optimizer schedule reacts to validation curves and warmup behaviour; the $\lambda$ cutover is a curriculum decision. Coupling them creates an opaque hyperparameter that no ablation can isolate.
• Disable MTP for eval. When you compute validation loss, you want to report the main next-token loss only — that is the comparable metric across runs and against the literature. Passing a separate $\lambda_{\text{eval}} = 0$ to the loss module (or just reading $L_{\text{main}}$ out of the return dict) keeps comparisons honest.
• The MTP heads are kept for inference. A naive read of "auxiliary loss" suggests the heads are thrown away after training. They are not — section 7.5 turns them into a speculative-decoding draft model and recovers a 1.8–2.0× decode speedup. The training objective above is what makes them good draft models. The fact that the same module serves both the training objective and the inference acceleration is the deepest reason MTP is worth its modest compute overhead.

End of the MTP training story. Architecture (7.3) gave us a sequential causal stack. This section gave us the loss that trains it: a depth-normalised, anneal-weighted sum of per-depth cross-entropies. Ablations confirmed the choices were not arbitrary — they sit on the corner of a diminishing-returns frontier. The next section, 7.5, cashes in the second dividend: those same heads, at inference time, become a tightly-coupled draft model and unlock speculative decoding without a second model to maintain.