@@ -80,9 +80,7 @@ Here, the reward function can be instantiated by different types of evaluators d
However, directly optimizing this objective with the RL formulation introduced in Section~\ref{sec:policy-gradient} is not straightforward. This is because that diffusion models generate samples through an iterative denoising process rather than an autoregressive generation process. Therefore, we need to redefine the RL formulation according to the characteristics of diffusion generation. Taking DPOK \citep{fan-etal:dpok} as an example, recent studies observe that the reverse diffusion process naturally forms a multi-step trajectory, where each denoising step can be viewed as an action conditioned on the current noisy state. Based on this observation, the diffusion generation process can be formulated as a MDP. Specifically, we can consider the denoising procedure as a multi-step MDP and applies a policy gradient-based RL algorithm to optimize the reward obtained from generated images. The state corresponds to the current noisy latent representation, while the action represents the next denoising step:
The policy is defined as the reverse diffusion transition:
...
...
@@ -93,14 +91,14 @@ The policy is defined as the reverse diffusion transition:
where the final generated image receives the reward signal from the reward model. Based on this formulation, the optimization objective in Eq.~(\ref{eq:optimization_objective}) can be optimized using a simple policy gradient loss function:
where $\tau$ denotes the denoising trajectory $\{\mathbf{x}_{T},\mathbf{x}_{T-1},\cdots,\mathbf{x}_{0}\}$. Here, many RL techniques developed for LLM alignment can be naturally extended to diffusion model optimization, such as incorporating KL regularization \citep{fan-etal:dpok}, introducing a reference diffusion model for importance sampling \citep{black-etal:training}, and applying GRPO-based optimization \citep{xue-etal:dancegrpo}.
where $S_{z}$ denotes the training set, $\tau_\mathrm{dm}$ denotes the denoising trajectory $\{\mathbf{x}_{T},\mathbf{x}_{T-1},\cdots,\mathbf{x}_{0}\}$. Here, many RL techniques developed for LLM alignment can be naturally extended to diffusion model optimization, such as incorporating KL regularization \citep{fan-etal:dpok}, introducing a reference diffusion model for importance sampling \citep{black-etal:training}, and applying GRPO-based optimization \citep{xue-etal:dancegrpo}.
\subsubsection{Flow Matching Models}
...
...
@@ -142,58 +140,26 @@ Since the flow model deterministically predicts the velocity field, the policy c
where $\delta(\cdot)$ denotes the Dirac delta function. Given the predicted velocity, the next state is obtained by discretizing the ODE:
where $\delta(\cdot)$ denotes the Dirac delta distribution, indicating that the next state is deterministically determined by the current state and the learned velocity field.
However, this deterministic formulation introduces challenges for RL optimization because policy gradient methods typically require stochastic sampling to explore different trajectories. To address this issue, Flow-GRPO converts the original ODE-based flow matching model into an equivalent stochastic differential equation (SDE), which preserves the marginal distribution while introducing sampling diversity \citep{liu-etal:flow}. Under this formulation, the transition probability becomes:
\begin{eqnarray}
\mathbf{x}_{t-\Delta t} = \mathbf{x}_t-\Delta t a_t
\mathrm{Pr}_\theta(a_t|s_t) = \mathcal{N} (a_t;
\mu_\theta(\mathbf{x}_t,t,\mathbf{z}),
\sigma_t^2\mathbf{I})
\end{eqnarray}
The initial state is sampled from the noise distribution:
where $\mu_\theta(\cdot)$ is determined by the learned velocity field and $\sigma_t$ controls the stochasticity during sampling. Based on this stochastic MDP formulation, RL algorithms can be applied to optimize flow matching models. Taking GRPO as an example, given a text condition $\mathbf{z}$, the flow model samples a group of generation trajectories $\{\tau^i_{\mathrm{fm}}\}_{i=1}^{G}$, where each trajectory represents a complete sampling process from the initial noise to the final generated sample. We evaluate the final generated sample of each trajectory using the reward function and obtain a group of rewards: $\mathbf{R}_{\mathrm{fm}}=\{R_{\mathrm{fm}}(\mathbf{x}_0^1,\mathbf{z}), R_{\mathrm{fm}}(\mathbf{x}_0^2,\mathbf{z}), \cdots, R_{\mathrm{fm}}(\mathbf{x}_0^G,\mathbf{z})\}$, where $R_{\mathrm{fm}}$ denotes the reward function for evaluating the generated sample. Similar to diffusion models, the reward function can be instantiated with different types of evaluators depending on the optimization objective, such as vision-language models for measuring text-image alignment. The advantage of each trajectory is then estimated based on its relative reward within the sampled group:
where $\mathcal{N}(0,I)$ denotes the standard Gaussian distribution and $I$ represents the identity covariance matrix.
% Similar to diffusion models, the reward is usually provided at the end of the generation process:
% \begin{eqnarray}
% R(s_t,a_t)=
% \begin{cases}
% r(\mathbf{x}_0,\mathbf{z}),&t=0,\\
% 0,&\text{otherwise}.
% \end{cases}
% \end{eqnarray}
Based on this MDP formulation, we can optimize flow matching models with policy gradient-based RL algorithms. Taking GRPO as an example, we sample multiple generation trajectories from the current flow model and compute their relative rewards to construct the advantage signals for policy optimization:
Based on these estimated advantages, we formulate the GRPO objective as follows:
where $A_i$ denotes the normalized advantage computed from the final rewards and $r_i(\theta)$ represents the policy ratio between the updated and reference flow models.
However, directly applying policy optimization to flow matching models introduces additional challenges. Unlike autoregressive models, flow matching models do not explicitly define the likelihood of each generation step, making the policy ratio difficult to compute. To address this issue, recent studies propose to construct surrogate objectives based on the flow matching loss. For example, Flow Policy Optimization (FPO) replaces the exact likelihood ratio with a flow-matching-based approximation:
\begin{eqnarray}
\hat{r}_{\mathrm{FPO}}(\theta)
=
\exp
(
\mathcal{L}_{\mathrm{CFM},\theta_{\mathrm{old}}}
-
\mathcal{L}_{\mathrm{CFM},\theta}
),
\end{eqnarray}
where $\mathcal{L}_{\mathrm{CFM}}$ denotes the conditional flow matching loss. This formulation avoids expensive likelihood estimation while maintaining compatibility with PPO-style optimization. Moreover, compared with methods that treat every integration step as an independent decision, FPO regards the sampling process as a whole and supports different numerical solvers and sampling strategies \citep{mcallister-etal:fpo}.
% Similar to diffusion-based RL, the reward is typically assigned at the terminal step based on the quality of the final generated sample:
