Flow-DPPO: Divergence Proximal Policy Optimization for Flow Matching Models

Abstract: Recent work has demonstrated that online reinforcement learning (RL) can substantially improve the quality and alignment of flow matching models for image and video generation. Methods such as Flow-GRPO and CPS cast the denoising process as a Markov Decision Process and apply PPO-style ratio clipping to enforce a trust region. However, we argue that ratio clipping is structurally ill-suited for flow models: the probability ratio between new and old policies is a noisy, single-sample estimate of the true policy divergence, leading to over-constraining in some regions of the trajectory and under-constraining in others. We propose Flow-DPPO (Flow Divergence Proximal Policy Optimization), which replaces ratio clipping with a divergence proximal constraint. A key observation is that the per-step policy in flow models is Gaussian, enabling exact and cheap computation of the KL divergence between old and new policies. Flow-DPPO employs an asymmetric divergence mask that blocks gradient updates only when they simultaneously move away from the trusted region and violate the divergence threshold. Experiments show that Flow-DPPO achieves higher rewards with better KL-proximal efficiency, alleviates catastrophic forgetting, promotes balanced multi-objective optimization, and enables stable multi-epoch training where ratio clipping degrades. Code and models are available at https://github.com/Tencent-Hunyuan/UniRL/tree/main/FlowDPPO.

• The paper introduces a KL-divergence based trust region mechanism that replaces noisy ratio clipping in flow matching models.
• It demonstrates superior reward optimization and sample efficiency in both single- and multi-reward settings on benchmarks like GenEval2 and FLUX2-9B.
• The study proposes an asymmetric divergence mask that mitigates catastrophic forgetting and stabilizes multi-epoch training.

Flow-DPPO: Divergence-Constrained Policy Optimization for Flow Matching Models

Introduction and Motivation

The "Flow-DPPO: Divergence Proximal Policy Optimization for Flow Matching Models" (2606.11025) paper addresses crucial limitations in existing reinforcement learning-based fine-tuning of flow matching models for visual generative tasks. Prior approaches such as Flow-GRPO and Flow-CPS adapt Proximal Policy Optimization (PPO) paradigms to this domain by casting the generative denoising process as a Markov Decision Process (MDP), enforcing trust regions via surrogate likelihood ratio clipping. However, under the Gaussian policy structure intrinsic to flow models, the ratio-based clipping becomes a significantly noisy, high-variance proxy for true policy divergence. This results in ill-calibrated regularization, hindering reward optimization and stability, especially under multi-objective reward schemes and in multi-epoch training scenarios.

The paper's central thesis is that, due to the analytic tractability of the KL divergence between per-step Gaussian policies in flow models, ratio clipping should be replaced by exact divergence-based masking. This leads to a more precise and theoretically grounded trust-region mechanism, enhancing sample efficiency, reward-optimality, and robustness to catastrophic forgetting.

Methodological Innovations

Critique of Ratio-Based Trust Regions

Ratio clipping, as implemented in previous RL fine-tuning of flow models (e.g., Flow-GRPO), utilizes single-sample importance ratios to constrain policy updates. The authors provide a detailed variance decomposition, showing that in the high-dimensional continuous action spaces of flow models, this results in asymmetric and frequently spurious clipping due to Monte Carlo noise. The consequence is a systematic bias that can either over- or under-constrain policy updates depending on stochasticity, hindering monotonic improvement mandates derived from canonical trust-region policy theory.

Divergence-Proximal Trust Region (Flow-DPPO)

Flow-DPPO leverages the Gaussianity of the policy at each denoising step, enabling the computation of the KL divergence in closed-form between old and updated policies at negligible computational overhead:

$$D_{\mathrm{KL}}(\mathcal{N}(\mu_1, \sigma^2 I) \| \mathcal{N}(\mu_2, \sigma^2 I)) = \frac{||\mu_1 - \mu_2||^2}{2\sigma^2}$$

