PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design

Jaemoo Choi, Yuchen Zhu, Wei Guo, Petr Molodyk, Bo Yuan, Jinbin Bai, Yi Xin, Molei Tao, Yongxin Chen

arXiv:2602.04663v2cs.LGcs.AI

TL;DR

Policy-gradient reinforcement learning is difficult to apply to diffusion models because their likelihoods are not directly tractable, while existing methods emphasize engineered objectives and ad hoc estimators. The paper disentangles objectives, likelihood estimators, and sampling schemes, finding that final-sample ELBO estimation dominates performance and efficiency across benchmarks. Its method raises GenEval from 0.24 to 0.95 in less than 90 GPU hours, with reported efficiency gains over FlowGRPO and DiffusionNFT.

  • Problem

    Diffusion models lack exact, efficiently computable likelihoods required by policy-gradient methods, and existing diffusion-RL studies do not systematically separate objective, likelihood-estimation, and sampling effects.

  • Method

    The paper systematically compares policy-gradient objectives, trajectory- and ELBO-based likelihood estimators, and SDE- versus ODE-based sampling strategies.

  • Results

    Across multiple reward benchmarks, ELBO-based likelihood estimation has greater impact than the policy-gradient objective or sampler; GenEval rises from 0.24 to 0.95 in less than 90 GPU hours.

  • Takeaways & Limitations

    The study identifies likelihood estimation, particularly a single-timestep ELBO estimator, as the practical focus for efficient diffusion-model RL.

  • Takeaways & Limitations

    Findings are validated mainly on SD3.5-Medium and a limited set of benchmarks, without systematic comparison of output diversity.

Abstract

from arXiv · show

Reinforcement learning has been widely applied to diffusion and flow models for visual tasks such as text-to-image generation. However, these tasks remain challenging because diffusion models have intractable likelihoods, which creates a barrier for directly applying popular policy-gradient type methods. Existing approaches primarily focus on crafting new objectives built on already heavily engineered LLM objectives, using ad hoc estimators for likelihood, without a thorough investigation into how such estimation affects overall algorithmic performance. In this work, we provide a systematic analysis of the RL design space by disentangling three factors: i) policy-gradient objectives, ii) likelihood estimators, and iii) rollout sampling schemes. We show that adopting an evidence lower bound (ELBO) based model likelihood estimator, computed only from the final generated sample, is the dominant factor enabling effective, efficient, and stable RL optimization, outweighing the impact of the specific policy-gradient loss functional. We validate our findings across multiple reward benchmarks using SD 3.5 Medium, and observe consistent trends across all tasks. Our method improves the GenEval score from 0.24 to 0.95 in 90 GPU hours, which is $4.6\times$ more efficient than FlowGRPO and $2\times$ more efficient than the SOTA method DiffusionNFT without reward hacking.

1. Introduction

The paper investigates why reinforcement learning for diffusion models is difficult and systematically separates policy-gradient objectives, likelihood estimators, and sampling strategies. Experiments identify ELBO-based likelihood estimation as the dominant factor for efficient and effective optimization.

  • Diffusion models lack exact, efficiently computable likelihoods required by policy-gradient methods, unlike autoregressive LLMs.
  • Existing diffusion-RL methods often modify PPO or GRPO objectives while relying on trajectory-based likelihood estimation that requires storing complete sampling paths.FlowGRPO therefore incurs substantial memory and compute costs, slowing convergence.
  • The study disentangles policy-gradient objectives, likelihood estimation recipes, and SDE- versus ODE-based sampling strategies through controlled experiments.It compares GRPO with three lightweight objectives and evaluates backward trajectory-based against forward ELBO-based estimators.
  • ELBO-based likelihood estimation has a substantially greater impact on optimization efficiency and performance than the specific policy-gradient objective or sampler.Its advantage persists across policy-gradient objectives, while ODE sampling adds efficiency and stability benefits.
  • 0.24 to 0.95 GenEval improvement was achieved in less than 90 GPU hours, with up to 4.6× greater efficiency than FlowGRPO and 2× greater efficiency than DiffusionNFT.The method reached state-of-the-art performance across GenEval, OCR, and DrawBench without complex designs.
  • Likelihood estimation is the paper’s central contribution, with the authors emphasizing its broader potential for advancing diffusion-model RL algorithms.

2. Background

