Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Know Your Step: Faster and Better Alignment for Flow Matching Models via Step-aware Advantages

Zhixiong Yue, Zixuan Ni, Feiyang Ye, Jinshan Zhang, Sheng Shen, Zhenpeng Mi

TL;DR

This paper tackles the latency and reward-signal sparsity of flow-matching Text-to-Image models by introducing TAFS-GRPO, a framework that enables efficient few-step generation while maintaining alignment with human preferences. It combines temperature-annealed few-step sampling, which injects adaptive temporal noise to generate semantically meaningful intermediate states, with a step-aware advantage integration of Group Relative Policy Optimization to produce dense, step-specific rewards without requiring differentiable reward functions. The approach yields faster convergence and robust human-preference alignment across composition-generation and evaluation on the Pick-a-Pic dataset, outperforming several baselines on both in-domain and out-of-domain metrics. Practically, TAFS-GRPO reduces inference latency and enhances deployment flexibility, making few-step, preference-aligned generation more scalable for real-world applications.

Abstract

Recent advances in flow matching models, particularly with reinforcement learning (RL), have significantly enhanced human preference alignment in few step text to image generators. However, existing RL based approaches for flow matching models typically rely on numerous denoising steps, while suffering from sparse and imprecise reward signals that often lead to suboptimal alignment. To address these limitations, we propose Temperature Annealed Few step Sampling with Group Relative Policy Optimization (TAFS GRPO), a novel framework for training flow matching text to image models into efficient few step generators well aligned with human preferences. Our method iteratively injects adaptive temporal noise onto the results of one step samples. By repeatedly annealing the model's sampled outputs, it introduces stochasticity into the sampling process while preserving the semantic integrity of each generated image. Moreover, its step aware advantage integration mechanism combines the GRPO to avoid the need for the differentiable of reward function and provide dense and step specific rewards for stable policy optimization. Extensive experiments demonstrate that TAFS GRPO achieves strong performance in few step text to image generation and significantly improves the alignment of generated images with human preferences. The code and models of this work will be available to facilitate further research.

Know Your Step: Faster and Better Alignment for Flow Matching Models via Step-aware Advantages

TL;DR

This paper tackles the latency and reward-signal sparsity of flow-matching Text-to-Image models by introducing TAFS-GRPO, a framework that enables efficient few-step generation while maintaining alignment with human preferences. It combines temperature-annealed few-step sampling, which injects adaptive temporal noise to generate semantically meaningful intermediate states, with a step-aware advantage integration of Group Relative Policy Optimization to produce dense, step-specific rewards without requiring differentiable reward functions. The approach yields faster convergence and robust human-preference alignment across composition-generation and evaluation on the Pick-a-Pic dataset, outperforming several baselines on both in-domain and out-of-domain metrics. Practically, TAFS-GRPO reduces inference latency and enhances deployment flexibility, making few-step, preference-aligned generation more scalable for real-world applications.

Abstract

Recent advances in flow matching models, particularly with reinforcement learning (RL), have significantly enhanced human preference alignment in few step text to image generators. However, existing RL based approaches for flow matching models typically rely on numerous denoising steps, while suffering from sparse and imprecise reward signals that often lead to suboptimal alignment. To address these limitations, we propose Temperature Annealed Few step Sampling with Group Relative Policy Optimization (TAFS GRPO), a novel framework for training flow matching text to image models into efficient few step generators well aligned with human preferences. Our method iteratively injects adaptive temporal noise onto the results of one step samples. By repeatedly annealing the model's sampled outputs, it introduces stochasticity into the sampling process while preserving the semantic integrity of each generated image. Moreover, its step aware advantage integration mechanism combines the GRPO to avoid the need for the differentiable of reward function and provide dense and step specific rewards for stable policy optimization. Extensive experiments demonstrate that TAFS GRPO achieves strong performance in few step text to image generation and significantly improves the alignment of generated images with human preferences. The code and models of this work will be available to facilitate further research.
Paper Structure (20 sections, 11 equations, 5 figures, 7 tables, 1 algorithm)

This paper contains 20 sections, 11 equations, 5 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preference Alignment for T2I Models
  4. Few-step Distillation of T2I models
  5. Method
  6. GRPO’s Robustness to Reward Objectives
  7. The Sparse Rewards of Flow-Based GRPO
  8. Temperature-Annealed Few-step Sampling with GRPO
  9. Experiments
  10. Experimental Setup
  11. Main Results
  12. Ablation Studies
  13. Conclusion
  14. Extended Experimental Results
  15. Additional Qualitative Results
  16. ...and 5 more sections

Figures (5)

  • Figure 1: Qualitative comparison of generations from 28 sampling steps of Flux.1-dev, MixGRPO, and 8 sampling steps of Flux-Hyper, Reward-Instruct, and the proposed TAFS-GRPO. TAFS-GRPO demonstrates superior performance in semantics and aesthetics alignment with fewer sampling steps than the baseline Flux.1-dev.
  • Figure 2: Comparison of sampling process between flow-based GRPO methods and TAFS-GRPO.
  • Figure 3: The architecture of the proposed TAFS-GRPO framework. Given a dataset of prompts, our temperature-annealed few-step sampling generates clean images in every sampling step. Then the step-aware advantage integration mechanism collects the rewards of these images and produces precise advantages of each sampling step. These precise feedback signals guide the few-step model updates.
  • Figure 4: The pick score of various few-step model based on Flux.1-dev under different number of inference steps. TAFS-GRPO sustains high performance across 4 to 8 inference steps, illustrating its flexibility under different computational budgets.
  • Figure 5: Qualitative comparison of generations on the composition image generation task from 28 sampling steps of Flux.1-Dev, and 8 sampling steps of Flux.1-Schnell, Flux-Hyper, and the proposed TAFS-GRPO.