Few-step flow models generate samples through a handful of large denoising transitions. Applying GRPO to such models requires introducing stochasticity into the denoising transitions, and when a single transition spans a wide range of noise levels, the choice of the resulting stochastic policy becomes a major factor in optimization. Recent constructions constrain the policy to preserve the target flow coefficients, but this constraint admits a family of policies and offers no criterion for choosing among them. We introduce Geo-FSRL (Geometry-Optimized Few-Step Reinforcement Learning), built on a selection principle: at a fixed conditioning state and timestep pair, choose the admissible policy that minimizes the conditional KL divergence induced by any given nonzero velocity-output update. Within a coefficient-preserving family of isotropic Gaussian transitions, this criterion reduces to minimizing a single scalar geometry factor that determines the conditional KL for every velocity-output update. The same factor governs velocity-space Fisher information and conditional transition log-likelihood-ratio variance. We derive the unique closed-form optimizer within this family, shared across all nonzero update magnitudes and directions, with transition coefficients determined by both the source and target noise levels. Because the minimized geometry factor still varies across timesteps, we further adopt gradient-scale matching from prior work and derive closed-form policy-gradient and KL weights that normalize explicit score scales and reference-KL curvature. Geo-FSRL trains on the original few-step schedule and retains the deterministic solver at inference. Our method achieves higher final reward scores than Flow-GRPO and Flow-CPS on both image generation and video generation models under the 4-step setting. Code is available at https://github.com/ModelTC/LightX2V/tree/main/lightx2v_rl
Figures & tables
Figure 1: Four-step image validation rewards versus optimizer updates: (a) HPSv3-only; (b,c) joint HPSv3 and CLIPScore. Flow-CPS curves use matched weights for their respective ϕ values. DiffusionNFT denotes the tested four-step adaptation.
Single-Reward
Multi-Reward
Method
HPSv3
CLIPScore
HPSv3
Qwen-Image-Lightning
12.17
0.971
12.17
Flow-GRPO ( Liu et al., 2025 )
13.62 (+1.45)
0.996 (+0.025)
13.45 (+1.28)
Flow-CPS ( Wang and Yu, 2025 )
13.90 (+1.73)
0.988 (+0.017)
13.58 (+1.41)
Geo-FSRL (Ours)
14.05 (+1.88)
1.008 (+0.037)
13.80 (+1.63)
Table 1: Final validation reward scores for single-reward setting and multi-reward setting text-to-image training.
Visual Quality (%)
Overall Preference (%)
Flow-CPS
46
42
Geo-FSRL (Ours)
54
58
Table 2: Human evaluation of 20 image pairs by five evaluators blinded to method identity. We report the percentage of preference votes for each method.
Method
Without policy reweighting
With policy reweighting
Gain
Flow-CPS
13.69
13.90
+0.21
Geo-FSRL (Ours)
13.61
14.05
+0.44
Table 3: Single-reward policy-weighting ablation on the validation set. Each method uses its own matched weights; KL weighting remains enabled throughout. Gain is the change from enabling policy weighting.
Figure 2: Four-step video validation rewards versus optimizer updates under joint optimization. Flow-CPS uses matched weights; DiffusionNFT denotes the tested four-step adaptation.
Method
HPSv3-General
HPSv3-Percentile
VideoAlign-MQ
Wan2.1-T2V-1.3B
4.13
8.85
0.25
Flow-GRPO ( Liu et al., 2025 )
7.29 (+3.16)
12.10 (+3.25)
1.05 (+0.80)
Flow-CPS ( Wang and Yu, 2025 )
8.15 (+4.02)
11.94 (+3.09)
1.26 (+1.01)
Geo-FSRL (Ours)
9.24 (+5.11)
12.90 (+4.05)
2.07 (+1.82)
Table 4: Final validation reward scores for joint text-to-video training.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Transition
αt,s
βt,s
αt,s2+βt,s2−s2
ODE Euler
0
s
0
Flow-GRPO (Euler–Maruyama)
σtΔt,s
s−2tΔt,sσt2
tΔt,s2σt2(1+4tσt2)
Flow-CPS
ssin(πϕ/2)
scos(πϕ/2)
0
Appendix
Table 5: Finite transitions from t to s ( 0<s<t<1 ) in the form of Eq. 4 . The last column is the deviation from the coefficient-preserving condition αt,s2+βt,s2=s2 . The Flow-GRPO row is the standard uncapped Euler–Maruyama construction at interior timesteps; endpoint and noise-cap variants are given in Appendix A.1 .
Setting
Text-to-image
Text-to-video
Optimizer updates
4,000
2,000
Optimizer
AdamW
AdamW
Learning rate
2×10−6
2×10−6
Adam (β1,β2)
(0.9,0.999)
(0.9,0.999)
Adam ϵ
10−8
10−8
Weight decay
10−4
10−4
Appendix
Table 6: Training hyperparameters for the image and video GRPO-based experiments. The geometry multiplier cgeo applies to the matched weights of Geo-FSRL and Flow-CPS.
Selected split ( ϕ⋆ )
Flow-CPS ( ϕ=0.3 )
Flow-CPS ( ϕ=0.7 )
(t,s)
α⋆
β⋆
G⋆
α
β
G
α
β
G
(1,15/16)
0.93750
0.00000
0.06667
0.42562
0.83532
0.14685
0.83532
0.42562
0.07482
(15/16,5/6)
0.78567
0.27778
0.17678
0.37833
0.74251
0.29034
0.74251
0.37833
0.17859
(5/6,5/8)
0.58926
0.20833
0.47140
0.28374
0.55688
0.77424
0.55688
0.28374
0.47624
Appendix
Table 7: Analytic Wan coefficients (five decimal places). Flow-CPS with ϕ=0.3 is the reported video baseline; ϕ=0.7 is shown for analytic comparison.
Figure 3: Validation HPSv3 with and without policy reweighting for (a) Geo-FSRL and (b) Flow-CPS. All curves use the same 256 validation prompts; horizontal axes count optimizer updates. KL reweighting remains enabled throughout. Scores are reported in Table 3 .
Figure 4: Qwen-Image-Lightning, Flow-CPS, and Geo-FSRL on the displayed prompts. Multi-reward training optimizes HPSv3 and CLIPScore; HPSv3-only training optimizes HPSv3 alone.
Figure 5: Additional text-to-image examples from the Qwen-Image-Lightning base model and Geo-FSRL under joint HPSv3 and CLIPScore training or HPSv3-only training. Each row uses the displayed prompt.
Figure 6: Selected frames from videos generated by the four-step distilled Wan2.1 base model and Geo-FSRL under joint reward optimization. For each prompt, the upper row shows the base model and the lower row (“After FSRL”) shows Geo-FSRL, with four frames per video. Some prompts are shortened for display.
Figure 7: Additional video comparisons under joint reward optimization. Each prompt has two rows of four selected frames: the four-step distilled Wan2.1 base model above and Geo-FSRL (“After FSRL”) below. Some prompts are shortened for display.
Few-step flow-map generators, such as consistency models and MeanFlow, accelerate sampling by learning long-range transport maps between noise and data. However, their deterministic transitions do not directly provide the stochastic trajectories and tractable likelihood ratios required by reinforcement learning (RL) post-training. Existing SDE-based stochasticization techniques target velocity-based samplers and do not directly extend to long-range flow-map transitions. We propose Flow-Map GRPO, an online RL post-training framework for deterministic few-step flow-map generators. Its key component, Anchored Stochastic Flow Map Composition (ASFMC), combines deterministic transport with anchor-based conditional resampling. We establish the conditions under which this construction preserves the marginal probability path and develop tractable local- and endpoint-anchor policies for two-time and single-time flow maps. These policies enable a unified GRPO training procedure. Experiments on FLUX-based MeanFlow and sCM generators demonstrate substantial improvements in OCR, PickScore, and GenEval at different numbers of inference steps, including joint OCR--PickScore gains with mixed rewards. Controlled ablations show that the stochastic transition design is essential for translating training rewards into generation quality. Flow-Map GRPO enables effective RL alignment of pretrained deterministic flow-map generators while retaining their original parameterization, without retraining them as native stochastic models.
Zhiqi Li, Wen Zhang, Bo Zhu
Georgia Institute of Technology, Atlanta, GA, USA.
Reinforcement learning is increasingly used to align image generators with reward signals, and Flow-GRPO recently extended this paradigm to flow-matching models by treating the denoising sampler as a stochastic policy that can be optimized from reward feedback. Training in this setting is unstable in a way specific to multi-step denoising: the policy update changes systematically across denoising steps, with importance ratios drifting below one, becoming increasingly dispersed, clipping at different rates, and leaving fewer usable samples late in training. Prior work treats these effects as separate failure modes and addresses each with a hand-tuned stabilizer. We show instead that they arise from a single per-step quantity, which we call path variance. This quantity is determined exactly by the sampler's Gaussian transition kernel and can be estimated cheaply during training. This reframes instability as a resource that can be measured and budgeted rather than a collection of symptoms to repair. Our method, λ-Controlled GRPO, calibrates importance-ratio behavior from this predicted law rather than from noisy empirical statistics, and allocates gradient effort across denoising steps according to their predicted cost. The two scales governing the update are fixed by standard policy choices rather than introduced as free tuning parameters. On a text-to-image model under two reward settings, rendering difficult target text scored by optical character recognition and matching human preferences scored by a preference model, λ-Controlled GRPO improves both text accuracy and preference reward over the strongest empirical stabilizer. It also keeps late-step path variance within its intended budget, precisely where the baseline systematically overshoots. The result is a Flow-GRPO update calibrated by its own transition law rather than stabilized after instability appears.
Yufeng Wang, Parivesh Priye, Meeshawn Marathe +1
Stony Brook University · Georgia Institute of Technology · Rivian and Volkswagen Group Technologies
We present AdvantageFlow, a forward-process reinforcement learning (RL) algorithm for rectified flow models. The algorithm minimizes an advantage-weighted prediction loss, which maximizes reward, regularized by the rollout policy, which convexifies the objective and makes its optimization stable. Our objective can be viewed as fitting a local reward-improving target distribution. The rollout regularization arises as a variance reduction step. We evaluate AdvantageFlow empirically on text-to-image generation with Stable Diffusion 3.5 Medium and FLUX.1, and compare it to both forward- and reverse-process RL algorithms.