The background formulates RL post-training as reward maximization with KL regularization and reviews policy-gradient methods. It then describes diffusion and flow models as noise-reversal generative processes trained through ELBO minimization and sampled through reverse dynamics.

  • RL post-training maximizes a reward function while KL regularization controls deviation from a pretrained reference model.The regularized optimum is proportional to the reference distribution multiplied by exp(R(x)/β).
  • KL-regularized REINFORCE, PPO, and GRPO are established policy-gradient approaches for likelihood-based generative models.GRPO additionally uses clipping, per-sample KL estimates, and group-derived reward advantages.
  • Reaching a GenEval score of 0.95 is compared using total GPU hours across fine-tuning methods, with ELBO estimation and ODE sampling reducing training cost.The comparison uses 8×H100 GPUs and holds target performance fixed.
  • 2.1. Diffusion and Flow Models: Diffusion and flow models gradually corrupt clean data with Gaussian noise and generate samples by reversing that process.Their forward process uses noise schedules α_t and σ_t, while flow models are a special case with α_t = 1 − t and σ_t = t.
  • 2.1. Diffusion and Flow Models: Velocity-parameterized diffusion and flow models are trained by minimizing an evidence lower bound and sampled through learned reverse dynamics.The velocity is defined from the time derivatives of the interpolation coefficients and Gaussian noise.

3. Rethinking RL for Diffusion Model with Policy-Gradient Methods

The section revisits policy-gradient RL for diffusion models by separating objective design from likelihood estimation and sampling choices. It argues that effective likelihood estimation makes multiple policy-gradient objectives viable, while estimator validity and sampler compatibility determine computational flexibility.

  • Diffusion models lack direct data likelihoods, making likelihood estimation essential for applying policy-gradient RL methods.
  • The likelihood-estimation design space includes forward- and backward-process formulas, sampler choice, and ELBO weighting.
  • Trajectory-based likelihood estimation requires SDE sampling because ODE sampling makes its Gaussian transition degenerate, constraining it to higher computational cost.
  • ELBO-based estimation works with any black-box sampler and stores only final samples, improving flexibility and reducing memory use.
  • With an effective likelihood estimator, different policy-gradient objectives achieve near-identical peak performance without classifier-free guidance during training.
  • PEPG uses a proximal KL penalty to limit updates relative to the old policy, while PAR instead uses an L2 regression between log-probability ratios and advantages.

4. Experiments

Experiments disentangle policy-gradient losses, likelihood estimators, and sampling schemes across SD3.5-M reward settings. ELBO-based estimation and ODE sampling provide the strongest efficiency advantages, while loss choices have comparatively limited impact.

  • Experimental Setup: Experiments evaluate policy-gradient objectives, likelihood estimators, and sampling schemes using SD3.5-M at 512 × 512 resolution.The study uses rule-based and model-based rewards across GenEval, OCR, PickScore, CLIPScore, HPSv2.1, Aesthetics, and ImageReward settings.
  • Policy-Gradient Loss Design: Policy-gradient objectives produce comparable GenEval performance, indicating limited sensitivity to the specific loss formulation once likelihood estimation and sampling are fixed.PEPG and PAR achieve slightly higher average performance, but differences remain small.
  • Likelihood Estimation: 4.68× faster convergence is achieved by ELBO with ODE sampling relative to FlowGRPO, while ELBO with SDE sampling achieves 1.24× faster convergence.The comparison measures prompts required to reach GenEval 0.95.
  • Sampling Strategy: ODE and SDE samplers reach comparable GenEval performance under ELBO estimation, but SDE sampling costs approximately 4 times more function evaluations.ODE sampling uses about 10 steps versus 40 for SDE sampling, yielding faster training.
  • Cross-Benchmark Results: The proposed recipe outperforms existing methods in most multi-task and multi-reward comparisons while using a unified training procedure.The evaluation compares PEPG, ELBO estimation, and ODE sampling against FlowGRPO, AWM, and DiffusionNFT.
  • Further Ablations: Single-timestep and all-timestep ELBO estimators achieve comparable performance, making the lower-cost single-timestep estimator the recommended practical default.The ablation evaluates path-KL weighted, simple weighted, and adaptive ELBO formulations with single- or multiple-timestep Monte Carlo schemes.
  • Further Ablations: Removing clipping and standard-deviation advantage normalization does not reduce performance, and may avoid prompt-level difficulty bias.GRPO with clipping reaches GenEval 0.94, whereas unclipped EPG, PEPG, and PAR reach approximately 0.96.
  • Further Ablations: CFG-free training achieves competitive performance and substantial efficiency gains when combined with ELBO-based likelihood estimation and ODE sampling.The experiments report this result across benchmarks.

5. Conclusion

The study finds that likelihood-estimator quality matters more than policy-gradient objective choice for RL fine-tuning of diffusion models, while noting limits to the current validation scope.

  • Likelihood-estimator quality contributes more to RL success than the choice of policy-gradient objective.
  • The findings were validated primarily on SD3.5-Medium and a limited set of benchmarks, including GenEval, OCR, PickScore, and DrawBench.
  • Broader architectures and tasks would strengthen the generality of the conclusions, including larger-scale models and text-to-video generation.

Impact Statement

The work advances generative-model training with reinforcement learning but may be misused to create harmful or offensive content.

  • The proposed reinforcement-learning approach for improved image generators may be abused to create harmful and offensive content.

A. Related Works

