cs.LG · 2605.26013 Copy arXiv ID · May 25, 2026 Save AdvantageFlow: Advantage-Weighted Least Squares for RL in Flow Models Authors: Branislav Kveton , Anup Rao , Subhojyoti Mukherjee , Krishna Kumar Singh , Viet Dac Lai
Organizations: Adobe Research
Abstract We introduce AdvantageFlow, a forward-process reinforcement learning algorithm for rectified flow models. Unlike Flow-GRPO, which optimizes the reverse process, we optimize an advantage-weighted forward-process prediction loss. This optimization problem is unstable when advantages are negative and the loss becomes non-convex. We stabilize it by rollout policy regularization, which reduces variance and arises from fitting a local reward-improving target distribution. We evaluate AdvantageFlow on image generation tasks with Stable Diffusion 3.5 Medium. It outperforms both Flow-GRPO and a state-of-the-art forward-process RL baseline based on negative-aware fine-tuning.
Explore similar work Jun 29, 2026 · Zheming Fu, Ruizhe He, Wei Shang +4 Generative Flow Networks Velocity Field
Sep 29, 2025 · Shuchen Xue, Chongjian Ge, Shilong Zhang +2 Diffusion Language Models Diffusion Policies
Jun 9, 2026 · Bowen Ping, Xiangxin Zhou, Penghui Qi +3 Flow-Based Policy Proximal Policy Optimization
Jun 29, 2026 · cs.LG J/K move · Enter open · S save
Zheming Fu, Ruizhe He, Wei Shang, Xiaoxiao Ma +3
Beihang University · Joy Future Academy · Zhongguancun Academy · Zhejiang University · University of Science and Technology of China
Aligning generative flow models on continuous spaces via online reinforcement learning is constrained by intractable trajectory likelihoods. Existing density-approximated policy gradient methods rely on stochastic SDE samplers to construct tractable transition kernels, which introduce training-inference inconsistencies and necessitates Classifier-Free Guidance (CFG). While implicit frameworks such as DiffusionNFT directly optimize forward-process velocity fields, its heuristic fixed-magnitude corrections prevent optimization strength from relative intra-group quality. We propose \textit{Flow Advantage-Weighted Rectification} (\textbf{FlowAWR}), a paradigm that recasts continuous generative policy optimization as supervised regression toward a theoretically optimal velocity field. Starting from the optimal policy of a KL-constrained reward maximization, FlowAWR derives the optimal velocity field that admits a magnitude-aware, advantage-weighted rectification form, yielding SDE-free optimization and CFG-free generation. In comparative evaluations on SD3.5-Medium, FlowAWR achieves improved alignment performance alongside a 2
× \times × to 5
× \times × convergence acceleration over DiffusionNFT (e.g., reaching a 24.12 PickScore in 1.2k steps, versus 23.82 in 2.0k steps for DiffusionNFT and 23.50 in
> > > 4k steps for FlowGRPO). Under multi-reward constraints, FlowAWR sustains generation quality, satisfying structural rules while maintaining stable out-of-domain performance.