TempFlow-GRPO: When Timing Matters for GRPO in Flow Models

TL;DR

TempFlow-GRPO addresses the mismatch between diffusion-based text-to-image optimization and human preferences by introducing temporally aware reinforcement learning for flow models. It achieves this through trajectory branching to allocate process rewards at specific timesteps, a noise-aware policy weighting that emphasizes high-impact early exploration, and a seed-group strategy to control initialization effects. The method delivers state-of-the-art human-preference alignment and compositional image generation across benchmarks, while significantly reducing training time and improving stability compared to Flow-GRPO. These contributions offer a practical, scalable approach to fine-grained reward optimization in flow-based generation with strong generalization across reward models.

Abstract

Recent flow matching models for text-to-image generation have achieved remarkable quality, yet their integration with reinforcement learning for human preference alignment remains suboptimal, hindering fine-grained reward-based optimization. We observe that the key impediment to effective GRPO training of flow models is the temporal uniformity assumption in existing approaches: sparse terminal rewards with uniform credit assignment fail to capture the varying criticality of decisions across generation timesteps, resulting in inefficient exploration and suboptimal convergence. To remedy this shortcoming, we introduce \textbf{TempFlow-GRPO} (Temporal Flow GRPO), a principled GRPO framework that captures and exploits the temporal structure inherent in flow-based generation. TempFlow-GRPO introduces three key innovations: (i) a trajectory branching mechanism that provides process rewards by concentrating stochasticity at designated branching points, enabling precise credit assignment without requiring specialized intermediate reward models; (ii) a noise-aware weighting scheme that modulates policy optimization according to the intrinsic exploration potential of each timestep, prioritizing learning during high-impact early stages while ensuring stable refinement in later phases; and (iii) a seed group strategy that controls for initialization effects to isolate exploration contributions. These innovations endow the model with temporally-aware optimization that respects the underlying generative dynamics, leading to state-of-the-art performance in human preference alignment and text-to-image benchmarks.

Figures (18)

• Figure 1: Images generated by our proposed TempFlow-GRPO with FLUX.1-dev. It substantially improves the baseline models, achieving superior photorealism and enhanced fine-grained detail.
• Figure 2: (Left) Reward Variance Analysis: We plot the standard deviation of PickScore at each denoising step for 200 prompts, per prompt group size is 24. The results, obtained via applying SDE at only one step, reveal that reward variance is highest in the initial steps, indicating that early-stage interventions are most impactful for exploration. (Right) Method Illustration: By branching a stochastic (SDE) exploration from a specific, known state on a deterministic (ODE) trajectory, the resulting difference in the final reward can be unambiguously attributed to the exploration action taken at that precise branching point.
• Figure 3: (Left) Performance comparison on the PickScore benchmark across training steps and GPU hours. Flow-GRPO (Prompt) represents an improved baseline with group-wise standard deviation stabilization. TempFlow-GRPO consistently outperforms both Flow-GRPO variants in both sample efficiency (steps) and computational efficiency (GPU hours), demonstrating superior training efficiency while achieving the best performance. (Right) On the Geneval benchmark, TempFlow-GRPO achieves the highest performance, significantly outperforming Flow-GRPO and surpassing state-of-the-art models including GPT-4o, FLUX.1 Dev, and SD3.5-Medium.
• Figure 4: Overview of TempFlow-GRPO Framework. Our method performs trajectory branching by switching from ODE to SDE sampling at selected timesteps (t=k, j, i), injecting noise $\sigma_t\sqrt{\Delta t}\bm{\epsilon}$ to create exploratory branches. Each branch generates a distinct outcome with reward $R_i$, enabling precise credit assignment. The framework applies noise-aware weighting where $\omega_i > \omega_j > \omega_k$, prioritizing optimization at high-noise early stages (larger circles) over low-noise refinement phases (smaller circles), aligning learning intensity with each timestep's intrinsic exploration capacity. We visualize the model's learning process as an astronaut exploring unknown planets: in early stages , the model explores vast possibility spaces with high uncertainty, while later stages involve focused navigation toward the final destination.
• Figure 5: (Left) Strong correlation between reward standard deviation and noise level across generative timesteps. (Right) Scale term analysis reveals a fundamental mismatch in standard GRPO: scale terms are inversely proportional to noise levels, causing low-noise refinement steps to dominate optimization despite minimal impact on image content.