PaperlayerView exact revision on arXiv

Source-linked AI summary

GrowMTP: Can RL Grow Its Own Draft Head?

Minghua He, Lingzhe Zhang, Yuan Liu, Xiao Zhou, Aiwei Liu

arXiv:2609.16648v1cs.LGcs.CL

TL;DR

RL rollout generation dominates training cost, while existing draft heads require training outside RL despite the task-specific rollout distribution and verification signals available during RL. GrowMTP trains a draft head online from those signals, achieving rollout speedups across models with no, weak, or strong pretrained heads.

  • Problem

    RL rollouts dominate training cost, while draft-head supervision is sourced outside RL despite its narrower rollout distribution and available verification signals.

    Source: §1 · pp. 1–2

  • Method

    GrowMTP trains a draft head online within the RL loop using rollout verification signals, rollout-consistent states, and detached updates to the policy backbone.

    Source: §1 · p. 2

  • Results

    2.13×, 1.93×, and 1.36× rollout speedups are reported for Qwen3-4B, MiMo-7B-SFT, and Qwen3.5-4B-Base, respectively.

    Source: §1 · p. 2

  • Takeaways & Limitations

    GrowMTP offers a from-scratch acceleration path for models without pretrained draft heads and also accelerates models with weak or strong heads.

    Source: §1 · p. 2, §6 · p. 9

    The fixed-record acceptance-surrogate objective is not assumed unbiased for deployment-time expected accepted length.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

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.13x, 1.93x, and 1.36x, and end-to-end speedups of 1.60x, 1.41x, and 1.20x, 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.

1 INTRODUCTION

RL rollout generation is a major training cost, yet existing draft heads require training before RL. GrowMTP trains a head online from rollout verification signals and accelerates rollouts across models with no, weak, or strong pretrained heads.

1 INTRODUCTION

  • 2.13× rollout and 1.60× end-to-end speedups on Qwen3-4B; MiMo-7B-SFT and Qwen3.5-4B-Base reach 1.93× and 1.36× rollout speedups, respectively.The Qwen3-4B run also uses 9.4% fewer total GPU-hours than an offline-trained head at comparable acceptance over 500 steps.Source: Figure 1 · p. 1, §1 · p. 2
  • Existing draft heads are trained before RL, adding offline training cost and leaving models without a head no route to speculative acceleration.Autoregressive decoding typically accounts for 70–90% of each RL training step.Source: §1 · pp. 1–2
  • RL rollouts provide task-specific distributions and continuous verification supervision, allowing a head to be trained from the rollouts RL already requires.This avoids an upfront head-training bill.Source: §1 · p. 2
  • Online multi-step training must preserve draft-conditioned states, account for earlier rejections limiting deeper contributions, and exclude drafts conditioned on rejected tokens.Otherwise, supervision can train on discarded suffixes that do not follow the executed sequence.Source: §1 · p. 2
  • GrowMTP trains a draft head online from verification signals under rollout-consistent states, couples supervision across accepted draft depths, and detaches head updates from the policy backbone.It can start from random initialization because the first draft position always supplies valid supervision, without offline pretraining or a depth schedule.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries describe RL post-training as rollout generation followed by policy updates, then explain speculative draft-then-verify decoding. They motivate online head training from rollout verification signals and identify the state, depth-dependence, and rejection challenges it must address.

2.1 REINFORCEMENT LEARNING FOR LARGE LANGUAGE MODELS

  • An RL step samples a group of responses, generates them token by token during rollout, then uses rewards to compute advantages and update the policy.The work targets the rollout bottleneck while retaining the original RL objective.Source: §2.1 · p. 3

2.2 SPECULATIVE DECODING AND MULTI-TOKEN PREDICTION

  • Speculative decoding uses a recurrent MTP head to propose K tokens, which the target verifies in one forward pass; acceptance length counts the average tokens committed per cycle.The committed tokens include the accepted prefix and a replacement or bonus token, so acceptance length is at least one.Source: §2.2 · p. 3

