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

---
license: apache-2.0
base_model: openbmb/MiniCPM5-1B
pipeline_tag: text-generation
library_name: transformers
tags:
  - 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
language:
  - en
datasets:
  - openai/gsm8k
  - openai/openai_humaneval
  - HuggingFaceH4/mt_bench_prompts
  - ewinregirgojr/minicpm5-stock-v2-forward-return
model-index:
  - name: MiniCPM5-1B MTP Head (FastMTP-style + LK losses)
    results:
      - task:
          type: text-generation
          name: Multi-token prediction draft acceptance (teacher-forced, held-out)
        dataset:
          name: GSM8K (held-out slice)
          type: openai/gsm8k
        metrics:
          - name: Acceptance @ k=1
            type: accuracy
            value: 60.9
          - name: Acceptance @ k=2
            type: accuracy
            value: 48.0
          - name: Acceptance @ k=3
            type: accuracy
            value: 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](https://arxiv.org/abs/2509.18362) (single shared MTP head, self-distillation on the model's own outputs) + [LK losses](https://arxiv.org/abs/2602.23881) (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](https://huggingface.co/openbmb/MiniCPM5-1B) with the [stock-analyst v2 LoRA](https://huggingface.co/ewinregirgojr/minicpm5-stock-analyst-gguf) 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](https://huggingface.co/ewinregirgojr/minicpm5-stock-analyst-gguf) (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.