PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiaqi Wang, Haoge Deng, Ting Pan, Yang Liu, Chengyuan Wang, Fan Zhang, Yonggang Qi, Xinlong Wang

arXiv:2604.18518v4cs.CVcs.LG

TL;DR

Integrating reinforcement learning with uniform discrete diffusion remains challenging for text-to-image generation because naive GRPO adaptation is unstable. UDM-GRPO treats the final clean sample as the action, reconstructs trajectories through forward diffusion, and improves performance across multiple benchmarks, including GenEval accuracy from 0.69 to 0.96.

  • Problem

    Uniform discrete diffusion models lack an established reinforcement-learning integration for improving preference alignment and compositional text-to-image generation.

  • Method

    UDM-GRPO uses the final clean sample as the action, reconstructs trajectories through forward diffusion, and adds Reduced-Step and CFG-Free training strategies.

  • Results

    UDM-GRPO improves the base model across multiple text-to-image tasks, raising GenEval accuracy from 0.69 to 0.96 and achieving state-of-the-art performance in continuous and discrete settings.

  • Takeaways & Limitations

    Stable reward growth without collapse and bounded KL divergence support robust training dynamics for reinforcement learning with uniform discrete diffusion.

  • Takeaways & Limitations

    Naive GRPO adaptation can collapse after transient gains, with severe reward oscillations and sharply increasing KL divergence.

Abstract

from arXiv · show

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.

1. Introduction

UDM-GRPO integrates GRPO with Uniform Discrete Diffusion for text-to-image generation, addressing instability through final-sample actions and forward-process trajectory reconstruction. Reduced-Step and CFG-Free strategies further improve training efficiency while enabling stable reward gains and bounded KL divergence.

  • Limitations: Intermediate predicted samples produce unstable optimization because early-step predictions are high-entropy and inaccurate.Direct adaptation also induces distribution shift because reverse-process optimization diverges from the forward process used during pretraining.
  • Method: UDM-GRPO is the first framework integrating Uniform Discrete Diffusion with GRPO for text-to-image generation.It addresses instability by defining the final clean sample as the policy action and reconstructing trajectories through the forward process.
  • Efficiency: Reduced-Step optimization concentrates training on critical high-noise timesteps to mitigate slow convergence from gradient dispersion across multiple steps.The strategy focuses optimization where it is most effective rather than distributing gradients across the full multi-step process.
  • Efficiency: CFG-Free training avoids jointly optimizing conditional and unconditional objectives, substantially reducing computational overhead.This strategy is introduced specifically to improve training efficiency.
  • Results: UDM-GRPO shows stable and sustained reward improvement without optimization collapse while maintaining low and bounded KL divergence.The baseline instead experiences collapse after 500 steps, with violent reward oscillation and exploding KL divergence.

2. Related work

Prior work established diffusion over categorical spaces and advanced discrete image generation through masked token prediction, while reinforcement-learning alignment developed from preference optimization and PPO-based diffusion methods toward GRPO extensions.

  • Discrete Diffusion Model: Discrete diffusion extends continuous diffusion to categorical variables, introducing challenges from their categorical nature.Multinomial transitions and discrete denoising objectives formalized diffusion over categorical spaces.
  • Discrete Diffusion Model: Masked image modeling enables discrete image generation through iterative masked-token prediction, with strong performance and efficient parallel decoding.This line builds on formalizations of diffusion over categorical spaces.
  • Reinforcement Learning for Text-to-Image Models: Text-to-image reinforcement-learning alignment commonly follows Direct Preference Optimization or policy-based reinforcement learning using human-preference feedback.Direct Preference Optimization treats alignment as preference classification over ranked output pairs and updates policies without explicit reward modeling.
  • Reinforcement Learning for Text-to-Image Models: DDPO formulates diffusion denoising as a multi-step Markov Decision Process, extending reinforcement learning beyond likelihood maximization.Earlier policy-based efforts primarily adopted Proximal Policy Optimization, while GRPO has since been extended to visual generation.
  • Reinforcement Learning for Text-to-Image Models: GRPO has been extended from large-language-model applications to autoregressive visual generation, mask-based diffusion, continuous diffusion models, and flow-matching.These extensions follow DDPO and earlier PPO-based efforts in diffusion reinforcement learning.

3. Initial Exploration

This section introduces the motivation and foundations for integrating GRPO with Uniform Discrete Diffusion, then describes a preliminary reverse-process formulation. The formulation redefines the action to preserve differentiability and enables policy optimization, with empirical improvement over the base model.

  • Motivation: Uniform Discrete Diffusion’s supervised cross-entropy training limits optimization of complex, non-differentiable objectives and intricate generation tasks.The authors therefore adopt GRPO to address these limitations.
  • Uniform Discrete Diffusion: Uniform discrete diffusion transports categorical noise toward data distributions by iteratively refining all tokens, unlike masking-based diffusion’s typically non-refinable local generation.The framework defines intermediate probability paths, continuous-time probability velocities, cross-entropy training, and Euler-based inference.
  • Trajectory Formulation: Flow-GRPO converts ODE-based denoising dynamics into an SDE, while UDM defines backward, pretraining, and forward-process-based trajectories for comparison.The forward-process-based trajectory perturbs the model’s estimated clean sample using the same forward diffusion process as pretraining.
  • Preliminary Integration: In preliminary UDM-GRPO integration, the non-differentiable intermediate sampling step blocks gradients, so the action is redefined as the intermediate prediction x_t1.Optimizing along Xbackward preserves differentiability, retains Euler sampling efficiency, and enables policy optimization.
  • Preliminary Integration: Optimizing the redefined action along Xbackward empirically improves the base model’s performance.The preliminary integration computes policy probabilities from discrete-token logits and sampled intermediate predictions.

