---
title: 'ESPO: ELBO-based Sequence-level Policy Optimization'
url: https://www.emergentmind.com/topics/elbo-based-sequence-level-policy-optimization-espo
type: topic
---

# ESPO: ELBO-based Sequence-level Policy Optimization

ELBO-based Sequence-level Policy Optimization (ESPO) is a reinforcement learning (RL) framework designed for fine-tuning diffusion language models (DLMs) by treating entire sequence generation as a single atomic decision, using the evidence lower bound (ELBO) as a tractable proxy for the intractable sequence-level likelihood. ESPO addresses key incompatibilities between RL algorithms designed for autoregressive language models and the non-autoregressive, iterative denoising characteristic of DLMs, enabling principled and stable policy optimization in domains such as mathematical reasoning, coding, and planning [2512.03759][2510.08554].

## 1. Formal Objective and Theoretical Foundation

Standard RL policy gradients optimize
$$
J(\pi_\theta) = \mathbb{E}_{x \sim D}\;\mathbb{E}_{y \sim \pi_\theta(\cdot|x)}[R(x, y)],
$$
where $R(x, y)$ is a reward for completion $y$ given prompt $x$. Token-level RL objectives are directly applicable to autoregressive models due to their explicit token conditional factorization. However, DLMs lack such a factorization: sequence likelihood $\pi_\theta(y | x)$ is defined implicitly through diffusion-style denoising, making per-token RL inapplicable.

ESPO overcomes this by directly optimizing sequence-level rewards. The policy is updated via an off-policy group-relative (GRPO-style) surrogate, but with the complete output sequence $y$ treated as a single decision:
$$
J_{\text{seq}}(\theta) = \mathbb{E}_{x,\,y^{(1:G)} \sim \pi_{\text{old}}} \left[ \frac{1}{G}\sum_{i=1}^G \min\left(\rho_{\text{seq}}(y^{(i)}) \cdot \hat{A}^{(i)},\, \text{clip}(\rho_{\text{seq}}(y^{(i)}), 1-\epsilon, 1+\epsilon)\cdot \hat{A}^{(i)}\right) \right],
$$
where the sequence-level importance ratio is
$$
\rho_{\text{seq}}(y) = \frac{\pi_\theta(y|x)}{\pi_{\text{old}}(y|x)}.
$$

Direct computation of $\pi_\theta(y|x)$ is intractable. Instead, ESPO uses the standard evidence lower bound (ELBO) in a $k$-masked form:
$$
\mathcal{L}_\theta(y|x) = \mathbb{E}_{l \sim \mathrm{Unif}(1 \ldots L)}\, \mathbb{E}_{y_l\sim q_l(y_l|y,x)}\left[ \left(\frac{L}{l}\right)\sum_{i=1}^L \mathbf{1}[y_l^i = M]\cdot \log p_\theta(y^i|y_l, x)\right],
$$
with the guarantee $\mathcal{L}_\theta(y|x) \leq \log \pi_\theta(y|x)$. Substituting $\exp\left(\mathcal{L}_\theta(y|x)\right)$ for $\pi_\theta(y|x)$ yields a practical, lower-bounding sequence likelihood proxy.

*Sequence-level RL with ELBO surrogates thus restores formal consistency (avoiding token-level decompositions) and enables non-autoregressive policy optimization*.

## 2. Importance-Ratio Normalization and Clipping

Because the difference $\mathcal{L}_\theta(y|x) - \mathcal{L}_{\theta_{\text{old}}}(y|x)$ grows with sequence length $L$, the naive sequence-ratio
$$
\rho_{\text{raw}}(y) = \exp(\mathcal{L}_\theta(y|x) - \mathcal{L}_{\theta_{\text{old}}}(y|x))
$$
can produce extreme values for long sequences. To remedy this, ESPO normalizes per-token:
$$
\rho_{\text{seq}}(y) = \exp\left(\frac{1}{L}\left(\mathcal{L}_\theta(y|x) - \mathcal{L}_{\theta_{\text{old}}}(y|x)\right)\right).
$$
This stabilization is critical for practical training. Further, PPO-style clipping is applied to limit $\rho_{\text{seq}}$ to $[1-\epsilon, 1+\epsilon]$ during the surrogate loss calculation, preventing instability from rare large-ratio outliers.

Ablations report that per-token normalization is necessary: without it, ratios explode or vanish, undermining optimization [2512.03759].

## 3. Efficient ELBO Estimation and Variance Reduction

Estimating $\mathcal{L}_\theta(y|x)$ for each sequence is computationally demanding. The ELBO is defined as an expectation over continuous mask-fractions $t \in [0, 1]$ and random masking patterns $y_t \sim \pi_t(\cdot | y)$, with
$$
\mathcal{L}_{\rm ELBO}(y|q) = \mathbb{E}_{t \sim U[0,1]} \mathbb{E}_{y_t \sim \pi_t(\cdot | y)} \left[ \frac{1}{t} \sum_{i=1}^L \mathbf{1}[y_t^i = M]\log \pi_\theta(y^i | y_t, q) \right].
$$

Naive double Monte Carlo estimation suffers from severe variance—*time-variance* from sampling $t$ dominates, requiring $N K \gg 100$ network forward passes per sequence for stable estimates [2510.08554].

