多 token 预测（MTP，让模型额外预测后面几个 token）已是常见的预训练配方，既提升下游表现，也能给投机解码当草稿头。作者把 MTP 的交叉熵梯度精确拆成两部分：一个把 MTP 头拉向模型自己下一词分布的 KL 自蒸馏项（有用、能提升投机解码），和一个信噪比很低、会干扰学习且随规模变严重的残差。Soft-MTP 用“直通梯度路由”只保留 KL 那一项，复用现有的交叉熵计算，不需要整词表的 logits 矩阵，兼容融合交叉熵内核。
点赞 9 · 代码星数 6
理由
- 改动极小：对标准 MTP 是一行代码，不加参数，作者称没有可测的额外开销；同时支持 DeepSeek 式和 Meta 式 MTP。
- 作者称在 167M 到 1.2B 的每个规模上，下一词困惑度更低、投机解码更快、零样本准确率更高，1.2B 微调后数学和代码推理更强，并在多数指标上不输直接最小化 KL。
- 可复现：四个规模 × 三种目标（NTP、MTP、Soft-MTP）共 12 个检查点放在抱脸，训练配置、投机解码评测脚本都在，tests/ 里能在 float64 下验证梯度与显式 KL 一致；MIT 许可。
用法
- 代码：git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5（预分词的 FineWeb-Edu，每个配置训练 31.5B token，至少 4 个分片，组装时约需 120 GB 磁盘）；Muon 优化器需要 PyTorch 2.9 以上。
- 训练：torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp（换成 hard_mtp / direct_kl / ntp 做对照）。
- 看现成模型：AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True)，分词器用原版 gpt2；投机解码评测用 python -m src.eval_spec --ckpts …。
- 要搬到自己的训练代码里，看 src/models/loss.py 里的 combined_mtp_loss()。
风险
- 放出的都是 1.2B 以下、只在 FineWeb-Edu 上训练的研究用小模型，不是能直接拿来用的对话模型；更大规模是否同样有效没有验证。
- 价值主要在预训练阶段，要自己从头训才能受益；训练评测默认接 Weights & Biases。
- 摘要没有给投机解码加速的具体倍数，效果为作者报告；README 提醒墙钟加速依赖 GPU，对比要在同一台机器上跑。