2.3 TRAINING DRAFT HEADS FROM RL ROLLOUTS

  • RL rollouts provide verification supervision aligned with the process’s rollout distribution, but online training must match draft-conditioned states, couple learning across depths, and avoid rejected suffixes.The first draft position always provides supervision, enabling training from random initialization without a separate dataset or offline head-training stage.Source: §2.3 · p. 3

3 GrowMTP: GROWING DRAFT HEADS WITHIN RL

GrowMTP trains the draft head online from rollout verification records alongside policy optimization, with head updates detached from the backbone. It reconstructs realized draft-conditioned states and weights acceptance-chain training to reflect the first-rejection boundary.

3 GrowMTP: GROWING DRAFT HEADS WITHIN RL

  • Rollout records train the draft head alongside policy optimization, using detached backbone states and no additional target-model forward pass.The head loss updates only head parameters; the backbone continues optimizing the original RL objective.Source: §3 · p. 3, §3.1 · p. 4

3.1 TRAINING PIPELINE AND DRAFT-PATH RECONSTRUCTION

  • Draft-Path Reconstruction replays recorded drafts from each cycle context, matching training states to realized prefixes that teacher forcing diverges from after rejection.Teacher-forcing states match drafting only when all preceding drafts were accepted, and the mismatch grows with draft depth.Source: §3.1 · p. 4

3.2 THE DEPTH-COUPLED ACCEPTANCE LOSS

  • DCA chains per-position acceptance probabilities into an accepted-length surrogate, then applies a logarithm to give smaller chain values greater relative gradient weight.The DCA objective is minimized with respect to the draft head alone; recorded target distributions stay constant.Source: §3.2 · p. 5

3.3 LEARNING UP TO THE FIRST REJECTION

  • VGM truncates the chain sum at and including the first rejected position, removing supervision for later drafts conditioned on discarded tokens.The first draft position is always included, and the first rejected position is included when one occurs.Source: §3.3 · pp. 5–6
  • Supervision is available from the first cycle, so training can start with a randomly initialized head without a shallow-depth warm-up or K schedule.The supervision depth follows the realized accepted prefix, while gradient weights reflect preceding acceptance overlaps.Source: §3.3 · pp. 5–6
  • The update uses a fixed-record acceptance surrogate and does not assume it is unbiased for deployment-time expected accepted length.Recorded tokens, target distributions, and the first-rejection position are held fixed while gradients flow through reconstructed head distributions.Source: §3.3 · p. 6

4 EXPERIMENTS

GrowMTP grows draft heads online during RL and evaluates it from scratch, on pretrained heads, and across component choices. The experiments assess rollout and end-to-end speed, acceptance, and held-out policy quality.

4.1 EXPERIMENTAL SETUP

  • The setup compares Qwen3-4B without a pretrained draft head, Qwen3.5-4B-Base with K=3, and MiMo-7B-SFT with K=1, measuring training speed, acceptance, and held-out quality.Evaluation spans mathematical and code reasoning; the shipped head configurations isolate from-scratch training and extension of shallow drafting capability.Source: §4.1 · p. 6

4.2 RQ1: CAN AN RL RUN GROW ITS OWN DRAFT HEAD FROM SCRATCH?

  • 2.13× and 1.93× rollout speedups, with 1.60× and 1.56× end-to-end step speedups, emerge from training a random K=5 head on Qwen3-4B.Acceptance length reaches 2.91 on mathematical reasoning and 2.64 on code reasoning; held-out benchmark scores are comparable.Source: §4.2 · p. 6, Table 1 · p. 7
  • 502.8 versus 555.3 GPU·h gives GrowMTP 9.4% lower total cost than offline training, while both heads reach nearly identical serving acceptance.The comparison covers 500 RL steps; offline training adds a separate trajectory-generation stage before RL.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 RQ2: HOW SHOULD THE DRAFT HEAD BE TRAINED ONLINE?

  • With gradients detached, DCA reaches end-to-end step speedups of 1.20× and 1.12× and held-out quality on par with the frozen head.Joint CE attains higher acceptance but zero held-out quality in both domains; detached CE also accelerates while avoiding that observed collapse.Source: Table 3 · p. 7, §4.3 · p. 7
  • At K=5, MiMo-7B-SFT reaches end-to-end step speedups of 1.41× and 1.35×; K=7 raises acceptance further but lowers speedups to 1.30× and 1.31×.Held-out scores are comparable across depths, and K=5 is highlighted because end-to-end speedups peak there.Source: §4.3 · p. 7, Table 4 · p. 8

