Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion Large Language Models

1 Introduction

• Although dLLMs have shown competitive language-modeling performance, research has focused mainly on pre-training and supervised fine-tuning, while applying RL remains challenging.Unlike autoregressive models, dLLMs iteratively unmask tokens in parallel.Source: §1 · p. 1
• ELBO-based RL estimates require retaining all MC computational graphs, restricting practical sample sizes and worsening likelihood and objective estimates.The resulting bias and variance in estimated objectives and gradients can degrade performance.Source: §1 · pp. 1–2
• BGPO maximizes a linear lower bound whose value and gradient equal those of the ELBO-based objective during on-policy training.Each term depends on one MC sample, enabling gradient accumulation with memory usage that does not depend on sample size.Source: §1 · p. 2
• BGPO significantly improves LLaDA-8B-Instruct across math, code, and planning tasks, outperforming previous dLLM RL algorithms with only marginal increases in average training step time.Experiments also find that increasing the MC sample size reduces gradient bias and variance and improves model performance.Source: §1 · p. 2

2.1 Masked Diffusion Language Models

• Masked dLLMs progressively mask and then recover response tokens while keeping the prompt unmasked during conditional generation.The forward process masks each token with probability t; the reverse process iteratively predicts masked tokens as time decreases.Source: §2.1 · p. 2

2.2 Challenges of Applying RL to dLLMs.

• Iterative denoising makes exact computation of log πθ(y|x) intractable, although likelihoods are needed for RL.Source: §2.2 · p. 3
• diffu-GRPO estimates likelihood in one pass but introduces notable bias, whereas VRPO approximates it with an ELBO using customized Monte Carlo sampling.Source: §2.2 · p. 3
• Retaining all nt computational graphs for gradient calculation requires substantial GPU memory, so practical sample sizes remain small, such as nt = 4.The resulting likelihood and objective approximations are inaccurate.Source: §2.2 · p. 3
• BGPO is proposed as a memory-efficient dLLM RL algorithm that supports large Monte Carlo sample sizes for more accurate approximations and better performance.Source: §2.2 · p. 3

3 BGPO

• BGPO maximizes a constructed lower bound of the ELBO-based objective, with each term tied to one MC sample for gradient accumulation and constant memory.The memory usage is irrelevant to the MC sample size.Source: §3 · p. 3

3.1 Linear Lower Bound Construction

• For nonnegative advantages, BGPO uses a Taylor-based construction; for negative advantages, it uses Jensen’s inequality, yielding a linear lower-bound sum.Separating each sample’s gradient backpropagation keeps memory usage constant and allows a larger sample size.Source: §3.1 · p. 4

3.2 Proof of Equivalence

• In on-policy training, the lower-bound and ELBO-based objectives match in value and gradient; off-policy, this equivalence no longer holds, though experiments show policy improvement.The on-policy equivalence also makes the lower-bound objective an effective approximation of the original RL objective.Source: §3.2 · p. 5

3.3 Final Loss of BGPO

• BGPO adopts group-based advantage estimation by sampling G responses per prompt and assigning each response a reward before defining its advantage.The paper then formulates the BGPO loss accordingly.Source: §3.3 · p. 5

4.1 Setup

• Experiments use LLaDA-8B-Instruct on math, coding, and planning, with MATH and GSM8K, DeepCoder, MBPP and HumanEval, and Countdown and Sudoku benchmarks.Math and coding models are evaluated on their respective test sets; planning uses the same splits as d1.Source: §4.1 · p. 5

4.2 Main Results

• BGPO beats diffu-GRPO and VRPO-OL across tasks, improving LLaDA-8B-Instruct by about 5.5% (math), 2.5% (coding), 14.9% (Sudoku), and 68.0% (Countdown).In the off-policy setting, BGPO also surpasses all baselines on math and Countdown and achieves comparable improvements on other tasks.Source: Table 1 · p. 6, §4.2 · p. 6
• BGPO’s training reward is higher than the two baselines in most steps, with particularly faster reward growth and higher reward on Countdown.The authors attribute these patterns to a more accurate optimization direction from larger MC sample sizes.Source: Figure 2 · p. 6, §4.2 · p. 7

4.3 Effect of Increasing MC Sample Sizes

• As MC sample size increases, BGPO’s gradient variance and bias decrease, while diffu-GRPO has comparatively large values and VRPO-OL exceeds H100 memory at nt = 8.On math tasks, model performance consistently improves as nt rises from 1 to 16.Source: §4.3 · p. 7, Table 2 · p. 7, Figure 3 · p. 7

4.4 Ablation of Lemma 1 and Lemma 2

• In the off-policy math ablation, using either Lemma 1 or Lemma 2 alone leads to significant performance degradation.The ablation uses an off-policy setting because on-policy equivalence does not reveal the lemmas’ individual roles as boundaries.Source: §4.4 · p. 7, Table 3 · p. 7

4.5 Out-of-domain Performance

• Models trained on math improve planning performance, while coding-trained models improve on both math and planning tasks.The authors describe these results as evidence of BGPO’s generalizability.Source: §4.5 · pp. 7–8, Table 4 · p. 8

4.6 Training Efficiency Comparison

• With four times VRPO-OL’s MC sample size, BGPO’s average math-training step time increases only slightly because response rollout dominates runtime.The comparison sets maximum response length to 512.Source: Figure 3 · p. 7, §4.6 · p. 8

5.1 Diffusion Large Language Models

• dLLMs iteratively unmask tokens in parallel, unlike left-to-right autoregressive generation; several open-source models adapt pretrained AR models, while LLaDA is trained from scratch.The passage also notes commercial dLLMs’ code-generation performance and faster inference.Source: §5.1 · p. 8

5.2 Reinforcement Learning for dLLMs

• For dLLM RL, diffu-GRPO’s single-pass likelihood estimate is efficient but biased, whereas ELBO-based Monte Carlo estimates can be accurate but are memory-limited.BGPO is presented as addressing that memory constraint to support larger samples and reduce approximation bias and variance.Source: §5.2 · p. 8

6 Conclusion

• BGPO uses large Monte Carlo samples to approximate sequence-level log-likelihoods and the final objective, with the authors reporting reduced approximation bias and variance and better performance.The paper also proves equivalence between the BGPO objective and the prior ELBO-based objective.Source: §6 · p. 8

7 Limitations

• Experiments cover only 8B-level models; no larger open-source dLLMs were available, and computational resources were limited.The authors state that BGPO may apply to larger dLLMs, but report no experiments at that scale.Source: §7 · p. 9

• All models and datasets used in the work are publicly published with permissible licenses.Source: §8 · p. 9

• Block-wise decoding is used in both training and evaluation, with response length, diffusion step, and block size following Zhu et al. (2025) and Zhao et al. (2025a).Source: A Detailed Hyperparameters · p. 11

• Performance generally peaks when inference length matches training length, then deteriorates as sequence length increases across multiple benchmarks.Table 6 presents the inference-time length scaling results.Source: B Length Extrapolation Analysis · p. 11, Table 6 · p. 12
• The authors attribute the decline to current open-source dLLMs' limited long-chain-of-thought reasoning and difficulty retaining and processing information over extended contexts.They note that this observation is consistent with Zhu et al. (2025) and Zhao et al. (2025a).Source: B Length Extrapolation Analysis · p. 11