4. Method

UDM-GRPO addresses instability in naive Uniform Diffusion–GRPO integration by using the final clean sample as the action and a forward-process trajectory. Reduced-Step and CFG-Free training further improve efficiency while preserving or surpassing conventional CFG-based performance.

  • Limitations: Naive Uniform Diffusion–GRPO training becomes unstable as reward fluctuations and sharply rising KL divergence degrade performance.Reward initially increases for 500 training steps before severe fluctuations emerge.
  • Limitations: Using unreliable intermediate predictions as actions teaches misleading information, potentially destabilizing training and causing collapse.Early predictions have high entropy and are incoherent, whereas the final clean prediction is more accurate.
  • UDM-GRPO: UDM-GRPO selects the accurate, reward-aligned clean sample as the action and replaces the generated backward trajectory with a forward diffusion trajectory.These changes align optimization with the final-sample reward and the pretraining process.
  • UDM-GRPO: Xforward achieves consistently lower FID than Xbackward across timesteps and better matches the pretraining distribution.Xbackward deviates particularly at steps 1 and 2 because of accumulated early-prediction errors.
  • Training acceleration: Reduced-Step training randomly selects three consecutive timesteps from the first half of diffusion, while CFG-Free training removes classifier-free guidance complexity.CFG-Free quality initially declines but later recovers and surpasses conventional CFG-based methods.

5. Experiments

Experiments show that UDM-GRPO substantially improves URSA-based text-to-image generation across compositional, preference-alignment, and OCR tasks. Ablations attribute these gains to clean-sample actions, forward trajectory reconstruction, few-step optimization, and CFG-free training.

  • Experimental Setup: Experiments cover compositional image generation, visual text rendering, and human preference alignment using pretrained URSA text-to-image models.Training uses group sampling with 16 groups of 8 image samples per batch and evaluates with 25 inference steps.
  • Main Results: 0.69 to 0.96: UDM-GRPO boosts the overall GenEval score and surpasses prior reinforcement-learning methods and pretrained baselines.The comparison establishes state-of-the-art performance across continuous and discrete settings.
  • Main Results: UDM-GRPO achieves state-of-the-art PickScore performance and substantially improves the pretrained model’s poor OCR performance.Qualitative results also show improved spatial arrangement, attribute binding, object counting, visual fidelity, and text-image alignment.
  • Few-Step Optimization: Optimizing early high-noise timesteps improves GenEval and PickScore over random consecutive or all-timestep strategies, with minor differences on OCR.This strategy demonstrates the efficiency and effectiveness of few-step training.
  • Ablation Studies: Using the final clean sample as the action consistently improves GenEval and PickScore while reducing KL divergence to the reference policy.The gain is less pronounced on OCR, but the lower KL divergence indicates more stable optimization.
  • Ablation Studies: Forward trajectory reconstruction and CFG-free training improve efficiency and stability, with CFG-free training eventually surpassing CFG-based training and outperforming it on OCR.CFG-free training performs poorly early because of lower sample quality, then converges faster with lower KL divergence; the final method also shows clear qualitative advantages.

6. Conclusion · Appendix

The paper concludes that UDM-GRPO integrates Uniform Discrete Diffusion Models with GRPO for text-to-image generation while addressing instability and improving training efficiency. The appendix organizes additional implementation details, distribution analysis, extended experiments, and qualitative results.

  • 6. Conclusion: UDM-GRPO is presented as the first method integrating Uniform Discrete Diffusion Models with GRPO for text-to-image generation.
  • 6. Conclusion: Treating the final clean sample as the action addresses instability caused by naive adaptation.
  • 6. Conclusion: Reconstructing trajectories through the forward diffusion process helps address instability from naive adaptation.
  • 6. Conclusion: Reduction-Step and CFG-free strategies are introduced to enhance training efficiency.
  • Appendix: The appendix organizes implementation details, experiments, and qualitative results.
  • Appendix: The appendix includes a Distribution Analysis section.
  • Appendix: The appendix includes Extended Experimental Results.

A. Training Details · A.1. Pseudo Code for UDM-GRPO

