WeChat AIGrowMTP Paper

GrowMTP: Can RL Grow
Its Own Draft Head?

Minghua He1,2,*Lingzhe Zhang2Yuan Liu1Xiao Zhou1Aiwei Liu1,†

1 WeChat AI, Tencent2 Peking University

* Work done during an internship at WeChat AI. † Corresponding author.

Paper Code

An RL run can train the draft head it needs.GrowMTP turns rollout verification into online supervision, growing a draft head from scratch while accelerating the same training run.

1.00 → 2.91Acceptance length τ
2.13×Rollout speedup
1.60×End-to-end speedup
0 GPU·hOffline head training
Qwen3-4B · Math reasoning · 500 RL steps · Head-update cost included. τ is the final-10-step mean; speedups use whole-run mean times.
Abstract

Reinforcement learning (RL) post-training drives the frontier capabilities of large language models, with its wall-clock dominated by autoregressive rollout generation. Speculative decoding is an established remedy for this bottleneck, but existing draft heads must be pretrained or warmed up before RL, introducing substantial training cost outside the RL run to be accelerated. We observe that RL training itself provides both conditions required for online draft-head training: its rollout distribution is far narrower than that of pretraining, and its verification step continuously produces supervision signals aligned with this distribution. Building on these observations, we propose GrowMTP, which uses this supervision to train a draft head from scratch entirely within the RL loop, with all head updates detached from the policy backbone. On Qwen3-4B (no draft head), MiMo-7B-SFT (weak head), and Qwen3.5-4B-Base (strong head), GrowMTP achieves rollout speedups of 2.13×, 1.93×, and 1.36×, and end-to-end speedups of 1.60×, 1.41×, and 1.20×, respectively. GrowMTP therefore serves existing RL training frameworks as a modular component, particularly offering a from-scratch acceleration path for models without pretrained draft heads.

THE GROWMTP LOOPInside the existing RL training step
01

Draft & verify

The head proposes. The policy checks.

hθ→d₁ ✓d₂ ✓d₃ ×d₄
02

Reuse the supervision

Keep the path that produced each signal.

Draft tokensTarget logitsCycle boundaries
03

Grow the draft head

Reconstruct, learn, and draft again.

DCA + VGM→Update φθ detached
01 MOTIVATION

Why train the draft head within RL?

Speculative decoding accelerates rollouts, but drafting capability usually has to be trained before RL begins.

Existing draft-head training routes require preparation before reinforcement learning can benefit.
Existing approach Where the head comes from Preparation before RL
Native MTPDeepSeek-V3 Joint backbone
pretraining 		14.8Tpretraining tokens†
EAGLE-style head70B target model Dedicated offline
head training 		~5,000GPU·h

Drafting capability enters RL as a prerequisite.

† 14.8T is the full backbone pretraining corpus, not the incremental cost of the MTP head. The two examples use different measures.

THE QUESTION

Can RL grow the draft head it needs?

02 OBSERVATION

RL already provides both ingredients.

A focused rollout distribution and valid verification supervision.

01

RL only needs a
domain-specific head.

A broad pretraining distribution and a focused Math reasoning RL run Pretraining spans Code, Knowledge, Math, Dialogue, and Writing. An arrow selects the Math domain for this example RL run. This is a conceptual distribution sketch, without measured coordinates. Broad pretraining This RL run Code Knowledge Dialogue Writing Math focus Math Task-specific rollouts
Schematic: a Math reasoning RL run.

Fit the current rollout distribution.

02

RL already produces
valid supervision.

Random headFirst draft rejected
Draftd₁ ×d₂d₃d₄ Valid
signalp₁ ✓———
Improving headLonger accepted prefix
Draftd₁ ✓d₂ ✓d₃ ×d₄ Valid
signalp₁ ✓p₂ ✓p₃ ✓—

× The first rejected position still provides a valid signal.

Valid supervision from the first cycle.

A focused rollout distributionValid verification supervision

Grow the draft head within RL.

RL already provides both the distribution to fit and the supervision to learn from.

03 CHALLENGE