4.4 RQ3: HOW DOES EACH COMPONENT CONTRIBUTE TO MULTI-STEP ADAPTATION?

  • Draft-path reconstruction lifts the end-to-end step speedup from 1.28× to 1.41× while aggregate acceptance ties within 0.03.Teacher forcing adds cost without raising aggregate acceptance, and a deeper-position gap remains despite the aggregate tie.Source: Table 5 · p. 8, §4.4 · p. 8
  • At K=5, DCA reaches τ=4.04 versus 3.86 for CE, 3.85 for KL, and 3.94 for TV; removing VGM also reduces step speedups.DCA's end-to-end speedup is 1.41× versus 1.36×–1.37× for alternatives; removing VGM lowers it to 1.35× under DCA and 1.28× under KL.Source: §4.4 · pp. 8–9, Figure 4 · p. 8

5 RELATED WORK

Related work covers draft-head designs and training approaches, then situates GrowMTP among methods for accelerating RL rollouts.

5.1 SPECULATIVE DECODING

  • Draft-head research spans lightweight heads that reuse backbone states and online adaptation, including state-alignment methods HASS [Zhang et al., 2025] and EAGLE-3 [Li et al., 2026b].GrowMTP shares their alignment motivation, reconstructs draft-path training states, and combines a log-transformed TV-chain objective with VGM.Source: §5.1 · p. 9

5.2 ACCELERATING REINFORCEMENT LEARNING FOR LLMS

  • RL rollout acceleration includes inference-engine and scheduling improvements, while RhymeRL [He et al., 2025] uses historical drafts and ReSpec uses RL supervision with offline warm-up.GrowMTP supports random initialization and integrates drafting, verification, and head updates within the RL run.Source: §5.2 · p. 9

6 CONCLUSION

GrowMTP trains draft heads from RL verification signals, including from random initialization, and reports acceleration across models with no, weak, and strong draft heads. The appendix specifies the tested model configurations and shared training setup.

6 CONCLUSION

  • GrowMTP trains draft heads online within RL, using verification signals to grow a random head without offline pretraining or warm-up.Across Qwen3-4B, MiMo-7B-SFT, and Qwen3.5-4B-Base, reported rollout speedups are 2.13×, 1.93×, and 1.36×, with comparable held-out benchmark scores.Source: §6 · p. 9

B.1 MODELS AND DRAFT HEADS

  • Qwen3-4B has no shipped draft head, while Qwen3.5-4B-Base and MiMo-7B-SFT ship native heads at K=3 and K=1, respectively.All use a single-layer recurrent MTP form; Qwen3-4B's newly built head alone is randomized, while the embedding and output head come from the backbone.Source: §B.1 · p. 15

B.2 RL TRAINING CONFIGURATION

  • Policy training uses GRPO with a fixed initial-checkpoint KL reference and shared optimization settings across models, domains, and ablation arms.The setup uses released math and code datasets with head supervision from RL rollouts, on a single node of 8×H800 GPUs.Source: §B.2 · p. 15

B.3 DRAFT-HEAD TRAINING CONFIGURATION

  • The head is updated once per RL step from that step's recorded verification signals, with detached gradients, DCA, and VGM throughout GrowMTP runs.Training has no freeze phase; head learning rates are 3 × 10−4 for random initialization and 1 × 10−4 for pretrained heads.Source: §B.3 · p. 16

