FlowAWR: Online Adaptive Flow Reinforcement via Advantage-Weighted Rectification

2026-06-29

AI Summary

Existing continuous generative flow models in online reinforcement learning suffer from intractable trajectory likelihoods, train-inference mismatch, and reliance on Classifier-Free Guidance (CFG). This work reframes generative policy optimization as supervised regression toward a theoretically optimal velocity field, deriving an advantage-weighted velocity correction mechanism grounded in KL-constrained reward maximization. The proposed approach eliminates the need for stochastic differential equation (SDE) sampling or CFG, enabling end-to-end optimization while naturally supporting stable, high-quality generation under multiple reward constraints. Experiments demonstrate that, on SD3.5-Medium, the method achieves a PickScore of 24.12 in only 1.2k training steps—accelerating convergence by 2–5× compared to baseline methods.

Abstract

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.

Problem

• generative flow models
• online reinforcement learning
• trajectory likelihoods
• training-inference inconsistency
• Classifier-Free Guidance

Innovation

• FlowAWR
• advantage-weighted rectification
• generative flow models
• SDE-free optimization
• CFG-free generation