What makes online draft-head training challenging?

RL provides the supervision. Using it for multi-step drafting requires resolving three challenges.

  1. State consistency

    Verification signals belong to specific draft states. Training on the final response can change cycle boundaries, recurrent states, and KV context.

  2. Depth dependence

    A deeper draft contributes only if all earlier drafts are accepted. Learning at different depths is therefore coupled.

  3. Supervision validity

    Verification also produces signals after the first rejection, where later drafts depend on a token the target discarded.

    p₂ still supervises a valid prefix; p₃ conditions on rejected d₂.

01 State consistency · Rollout vs. teacher forcing

Rollout above, teacher forcing below, connected by the same committed sequence Two rollout cycles produce d tilde 11, d21 and d tilde 22. Teacher forcing reads these committed tokens and recomputes its own hidden states and logits. It writes KV for each input token. Entries for d tilde 11 and d21 are red because their recursively computed features differ from those used at the corresponding rollout positions. The first hidden output matches rollout. When teacher forcing carries it across a cycle boundary instead of restarting from the target's feature, that hidden input, the feedback path, and subsequent hidden outputs are also marked red. Tokens keep their original colors. Rollout Cycle 1 · draft 1 / 3 h₀ yₜ hidden token Drafter φ — — hidden token Input Output KV Prefix — — — Drafts — — — Target wait for drafts Committed tokens Prefix — — — Teacher forcing Uses the committed tokens — — hidden token Drafter φ — — hidden logits Input Output KV writes Prefix — — — Logits q₁ ↔ p₁₁ q₂ ↔ p₂₁ q₃ ↔ p₂₂ h₁ · d₁ d₁₁ KV

Draft tokens first; train on the committed sequence afterward.

Target-derived Draft-derived Hidden / KV inconsistent with rollout
04 METHOD

How does the draft head learn?

Reconstruct draft states, optimize the acceptance chain, and learn up to the first rejection.

  1. Draft-Path Reconstruction

    Replay each cycle from its recorded boundary using the original drafts. Recompute hidden states and KV to pair each head prediction with the verification signal under the same context.

    Recorded drafts become inputs. No resampling.

  2. Depth-Coupled Acceptance

    Optimize an acceptance-chain surrogate so deeper predictions contribute through earlier acceptance.

    αi = 1 − TV(pi, qi) · K = drafting depth

  3. Verify-Gated Masking

    Include the first rejected position and exclude the later suffix. DCA’s soft weights alone do not enforce this rejection boundary.

    The first position always teaches the head, even from a random start.

Recorded rollout → draft-path reconstruction

Rollout records each cycle; GrowMTP reconstructs its draft path The same two rollout cycles as the preceding animation: cycle 1 rejects its first draft, and cycle 2 accepts its first draft and rejects its second. Each cycle records its target boundary, all three draft tokens, and verification distributions. Reconstruction starts each cycle from that recorded boundary, rebuilds its prefix KV, and feeds the recorded drafts back in order. The head recomputes hidden states, logits, and KV; the draft tokens are not sampled again. Green features follow the same computation as rollout. Logit pairs after the first rejection are gray because VGM excludes them from the loss. Rollout Cycle 1 · draft 1 / 3 h₀ yₜ hidden token Drafter φ — — hidden token Input Output KV Prefix — — — Drafts — — — Target wait for drafts Recorded cycles Cycle 1 — — — — Cycle 2 — — — — GrowMTP reconstruction Uses the recorded cycles — — hidden token Drafter φ — — hidden logits Input Output KV writes Prefix — — — Loss pairs q₁₁ ↔ p₁₁ q₁₂ ↔ p₁₂ q₁₃ ↔ p₁₃ h₁ · d₁ d₁₁ KV

Record each cycle, then reconstruct its draft path.

Target-derived Reconstructed Gray pairs: VGM masked
Detached updates
Head lossDraft head φ
RL lossPolicy θ

Backbone features are detached from head training.

05 From Scratch Evaluation

Can RL grow its own draft head from scratch?

A randomly initialized head accelerates the same RL run, with head-update costs included.