Algorithm 1 specifies the UDM-GRPO training procedure, from sampling trajectories and clean samples to reconstructing noisy states and updating the policy. It includes CFG-Free sampling, Reduced-Step timestep selection, an Accurate Action Strategy, and iterative policy optimization.

  • A. Training Details: Algorithm 1 takes KL weight β, clip parameter ϵ, reference policy πref, and candidate timestep groups Tgroup as inputs.The procedure initializes θ ← θold before iterations begin.
  • A.1. Pseudo Code for UDM-GRPO: For each prompt c ∼ C, the method samples G trajectories {τi}G_i=1 from πθold(· | c) using CFG-Free sampling.Trajectory sampling occurs within each training iteration.
  • A.1. Pseudo Code for UDM-GRPO: The procedure extracts clean samples {x̂1} from the sampled trajectories before subsequent loss computation.The supplied pseudocode identifies this extraction step but truncates the displayed sample notation.
  • A.1. Pseudo Code for UDM-GRPO: For each trajectory, the algorithm initializes total loss L ← 0 and samples noisy states across j = 1, 2, 3.The noisy-state sampling step is part of the per-trajectory loss-accumulation loop.
  • A.1. Pseudo Code for UDM-GRPO: The noisy-state computation applies the Accurate Action Strategy before completing the nested trajectory and timestep loops.The provided pseudocode fragment ends with the strategy annotation after the noisy-state expression.
  • A.1. Pseudo Code for UDM-GRPO: After loop completion, policy parameters update as θ ← θ − λ∇θL, then θold ← θ for the next iteration.This is the final policy-optimization and reference-policy refresh sequence shown in Algorithm 1.

B. Distribution Analysis

This section evaluates whether Xforward more closely follows the forward trajectory than Xbackward by comparing each trajectory with Xpretrain using FID and visual analysis. The experiment samples caption-image pairs from URSA and compares predicted-image distributions at matching diffusion timesteps.

  • Distribution Analysis: The analysis compares Xforward and Xbackward against Xpretrain using FID and visual comparisons to assess trajectory alignment.The stated goal is to validate whether Xforward is closer to the forward trajectory than Xbackward.
  • Distribution Analysis: 2,048 caption-image pairs are sampled from the URSA pretraining dataset for constructing Xforward, Xbackward, and Xpretrain.Each pair consists of captions and corresponding images (c, x1).
  • Distribution Analysis: FID is computed between predicted-image distributions from Xforward and Xbackward at the same timestep to measure alignment with Xpretrain.Predictions are converted to image space using the model’s standard decoding procedure before comparison.

C. Extended Experimental Results · C.1. Generalized Validation · C.2. Different Model Performance Comparison

The extended experiments evaluate UDM-GRPO across additional benchmarks and model comparisons. Results indicate consistent improvements on FUDOKI and stronger qualitative generation than several baselines and reference models.

  • C. Extended Experimental Results: The section presents additional experiments designed to assess UDM-GRPO from multiple perspectives.These experiments comprise generalized validation and comparisons across different models.
  • C.1. Generalized Validation: FUDOKI extends validation to a UDM-based multimodal large language model that unifies visual understanding and image generation.The method’s generality is tested beyond the URSA setting.
  • C.1. Generalized Validation: The FUDOKI experiments use the URSA experimental setup across GenEval, PickScore, and OCR.Group sampling uses 15 inference steps, while evaluation uses 32 inference steps.
  • C.1. Generalized Validation: UDM-GRPO consistently and significantly improves FUDOKI performance across the three evaluated benchmarks.The supplied passage introduces Table 4 as evidence for these benchmark improvements, but does not include its numerical values.
  • C.2. Different Model Performance Comparison: Qualitative comparisons include the baseline, UDM-GRPO, SD3.5-L, and Flux.1 Dev.Figure 8 presents visual comparisons using prompts from GenEval and PickScore.
  • C.2. Different Model Performance Comparison: Compared with the reference models, UDM-GRPO demonstrates stronger generative capacity for complex scenes containing more objects.The reported visualizations also compare SD3.5-L and Flux.1 Dev with the proposed model.
  • C.2. Different Model Performance Comparison: The visual results show outputs that better align with human perceptual quality metrics.The passage presents these observations as further evidence of UDM-GRPO’s potential to augment model capabilities.

C.3. Different Method Qualitative Comparision

Qualitative comparisons show that integrating GRPO improves model capabilities, while Accurate Action and the forward strategy provide further gains. CFG-free preserves generative quality, indicating strong generative capacity.

  • C.3. Different Method Qualitative Comparision: The compared methods are backward + x_t, backward + x_1, forward + ˆx_1, and forward + ˆx_1 + CFG-free, in that order.These methods correspond to the visualizations from left to right in Figure 9.
  • C.3. Different Method Qualitative Comparision: Initial GRPO integration improves model capabilities, while Accurate Action and the forward strategy further improve performance.The comparison evaluates models trained with different methods discussed in Section 5.
  • C.3. Different Method Qualitative Comparision: CFG does not compromise generative quality, demonstrating the model’s strong generative capacity.This observation comes from the qualitative comparison of the different training methods.

C.4. Training Process

The training process shows progressively improving generated-sample quality and steadily increasing accuracy across optimization iterations.

  • Training Dynamics: Generated samples for fixed evaluation prompts improve progressively throughout training, with accuracy steadily increasing.The evolution is visualized at regular intervals during training.
  • Training Dynamics: Generated samples are visualized across successive training iterations during optimization.
Loading 2604.18518v4…