B.4 EVALUATION PROTOCOL

The evaluation protocol measures training-time speed and acceptance, then assesses held-out policy quality and drafting capability on math and code benchmarks.

B.4 EVALUATION PROTOCOL

  • Training evaluation reports rollout time and end-to-end step time, including head updates, against a frozen-head run under the same script, data, and hyperparameters.Training-side τ is averaged over the final 10 steps, while time averages cover the full run.Source: §B.4 · p. 16
  • Held-out policy quality uses Mean@16 on AMC23, AIME24, and AIME25, and Pass@4 on LiveCodeBench v6; drafting capability is held-out τ.τ is one plus the ratio of accepted tokens to verifications, measured on the same generations used to assess quality.Source: §B.4 · p. 16

C.1 ADDITIONAL RESULTS FOR GROWING FROM SCRATCH (RQ1)

Online draft-head training improves from-scratch acceptance across math and code, while detached training preserves policy quality and DCA leads among quality-preserving strategies. The additional results examine drafting depth, component ablations, model scope, RL-scale projections, domain specialization, and reconstruction implementation.

C.1 ADDITIONAL RESULTS FOR GROWING FROM SCRATCH (RQ1)

  • τ reaches 2.91 in math and 2.64 in code, with step time falling below the no-head baseline by steps 30 and 22, respectively.On held-out benchmarks, policy quality stays on par with the autoregressive baseline and τ reaches 2.30–2.34 for math and 2.11–2.27 for code.Source: §C.1 · p. 17, Figure 5 · p. 17, Figure 6 · p. 17

C.2 ADDITIONAL RESULTS FOR TRAINING STRATEGY (RQ2)

  • DCA leads among quality-preserving strategies on all six held-out benchmarks, while joint CE reaches high τ only with policy quality collapsed to zero.The code-domain comparison also shows that joint CE's step-time outcome reverses relative to math, so speed alone does not establish preserved quality.Source: §C.2 · p. 18, Table 8 · p. 18

C.3 ADDITIONAL RESULTS FOR DRAFTING DEPTH (RQ2)

  • At K=3, 5, and 7, deeper heads achieve higher τ, but K=5 is highlighted because end-to-end speedups, including head-update cost, determine depth selection.Held-out τ rises with K while policy quality stays on par with the frozen head; K=7 has the highest τ in every benchmark column.Source: §C.3 · p. 19, Figure 9 · p. 19, Table 9 · p. 19

C.4 ADDITIONAL RESULTS FOR COMPONENT ABLATIONS (RQ3)

  • Reconstruction avoids teacher forcing's roughly 30-second extra actor-update cost; at K=5 DCA leads held-out τ, while removing VGM lowers end-to-end speedups.With VGM removed, speedup falls from 1.37× to 1.28× under KL and from 1.41× to 1.35× under DCA.Source: §C.4 · p. 19, Figure 10 · p. 20, Table 10 · p. 20

D.1 WHICH MODELS CAN GrowMTP ACCELERATE?

  • Across Qwen3-4B, MiMo-7B-SFT, and Qwen3.5-4B-Base, mathematical-reasoning end-to-end speedups are 1.60×, 1.41×, and 1.20×, respectively.The authors describe GrowMTP as model-agnostic and a modular RL-framework component, particularly for base models without MTP heads.Source: §D.1 · p. 21

D.2 DOES THE ACCELERATION GROW WITH RL SCALE?

  • Under stated assumptions, projected end-to-end speedup grows with RL scale, reaching 1.91× at 16K rollout length and 1,000 steps and 2.17× at 32K and 2,000 steps.These are illustrative projections from the measured Qwen3-4B run; the assumptions include linear rollout-time scaling and unchanged acceptance length.Source: §D.2 · p. 21, Table 12 · p. 21