Qwen3-4B · Random draft head, K=5 · 500 RL steps · 8 × H800

A draft head grows inside RL

Math: DAPO-Math-17K · Code: TACO-Verified

RL TRAINING STEP500 / 500

Math

Acceptance length τ

Math

End-to-end step time (s)

Code

Acceptance length τ

Code

End-to-end step time (s)
GrowMTP AR RL Solid: plotted trend · Faint: plotted variability
500 / 500
500 steps: a useful head, grown from scratch.

Single-step time crosses below AR at step 30 on Math and step 22 on Code. These are per-step crossovers.

06 Training Strategy

How should the draft head be trained online?

Detach the head's gradients from the policy, and train it on an objective aligned with acceptance. A rising τ on its own is not evidence of a healthy run.

Qwen3.5-4B-Base · Pretrained head, K=3 · 200 RL steps · Baseline: the same head frozen

Acceptance and policy quality on the same 200-step timeline. Read them together: the arm with the highest τ is the one whose quality goes to zero.

Math reasoning

DAPO-Math-17K

Code reasoning

TACO-Verified
06 · DRAFTING DEPTH

How far should the draft head look ahead?

GrowMTP extends a single-step head to multi-step drafting. Acceptance increases with depth, while end-to-end speedup peaks at K=5 in this setting.

MiMo-7B-SFT · Pretrained single-step head · K=3, 5, 7 · 200 RL steps · Baseline: frozen K=1

Deeper drafting yields higher acceptance throughout training. Step time improves across all trained depths, but its ranking does not follow τ.

Scaling drafting depth

07 · COMPONENT ABLATIONS

How does each component contribute to multi-step adaptation?

Reconstruction lowers training cost, DCA improves multi-step acceptance, and VGM adds further efficiency by enforcing the rejection boundary.

MiMo-7B-SFT · Math reasoning · Pretrained single-step head · 200 RL steps · K=5, with K=1 also evaluated for DCA

Similar aggregate acceptance, lower training cost. Reconstruction reaches τ=4.04 versus 4.02 for teacher forcing, while reducing end-to-end step time from 408.09s to 371.02s.

Full reconstruction ablation

Training evaluation

DAPO-Math-17K · 200 RL steps · Speedups relative to frozen K=1
ConfigurationτRolloutEnd-to-end step
Time (s)SpeedupTime (s)Speedup
Teacher forcing4.02169.051.89×408.091.28×
GrowMTP4.04164.841.93×371.021.41×

Inference Evaluation

Math reasoning · Mean@16 (%) and held-out acceptance τ
ConfigurationAMC23AIME24AIME25
Mean@16τMean@16τMean@16τ
Teacher forcing91.563.9354.373.8243.963.79
GrowMTP91.253.9556.873.8446.253.82
08 DISCUSSION

Where does GrowMTP help, and how far can it scale?

Three model starting points, the potential of larger RL runs, and the domain specialization of an online-trained head.

GrowMTP adapts to different starting points, from a randomly initialized head to pretrained draft heads. The three evaluated models all gain end-to-end acceleration on mathematical reasoning.

Random initialization

Qwen3-4B

1.60×

End-to-end step speedup

K=5 · 500 RL steps

Baseline
Autoregressive decoding

Pretrained single-step head

MiMo-7B-SFT

1.41×

End-to-end step speedup

K=5 · 200 RL steps

Baseline
Frozen pretrained head, K=1

Pretrained head

Qwen3.5-4B-Base

1.20×

End-to-end step speedup

K=3 · 200 RL steps

Baseline
Frozen pretrained head, K=3

09 CONCLUSION

Conclusion

Grow From Scratch during RL

A randomly initialized draft head learns entirely within RL, without offline pretraining or warm-up.

Accelerate the Same Run

On Qwen3-4B, end-to-end training speeds up by 1.60× on Math and 1.56× on Code, including head updates.

Adapt Across Starting Points

Online adaptation improves both random and pretrained heads, with comparable held-out quality in the evaluated runs.

Read the Paper ↗

Paper figure

Image