ELBO-Based Block-Level Policy Optimization

• The paper introduces a novel ELBO-based block-level policy optimization method that decomposes Monte Carlo estimations to tackle intractable sequence-level likelihoods in RL for diffusion models.
• It employs block-wise gradient accumulation and variance reduction techniques to mitigate memory bottlenecks and enhance computational efficiency.
• Empirical results demonstrate improved performance on benchmarks such as MATH500, Countdown, and HumanEval with increased MC sample sizes compared to prior methods.

ELBO-based block-level policy optimization encompasses a set of methodologies for reinforcement learning (RL) with diffusion-based LLMs (dLLMs) and discrete diffusion LLMs (DLMs), where the intractability of the sequence-level log-likelihood necessitates evidence lower bound (ELBO) estimation. These algorithms address the computational and statistical challenges of optimizing RL objectives that incorporate such ELBO-based surrogates, combining strategies for memory efficiency, scalable Monte Carlo (MC) estimation, and variance reduction by leveraging block-level decompositions and novel estimator constructions (Lin et al., 13 Oct 2025, Rojas et al., 9 Oct 2025).

1. The ELBO Surrogate and Block-Level Estimation

For dLLMs and DLMs, the log-likelihood $\log\pi_\theta(y\mid x)$ of generating target $y$ from prompt $x$ is intractable. The standard approximation is the ELBO,

$$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$

which is estimated via MC samples over both $t$ and $y_t$ (for a diffusion time $t$ and a partially masked sequence $y_t$) (Rojas et al., 9 Oct 2025). The RL policy objective, as used in PPO-style surrogates, composes these likelihood estimates in the policy-ratio: $$\mathcal{J}(\theta) = \mathbb{E}_{x,y}\left[e^{\log\pi_\theta(y\mid x) - \log\pi_{\theta_{\mathrm{old}}}(y\mid x)}\; A(x,y)\right],$$ and, with ELBO approximations, yields a surrogate objective involving nonlinear functions of averages over MC samples.

Block-level estimators refer to the strategy of decomposing this ELBO estimation into computational "blocks"—fixed quadrature points or independent MC samples—enabling fine-grained control over bias and variance while maintaining computational tractability.

2. Monte Carlo Bias, Variance, and Computational Limitations

A central challenge is that direct MC estimation of the ELBO and its appearance within nonlinearities exacerbates both variance and memory usage. Specifically, naively sampling $n_t$ MC pairs $y$0 and forming

$y$1

leads to gradients that depend on all $y$2 forward computational paths due to the nonlinear exponential in the RL objective. Consequently, memory usage grows linearly with $y$3, restricting feasible sample sizes (as low as $y$4–$y$5 in practice), increasing estimation bias and variance in both objectives and gradients (Lin et al., 13 Oct 2025).

Analysis of estimator variance reveals two distinct contributions: (1) masking variance from $y$6 at each $y$7, and (2) time-sampling variance from $y$8 itself. With small $y$9, the time-sampling variance typically dominates, deteriorating the statistical efficiency of the estimator (Rojas et al., 9 Oct 2025).

3. Boundary-Guided Policy Optimization (BGPO): Linear Lower Bound

Boundary-Guided Policy Optimization (BGPO) introduces a linear lower bound to the original ELBO-based surrogate to overcome memory bottlenecks by enabling block-wise (per-sample) gradient accumulation. The core algebraic insight is to lower-bound the exponential of the mean per-sample difference $x$0 between new and old policies: $x$1 so that

$x$2

For $x$3, first-order Taylor gives $x$4; for $x$5, Jensen's inequality provides $x$6. This results in the lower bound: $x$7 where

$x$8

A key property is that each $x$9 depends only on the $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$0-th sample, permitting gradients to be accumulated sequentially (block-wise) rather than requiring all forward graphs to be in memory, so memory no longer scales with $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$1 (Lin et al., 13 Oct 2025).

In the on-policy regime ($$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$2), $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$3 for all $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$4, and both the original surrogate and the lower bound coincide in value and gradient: $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$5

4. Block-Wise Decomposition and Algorithmic Structure