**Semi-deterministic Monte Carlo (SDMC) / Quadrature Scheme**: ESPO and its extensions (e.g., Group Diffusion Policy Optimization, GDPO) employ deterministic quadrature for the integral over $t$ and minimal Monte Carlo sampling over masking, so that
$$
\int_0^1 g(t)dt \approx \sum_{n=1}^N w_n \frac{1}{K}\sum_{k=1}^K Z(t_n, y_t^{[k]}),
$$
with $N=2$ or $3$ quadrature points and $K=1$ inner mask sample, leading to dramatically lower estimator variance without additional compute.

**Variance Decomposition and Theoretical Guarantees**: The total mean squared error for SDMC decomposes into a Monte Carlo variance scaling as $O(1/(N K))$ and a quadrature bias term that rapidly decays with $N$ (midpoint/trapezoidal: $O(1/N^4)$, Simpson: $O(1/N^8)$). This ensures practical, low-variance estimates with few network evaluations [2510.08554].

| Estimator                     | Variance Term       | Bias$^2$ Term      |
|-------------------------------|--------------------|--------------------|
| Riemann (generic)             | $O(1/(N K))$       | $O(1/N^2)$         |
| Smooth/Quadrature             | $O(1/(N^2 K))$     | $O(1/N^4)$         |

## 4. Robust KL Regularization

A KL regularizer penalizes deviation of $\pi_\theta$ from a reference $\pi_{\text{ref}}$ to maintain conservative policy updates. While exponential-based KL estimators (e.g., $k_3$) can be unstable due to $\exp(\cdot)$ effects for long sequences, ESPO instead applies the quadratic $k_2$ estimator:
$$
\widehat{\text{KL}}_{k_2} = \mathbb{E}_{y \sim \pi_\theta}\left[\frac{1}{2}(\mathcal{L}_\theta(y|x) - \mathcal{L}_{\text{ref}}(y|x))^2\right].
$$
This provides unbiased gradients for $\mathrm{KL}(\pi_\theta \| \pi_\text{ref})$, is free of exponentials, and remains robust for extended outputs [2512.03759].

## 5. Algorithmic Workflow and Practical Implementation

A typical ESPO (or GDPO) RL iteration proceeds as follows:

- For each prompt $x$, sample $G$ completions $\{y^{(i)}\}$ from the behavior policy via a diffusion sampler.
- Compute rewards $R(x, y^{(i)})$ and centered advantages $\hat{A}^{(i)}$.
- For each $y^{(i)}$, estimate $\mathcal{L}_\theta(y^{(i)}|x)$ and $\mathcal{L}_{\theta_{\text{old}}}(y^{(i)}|x)$ using SDMC quadrature.
- Compute length-normalized, clipped importance ratios and robust quadratic KL divergence.
- Formulate the surrogate loss as the mean of clipped-importance-weighted advantages, add the KL penalty, and take a gradient step.

Key training hyperparameters include group size $G$ (e.g., $6$–$16$), number of quadrature points ($N=2$–$3$), Monte Carlo masks per $t$ ($K=1$–$2$), GRPO clipping parameter $\epsilon=0.2$, and KL penalty weight $\beta=0.02$.

## 6. Empirical Results and Benchmark Performance

On arithmetic, planning, and code-completion tasks, ESPO demonstrates consistent, often dramatic, improvements over token-level RL baselines and one-step unmasking methods (diffu-GRPO).

| Method        | GSM8K$\uparrow$ | MATH$\uparrow$ | Countdown$\uparrow$ | Sudoku$\uparrow$ | HumanEval-avg$\uparrow$ | MBPP-avg$\uparrow$ |
|---------------|-----------------|----------------|---------------------|------------------|-------------------------|--------------------|
| Base          | 75.9            | 37.0           | 18.7                | 15.7             | 37.8                    | 37.8               |
| +d1 (GRPO)    | 78.0            | 37.7           | 33.9                | 22.2             | 37.2                    | 36.5               |
| +wd1          | 80.1            | 36.9           | 48.3                | 23.1             | 37.2                    | 36.5               |
| **+ESPO**     | **82.0**        | **39.5**       | **81.0**            | **86.0**         | **40.1**                | **45.4**           |
| Δ vs base     | +6.1            | +2.5           | +62.3               | +70.3            | +2.3                    | +7.6               |

These results show especially large gains in planning (Countdown, Sudoku: $+62$ to $+70$ points), and consistent improvements in math and coding, with evaluation at sequence lengths up to $L=512$, confirming scalability [2512.03759][2510.08554]. Sequence-level+ELBO consistently outperforms any token-level proxy (mean-field or token-ELBO).

## 7. Limitations, Extensions, and Future Directions

- The practical efficacy of ESPO is closely tied to the tightness of the ELBO as a likelihood surrogate—*a plausible implication is that further variance or bias outliers could disrupt optimization for some distributions, though no such breakdown was observed empirically*.
- While ESPO’s convergence is empirically robust, no formal convergence guarantees are established beyond those of standard PPO.
- Future directions include integrating learned value functions, extending to multimodal diffusion LLMs, and exploring cost reductions via distillation or adaptive masking.
- For tasks demanding long-range, sequence-level coherence (e.g., planning, full-program synthesis), ESPO’s atomic sequence optimization provides a distinct practical advantage over denoising-token RL methods.

*This framework establishes sequence-level RL based on ELBO surrogates as a new paradigm for RL in diffusion LLMs, combining theoretical justification, stable optimization, and large empirical gains across challenging domains* [2512.03759][2510.08554].

Source: https://www.emergentmind.com/topics/elbo-based-sequence-level-policy-optimization-espo