Reinforcement learning grows its own faster draft head online
GrowMTP: Can RL Grow Its Own Draft Head?
Machine LearningComputation and Language
Summary
Training large language models with reinforcement learning (RL) often takes a long time because generating text step-by-step is slow. The authors noticed that the process of RL itself can be used to train a small helper model, called a draft head, that speeds up text generation without extra pretraining. They created GrowMTP, a way to train this draft head during RL training, making the whole system faster. This approach works well on different models, speeding up both the text generation step and the overall training.
What this means in practice
- •For machine learning engineers: Accelerate reinforcement learning training of language models by integrating draft head training within the RL process, reducing the need for separate pretraining.
- •For natural language processing developers: Speed up the generation of text from language models during reinforcement learning by using an automatically trained draft head that adapts online.
Authors
Minghua He, Lingzhe Zhang, Yuan Liu, Xiao Zhou, Aiwei Liu
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.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.