UDM-GRPO: Stable and Efficient Group Relative Policy Optimization for Uniform Discrete Diffusion Models

Jiaqi Wang, Haoge Deng, Ting Pan, Yang Liu, Chengyuan Wang +3 more
4/20/2026
cs.CVcs.LG

Official code: Yovecent/UDM-GRPO (29 stars, Apache-2.0, last commit 2026-05-01).

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 UDM-GRPO, 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. UDM-GRPO 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.

View on arXivView PDF

Code Implementations(1)

Yovecent/UDM-GRPOOfficial100%
Apache-2.0

[ICML 2026 Spotlight] UDM-GRPO: Stable and Efficient Group Relative Policy Optimization for Uniform Discrete Diffusion Models

291Shell, PythonApr 13, 20265 months agoApache-2.0

Similar Papers

Cite this paper

@article{wang2026udmgrpo,
  title  = {UDM-GRPO: Stable and Efficient Group Relative Policy Optimization for Uniform Discrete Diffusion Models},
  author = {Jiaqi Wang and Haoge Deng and Ting Pan and Yang Liu and Chengyuan Wang and Fan Zhang and Yonggang Qi and Xinlong Wang},
  year   = {2026},
  eprint = {2604.18518},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2604.18518v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2604.18518.svg)](https://codeofpaper.com/paper/2604.18518)

Discussion