ewinregirgojr

minicpm5-stock-analyst-v2-mtp

README

License: apache-2.0

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

Table with columns: benchmark, k=1, k=2, k=3
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

Table with columns: file, what it is
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

python
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

Table with columns: version, data, GSM8K k=2, note
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.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider

ewinregirgojr

Model Tree

Base

openbmb/MiniCPM5-1B

Fine-tuned

this model

Input Modalities

Text

Output Modalities