Technical Summary: dFlowGRPO

Problem Statement

Discrete Flow Models (DFMs) represent a flexible class of generative models for discrete data, offering a broader design space than Discrete Diffusion Large Language Models (dLLMs) by supporting general probability paths and non-masked source distributions. While Reinforcement Learning (RL) has been successfully applied to dLLMs to improve generation quality, existing methods face significant challenges when extended to general DFMs. The primary obstacle is the difficulty in estimating the log-likelihood of DFMs, which cannot be naturally decomposed as in autoregressive models. Current RL approaches for dLLMs rely on approximations (e.g., mean-field or ELBO-based) that often fail to account for the specific conditional transition rates and posterior dynamics inherent to general DFMs. Consequently, there is a lack of a unified RL framework that effectively leverages the full trajectory information of DFMs, including both conditional transition rates and posterior models.

Methodology

The authors propose dFlowGRPO (Discrete Flow-GRPO), a unified rate-aware policy optimization framework designed specifically for DFMs. The methodology proceeds through the following key steps:

1. MDP Formulation: The denoising process of DFMs is formulated as a Markov Decision Process (MDP), where the state includes the current token sequence and time step, and the action is the transition to the next state.
2. Trajectory Probability Derivation: A core theoretical contribution is the derivation of the full trajectory probability for DFMs using an Euler sampler. The authors prove that the transition probability ratio between the current policy and the old policy can be expressed as a product of token-wise expected posterior ratios, reweighted by a rate-dependent term.
   • Specifically, the ratio $r_k(\theta)$ is derived as:
     $$r_k(\theta) = \prod_{d \in [D]} \mathbb{E}_{X^d_1 \sim p^{\theta_{old}}_{1|t_k}} \left[ w^{\theta_{old}, d}_{Q,k} \cdot \frac{p^\theta_{1|t_k}(X^d_1|x_{t_k})}{p^{\theta_{old}}_{1|t_k}(X^d_1|x_{t_k})} \right]$$
   • Here, $w^{\theta_{old}, d}_{Q,k}$ is a rate-dependent weight derived from the conditional transition rates, and the second term is the posterior ratio.
3. Efficient Estimation: The framework enables the estimation of this transition probability ratio using Monte Carlo (MC) samples from the posterior model without requiring additional forward passes. This allows the incorporation of information from both the conditional transition rates and the posterior model into the RL objective.
4. Optimization Objective: Building on the GRPO (Group Relative Policy Optimization) framework, dFlowGRPO optimizes the policy using a clipped objective that incorporates the derived transition probability ratio and a group-normalized advantage estimator. The objective includes a KL regularization term to prevent policy collapse.

Key Contributions

The paper outlines three primary contributions:

1. Unified Framework: The proposal of dFlowGRPO, a unified rate-aware policy optimization framework for DFMs that supports general probability paths and conditional rates by utilizing information from both the conditional rate and the posterior model.
2. Theoretical Derivation: The derivation of the transition probability ratio for the denoising trajectory, showing it can be decomposed into a product of token-wise expected posterior ratios reweighted by a rate-dependent term. This term is shown to be estimable via MC samples without extra forward passes.
3. Empirical Validation: The application of dFlowGRPO to a multimodal DFM (FUDOKI), demonstrating significant improvements in both text-to-image generation and multimodal understanding tasks compared to existing RL methods.

Experimental Results

The authors evaluated dFlowGRPO using FUDOKI as the base model across two domains:

• Text-to-Image Generation:

  • Evaluated on GenEval, PickScore, and DrawBench.
  • Without classifier-free guidance (CFG), dFlowGRPO improved FUDOKI's PickScore from 20.87 to 23.00 and GenEval score from 61% to 93%.
  • On the GenEval benchmark, dFlowGRPO achieved an overall score of 0.93, outperforming other unified multimodal models and matching or exceeding specialized generation-only models, despite using fewer denoising steps (8 NFEs) and no CFG.
  • The method demonstrated robustness against reward hacking, with consistent improvements across multiple metrics (HPSv3, Aesthetic Score, etc.).

• Multimodal Understanding:

  • Evaluated on ScienceQA, POPE, MME-P, SEED, MMB, and GQA.
  • On ScienceQA, dFlowGRPO boosted accuracy from 75.1% to 81.1% (with LLM extraction) and from 54.1% to 59.1% (without LLM extraction), with minimal reward hacking.
  • Performance improvements were also observed on GQA, SEED, and MMB benchmarks.

• Ablation Studies:

  • MC Samples: Increasing the number of MC samples (from $n=1$ to $n=24$) accelerated training convergence and improved final rewards.
  • KL Regularization: While training without KL regularization achieved high rewards, adding KL regularization ($\beta=0.01$) significantly mitigated reward hacking and improved performance on auxiliary metrics.
  • Comparison: dFlowGRPO consistently outperformed mean-field approximation methods (diffu-GRPO, diffu-GSPO) and online DPO variants, which were found to be less stable or less effective for general probability paths.

Significance and Claims

The paper claims that dFlowGRPO provides a unified and effective solution for applying reinforcement learning to general discrete flow models, addressing the limitations of previous methods that were restricted to specific dLLM configurations or relied on inaccurate approximations. By deriving the exact transition probability ratio and incorporating rate-dependent weights, the framework leverages the full structural information of DFMs.

The authors position dFlowGRPO as a significant step forward in aligning discrete flow models with human preferences and task-specific objectives, demonstrating that RL can be successfully applied to the broader class of DFMs beyond just masked diffusion models. The results suggest that dFlowGRPO achieves performance competitive with continuous flow-based models trained using Flow-GRPO, while maintaining strong capabilities in understanding tasks. The paper modestly notes that its evaluation is currently limited to a single DFM (FUDOKI) and suggests that future work could explore applying these rate-aware techniques to other dLLMs with alternative conditional rates.