Fine-tuning Flow Matching Generative Models with Intermediate Feedback

TL;DR

AC-Flow tackles the credit-assignment and stability problems in online fine-tuning of continuous-time flow matching models by learning intermediate-state values with a lightweight critic and weighting policy updates via generalized critic weights. It introduces reward shaping, a dual-stability mechanism (advantage clipping and critic warm-up), and Wasserstein-2 regularization to maintain diversity while refining performance. Empirical results on Stable Diffusion 3 show state-of-the-art text-to-image alignment (CLIP Score, HPS v2, ImageReward) and strong generalization to unseen human preferences, with manageable computational overhead. Together, these contributions enable robust, data-efficient online refinement of flow-based generators without sacrificing quality or diversity.

Abstract

Flow-based generative models have shown remarkable success in text-to-image generation, yet fine-tuning them with intermediate feedback remains challenging, especially for continuous-time flow matching models. Most existing approaches solely learn from outcome rewards, struggling with the credit assignment problem. Alternative methods that attempt to learn a critic via direct regression on cumulative rewards often face training instabilities and model collapse in online settings. We present AC-Flow, a robust actor-critic framework that addresses these challenges through three key innovations: (1) reward shaping that provides well-normalized learning signals to enable stable intermediate value learning and gradient control, (2) a novel dual-stability mechanism that combines advantage clipping to prevent destructive policy updates with a warm-up phase that allows the critic to mature before influencing the actor, and (3) a scalable generalized critic weighting scheme that extends traditional reward-weighted methods while preserving model diversity through Wasserstein regularization. Through extensive experiments on Stable Diffusion 3, we demonstrate that AC-Flow achieves state-of-the-art performance in text-to-image alignment tasks and generalization to unseen human preference models. Our results demonstrate that even with a computationally efficient critic model, we can robustly finetune flow models without compromising generative quality, diversity, or stability.

Figures (5)

• Figure 1: General Framework of Our AC-Flow including our Model Architecture (Left) and Training Techniques (Right). We also include pipeline of outcome-driven methods like RWR iclr_rwr and AWR awr in actor part (Right), to highlight the differences and our unique contribution.
• Figure 2: Ablation Studies of Proposed Stabilization Techniques.
• Figure 3: Qualitative Comparison. Our AC-Flow demonstrates superior performance across Spatial Positioning, Coloring, Attribute Binding, Object Counting, and Text Rendering. More generative results can be found in App. \ref{['app: Additional Experimental Results']}.
• Figure 4: Reward Curve in Case Study of Proposed Stabilization Techniques.
• Figure 5: Complex Semantic Alignment Results. Comparison between original base model SD3 (left column pairs) and our fine-tuned model (right column pairs) on challenging text-to-image generation tasks. Our method demonstrates superior performance in handling spatial relationships (banana-apple positioning), attribute binding (black apple, green backpack), environmental context (New England winter scene), and unconventional concepts (blue pizza, underwater elephant) while maintaining visual quality and semantic coherence. Each test case evaluates different aspects of semantic understanding and compositional generation capabilities.

Theorems & Definitions (17)

• Theorem 3.1: Policy Update in Online Reward-Weighting
• Remark 3.2
• Theorem 3.3: Bound on $W_2$ for Flow Matching
• Theorem B.1: Intractability of Exact Likelihood
• proof : Proof of Theorem \ref{['theorem: Intractability of Exact Likelihood']}
• Theorem B.2: Policy Gradient Theorem
• Remark B.3
• Theorem B.4: Value Regression Consistency
• Definition B.5: Admissible Value Function
• Theorem B.6: RWR Policy Update