Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling

Paper: 2606.12370
Authors: Yucheng Li, Huiqiang Jiang, Yang Xu, Jianxin Yang, Yi Zhang, Yizhong Cao, Yuhao Shen, Fan Zhou, Rui Men, Jianwei Zhang
Categories: cs.LG, cs.CL

The Gap

Existing work uses Multi-Token Prediction (MTP) to accelerate RL rollouts via speculative decoding: a draft model predicts K tokens, and the target model validates them. This works well during pretraining, but during post-training RL, the acceptance rate collapses (often below 50%), killing the speedup. Prior attempts treat this as a training stability issue or simply avoid MTP in RL. The paper pinpoints the root cause: the acceptance rate is bounded by the model’s entropy—as entropy rises during RL (new tasks, reward shaping), the agreement between the draft and target drops linearly. This is an information-theoretic barrier, not just an engineering bug.

[Problem: MTP acceptance rate drops in RL]
       |
       v
[Assumption: Entropy fluctuation causes this]
       |
       v
[Method: Rejection sampling + TV loss + offline MTP training]
       |
       v
[Evidence: ~95% acceptance rate, 1.8x end-to-end acceleration across 3 tasks]
       |
       v
[Conclusion: The entropy bound can be broken by changing the sampling strategy and training objective]

The Increment

One sentence: Before this paper, MTP acceptance rates collapsed in RL, making speculative decoding impractical; after this paper, by switching to probabilistic rejection sampling and directly optimizing acceptance via TV loss, the rate stays high (95%) and the wall-clock acceleration is real (1.8x).

Core Mechanism

The method, called Bebop, has three components that work together.

First, replace the standard greedy top-1 draft sampling with probabilistic rejection sampling. Instead of generating the most likely draft token, the draft model samples from its full distribution, and the target model accepts or rejects based on a ratio of probabilities. This accounts for the fact that during RL, the target’s distribution becomes more uniform (higher entropy), so greedy drafts disagree more often. Probabilistic sampling aligns the two distributions better by design.

Second, the training objective for the MTP head is changed from cross-entropy or KL divergence to an end-to-end Total Variation (TV) loss. The KL divergence penalizes mismatched probabilities in expectation but does not directly maximize the expected number of accepted tokens under the rejection sampling procedure. The TV loss directly maximizes the acceptance rate because, for a given draft distribution p and target q, the optimal rejection rate is proportional to the TV distance $TV(p,q)$. By minimizing TV, Bebop achieves ~10% higher acceptance.

Third, the paper shows that online training of the MTP head during RL is unnecessary. A single pre-RL training session with the TV loss and rejection sampling produces an MTP head that maintains a constant acceptance rate throughout the entire RL training. This eliminates the cost of updating the draft model during RL, a major practical hurdle.

[ASCII diagram of method internals]

                RL rollout loop:
                    |
   +---

![Hero diagram](/arxiv-visuals/breaking-entropy-bounds-accelerating-rl-training/HeroScene.png)

-------------v------------------------------+
   |    Pre-RL: train MTP head with TV loss       |
   |        + rejection sampling in data gen      |
   +----------------+-----------------------------+
                    |
                    v (fixed MTP head, no further update)
   +----------------v------------------------------+
   | During RL:                                    |
   |  (a) Draft model (MTP head) samples tokens    |
   |      probabilistically from its distribution  |
   |  (b) Target model (RL policy) accepts/rejects |
   |      via rejection sampling                   |
   |  (c) Accepted tokens executed, skipped steps  |
   +----------------+-----------------------------+
                    |
                    v
              Rollout speedup (1.8x)

Structural metaphor: Think of the RL rollout as a live theater improvisation where the main actor (the RL policy) must deliver lines, but each line takes time to think. A script prompter (the MTP head) whispers multiple future lines at once. In the old way, the prompter always suggests the most likely next line (greedy). But during a challenging scene (high entropy RL task), the actor often deviates, making the prompts useless. Bebop changes the prompter to occasionally give alternative suggestions (probabilistic sampling), and trains the prompter specifically for this back-and-forth (TV loss). Crucially, the prompter rehearses before the show starts (pre-RL training) and never changes again, because its skill transfers perfectly to any scene in the play. The actor now accepts most prompts, finishing lines much faster.

Key Concepts

• Multi-Token Prediction (MTP): Instead of predicting only the next token, the model predicts the next K tokens in one forward pass. These are usually handled by a lightweight “draft head” appended to the model. During generation, the draft provides a block of tokens that can be validated in parallel (speculative decoding), saving K-1 sequential steps if the block is accepted. The challenge is that the draft must be close to the target’s output, otherwise the rejection rate kills the benefit.

• Rejection Sampling (in speculative decoding): Given a draft token from distribution p and the target’s distribution q, the token is accepted with probability min(1, q(x)/p(x)). If accepted, we move to the next draft token; if rejected, we sample from a corrected distribution (q - p)+. This mechanism ensures the final output distribution matches q exactly, regardless of p’s quality. The expected number of accepted tokens per step is related to the similarity between p and q. The paper’s insight is that during RL, q’s entropy increases, making greedy p far from q, but probabilistic p (sampling from p itself) remains closer.

