Multi-Token Prediction (MTP) Head for a 1B LLM — FastMTP-style Speculative Decoding, Trained for ~$0 on Kaggle
A single 33M-parameter multi-token prediction head bolted onto a frozen MiniCPM5-1B, trained with self-distillation (FastMTP recipe) + a direct acceptance-rate loss (LK losses) — with honestly reported, per-benchmark acceptance rates instead of a single cherry-picked number. If you searched for multi-token prediction, MTP head, speculative decoding for small LLMs, or how to speed up LLM inference without a separate draft model — this is a complete, reproducible worked example, including the failure modes.
TL;DR
- What: one extra Llama decoder layer (+ input projection, 33M params ≈ 2% of the backbone) that recursively drafts k=1..3 future tokens from the frozen base model's own hidden states, embeddings, and LM head. No separate draft model. Backbone weights untouched.
- Recipe: FastMTP (single shared MTP head, self-distillation on the model's own outputs) + LK losses (hybrid KL/Total-Variation objective that directly optimizes acceptance rate instead of plain cross-entropy).
- Cost: trained entirely on free Kaggle T4x2 GPUs, ~8,400 self-distilled sequences total.
- Result (held-out, teacher-forced argmax acceptance vs. the backbone's own predictions):
| benchmark | k=1 | k=2 | k=3 |
|---|
| GSM8K (math) | 60.9% | 48.0% | 42.1% |
| HumanEval (code) | 40.2% | 28.3% | 24.0% |
| MT-Bench (chat) | 38.1% | 29.0% | 25.0% |
- Not: a claim of production speedup. These are draft-acceptance rates, not wall-clock numbers; integration into an inference engine (vLLM/SGLang-style verify step) is up to you.
Why this exists
Vanilla reuse of a single MTP module collapses at depth — the FastMTP paper measured ~70% acceptance at k=1 dropping to ~10% at k=2 and ~0% at k=3 when a stock MTP head is recursed naively. This repo demonstrates the two published fixes that actually move that number, combined and reproduced at 1B scale on free hardware:
- Self-distillation — the head trains on the model's own greedy completions (MT-Bench / GSM8K / HumanEval prompt mixes), so it learns the distribution it will actually be drafting for.
- Direct acceptance-rate optimization — the LK hybrid loss
λ·KL + (1−λ)·TV with an adaptive λ = exp(−η·α) schedule, which their ablation showed beats plain KL by ~5.6% acceptance length when fine-tuning a native MTP module.
What's in the repo
| file | what it is |
|---|
model.safetensors + configs | the backbone: MiniCPM5-1B with the stock-analyst v2 LoRA merged in. Frozen during MTP training. |
mtp_head.pt | the trained 33M-param MTP head (input_proj + 1 LlamaDecoderLayer + norms). |
| tokenizer files | standard MiniCPM tokenizer. |
The head projects through the backbone's own LM head matrix (no new output weights).
How to load
1import torch, torch.nn as nn
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRMSNorm
4from huggingface_hub import hf_hub_download
5
6REPO = "ewinregirgojr/minicpm5-stock-analyst-v2-mtp"
7backbone = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16).eval()
8tokenizer = AutoTokenizer.from_pretrained(REPO)
9config = backbone.config
10
11class MTPHead(nn.Module):
12 def __init__(self, config):
13 super().__init__()
14 self.input_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
15 self.pre_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
16 self.decoder_layer = LlamaDecoderLayer(config, layer_idx=config.num_hidden_layers)
17 self.out_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
18 def forward(self, prev_hidden, token_embed, position_ids, position_embeddings):
19 x = self.pre_norm(self.input_proj(torch.cat([prev_hidden, token_embed], dim=-1)))
20 out = self.decoder_layer(x.unsqueeze(1), position_ids=position_ids.unsqueeze(1),
21 position_embeddings=position_embeddings)
22 if isinstance(out, tuple):
23 out = out[0]
24 return self.out_norm(out.squeeze(1))
25
26head = MTPHead(config).to(torch.bfloat16)
27head.load_state_dict(torch.load(hf_hub_download(REPO, "mtp_head.pt"), map_location="cpu"), strict=True)
28head.eval()
29# drafting: feed h_i (backbone hidden at position i) + embedding of token i+1 -> logits for token i+2,
30# then recurse with the head's own output hidden state. Verify drafts with the backbone as usual.
FAQ
Is there a multi-token prediction head I can train myself without a big cluster?
Yes — that's this repo's point. One epoch + a continuation round on ~8.4k self-distilled sequences, free Kaggle T4x2, a few hours total. The full recipe (FastMTP-style shared head + LK loss) is described above and in the linked papers.
Why is GSM8K so much higher than MT-Bench and HumanEval?
Data mix. The self-distillation rounds were math/general-instruction heavy; code got only ~120 prompts. Acceptance tracks training-data volume per domain almost linearly at this scale — that's the main practical lesson from this run. If you continue training with a code-heavy mix, expect HumanEval to move the same way GSM8K did (it roughly doubled from a 7k-sequence round).
Is this the same as DeepSeek's MTP or Qwen's MTP-native models?
Same idea (a small head drafts future tokens, the base model verifies), different provenance: those models pretrain the MTP module jointly; this one is retrofitted onto a finished model post-hoc, which is the situation you're in if your model didn't ship with an MTP head.
What's the catch?
Three honest ones. (1) These are teacher-forced acceptance rates, an upper-bound proxy — real speculative-decoding speedup depends on your verify-step implementation and batch regime. (2) Chat (MT-Bench) is the hardest domain for a 33M head and is still under 30% at k=2 here. (3) The backbone includes a stock-prediction LoRA (see lineage below); for a general-purpose backbone you may want to re-run the recipe on plain MiniCPM5-1B.
Where did the backbone come from?
It's the
MiniCPM5-1B Stock Analyst v2 (a LoRA fine-tune for BUY/SELL stock-direction prediction with honestly-reported ~53% held-out accuracy at the feature set's measured ceiling), merged to plain weights. On the stock task itself the model answers in one token, so MTP does nothing there — the head is trained and evaluated on general text.
Reproduce / extend
- Levers, in order of measured impact: more self-distillation data (dominant — one 7k round roughly doubled every benchmark), domain mix (train on what you want accepted), epochs, then loss hyperparameters (
η, γ).
- Eval protocol: held-out prompts (seeded slices disjoint from training), backbone generates completions greedily, head predicts each next token teacher-forced, acceptance = argmax match with the backbone's own logits.
Version history
| version | data | GSM8K k=2 | note |
|---|
| round 1 | 1.4k seqs (incl. stock) | 26.4% | stock EOS-patterns polluted training |
| round 2 (this) | +7k general seqs, stock removed | 48.0% | every benchmark ~doubled |
License: Apache 2.0 (matches base model and LoRA lineage). Last updated: 2026-07.
