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):

• 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:

1. 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.
2. 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

The head projects through the backbone's own LM head matrix (no new output weights).

How to load

import torch, torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRMSNorm
from huggingface_hub import hf_hub_download

REPO = "ewinregirgojr/minicpm5-stock-analyst-v2-mtp"
backbone = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16).eval()
tokenizer = AutoTokenizer.from_pretrained(REPO)
config = backbone.config

class MTPHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
        self.pre_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.decoder_layer = LlamaDecoderLayer(config, layer_idx=config.num_hidden_layers)
        self.out_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
    def forward(self, prev_hidden, token_embed, position_ids, position_embeddings):
        x = self.pre_norm(self.input_proj(torch.cat([prev_hidden, token_embed], dim=-1)))
        out = self.decoder_layer(x.unsqueeze(1), position_ids=position_ids.unsqueeze(1),
                                 position_embeddings=position_embeddings)
        if isinstance(out, tuple):
            out = out[0]
        return self.out_norm(out.squeeze(1))

head = MTPHead(config).to(torch.bfloat16)
head.load_state_dict(torch.load(hf_hub_download(REPO, "mtp_head.pt"), map_location="cpu"), strict=True)
head.eval()
# drafting: feed h_i (backbone hidden at position i) + embedding of token i+1 -> logits for token i+2,
# 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

License: Apache 2.0 (matches base model and LoRA lineage). Last updated: 2026-07.