Working paper 01 Experiments active

Small-vocabulary MTP heads for memory-bound speculative decoding

Can the drafter lose most of its compute without losing the acceptance that makes it useful?

Abstract

Multi-token-prediction heads inherit an expensive shape: a model-width transformer block followed by an output projection over the target model's full vocabulary. On Qwen3.5-9B that means scoring 248,320 vocabulary rows at every draft step. Vocabulary trimming reduces the projection cost but leaves the draft block unchanged. This study asks whether a separately trained, half-width MTP block with its own 32,768-token output head can preserve enough target acceptance to improve the full speculative-decoding path. The current StudentSV arm reaches 0.527 slot-0 acceptance on held-out code versus 0.675 for the windowed co-trained teacher, while using an estimated 0.13× of its draft FLOPs. The result is preliminary, below the study's acceptance-retention target, and substantially weaker out of domain. The main finding so far is directional: jointly shrinking the block and vocabulary is viable enough to scale, but the current training and domain mixture are not yet sufficient.

Preliminary result · code replay
0.527StudentSV acceptance
0.675windowed teacher
~0.13×draft FLOPs

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 Qwen3.5-9B measurements, 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 vocabulary coverage capped acceptance.

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. That assumption did not survive the first ablations. Equal-budget knowledge distillation lost to plain cross-entropy twice, so the current StudentSV arm is trained with cross-entropy only. The teacher remains an evaluation reference, not a mandatory training path.

Figure 1

Two compression axes in one drafter

token embeddingtrunk hidden
Teacher full-width MTP block 248,320-row head
StudentSV half-width MTP block 32,768-row head
draft tokensfull target verificationtarget distribution retained
The target model still verifies against its full vocabulary. The study changes the proposal path, not the verifier or the exactness contract.

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 uses 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.

Primary metric

The current page reports slot-0 teacher-forced replay acceptance: how often the draft head's first proposal matches the target model's greedy token on a fixed hidden-state track. It is a useful training metric, not the final systems metric. A publishable conclusion requires recursive K-step acceptance, end-to-end tokens per second, memory, multiple seeds, and a matched target-verification cost.

3. Preliminary results

Draft headHeld-out codeGeneration distributionEst. draft FLOPs
Engine teacher · full context0.6780.3481.00×
PyTorch teacher · 2k window0.6750.3711.00×
StudentSV · 960k samples0.5270.166~0.13×
StudentSV · 60k samples0.3370.105~0.13×
Half-width · full vocabulary · 60k0.373~0.35×
Figure 2

Acceptance is still buying scale

0.3 0.4 0.5 0.6 0.7 60k samples 960k samples co-trained teacher · 0.675 60k samples · 0.337 acceptance 960k samples · 0.527 acceptance 0.337 0.527
Slot-0 code acceptance on held-out replay. A 16× increase in training samples closed half of the gap to the teacher; the loss was still falling when the run stopped.

The scaled StudentSV run retains 78% of the windowed teacher's code acceptance at roughly one eighth of the draft FLOPs. The improvement from 60k to 960k samples is large and the training loss was still falling, which supports another scale step. It does not establish a speedup by itself.

The generation-distribution score is much weaker. Two likely causes are already measurable: only about 20% of the current corpus is conversational, and the chat shortlist covers 94.6% of target tokens versus 98.9% on code. The next run must change both data mixture and vocabulary allocation rather than merely train longer on the same distribution.

4. Negative results and changed decisions

RefutedQuantization alone breaks MTP agreement.

Clean NVFP4-versus-BF16 controls found no degradation, so the original “heal the head” premise was closed.

RejectedKnowledge distillation is required.

KD lost to cross-entropy twice at equal budget. Teacher-free CE is both better so far and materially faster to train.

InsufficientTrim the inherited head and stop.

Trim-only saved most LM-head kernel work but only 1–3% end to end in the measured 9B path because the transformer block remained.

5. What this does not establish

  • The current metric is first-proposal replay acceptance, not recursive draft-chain acceptance.
  • The FLOP ratio is an architectural estimate, not a measured end-to-end speedup.
  • 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 training harness, extracted datasets, checkpoints, and full run ledger are currently being cleaned for release. The public systems context already lives in bw24, and the trim-only motivation and implementation discussion is public in llama.cpp issue #25187.

  1. Available now Study question, protocol, preliminary table, and changed decisions.
  2. Next release Training code, data manifest, fixed evaluation split, and checkpoint provenance.
  3. Paper v1.0 Recursive acceptance, wall-time and memory measurements, ablations, multi-seed replication, and PDF/BibTeX.

References

  1. W. Zhao et al. FR-Spec: Accelerating Large-Vocabulary Language Models via Frequency-Ranked Speculative Sampling. ACL 2025.
  2. R. Goel et al. VocabTrim: Vocabulary Pruning for Efficient Speculative Decoding in LLMs. 2025.
  3. H. Cai et al. FastMTP: Accelerating LLM Inference with Enhanced Multi-Token Prediction. 2025.
Update discipline

Substantive result changes will increment the version and remain visible here. The page will not silently turn a preliminary number into a final one.