Source-linked AI summary
Flow-DPPO: Divergence Proximal Policy Optimization for Flow Matching Models
Bowen Ping, Xiangxin Zhou, Penghui Qi, Minnan Luo, Liefeng Bo, Tianyu Pang
TL;DR
PPO-style ratio clipping provides a noisy trust-region proxy for flow-model policies, whose continuous high-dimensional actions amplify clipping errors. Flow-DPPO uses exact Gaussian-policy KL with an asymmetric divergence mask, and experiments report stronger optimization, reduced forgetting, balanced objectives, and stable multi-epoch training.
Problem
PPO-style ratio clipping uses a noisy single-sample proxy for true policy divergence, making trust-region enforcement unreliable for flow models.
Method
Flow-DPPO replaces ratio clipping with an asymmetric divergence mask based on exact per-step KL computation for Gaussian flow policies.
Results
Flow-DPPO consistently achieves superior reward optimization and catastrophic-forgetting performance across base models, while enabling stable multi-epoch training where ratio clipping degrades.
Takeaways & Limitations
Exact divergence-based trust regions support efficient and stable flow-model RL fine-tuning across reward objectives and sampling settings.
Takeaways & Limitations
The theoretical setting is an undiscounted finite-horizon task with trajectory-level advantages induced by terminal rewards.
Abstract
from arXiv · showhide
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.
1 Introduction
Flow matching models have become a major target for online RL alignment, but PPO-style ratio clipping can poorly enforce trust regions because its divergence estimate is noisy. Flow-DPPO replaces ratio clipping with an exact divergence-based constraint enabled by Gaussian per-step policies.
- Motivation: Flow-GRPO and related methods formulate flow-model denoising as stochastic trajectories and use PPO-style ratio clipping for trust-region optimization.These approaches transform deterministic ODE sampling into stochastic SDE trajectories and apply clipped policy optimization.
- Motivation: Ratio clipping relies on a single-sample proxy for policy divergence, which becomes especially unreliable in continuous, high-dimensional flow-model action spaces.The resulting ratio distribution can shift below one, making clipping effectively asymmetric.
- Key observation: Flow models provide exact per-step KL computation because their policies are Gaussian with velocity-network means and fixed schedule-dependent variance.The KL reduces to ∥µθold −µθ∥2/(2σ2) and can be computed from forward passes already used during training.
- Contribution: Flow-DPPO replaces ratio clipping with a divergence-based mask that blocks only away-from-trust-region updates exceeding the KL threshold.Updates moving toward the old policy remain unblocked, preserving PPO’s beneficial asymmetry.
- Contribution: Experiments report higher rewards, better KL-proximal efficiency, reduced catastrophic forgetting, balanced multi-objective optimization, and stable multi-epoch training.Qualitative results show competitive compositional accuracy with less image-quality degradation than compared baselines.
2 Preliminaries
The paper models flow-generation denoising as a finite-horizon MDP with Gaussian sampler-induced actions and terminal rewards. RL fine-tuning optimizes reward while KL regularization limits deviation from the pretrained reference policy.
- Flow matching: Flow matching learns a time-dependent velocity field that transports samples along an interpolating path from noise toward the data distribution.The paper adopts the rectified-flow linear schedule throughout.
- Stochastic sampling: Sampling solves the learned flow ODE, while stochastic samplers such as Flow-SDE and CPS induce Gaussian policies at each denoising step.CPS reduces noise injection relative to Flow-SDE while preserving more of the scheduler’s interpolation structure.
- MDP formulation: The denoising process becomes a finite-horizon MDP whose actions are next latent samples drawn from sampler-induced per-step policies.States include the prompt, flow time, and current latent, with deterministic transitions after each sampled action.
- RL objective: RL fine-tuning maximizes terminal reward with KL regularization penalizing deviation from the pretrained reference policy.The KL penalty discourages reward hacking and mitigates catastrophic forgetting of pretrained capabilities.
- Flow-GRPO: Flow-GRPO applies GRPO by generating groups of samples, normalizing their rewards into relative advantages, and optimizing with a clipped surrogate objective.The clipped objective controls policy staleness caused by multiple updates on rollout data.
3 Methodology
The methodology derives a trust-region guarantee for finite-horizon flow-model RL, identifies ratio clipping as a noise-dependent divergence proxy, and replaces it with an exact-KL asymmetric mask.
- Trust-region theory: The flow-model policy improvement analysis accounts for finite-horizon undiscounted trajectories and trajectory-level advantages induced by terminal rewards.A performance-difference identity separates first-order improvement from higher-order interactions between per-step policy changes.
- Trust-region theory: Bounding the higher-order error yields a policy improvement guarantee in which constraining per-step divergence supports stable learning and monotonic improvement.The paper formulates a trust-region optimization problem analogous to TRPO.
- Exact divergence: For equal-covariance Gaussian policies, KL and TV are monotone functions of mean displacement, making the exact KL a deterministic trust-region measure.This divergence is available at negligible or zero additional cost from velocity-network forward passes.
- Ratio-clipping pitfalls: Ratio clipping constrains a noisy single-sample surrogate rather than the true divergence, and its decisions depend heavily on sampling noise.The log-ratio combines a random term with variance ∥d∥2/σ2 and a deterministic term equal to negative KL.
- Flow-DPPO mask: Flow-DPPO blocks gradients only when the update moves away from the old policy and the divergence exceeds the threshold.Updates that move toward the old policy remain unblocked, allowing corrective recovery after policy drift.
4 Experiments
Across base models and reward settings, Flow-DPPO variants outperform baselines while preserving out-of-domain behavior. They also support stable multi-epoch optimization through asymmetric masking and divergence-based trust-region control.
- Out-of-domain Behavior and Catastrophic Forgetting: Flow-DPPO variants achieve state-of-the-art performance with less catastrophic forgetting on out-of-domain rewards in FLUX2-9B single-reward training.
- Performance and Generalization: Flow-DPPO variants consistently outperform baselines across both base models and all evaluation metrics, with particularly substantial gains in GenEval2 reward.
- Performance and Generalization: In multi-reward SD3.5 training, Flow-DPPO variants consistently outperform all baselines across most metrics without sacrificing any individual objective.
- Asymmetric Masking and Divergence Threshold: Without asymmetric masking, training collapses because samples outside the trust region are largely ignored, whereas asymmetric masking constrains them back within the trust region.
- Asymmetric Masking and Divergence Threshold: A looser divergence threshold of 10^-5 diminishes stability, while a tighter threshold of 10^-7 slows initial learning but yields superior stability and slightly better final performance.
- Multi-epoch Training and Sample Efficiency: Under multi-epoch training, Flow-DPPO variants sustain long-term gains with G64-I2 and G32-I2, while baselines plateau or degrade.Two inner loops reuse each rollout batch for two consecutive gradient steps; the divergence mask supports this sample reuse within the trust region.
5 Conclusion
The paper identifies ratio clipping as a noisy, biased proxy for divergence and proposes Flow-DPPO, which uses an exact-KL divergence mask for policy optimization. The method is presented within a finite-horizon flow-model MDP and uses the CPS sampling framework with group-relative advantages.
- Ratio clipping in flow models is a noisy, biased proxy for the true policy divergence.
- Flow-DPPO replaces ratio clipping with a divergence-based mask computed from the exact KL at zero extra cost.
- Across multiple base models, sampling schedules, and reward objectives, Flow-DPPO achieves superior reward optimization and less catastrophic forgetting while enabling stable multi-epoch training.
- The training procedure adopts CPS sampling, group-relative advantage estimation, and divergence-based masking during policy optimization.
- Flow-DPPO is formulated for an undiscounted finite-horizon task with K − 1 decision steps, trajectory-level advantages, and terminal reward R(x0, c).
B.1 Proof of Performance Difference Identity
The section derives trust-region guarantees for flow models and connects per-step Gaussian mean distance to KL and TV divergence constraints. It shows that Flow-DPPO’s mask enforces an equivalent divergence-based trust region, with bounds that adapt to policy change and horizon length.
- A telescoping decomposition of trajectory probabilities underlies the policy-improvement analysis for flow-model MDPs.
- The trajectory-level analysis bounds policy improvement using per-step TV divergences and the maximum absolute terminal reward.
- A tighter guarantee is linear rather than quadratic in horizon K when using the universal bound DTV ≤ 1.
- The composite bound is quadratic for small policy changes but linear for larger updates or longer horizons.
- Therefore, Flow-DPPO’s mean-distance mask implements a per-step TV trust region and supports a monotonic policy-improvement guarantee.
- For equal-covariance Gaussian policies, squared mean distance is equivalent up to monotone transformations to both KL and TV divergence thresholds.
Appendix D. Ratio Variance Analysis
The appendix analyzes why ratio clipping is unreliable in flow models by quantifying the sampling variance of log importance ratios. Even moderate KL divergence creates fluctuations large enough to trigger clipping for noise-related reasons.
- The log-ratio decomposes into a noise-dependent inner product and a deterministic negative KL term.
- At DKL = 0.5, the log-ratio standard deviation is 1.0, producing roughly a 3× multiplicative spread in the ratio.
- With clip range [0.8, 1.2], many samples can be clipped because of sampling noise rather than excessive true divergence.
Appendix E. Towards a Predictive Divergence Mask
This appendix examines the directional logic of Flow-DPPO’s divergence mask and explains why ratio-based signs are only heuristic proxies for distributional divergence changes. It motivates predictive refinements based on the Gaussian structure.
- The current mask blocks updates only when divergence exceeds δ and the optimization direction moves farther from the old policy.
- The ratio sign is less reliable for KL constraints because a single sampled action yields a stochastic ratio, whereas KL integrates over all actions.
- Gaussian structure enables a closed-form divergence-change analysis that separates first-order direction from second-order step-size and gradient-magnitude effects.
E.1 Predicting Post-Update Divergence
The analysis predicts how one gradient step changes per-step divergence and derives increasingly precise masks. These masks range from the cheap current criterion to a full prediction of whether the updated divergence crosses δ.
- E.1 Predicting Post-Update Divergence: The post-update divergence consists of current divergence, a signed first-order change, and a non-negative second-order term.
- E.1 Predicting Post-Update Divergence: The first-order direction is determined by the inner product between the surrogate gradient and divergence gradient.
- E.1 Predicting Post-Update Divergence: When divergence is small, the predictive directional criterion recovers the existing sign test used by Flow-DPPO.
- E.2 The Predictive Mask: The full predictive mask blocks an update when its predicted post-update divergence exceeds δ, accounting for both direction and step magnitude.
- E.2 The Predictive Mask: The current mask is cheapest, the first-order mask adds no hyperparameters, and the full mask requires an effective learning-rate estimate.
E.3 Discussion on Mask Variants
The discussion compares three increasingly faithful divergence masks and identifies their computational and approximation trade-offs. It also frames Gaussian flow policies as enabling principled trust-region enforcement, while noting limits of the predictive analysis.
- The three masks form a hierarchy of increasing fidelity, from the current mask to first-order and full predictive variants.
- The current mask is cheapest, requiring no additional computation, and suffices when the trust region keeps D_t small throughout training.
- The first-order mask refines the directional decision without introducing additional hyperparameters, whereas the predictive mask estimates divergence quantitatively using an effective learning rate.
- The predictive mask is a local approximation because shared network parameters couple changes in µ across inputs; its accuracy depends on a small effective learning rate and a stable Jacobian.
- The analysis theoretically justifies the existing asymmetric condition and outlines a principled path toward refined trust-region enforcement using Gaussian flow-policy structure.
Appendix F. Experimental Details
The appendix describes the experimental hardware, adaptation and training settings, baselines, and clipping thresholds used across the reported configurations. It also identifies the substantial computational cost of the study.
- All experiments use NVIDIA H20 96GB GPUs, with approximately 90K GPU hours for the main results and 140K GPU hours overall.
- LoRA is used for every model, with model-specific ranks and scaling factors, a 3 × 10−4 learning rate, 512 × 512 resolution, and model-dependent denoising steps.
- The GRPO experiments use group size 16 and 64 groups per epoch, while PPO and KL-clipping thresholds vary by method and KL-scaling factor.
G.1 Additional Training Curves
Additional curves and tables evaluate Flow-DPPO across models, reward settings, CFG, regularization strengths, and KL behavior. Across these views, the method shows strong performance, reduced forgetting, and improved multi-reward balance under supported settings.
- Additional training curves: Flow-DPPO achieves state-of-the-art SD3.5 single-reward performance while reducing catastrophic forgetting on out-of-domain rewards.
- Additional training curves: Flow-DPPO variants consistently outperform baselines across FLUX2-9B multi-reward metrics, with a notable improvement on GenEval2.
- KL behavior: Figure 12 compares per-step KL divergence from the pretrained model across six training settings and two SDE schedules, with corresponding end-of-training values reported in Table 2.
- Tabulated results: Tables 3 and 4 report end-of-training Soft TIFAGM and ancillary CLIP, PickScore, and HPSv2 rewards across in-domain and out-of-domain prompt sets.
- CFG robustness: With CFG scale 4.0, Flow-DPPO remains robust on FLUX2-9B, achieving state-of-the-art GenEval2 training reward and mitigating forgetting on out-of-domain prompts.
- Regularization ablation: A moderate β=10−3 regularization strength suppresses early reward hacking, balances cross-reward gradients, and improves final GenEval2 without degrading any individual end-of-training reward.