BGPO proceeds by sequentially iterating over MC samples (blocks), backpropagating each $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$6 independently and accumulating gradients. The following table summarizes the key computational properties afforded by block-level decomposition in BGPO:

Property	Description	Implication
Linearity	Objective decomposes into sum over $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$7 samples	Constant memory usage
Block-independence	Each term depends on only one MC sample	Enables gradient accumulation
On-policy equivalence	Value/gradient match to full non-linear objective	Guarantees tight approximation

The corresponding pseudocode applies to dataset minibatches, group sizes $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$8, and MC sample sizes $$\mathcal{L}_{\mathrm{ELBO}}(y\mid x) = \mathbb{E}_{t\sim U[0,1]}\; \mathbb{E}_{y_t\sim q(\cdot\mid t,y,x)} \left[\frac{1}{t}\sum_{i}\mathbf{1}[y_t^i=\mathrm{mask}]\; \log\pi_\theta(y^i \mid y_t, x)\right] \leq \log \pi_\theta(y\mid x),$$9, updating parameters via gradient accumulation with constant memory (Lin et al., 13 Oct 2025).

5. Group Diffusion Policy Optimization (GDPO) and Deterministic Block Selection

An alternative block-level methodology is exemplified by Group Diffusion Policy Optimization (GDPO) (Rojas et al., 9 Oct 2025). GDPO reduces variance by partitioning the ELBO integral into $t$0 deterministic "blocks" (quadrature points $t$1 with weights $t$2), at each of which a small $t$3 number of MC samples are drawn over $t$4. The estimator is

$t$5

where $t$6 is defined as in the ELBO.

This semi-deterministic MC (SDMC) block-wise approach shrinks the "time-sampling variance" by fixing quadrature points, yielding MSE bounds of $t$7 with optimal $t$8 versus $t$9 for double-MC, where $y_t$0 is the total computational budget (Rojas et al., 9 Oct 2025).

6. Empirical Results and Performance Analysis

BGPO achieves substantial improvements in downstream tasks by exploiting its block-level linear lower bound and scalable $y_t$1. Experimental comparisons using LLaDA-8B-Instruct as the backbone show that increasing the MC sample size $y_t$2 yields higher exact match and code pass@1 accuracy than previous VRPO-OL and diffu-GRPO baselines at comparable or lower memory cost. For instance, on the MATH500 mathematics benchmark, BGPO at $y_t$3 attains 45.7% accuracy versus 44.1% (VRPO-OL, $y_t$4), while on Countdown (planning), BGPO with $y_t$5 achieves 87.5% accuracy versus 84.8% (VRPO-OL). Code generation on HumanEval also improves to 47.6% (BGPO, $y_t$6) over 44.8% (VRPO-OL, $y_t$7). An ablation on $y_t$8 demonstrates monotonic gains up to $y_t$9 (Lin et al., 13 Oct 2025).

GDPO, using $t$0 and $t$1, reports improvements of +4.5 percentage points on GSM8K, +5.6 on MATH500, and +47.7 on Countdown over base models, confirming that low-variance block-level ELBO estimators reliably improve policy optimization for DLMs (Rojas et al., 9 Oct 2025).

7. Connections, Limitations, and Prospects

ELBO-based block-level policy optimization exploits the structure of sequence-level objectives to circumvent computational obstacles inherent in diffusion-based generative modeling. BGPO's linear lower bound allows memory-efficient exploitation of large MC sizes, reducing approximation and gradient variance, and is provably tight in the on-policy regime. Alternative methodologies such as GDPO use deterministic block selection to minimize variance while maintaining computational efficiency.

A plausible implication is that future extensions may further improve estimator tightness, variance–bias tradeoff, or adaptively select block partitions for maximal sample efficiency. However, the necessity of ELBO-based surrogacy for intractable likelihoods enforces a fundamental statistical ceiling; thus, advances hinge on balancing approximation fidelity, computational resource constraints, and estimator design for highly expressive diffusion architectures.

References:

• "Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion LLMs" (Lin et al., 13 Oct 2025)
• "Improving Reasoning for Diffusion LLMs via Group Diffusion Policy Optimization" (Rojas et al., 9 Oct 2025)