Grow While You Roll: Training the Draft Head Inside the RL Run

TL;DR for operators

LLM reinforcement learning spends a large share of its wall-clock budget generating responses. Speculative decoding can reduce that cost, but it normally assumes that a smaller predictor is already capable of proposing useful tokens before the main model verifies them.

GrowMTP¹ tests whether that capability can instead be learned during the RL workload itself. On Qwen3-4B, a randomly initialized five-token draft head reached an acceptance length of 2.91 after 500 mathematical-reasoning RL steps. Rollout generation became 2.13x faster, and the complete RL step—including draft-head training—became 1.60x faster.

The more consequential result is about control. Higher speculative-token acceptance did not consistently mean better economics. On MiMo-7B-SFT, increasing drafting depth from $K=5$ to $K=7$ raised mathematical-reasoning acceptance from 4.04 to 4.38 while reducing end-to-end speedup from 1.41x to 1.30x. In another experiment, joint cross-entropy training drove acceptance still higher while held-out policy quality fell to zero across all six reported evaluations.

For teams running costly RL post-training, GrowMTP supports a specific possibility: treat the accelerator as a workload-adaptive component that can learn from verification computation the pipeline already performs. But govern it using total step time and policy quality, not acceptance alone.

The accelerator does not necessarily need to exist before the workload

Speculative generation uses a cheaper predictor to propose several future tokens, then asks the main model to verify them. If enough proposals are accepted, the system avoids some sequential calls to the expensive model. The usual operational complication is preparation: a useful predictor may need to be pretrained or warmed up before RL begins.

GrowMTP changes that sequencing. During speculative rollouts, the system already produces target-model verification distributions and records which proposed tokens survive. The paper uses those artifacts as supervision for the draft head during the same RL run.

The Qwen3-4B experiment is the clearest test. Its newly attached recurrent draft module starts randomly initialized, although its embedding and output projection are inherited from the policy backbone. After 500 RL steps, its acceptance length reaches 2.91 on mathematical reasoning and 2.64 on code reasoning. Relative to autoregressive RL, rollout speed improves by 2.13x and 1.93x respectively; after adding the cost of updating the head, end-to-end step speed still improves by 1.60x and 1.56x.

A matched mathematical-reasoning comparison makes the compute tradeoff more concrete. Offline head training reaches essentially the same acceptance length—2.92 versus 2.91—but consumes 555.3 total GPU-hours compared with 502.8 for GrowMTP. The reported 9.4% difference belongs to this 500-step Qwen3-4B configuration, not to speculative training in general.

The business inference is narrower than “online training is cheaper.” For workloads long enough to amortize the learning period, verification computation that would otherwise be transient can become supervision for an adaptive acceleration module. That can remove a separate preparation stage for models that do not already ship with a suitable draft head.

Online draft learning needs rollout-consistent supervision

The head cannot simply be trained as an ordinary next-token predictor and expected to behave correctly several speculative steps deep.

GrowMTP addresses three distinct problems.

First, it reconstructs the actual draft path. During rollout, later draft predictions are conditioned on earlier drafted tokens. Conventional teacher forcing would instead condition training on the committed response, creating different internal states from those encountered during speculation. GrowMTP records the starting state and drafted tokens, then replays those drafts so that each reconstructed head distribution is paired with the verification distribution generated under the same prefix.

The reconstruction ablation is informative because aggregate acceptance barely changes. At $K=5$ on MiMo-7B-SFT mathematical reasoning, teacher forcing reaches $\tau=4.02$ and GrowMTP reaches 4.04. Yet end-to-end speedup rises from 1.28x to 1.41x with reconstruction, partly because the reconstruction route reduces actor-update cost while removing the deeper-position state mismatch.

Second, GrowMTP trains for the dependency across draft depth. At position $k$, expected rejection-sampling acceptance is

where $p_k$ is the target distribution and $q_k$ is the draft distribution under the same state.

A token proposed five steps ahead has no operational value if the speculative prefix was already rejected at step two. Depth-Coupled Acceptance therefore optimizes chains of acceptance rather than treating positions independently. On MiMo-7B-SFT at $K=5$, DCA reaches $\tau=4.04$ versus 3.86 for cross-entropy, 3.85 for KL, and 3.94 for the TV objective, while producing the strongest reported end-to-end speedup among those variants.

Third, Verify-Gated Masking stops using speculative states after the first rejection. The rejected position itself still contains valid verification information; later positions depend on a branch the verifier has already discarded. Removing this masking reduces mathematical-reasoning step speedup from 1.41x to 1.35x under DCA.

Together, these components make online learning follow the states, sequential dependencies, and rejection boundary that the accelerator encounters during actual rollout.

Acceptance is a local metric, not a system objective

The paper’s strongest warning comes from experiments where acceptance improves but the system does not.

On MiMo-7B-SFT, deeper drafting illustrates the cost problem:

The additional accepted tokens at $K=7$ do not compensate for the extra drafting and verification work. A team tuning only $\tau$ would select a configuration that is measurably slower end to end.

The Qwen3.5-4B-Base training-strategy comparison exposes a more serious failure. Joint cross-entropy training reaches acceptance lengths of 3.99 on math and 3.42 on code—higher than the detached DCA configuration’s 3.22 and 3.11. Yet the joint-CE policy scores zero across all six reported held-out benchmarks.

GrowMTP avoids that direct interference by detaching the draft-head update from the policy backbone. The accelerator learns from rollout verification while the backbone continues to follow its original RL objective.

For infrastructure owners, this changes the monitoring requirement. At minimum, an adaptive speculative system needs separate measures for accelerator behavior, total compute cost, and underlying model quality. None can safely stand in for the others.

The learned accelerator belongs to its workload

GrowMTP benefits from specialization: the draft head only needs to approximate the distribution being generated by the current RL workload. The same property limits reuse.

The Qwen3-4B head grown on mathematics averages $\tau=2.32$ across held-out math benchmarks but only 1.45 across code benchmarks. The code-grown head averages 2.20 on code and 1.75 on math.

That is evidence for workload adaptation, not a broadly interchangeable speculative model. A team changing domains should re-evaluate the head rather than assuming that acceleration learned on the previous rollout distribution will transfer.

Other boundaries remain material. The measured experiments cover 4B-7B models, mathematical and code reasoning, and synchronous execution on one node with eight H800 GPUs. Multi-node systems, asynchronous RL, larger models, and different ratios of rollout to optimization cost could change the economics. Short RL runs may also end before the cost of learning the head is recovered.

Treat acceleration as part of the training system

GrowMTP’s main contribution is not simply another way to raise speculative acceptance. It demonstrates that an RL workload can, under the tested conditions, train an accelerator from within its own execution and still finish faster overall.

That creates a different engineering decision. Instead of asking only whether a model has a pretrained draft head, teams can ask whether the workload itself generates enough aligned supervision—and lasts long enough—to grow one economically.

The paper also supplies the governance constraint for that decision. The accelerator should be optimized as part of the complete training system. Acceptance indicates whether drafts resemble what the verifier will accept. End-to-end time determines whether that resemblance saves compute. Held-out policy evaluation checks whether the mechanism used to obtain the speed has damaged the model being trained.

Those are three different questions. GrowMTP’s experiments show why an RL platform needs answers to all three before declaring its accelerator successful.

Cognaptus: Automate the Present, Incubate the Future.