Soft-MTP — 332m

Multi-token prediction against the model's own next-token verifier. One of twelve checkpoints released with Multi-Token Prediction Done Right (arXiv:2610.38331): three objectives x four sizes, trained under an identical recipe so they are directly comparable.

Code: JaydenTeoh/Soft-MTP — configs/332m_fineweb.yaml.

Objective

Each MTP head is trained against the model's own next-token distribution q rather than the realized token e_y, via a straight-through KL-recovery coupling that costs one backward pass and no vocabulary softmax. See the paper for the gradient decomposition this removes.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-332m-model", trust_remote_code=True)
tok = AutoTokenizer.from_pretrained("gpt2")          # the corpus tokenizer

ids = tok("The capital of France is", return_tensors="pt").input_ids
print(tok.decode(model.generate(ids, max_new_tokens=20)[0]))

trust_remote_code=True is required: the architecture (model_type: hf_mtp) is defined by the modelling files in this repo, not by transformers.

The tokenizer is stock GPT-2 — this repo ships no tokenizer files. vocab_size is padded from 50257 to 50304 for alignment; the extra rows are unused.

Model

MTP heads are auxiliary: they are discarded at inference, so what you deploy is an ordinary next-token transformer. They can also serve as a draft model for self-speculative decoding.

Training

Single run per cell, one seed. The twelve checkpoints share seed, data order and schedule, so differences between them are attributable to the objective.

Citation

@article{teoh2026multitokenpredictionright,
  title={Multi-Token Prediction Done Right},
  author={Jayden Teoh},
  journal={arXiv preprint arXiv:2610.38331},
  year={2026},
}