minicpm5-stock-analyst-v2-mtp/README.md

license, base_model, pipeline_tag, library_name, tags, language, datasets, model-index

license	base_model	pipeline_tag	library_name	tags	language	datasets	model-index
apache-2.0	openbmb/MiniCPM5-1B	text-generation	transformers	multi-token-prediction mtp speculative-decoding self-speculative-decoding inference-acceleration draft-model mtp-head fastmtp eagle lk-losses acceptance-rate llm-inference minicpm small-llm 1b self-distillation finance stock-prediction	en	openai/gsm8k openai/openai_humaneval HuggingFaceH4/mt_bench_prompts ewinregirgojr/minicpm5-stock-v2-forward-return	name results MiniCPM5-1B MTP Head (FastMTP-style + LK losses) task dataset metrics type name text-generation Multi-token prediction draft acceptance (teacher-forced, held-out) name type GSM8K (held-out slice) openai/gsm8k name type value Acceptance @ k=1 accuracy 60.9 name type value Acceptance @ k=2 accuracy 48.0 name type value Acceptance @ k=3 accuracy 42.1

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:

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

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

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

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.