D.3 DOES GrowMTP REPLACE PRETRAINED DRAFT HEADS?

  • From-scratch heads specialize to their training domain: τ falls from 2.91 to 1.45 for the math-trained head evaluated on code, and from 2.64 to 1.75 for the code-trained head evaluated on math.The authors distinguish this narrow-distribution RL setting from general serving across unforeseen deployment distributions, where pretrained MTP heads remain relevant.Source: §D.3 · p. 21, Table 13 · p. 21

E.1 DRAFT-PATH RECONSTRUCTION / E.2 IMPLEMENTING DRAFT-PATH RECONSTRUCTION IN VERL

  • Reconstruction replays recorded drafts with recorded verification distributions at matching draft-conditioned states; implementation stores top-64 target log-probabilities and integrates recording and head updates into the RL loop.Top-64 sparsification yields mean absolute overlap errors of 1.8 × 10^-6 to 5.7 × 10^-6 and per-batch gradient cosine similarity of at least 0.99996; rollout verification still uses the full distribution.Source: §E.1 · p. 22, §E.2 · p. 22

E.3 PARALLELIZING GrowMTP TRAINING

GrowMTP parallelizes reconstruction across independent rollout cycles while retaining serial computation across drafting depth. Batched attention and chunked backward passes preserve chain-by-chain context and bound memory.

Parallelism across Cycles

  • Reconstruction runs K batched forwards, advancing all recorded cycles together at each depth instead of executing each cycle chain by chain.Across a response of three thousand cycles at K=5, the serial forward count falls from fifteen thousand to five.Source: §E.3 · p. 23, Figure 11 · p. 23
  • The shared-cache attention mask gives each query the causal committed prefix and only its own chain's draft entries, preserving chain-by-chain computation.Draft entries are recomputed from each cycle's recorded tokens, and positions come from the rollout record.Source: §E.3 · p. 23, Figure 11 · p. 23

Bounding the Memory

  • Chunks of 1,024 cycles run their own K forwards and backward, releasing activations before the next chunk while accumulated gradients match one response-wide backward.Chunking bounds resident memory; it does not add parallelism across chunks.Source: §E.3 · p. 23

F GrowMTP TRAINING ALGORITHM

The training algorithm records rollout-cycle inputs and verification signals, updates the policy with the RL objective, then updates only the draft head. The accompanying analysis explains why reconstruction matches inference states whereas teacher forcing diverges from them, especially at deeper positions.

F GrowMTP TRAINING ALGORITHM

  • Each training step records speculative rollout cycles, applies the RL objective to responses, updates the head from recorded cycles, and synchronizes the updated weights for the next step's drafting.The head gradient reaches only ϕ; the recorded inputs include the cycle-start hidden state, drafts, verification distributions, and accepted length.Source: F GrowMTP TRAINING ALGORITHM · p. 24

G TRAIN–INFERENCE INCONSISTENCY UNDER TEACHER FORCING

  • Teacher forcing starts a chain at every response position, but inference restarts only at cycle boundaries, so only one of its trained states per cycle matches inference.At greater depths, teacher-forced chains cross cycle boundaries more often and depart further from inference states.Source: G TRAIN–INFERENCE INCONSISTENCY UNDER… · pp. 25–26
  • GrowMTP reconstructs each recorded cycle from its boundary feature using recorded drafts, making all K training states coincide with the states used during inference.The reconstruction rebuilds the inference cache; a teacher-forcing mask cannot substitute for computing those drafted states.Source: G TRAIN–INFERENCE INCONSISTENCY UNDER… · p. 26

H ONE-STEP LAG BOUND UNDER REJECTION SAMPLING

A one-update lag separates the target that supplies draft-head training signals from the target that verifies the next rollout's drafts. The analysis bounds the resulting acceptance shift and relates it to policy movement, while distinguishing a fixed-reference KL measure from the consecutive-policy shift.

