MTP head for speculative decoding

#2
by burkap - opened

Hello, thanks for sharing this model.

Qwen3.5 includes an MTP head that enables speculative decoding in vLLM, but your repo doesn’t have it.
I first tried to use the original Qwen3.5-2B MTP head with the model as is, and it gave about a 22% speedup. Then I retrained the head on the model’s own outputs (self-distillation). The result is about a 57% speedup on an RTX 3090 Ti (156 -> 245 tok/s, 1.57x).

Training data: DocLayNet, PubTabNet, IAM, SROIE, invoice and Turkish OCR images, plus your turkish-synthetic-documents set. No TOB pages were used for training.

The main model still decides every token, so there is no quality loss; the TOB score is the same without a head, with the original head, and with the new head (89.2). It only needs one extra 122 MB file.

If you’d like to take a look, here is the link: https://huggingface.co/burkap/MoganOCR-TR-MTP

If you want, I can also open this as a PR to your repo.

Mogan AI org

Hi, thanks a lot for this, and for showing the MTP head was worth adding.

Your post prompted us to train our own head and add it to the main repo (mtp.safetensors + model.safetensors.index.json; model.safetensors is unchanged). We trained it with self-distillation on 50k pages of the model's original training data. The main model is frozen, and the head learns its next-token distributions over 3 chained draft steps (782 steps on an RTX PRO 6000). On our eval set, the first draft token matches the main model 96.1% of the time, the second 94.1% and the third 92.9%.

We measured speed with your setup: pages read one at a time, 50 TurkishOCR-Bench pages, vLLM 0.26, RTX PRO 6000.

Setup tok/s Speedup
No MTP 265 1.00×
Original Qwen3.5-2B head, 3 drafts 310 1.17×
Your head, 3 drafts 407 1.53×
Our head, 2 drafts 398 1.50×
Our head, 3 drafts 458 1.73×
With 64 concurrent requests over the full 488-page benchmark, the speedup is 1.22×.

These numbers are from our GPU. It would be great to see how our head compares with yours on your 3090 Ti.

You're welcome, and thanks for training a proper head and adding it to the main repo so quickly.

Here is the comparison on RTX 3090 Ti: 49 TurkishOCR-Bench pages (every 10th page, so all categories are covered), vLLM 0.33.
Results are slower than usual because I have some other stuff running on the machine, but speedup ranking should be relevant in any case.

Setup tok/s Speedup
No MTP 139 1.00×
Original Qwen3.5-2B head, 3 drafts 151 1.08×
My head, 3 drafts 197 1.42×
Your head, 2 drafts 203 1.46×
Your head, 3 drafts 220 1.58×
Your head, 4 drafts 237 1.70×
Your head, 5 drafts 245 1.76×
Your head, 6 drafts 249 1.79×
Your head, 7 drafts 249 1.79×

So the ranking matches yours, your head is clearly better.

Another thing you may find useful:

  • Your head keeps paying off beyond 3 drafts. Although it was trained with 3 chained steps, acceptance stays high at deeper positions: 98 / 95 / 91 / 84 / 77 / 68 / 58% for positions 1–7 (share of drafts accepted up to that position). After 6 it stops improving. I have not measured this under concurrency, where rejected drafts cost more, and I have not re-scored TOB at 5 drafts, so I'd treat it as a single-stream recommendation only.

Two questions:

  • Are you planning to release the model's original training dataset, or at least the 50k pages used for the head?
  • Was 3 drafts chosen simply because the head was trained with 3 steps?
Mogan AI org

Thanks for running this. We hadn't tested beyond 3 drafts, so that's very useful.

Dataset: We'd like to release it, but copyright restrictions on parts of the source material mean we can't publish the full training set. Our arXiv preprint comes out on Monday and describes the data in detail: sources, composition and filtering.

Why 3 drafts: Mostly because the head was trained with 3 chained steps. Each extra step adds another head pass and loss term during training, so 3 kept the cost reasonable. It also matched your setup, which made the comparison direct. And 3 seemed a safe default under concurrency, where rejected drafts cost more.

Sign up or log in to comment