Small-vocabulary MTP heads for memory-bound speculative decoding

1. Research question

Large-vocabulary speculative drafters can spend a surprising amount of time moving and multiplying the output head. FR-Spec and VocabTrim established that a static high-frequency subset can cut that cost while the full target still protects output correctness. In my first Qwen3.5-9B measurements (llama.cpp path), trimming removed roughly 85% of the draft LM-head kernel cost but produced only a 1–3% end-to-end gain: the MTP transformer block remained, and that serving path wasn't draft-bound. Section 4 records how this number reversed once the same recipe ran in an engine whose critical loop is the draft path.

The question is therefore not only which logits should the drafter score? It is:

Can a small block and a small vocabulary be trained as one deployment-specific drafter, preserving enough acceptance to beat the inherited co-trained head on wall time and memory?

The initial project framed this as teacher-to-student distillation, and the first ablations appeared to reject it: equal-budget KD lost to plain cross-entropy twice. That verdict turned out to be a loss bug rather than a finding—section 4 documents the reversal. The current recipe is distillation-first: soft cross-entropy against precomputed teacher logits, then a short CE reinforcement.

2. Experimental design

Subject and interface

The subject is Qwen3.5-9B with its natural co-trained MTP block. Each prediction receives the current token embedding and the predecessor trunk hidden state, projects the pair into the draft block, and scores the next token. The experiment holds the target trunk fixed and replaces only this proposal path.

Current student

StudentSV is a 212M-parameter drafter: a half-width draft block and an independent 32,768-row LM head initialized from frequently used rows of the target head. The shortlist is calibrated separately for code and conversational distributions; the current production run is code-weighted.

Training and controls

• 960,000 training samples drawn from a roughly seven-million-token mixed corpus.
• Cross-entropy training without a teacher forward pass in the current arm.
• A full-vocabulary half-width baseline to separate block compression from vocabulary compression.
• The co-trained MTP head measured both through the engine and through the windowed PyTorch reproduction.
• Held-out code and generation-distribution replay sets; no evaluation text in training.
• Training on the model's own generated outputs beat a generic corpus by 17 points at one third of the steps (0.548 versus 0.376 held-out acceptance)—the training distribution should match the serving distribution.
• Chain training: four recursive draft slots against precomputed teacher top-64 logits, soft cross-entropy over the full 32,768-row draft vocabulary.

Primary metric

Two metrics appear below. Slot-0 teacher-forced replay acceptance—how often the drafter's first proposal matches the target's greedy token on a fixed hidden-state track—is the training diagnostic. Self-drafting chain acceptance over four recursive slots is the deployment-shaped metric: each slot drafts from the previous slot's own output, the regime speculation actually runs in. End-to-end wall time was later measured in the memra engine (section 5); memory accounting and multi-seed replication stayed estimates.

3. Results

The scaled StudentSV run retained 78% of the windowed teacher's code acceptance at roughly one eighth of the draft FLOPs. The improvement from 60k to 960k samples was large and the training loss was still falling—the scaling lever was real. It did not establish a speedup by itself.

The generation-distribution score was much weaker, and the audit later explained why this gap was structural rather than fixable by another run: only about 20% of the corpus was conversational, the chat shortlist covered 94.6% of target tokens versus 98.9% on code—and no finite calibration sample accumulates mass on parameters serving rare content. That mechanism is what closed the student direction.

Chain distillation

Distilling directly from precomputed teacher logits—no CE warmup—matches the staged recipe at half the steps, and doubling the distillation schedule before the CE reinforcement sets the current best: 0.507 chain acceptance, 85.5% of the teacher's 0.593. Slots 1–3 individually exceeded the teacher (+0.013, +0.028, +0.007); the whole remaining gap was slot 0, which chain acceptance compounds. The 24k distillation curve was near-saturated, so the steps lever was spent. The remaining levers—more own-output data, a wider student block—were never pulled: the off-distribution audit closed the direction first.

4. Negative results and changed decisions

5. What this does not establish

• End-to-end wall time is now measured in the memra engine (1.82× MT-Bench at 9B; 2.88× best composed cell at 27B); memory accounting is still an estimate.
• The FLOP ratio remains an architectural estimate; the measured speedups above are the empirical counterpart.
• The result is one model family and one dominant domain; cross-model generality is unknown.
• The student is undertrained and below the desired acceptance-retention threshold.
• The current 2k training/evaluation window does not establish long-context behavior.

6. Artifacts and release plan

The harness is public: github.com/avifenesh/hqmtp — training and evaluation code, fixed corpora, the run ledgers behind every number here, and the negative-result verdicts with their artifact forensics. The engine-side context lives in memra, and the trim-only motivation is public in llama.cpp issue #25187.

1. Available now Harness, protocol, run ledgers, verdicts, and changed decisions.
2. Next release Checkpoints and extracted training data with provenance manifest.
3. Paper v1.0 Wall-time and memory measurements, ablations, multi-seed replication, and PDF/BibTeX.