• Total Variation (TV) Loss: TV distance between two distributions p and q is sup_A |p(A) - q(A)|, or equivalently 1/2 sum_x |p(x) - q(x)|. It directly measures the probability that optimal rejection sampling would reject a sample from p. The paper proposes to train the draft head by minimizing TV(p||q) on the RL-specific data distribution, which is equivalent to maximizing the expected acceptance rate. This is a direct optimization, unlike cross-entropy or KL which are surrogates.

Framework Shift

Before (mainstream approach):                After (this paper):
                                              (Pre-RL phase)
   [MTP head trained with CE/KL]              [MTP head trained with TV loss]
          |                                             |
          v                                             v
   (During RL: rollout)                      (During RL: rollout)
   [Draft head = greedy top-1]               [Draft head = probabilistic sample]
   [Target accepts/rejects]                  [Target accepts/rejects via rejection sampling]
          |                                             |
          v                                             v
   Acceptance rate drops (40-60%)            Acceptance rate stays ~95%
   Speedup: 1.0x - 1.2x                      Speedup: up to 1.8x

One sentence: From greedy draft + CE training to probabilistic draft + end-to-end TV training plus offline pre-RL training, the core shift is recognizing that the conventional training objective and sampling strategy fundamentally limit acceptance during RL, and that these can be replaced with a principled alignment method.

Expert Assessment

Problem choice: Real gap. RL training for LLMs is expensive, and rollout is the bottleneck. Prior work on speculative decoding implicitly avoided the post-training setting because “everyone knows” acceptance rates drop. This paper formalizes that intuition and solves it. It sits at the intersection of RL for LLMs and efficient generation—both hot areas.

Method maturity: Clever insight, not brute force. The key realization (entropy bound) is elegant and backed by a clear negative linear relationship. The proposed fixes (rejection sampling, TV loss) are principled and light-touch. However, the offline training trick is partly empirical luck; it might not hold for all tasks or model scales. The method is mature enough for deployment.

Experimental integrity: Baselines are fair—they compare against no-MTP, MTP with CE, MTP with KL, and a naive online update. The speedup measurements account for the overhead of rejection sampling. The paper tests three tasks (math, code, agent) and three model families. Biggest red flag: the TV loss is reported to give ~10% acceptance improvement over KL, but the variance and statistical significance are not fully discussed. Also, the offline training claim needs more ablation: what if RL data distribution drifts heavily? Still, the evidence is strong enough.

Writing quality: The paper is well-structured, with clear problem statement and logical flow. The “entropy bound” section is the highlight. Where it cuts corners: the ablation of TV loss vs. other alternatives (e.g., Jensen-Shannon) is missing. The “related work” section is thin—readers benefit from a longer discussion on speculative decoding in non-stationary settings. If the authors added a 2-page appendix on failure cases (e.g., tasks where offline MTP fails), the paper would be significantly stronger.

Verdict: weak accept — The paper solves a real engineering problem with a sound theoretical insight and produces convincing results. It deserves attention from anyone building RL pipelines for LLMs, though the offline training claim warrants caution in production.

Takeaways

1. If you use speculative decoding in RL, switch from greedy to probabilistic draft sampling. It costs nothing (the distribution is already computed) and makes acceptance robust to entropy changes.
2. Optimize your draft head with a loss that directly targets your acceptance metric. CE/KL are fine for pretraining but suboptimal for rejection sampling. The TV loss is a drop-in replacement that yields consistent gains.
3. Do not bother updating the draft head online during RL. Train it once on pre-RL data using the TV loss, then freeze it. This eliminates a major source of complexity and instability.
4. Monitor entropy of the RL policy during training. If it rises, your speculative decoding will degrade. The paper’s linear model can even predict the acceptance rate drop and schedule MTP head refreshes (though they show it’s unnecessary here).

[问题: 强化学习中 MTP 接受率下降]
       |
       v
[假设: 熵波动是罪魁祸首]
       |
       v
[方法: 拒绝采样 + 全变差损失 + 离线 MTP 训练]
       |
       v
[证据: 约 95% 接受率，三个任务上 1.8 倍加速]
       |
       v
[结论: 改变采样策略和训练目标可以打破熵界]

[方法内部 ASCII 图]

                  强化学习推演循环:
                      |
   +------------------v----------------------------+
   |  预训练阶段: 用 TV 损失训练 MTP 头           |
   |   + 用拒绝采样生成训练数据                    |
   +------------------+---------------------------+
                      |
                      v (MTP 头固定，不再更新)
   +------------------v---------------------------+
   |  强化学习阶段:                               |
   |  (a) MTP 头从分布中概率采样草稿 token        |
   |  (b) 目标模型用拒绝采样接受/拒绝             |
   |  (c) 接受的 token 直接执行，跳过顺序步骤     |
   +------------------+---------------------------+
                      |
                      v
                 推演加速 (1.8 倍)

之前（主流方法）:                          之后（本文方法）:
                                            （预训练阶段）
   [MTP 头用 CE/KL 训练]                    [MTP 头用 TV 损失训练]
          |                                            |
          v                                            v
   （强化学习推演中）                        （强化学习推演中）
   [草稿头 = 贪婪 top-1]                    [草稿头 = 概率采样]
   [目标接受/拒绝]                          [目标通过拒绝采样接受/拒绝]
          |                                            |
          v                                            v
   接受率下降 (40-60%)                      接受率保持约 95%
   加速: 1.0x - 1.2x                        加速: 最高 1.8x