One-step lag

  • The head used at step t+1 is trained on verification signals from step t, so its training target is one policy update old.The authors identify this as a possible source of degraded acceptance against the updated target.Source: H ONE-STEP LAG BOUND UNDER REJECTION… · p. 27

Single-position bound

  • At a fixed state, deployment acceptance differs from training-time acceptance by at most the total-variation displacement between the successive target distributions.The bound is uniform in the draft distribution and applies whether acceptance rises or falls.Source: H ONE-STEP LAG BOUND UNDER REJECTION… · p. 27

Magnitude of the target shift

  • To leading order, the target shift is bounded by half the policy update's length in the Fisher metric and vanishes linearly with learning rate.The Fisher matrix is averaged over conditioning states at the maximizing draft position.Source: H ONE-STEP LAG BOUND UNDER REJECTION… · p. 28

Measuring the shift

  • Measured per-token KL stays below 0.007 nats, but it is contextual evidence of a localized policy trajectory, not a direct estimate of the consecutive-policy shift in the bound.Refreshing the head after each RL step makes the acceptance-chain bound depend on the current update rather than accumulated training drift.Source: H ONE-STEP LAG BOUND UNDER REJECTION… · p. 28

I OPTIMIZATION ANALYSIS OF TRAINING OBJECTIVE

The analysis compares the TV-chain and DCA objectives, which use the same acceptance probabilities but differ through a logarithmic transformation. Their gradients behave alike on a single draft cycle, while cycle reweighting can make their batch objectives differ.

Gradient Analysis

  • On a single cycle, TV and DCA share descent direction, stationary points, minimizers, and relative depth weights; DCA rescales gradient magnitude by 1/M.The objectives use the same cycle acceptance probabilities, with DCA taking the logarithm of the shared chain value.Source: I OPTIMIZATION ANALYSIS OF TRAINING… · p. 29

Optimization on a Draft Step

  • Under SGD, DCA applies the scheduled learning rate, while TV’s effective rate is ηM and depends on that cycle’s acceptance.For a chain value M, the respective updates are −η∇ϕLDCA and −ηM∇ϕLDCA.Source: I OPTIMIZATION ANALYSIS OF TRAINING… · p. 30
  • With Adam, a positive factor shared across the gradient history cancels in the moment ratio, but factors varying across draft cycles generally do not.The appendix compares TV and DCA gradients for draft steps and batches under SGD and Adam.Source: I OPTIMIZATION ANALYSIS OF TRAINING… · pp. 29–30

Optimization on a Batch

  • Across a batch, TV weights each DCA cycle gradient gr by Mr, so the objectives can have different stationary conditions without guaranteeing faster convergence.DCA averages gr without those factors, giving smaller-M cycles greater relative weight under DCA.Source: I OPTIMIZATION ANALYSIS OF TRAINING… · p. 30

J LIMITATIONS

The reported savings depend on training-run length and rollout compute share, and the grown head specializes to its RL rollout distribution. Evaluation is limited to 4B–7B models on mathematical and code reasoning using one eight-GPU node.

Run length and compute allocation

  • Per-step acceleration begins within 30 steps in both from-scratch experiments, and 500-step runs recover the initial cost with net savings.Typical RL training spans hundreds or thousands of steps, which the authors say provides time to amortize the initial training overhead.Source: J LIMITATIONS · p. 30

Distribution specialization

  • The grown head retains higher acceptance in its training domain than across domains, so reuse on unrelated tasks requires separate evaluation.Broad deployment is distinct from the within-run acceleration setting studied.Source: J LIMITATIONS · p. 30

Evaluation scope

  • Measurements cover 4B–7B models on mathematical and code reasoning with a single eight-GPU node; larger models, multi-node systems, and asynchronous RL require direct measurement.The paper describes scaling projections as assumption-based rather than measured performance.Source: J LIMITATIONS · p. 30

Related papers

Loading 2609.16648v1…