Multi-Token Prediction · LLM Glossary

必须分开的三层

1. 训练目标

普通 next-token loss 让位置 $t$ 的表示预测 $u_{t+1}$ 。MTP 增加未来 offsets，例如 $k=2,3$ ：

\mathcal L_k=-\sum_t\log p_k(u_{t+k}\mid h_t,\text{auxiliary context})

\mathcal L=\mathcal L_1+\lambda_2\mathcal L_2+\lambda_3\mathcal L_3

具体架构可能使用独立 output heads、sequential MTP modules、共享 embedding/output projection，或把前一个未来 token embedding 融入下一个辅助 block。仅有“多个 heads”不是统一 MTP 架构定义。

2. 候选生成模块

推理时可以用这些 heads/blocks 从当前 target hidden state 提出多个未来 candidates。候选可能是单链、多个 branches 或 tree；它们的 logits 是 proposal distribution/score，不因来自同一模型训练就自动成为 target 在对应自回归 prefix 下的精确分布。

3. Target verification 与接受协议

候选必须交给权威 target model 验证。Greedy 可接受与 target argmax 匹配的连续前缀；随机 exact sampling 需采用投机解码的 $p/q$ 接受与校正。训练时同时预测 $t+1,t+2,t+3$ ，不等于推理时一次无条件提交三个 token。

shape 与一个 worked example

设 hidden states $H\in\mathbb{R}^{B\times T\times d}$ ，词表 $V$ ，3 个 future heads：

Z^{(k)}=HW^{(k)}\in\mathbb{R}^{B\times T\times V},\qquad k=1,2,3

计算 $\mathcal L_k$ 时只有前 $T-k$ 个 positions 有 target，labels 是 token_ids[:, k:]。因此不能让所有 heads 都与同一 token_ids[:, 1:] 对齐。

推理例：proposal 输出 [1,4,3]，target 的 greedy 路径是 0→1→2→3。Verifier 接受候选 1，在第二位置发现 target 应为 2 而不是 4，于是提交 [1,2] 并丢弃后缀 3。这展示候选数与接受数不同。

可运行三层实验

import platform,torch
import torch.nn.functional as F

torch.manual_seed(113)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float32

def shifted_cross_entropy(logits,token_ids,offset):
    usable=token_ids.shape[1]-offset
    return F.cross_entropy(logits[:,:usable].reshape(-1,logits.shape[-1]),
                           token_ids[:,offset:].reshape(-1))

def greedy_verify(current,candidates,target_logits):
    accepted=[]; state=current
    for candidate in candidates:
        target_token=target_logits[state].argmax().item()
        if candidate!=target_token: return accepted,target_token
        accepted.append(candidate); state=candidate
    return accepted,target_logits[state].argmax().item()

def main():
    b,t,d,v,k=2,7,8,5,3
    token_ids=torch.randint(0,v,(b,t),device=DEVICE)
    hidden=torch.randn(b,t,d,device=DEVICE,dtype=DTYPE)
    heads=[torch.randn(d,v,device=DEVICE,dtype=DTYPE) for _ in range(k)]
    logits=[hidden@head for head in heads]
    losses=[shifted_cross_entropy(logits[i-1],token_ids,i) for i in range(1,k+1)]
    total=losses[0]+.3*losses[1]+.1*losses[2]
    assert torch.isfinite(total)

    proposal_logits=torch.tensor([[0.,5.,1.,0.,-1.],
                                  [0.,1.,2.,0.,5.],
                                  [0.,0.,1.,5.,2.]],device=DEVICE)
    candidates=proposal_logits.argmax(-1).tolist()  # [1,4,3]
    target_logits=torch.tensor([[0.,5.,1.,0.,0.],
                                [0.,0.,5.,1.,0.],
                                [0.,0.,0.,5.,1.],
                                [0.,0.,0.,0.,5.],
                                [5.,0.,0.,0.,0.]],device=DEVICE)
    accepted,replacement=greedy_verify(0,candidates,target_logits)
    assert candidates==[1,4,3] and accepted==[1] and replacement==2
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} hidden={tuple(hidden.shape)} vocab={v}")
    print(f"loss_t+1={losses[0].item():.6f} loss_t+2={losses[1].item():.6f} loss_t+3={losses[2].item():.6f}")
    print(f"weighted_mtp_loss={total.item():.6f}")
    print(f"proposal_candidates={candidates}")
    print(f"target_accepted_prefix={accepted} target_replacement={replacement}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/mtp_reference.py。实际输出未来 1/2/3 token losses 为 2.650753 / 3.766954 / 3.823235，候选 [1,4,3]，target 只接受 [1] 并以 2 替换拒绝位置。

教学 loss heads 没有复现 DeepSeek 的 sequential MTP modules、共享参数或训练优化；greedy verifier 也不代表 exact stochastic sampling。它证明的是 offset label 对齐和三层边界。

成本与收益条件

训练成本包括 auxiliary modules/heads 的参数、激活、额外 logits projection 与交叉熵；推理候选生成增加计算与临时 KV。收益只有在候选质量高、verification 能并行、target 串行步数下降且 proposal/rollback 开销可控时成立。

需要分别报告：各 offset loss/accuracy、proposal branch/length、accepted tokens、target calls/token、verification shape、端到端 TPOT/吞吐。仅报告“预测 3 个 token”不能说明接受了 3 个，也不能说明 3× 加速。

上下游关系

Logits 与基础采样定义每个 head 的 vocabulary distribution 与采样边界。
投机解码提供 target verifier 与 exact acceptance；MTP 可作为 proposal 来源，但不是协议本身。
Tree/Self Speculation 总览说明多分支候选的 parent mask、tentative KV 与 rollback；MTP heads 可供给候选，但仍不定义 tree commit 协议。
KV Cache / PagedAttention要管理候选临时状态、接受提交和拒绝回滚。
Prefill/Decode是普通串行基线，MTP-assisted speculation 尝试减少 target 串行 decode calls。

参考资料

DeepSeek-V3 Technical Report — sequential Multi-Token Prediction modules 与推理用途。
DeepSeek-V2 — MLA/MoE 上下文；不要把 MLA 与 MTP 混成同一机制。