Prior work applies reinforcement learning to language, diffusion, and flow models, with diffusion methods adapting policy-gradient approaches such as GRPO through approximate likelihood estimation.

  • Reinforcement learning methods for language models include PPO, GRPO, and related objective-design variants.
  • Diffusion and flow models have been post-trained with reinforcement learning to align outputs with human preferences.
  • FlowGRPO adapted GRPO to diffusion models using trajectory-based likelihood estimation and motivated subsequent improvements.

B. Proof of Thm. 3.1

The appendix derives several policy-gradient objectives and shows that their iterative updates share a common optimal target distribution, while ELBO variants provide alternative likelihood objectives for diffusion and flow models.

  • Policy-gradient objectives: EPG, PEPG, and PAR are Monte Carlo estimators of their corresponding policy-gradient objectives using group rewards and a mean-reward baseline.
  • Policy-gradient objectives: PEPG and PAR use iterative policy optimization with a sampling policy updated across stages.
  • Policy-gradient objectives: PEPG converges to the target path measure proportional to exp(R(x0 | c)/β)πref(x0 | c), providing an exact alternative to GRPO.
  • Policy-gradient objectives: PAR penalizes deviations of the log-likelihood ratio from the advantage function through least-squares regression.
  • Policy-gradient objectives: PAR shares PEPG’s first-order optimum and converges to the same target path measure under the iterative update scheme.
  • ELBO likelihood objectives: The ELBO likelihood objectives considered in the study differ mainly in their weighting functions and are expressed using a unified v-loss form.

C.2. Connection between methods

AWM and FlowGRPO can be viewed through their likelihood-estimation choices, which distinguish their adaptations of GRPO for diffusion models.

  • AWM uses an ELBO estimator with simple weighting and one Monte Carlo sample.
  • FlowGRPO uses a trajectory estimator with a sum-of-exp approximation for the probability ratio.
  • FlowGRPO additionally clips each term in the sum individually, distinguishing it from the adapted GRPO version in Table 1.

C.3. Details on ELBO and Loss computation

The paper approximates ELBO likelihoods with Monte Carlo estimates over diffusion time and noise, then compares single- and all-timestep strategies for efficiency and fidelity.

  • ELBO likelihood estimation approximates expectations over diffusion time and noise for each clean sample x_0.
  • A straightforward Monte Carlo estimator draws M independent timestep-noise samples and constructs x_t from each pair.
  • Multiple samples are unbiased but computationally expensive, motivating simplified schemes that trade variance for efficiency.
  • Single-timestep estimation: Single-timestep estimation approximates the ELBO expectation using one discretized timestep and one Gaussian noise sample.
  • All-timestep estimation: All-timestep estimation aggregates contributions across discretized timesteps, while reusing the likelihood ratio across trajectory contributions.
  • Both estimators achieve comparable performance, so the paper recommends single-timestep estimation as the lower-cost, simpler default.

D. Implementation Details

The implementation uses a unified reward-based fine-tuning framework whose objectives, likelihood estimators, samplers, and experimental settings are varied within a common procedure.

  • The unified algorithm samples images from an old policy, estimates likelihoods, updates the current policy, and refreshes the old policy with EMA.
  • Algorithm 1 allows policy-gradient losses, likelihood estimators, and sampling strategies to be selected independently.
  • The training loop computes old- and current-policy likelihoods for sampled data before applying the chosen loss.
  • Experiments use 48 prompts per epoch, 24 rollouts per prompt, LoRA with α = 64 and r = 32, and a 3 × 10^-4 learning rate.
  • The study evaluates GenEval, OCR, and multiple reward combinations, comparing against FlowGRPO, DiffusionNFT, and AWM implementations.
  • ODE sampling uses DPM with 10 steps, while SDE sampling uses Euler–Maruyama with 40 steps.
  • Evaluation includes GenEval, PickScore, CLIPScore, HPSv2.1, Aesthetic Score, and ImageReward alongside several baseline models.

E. Additional Experimental Results

Additional evaluations report strong GenEval performance, substantial efficiency gains, low seed variance, and comparisons across benchmarks, models, and qualitative outputs.

  • 0.96 is the highest overall GenEval score, matching or exceeding baselines across individual subtasks.The method reaches 1.00 on Single Object, 0.99 on Two Objects, 0.99 on Position, and 0.97 on Color.
  • The method is evaluated on unseen GenEval categories and DrawBench prompts outside the GenEval training distribution.
  • 90–100 GPU hours are required for ELBO+ODE to reach GenEval ≥0.95, versus 423 hours for FlowGRPO and 750.3 ± 69.2 hours for EPG (Traj.).
  • Figures provide qualitative comparisons on GenEval, OCR, PickScore, and ELBO-based likelihood estimation across losses and samplers.
  • The approach is also evaluated on FLUX.1-Dev using PEPG, ELBO-based likelihood estimation, and ODE sampling.
  • Three-run results for PEPG and PAR in the ELBO+ODE setting show low variance and overlapping confidence intervals.
Loading 2602.04663v2…