JaydenTeoh/hard-mtp-167m-model

MTP — 167m

Standard multi-token prediction. 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/167m_fineweb.yaml.

Objective

Each MTP head is trained with cross-entropy against the realized token k ahead. This is the usual MTP objective, as in DeepSeek-V3.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("JaydenTeoh/hard-mtp-167m-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

hidden size x layers	512 x 36
attention heads	8
context length	2048
MTP heads	3 (DeepSeek-style cascade, 1 block each)
parameters	178M total
deployed (NTP path)	178M

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

data	FineWeb-Edu, GPT-2 BPE
tokens	31.5B (60,000 steps x 256 x 2048)
optimizer	Muon on 2D hidden matrices + AdamW on embeddings/head/norms
peak LR	0.002, cosine to 1e-05, 2,000 warmup steps
weight decay	0.1
precision	bf16 autocast over fp32 master weights
seed	13

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},
}