Multi-Token Prediction (MTP, where the model additionally predicts several subsequent tokens) has become a common pre-training recipe, improving downstream performance and serving as a draft head for speculative decoding. The authors precisely decompose the cross-entropy gradient of MTP into two parts: a KL self-distillation term that pulls the MTP head toward the model's own next-token distribution (useful and enhancing speculative decoding), and a residual with a low signal-to-noise ratio that interferes with learning and worsens with scale. Soft-MTP uses "straight-through gradient routing" to retain only the KL term, reusing existing cross-entropy computations, requiring no full-vocabulary logits matrix, and remaining compatible with fused cross-entropy kernels.
Likes 9 · Code Stars 6
Rationale
- Minimal changes: It requires only a one-line code change to standard MTP, adds no parameters, and the authors claim no measurable additional overhead; it supports both DeepSeek-style and Meta-style MTP.
- The authors claim that at every scale from 167M to 1.2B, next-token perplexity is lower, speculative decoding is faster, and zero-shot accuracy is higher. After fine-tuning at 1.2B, mathematical and code reasoning are stronger, and it performs no worse than directly minimizing KL divergence on most metrics.
- Reproducible: Checkpoints for four scales × three objectives (NTP, MTP, Soft-MTP), totaling 12 checkpoints, are hosted on Hugging Face. Training configurations and speculative decoding evaluation scripts are included. The tests/ directory verifies that gradients match explicit KL divergence under float64. Licensed under MIT.
Usage
- Code: git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5 (pre-tokenized FineWeb-Edu, training 31.5B tokens per configuration, requiring at least 4 shards and approximately 120 GB of disk space during assembly); the Muon optimizer requires PyTorch 2.9 or higher.
- Training: torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp (switch to hard_mtp / direct_kl / ntp for comparison).
- View existing models: AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True), using the original gpt2 tokenizer; for speculative decoding evaluation, use python -m src.eval_spec --ckpts ….
- To integrate into your own training code, refer to combined_mtp_loss() in src/models/loss.py.
Risks
- The released models are small research models under 1.2B parameters, trained only on FineWeb-Edu, and are not ready-to-use conversational models; effectiveness at larger scales has not been verified.
- The value lies primarily in the pre-training stage, requiring training from scratch to benefit; training evaluation defaults to Weights & Biases.
- The abstract does not provide specific speedup multiples for speculative decoding, and results are as reported by the authors; the README notes that wall-clock speedup depends on the GPU, and comparisons must be run on the same machine.