Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective

2025-12-03

AI Summary

To address the challenge that diffusion-based large language models (dLLMs) lack token-level conditional probability decomposition—rendering them incompatible with standard token-level reinforcement learning (RL) methods such as GRPO—this paper proposes the first principled, end-to-end sequence-level RL framework. Methodologically, it introduces a differentiable, stable, and scalable policy optimization objective by leveraging the evidence lower bound (ELBO), derived via variational inference, as a sequence-level likelihood surrogate; further incorporating per-token importance-weight normalization and robust KL divergence estimation. Crucially, the framework supports large-scale training without requiring autoregressive assumptions. Empirically, it achieves significant improvements over baselines across mathematical reasoning, code generation, and planning tasks: +20–40 points on Countdown, and consistent gains on mainstream mathematical and coding benchmarks.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Reinforcement Learning (RL) has proven highly effective for autoregressive language models, but adapting these methods to diffusion large language models (dLLMs) presents fundamental challenges. The core difficulty lies in likelihood approximation: while autoregressive models naturally provide token-level conditional probabilities essential for token-level RL objectives (e.g., GRPO), dLLMs generate sequences through iterative non-autoregressive denoising steps that lack this factorization. To address this fundamental mismatch, we propose ELBO-based Sequence-level Policy Optimization (ESPO), a principled RL framework that treats entire sequence generation as a single action and uses the ELBO as a tractable sequence-level likelihood proxy. Our method incorporates per-token normalization of importance ratios and robust KL-divergence estimation to ensure stable large-scale training. Extensive experiments on mathematical reasoning, coding, and planning tasks demonstrate that ESPO significantly outperforms token-level baselines, achieving dramatic improvements of 20-40 points on the Countdown task, while maintaining consistent gains on math and coding benchmarks. Our approach establishes sequence-level optimization as a principled and empirically effective paradigm for RL in dLLMs. Our code is available at https://github.com/ML-GSAI/ESPO.

Problem

• Adapts RL to diffusion LLMs lacking token-level probabilities
• Proposes sequence-level RL using ELBO as likelihood proxy
• Enhances performance on reasoning, coding, and planning tasks

Innovation

• Sequence-level policy optimization using ELBO proxy
• Per-token normalization for stable importance ratios
• Robust KL-divergence estimation in large-scale training