Thus, the true divergence is substituted directly into the trust-region update rule.

A novel asymmetric divergence mask is introduced: gradients are masked only if the update moves the policy away from its prior and the divergence exceeds a fixed threshold. The asymmetric nature ensures that updates bringing the policy "back" towards the reference are never blocked, allowing recovery from detrimental drift while rigorously enforcing the trust region when diverging.

Theoretical Foundations

The paper derives a finite-horizon, undiscounted policy improvement theorem specialized to flow matching MDPs, where only terminal rewards are defined. The authors show that constraining the per-step (max) KL or Total Variation divergence provides a computable bound on policy improvement, generalizing TRPO-style guarantees to this setting. In the Gaussian policy case, the equivalence between thresholding the mean square displacement and constraining TV/KL enables exact control over distributional drift.

Empirical Evaluation

Experiments span large-scale text-conditioned flow matching models, including SD3.5-Medium and FLUX2-9B, evaluated on both in-domain (GenEval2) and out-of-domain (PickScore, CLIP, HPSv2) benchmarks.

Key empirical findings:

• Reward Optimization: Flow-DPPO variants achieve state-of-the-art GenEval2 rewards. On FLUX2-9B, Flow-DPPO+CPS achieves 92.6% (Soft TIFAGM) in single-reward settings and 55.2% in multi-reward setups, outperforming all baselines on both in-domain generation quality and compositional accuracy.
• KL-proximal efficiency: The divergence-constrained trust region controls off-policy drift more tightly, enabling stable multi-epoch training. Baselines such as Flow-GRPO plateau or diverge under sample-efficient multi-epoch protocols, whereas Flow-DPPO continues to improve, demonstrating robust sample reuse.
• Robustness and Forgetting: OOD metrics and KL divergence to the pretrained model show that Flow-DPPO mitigates catastrophic forgetting and reward hacking more effectively, maintaining higher visual fidelity and compositional generalization under distribution shift.
• Balanced Multi-Objective Optimization: In multi-reward fine-tuning, Flow-DPPO achieves superior trade-off across all reward signals, preventing domination by any single reward (a common failure mode in ratio-clipping-based methods).

Analysis and Ablation

The ablation studies confirm:

• Asymmetric Masking is Crucial: Removing asymmetry from the mask collapses stability, preventing meaningful optimization.
• Divergence Threshold Tuning: Tighter divergence thresholds improve stability and final reward but slow initial learning.
• Multi-Epoch/Sample Efficiency: The divergence mask is essential for enabling high sample efficiency; under multi-epoch regimes, only Flow-DPPO variants are able to exploit sample reuse without catastrophic drift.
• Classifier-Free Guidance: Flow-DPPO remains robust under different guidance scales, with performance benefits persisting.

Theoretical analysis in the appendices justifies the choice of the divergence-proximal mask, and outlines avenues for further refinement via predictive (forward-looking) divergence constraints.

Implications and Future Directions

Flow-DPPO advances RL fine-tuning for flow matching models by introducing a theoretically rigorous, practically stable trust-region paradigm tailored for Gaussian policies. This contribution is particularly impactful for high-value tasks such as long-sequence video generation, where rollout costs are prohibitive and sample efficiency is critical. It also addresses failure modes in distributional generalization, providing a pathway towards more robust, reward-aligned visual generative models.

Potential extensions include adaptation of the predictive divergence masks, scaling to even larger models and more complex reward schemas, and exploration of divergence-proximal RL in non-Gaussian settings or other domains where closed-form divergences are tractable.

Conclusion

Flow-DPPO establishes that ratio-based trust regions are suboptimal for flow matching models due to high-variance, biased clipping. By employing exact KL-based masking, Flow-DPPO delivers superior reward optimization, stable sample-efficient training, and robust resistance to catastrophic forgetting. Its methodological and theoretical contributions suggest a general strategy for RL fine-tuning in domains where policy divergence is efficiently computable (2606.11025).