English

AdvantageFlow: Advantage-Weighted Least Squares for RL in Flow Models

Machine Learning 2026-05-26 v1 Artificial Intelligence Computer Vision and Pattern Recognition

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.

Keywords

Cite

@article{arxiv.2605.26013,
  title  = {AdvantageFlow: Advantage-Weighted Least Squares for RL in Flow Models},
  author = {Branislav Kveton and Anup Rao and Subhojyoti Mukherjee and Krishna Kumar Singh and Viet Dac Lai},
  journal= {arXiv preprint arXiv:2605.26013},
  year   = {2026}
}
R2 v1 2026-07-22T07:32:49.345Z