UDM-GRPO: Stable and Efficient Group Relative Policy Optimization for Uniform Discrete Diffusion Models
| Authors | Jiaqi Wang et al. |
| Year | 2026 |
| HF Upvotes | 6 |
| arXiv | 2604.18518 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Uniform Discrete Diffusion Model (UDM) has recently emerged as a promising paradigm for discrete generative modeling; however, its integration with reinforcement learning remains largely unexplored. We observe that naively applying GRPO to UDM leads to training instability and marginal performance gains. To address this, we propose \Ours, the first framework to integrate UDM with RL. Our method is guided by two key insights: (i) treating the final clean sample as the action provides more accurate and stable optimization signals; and (ii) reconstructing trajectories via the diffusion forward process better aligns probability paths with the pretraining distribution. Additionally, we introduce two strategies, Reduced-Step and CFG-Free, to further improve training efficiency. \Ours significantly improves base model performance across multiple T2I tasks. Notably, GenEval accuracy improves from 69% to 96% and PickScore increases from 20.46 to 23.81, achieving state-of-the-art performance in both continuous and discrete settings. On the OCR benchmark, accuracy rises from 8% to 57%, further validating the generalization ability of our method. Code is available at https://github.com/Yovecent/UDM-GRPO{https://github.com/Yovecent/UDM-GRPO}.
Engineering Breakdown
Plain English
This paper addresses a critical gap in generative modeling by being the first to successfully integrate Uniform Discrete Diffusion Models (UDM) with reinforcement learning for optimization. The authors discovered that naively applying existing RL methods like GRPO to UDM causes training instability and minimal improvements. They propose a framework with two key innovations: treating the final clean sample as the action (rather than intermediate diffusion states) for clearer reward signals, and reconstructing trajectories through the diffusion forward process to maintain alignment with pretraining distributions. They further introduce Reduced-Step and CFG-Free strategies to improve training efficiency.
Core Technical Contribution
The paper's core novelty is identifying and solving the fundamental mismatch between how UDM operates and how standard RL algorithms expect to interact with a policy. Rather than treating each diffusion denoising step as an action in a sequential decision-making framework, the authors reframe the problem by using the final discrete output as the action space—this creates a well-defined reward signal that doesn't degrade through the diffusion process. The trajectory reconstruction via forward diffusion ensures that probability distributions remain anchored to the pretraining distribution, preventing the RL optimization from diverging into out-of-distribution territory. This represents the first principled integration of discrete diffusion with RL, moving beyond naive application of existing RL-to-diffusion techniques.
How It Works
The method operates by first taking a discrete diffusion model pretrained on a task (like text or code generation) and treating the final clean sample—the output after all denoising steps—as the action that receives a reward signal from an external evaluator or preference model. During RL training, the model generates trajectories by: (1) sampling from the diffusion process, (2) computing rewards on the final output, (3) reconstructing the probability path backward using the diffusion forward process to ensure alignment with pretraining. The key architectural insight is that instead of assigning rewards to intermediate denoising states (which causes instability), all credit flows from the final sample backward through the trajectory. The Reduced-Step strategy shortens the diffusion chain during RL fine-tuning to reduce compute, while CFG-Free removes classifier-free guidance during training to eliminate another source of instability. This creates a cleaner optimization landscape where gradients flow consistently from a single reward signal.
Production Impact
For engineers deploying discrete generative models (text, code, molecule generation) with preference optimization, this framework enables significantly more stable training compared to naive RL approaches. You would integrate this by: (1) pretraining a UDM on your task as usual, (2) swapping your RL training routine to use the final-sample-as-action formulation, and (3) optionally reducing diffusion steps during fine-tuning. The production benefits include better convergence (fewer training divergences), faster convergence through Reduced-Step (potentially 3-4x fewer denoising calls per sample), and lower memory overhead by avoiding intermediate state tracking. The trade-off is that you lose fine-grained control over intermediate generation steps—the entire trajectory is optimized as a single unit toward the final reward. Integration complexity is moderate: you need to restructure how you compute rewards and gradients, and you'll need to tune hyperparameters differently than standard GRPO, but the core diffusion model code remains unchanged.
Limitations and When Not to Use This
The paper doesn't address what happens when reward functions are sparse, noisy, or misaligned with the pretraining distribution—these are common in real-world preference data, and the trajectory reconstruction mechanism may struggle if the reward signal contradicts pretraining. The method assumes access to differentiable rewards or preference models during training, which isn't always available in production (e.g., human-in-the-loop evaluation). Computational cost during RL training may still be prohibitive compared to non-diffusion alternatives, since you're still running full diffusion trajectories even with Reduced-Step optimization. The paper doesn't provide theoretical analysis of convergence properties or guidance on when this approach is preferable to simpler alternatives like supervised fine-tuning with rejection sampling. Additionally, scaling to very large models (billion+ parameters) and long sequences is not demonstrated, leaving open questions about practical limits.
Research Context
This work builds directly on two recent research threads: (1) the emergence of discrete diffusion models as competitive alternatives to autoregressive generation, and (2) the growing integration of RL with generative models (exemplified by RLHF for language models and GRPO for diffusion models). The paper positions itself as solving the specific failure mode of GRPO when applied to discrete diffusion—prior work showed GRPO works well for continuous diffusion but breaks down in discrete settings. It advances the state-of-the-art in preference optimization for discrete generation tasks (text, code, categorical data) and opens a research direction toward better RL-diffusion alignment. The insights about trajectory reconstruction and final-sample-as-action could inform future work on other discrete generative architectures, and the efficiency improvements (Reduced-Step, CFG-Free) extend a broader push in the field toward cheaper RL training loops.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::