FewStepRL: Fine-tuning Few-step Diffusion Models with Reinforcement Learning
Organizations: Tencent Hunyuan · Westlake University · University of Chinese Academy of Sciences
Abstract
Reinforcement learning (RL) has become a standard post-training method for diffusion models, but most RL objectives roll out the current policy at every iteration, and with a many-step sampler these rollouts typically dominate the cost of training. Few-step models cut this cost, yet they model an average velocity over an interval, whereas existing RL objectives are built for an instantaneous velocity with a many-step sampler. We adapt RL objectives to few-step models systematically in three phases, namely a transformation, a decomposition and an adaptation. The transformation to -space expresses each objective we consider through a common policy prediction, at the clean end by default and at an explicitly marked level for transition or flow-map objectives, so that objectives stated over densities, velocities or rewards become comparable. The decomposition writes the gradient of each of these objectives as the sum of a self-regression toward a reference and a reward augmentation, so that a method reduces, at the gradient level, to a choice of anchors and weights. The adaptation computes each anchor from a few-step model at the end of its transport or along its tangent, with the instantaneous velocity recovered from the network itself or from an auxiliary network. This yields few-step versions of FlowGRPO, AWM, DiffusionNFT, DiffusionOPSD and ReFL in one modular codebase. In our experiments, reward-gradient-based methods are the most effective when the reward is differentiable and its gradient is affordable; otherwise DiffusionNFT-style methods generally reach high rewards in relatively few iterations despite a sometimes higher per-iteration cost. Extra inference steps show no systematic aggregate benefit in our experiments, although selected outputs visibly improve with two to four times as many sampling steps as in training.
Figures & tables
| u-form | v-form, shared- ( v-ind ) | v-form, auxiliary ( v-aux ) | |
|---|---|---|---|
| (rollout) | |||
| trained | with | ; queried at only |
| FewStepGRPO | FewStepAWM | FewStepNFT | FewStepOPSD | |||||
| stage | quantity | MeanFlowAWM | DMDAWM | MeanFlowNFT | DMDNFT | -general † | ReFL | |
| rollout | SDE | ODE | ODE | ODE | ODE | ODE with | ODE with | |
| advantage | — | — | — | — | — | — | ||
| queries | , drawn as | |||||||
| anchors | ‡ | ‡ | ‡ | |||||
| ‡ | ‡ | |||||||
| anchors | , | |
|---|---|---|
| v-ind | ||
| v-aux | , |
| Trainers | |||||||
|---|---|---|---|---|---|---|---|
| § 5.3 , FewStepGRPO | 1 | 1 | 2 | — | 1 | — | |
| § 5.4 , MeanFlowAWM | 1 | 1 | 2 | — | 1 | — | |
| § 5.4 , DMDAWM | 1 | 1 | 0 | 1 | 1 | forwards and backwards of | |
| § 5.5 , MeanFlowNFT | 1 | 1 | 2 | — | 1 | — | |
| § 5.5 , DMDNFT | 1 | 1 | 0 | 1 | 1 | forwards and backwards of | |
| § 5.6 , FewStepOPSD | 1 | 1 | 1 | — | 1 | reward backwards, 2 no-grad forwards of |
| class | method | reward signal | query of | instantaneous field | |
| Group-adv. | FewStepGRPO, § 5.3 | , replayed | Eq. ( 109 ) | ||
| MeanFlowAWM, § 5.4 | Eq. ( 87 ) | ||||
| DMDAWM, § 5.4 | baseline; | ||||
| MeanFlowNFT, § 5.5 | Eq. ( 87 ) | ||||
| DMDNFT, § 5.5 | |||||
| Reward-grad. | FewStepOPSD, § 5.6 | , | (App. C ) |
| OCR accuracy at iteration | first it. with | at the first saved checkpoint with OCR accuracy | ||||||||||
| trainer | s/iter | 50 | 100 | 150 | 200 | OCR (h) | it. | OCR | aesthetic | HPSv3 | ImageReward | GenEval2 |
| AnyFlow Init. | — | — | — | — | — | —– | 0 | 0.498 | 5.24 | 12.86 | 1.233 | 0.363 |
| FewStepGRPO | 451 | 0.599 | 0.645 | 0.708 | 0.776 | — a | — | — | — | — | — | — |
| MeanFlowAWM | 496 | 0.575 | 0.661 | 0.755 | 0.802 | — a | — | — | — | — | — | — |
| DMDAWM | 342 | 0.698 | 0.820 | 0.902 | 0.936 | 150 (14.3) | 200 | 0.936 | 4.49 | 5.05 | 0.092 | 0.244 |
| MeanFlowNFT | 483 | 0.785 | 0.850 | 0.945 | 0.934 | 100 (13.4) | 100 | 0.850 | 4.99 | 10.64 | 0.887 | 0.306 |
| trainer | s/iter | sampling (s) | update (s) | forwards for | |||
| FewStepGRPO | 423 | 78.7 | 344.2 | 4 | 8 | 0 | 0 |
| MeanFlowAWM | 425 | 78.3 | 344.5 | 4 | 8 | 0 | 0 |
| DMDAWM | 304 | 78.2 | 223.4 | 4 | 0 | 1 | 5 |
| MeanFlowNFT | 435 | 78.8 | 354.6 | 4 | 8 | 0 | 0 |
| DMDNFT | 326 | 78.4 | 244.7 | 4 | 0 | 1 | 5 |
| FewStepOPSD | 268 | 224.7 | 43.3 | 1 | 0 | 0 | 0 |
| training reward | VBench (held-out) | ||||||||||
| method | GPUs | s/iter | GPU-h | sum | HPSv3-G | HPSv3-P | MQ | TA | overall | quality | semantic |
| AnyFlow Init. | — | — | — | 16.54 | 6.42 | 9.43 | 2.64 | 1.85 | 83.68 | 85.36 | 76.98 |
| FewStepGRPO | 32 | 735 | 1633 | 15.89 | 6.24 | 9.38 | 2.40 | 1.23 | 83.62 | 85.27 | 76.98 |
| MeanFlowAWM | 32 | 737 | 1637 | 16.66 | 6.49 | 9.46 | 2.68 | 1.84 | 83.70 | 85.36 | 77.04 |
| DMDAWM | 32 | 821 | 1824 | 16.78 | 6.51 | 9.53 | 2.74 | 1.82 | 83.64 | 85.35 | 76.76 |
| MeanFlowNFT | 32 | 734 | 1631 | 16.80 | 6.53 | 9.51 | 2.76 | 1.80 | 83.62 | 85.34 | 76.75 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Symbol | Meaning |
|---|---|
| text prompt (conditioning); omitted from an argument list when it is fixed or irrelevant. A bare always has this meaning; subscripted quantities such as and are distinct | |
| standard Gaussian noise, ; a prime marks a fresh rollout-time renoise draw | |
| flow-map endpoints, ; generation runs from noise at to data at . The unhatted letter always names a noise level and never a reward; the hatted below is the mixing weight. Individual objectives state any stricter query domain required by divisions by or | |
| a generic level on the same scale, , used as the time of the SDEs of Appendix B | |
| positive-integer group size : the number of rollouts sharing one prompt, ; members carry the index | |
| a clean sample . Without a noise-level subscript, an -symbol denotes a quantity at ; a subscript , or indicates its level, while , , or a role label indicates its source. Method-specific quantities have the level of the object they represent; in particular, the MeanFlow quantities of Equation ( 43 ) and the FewStepGRPO prediction and anchors of Equation ( 97 ) are defined at level |
| trainer | it. | training reward | HPSv2 | PickScore | CLIPScore | aesthetic | OCR |
|---|---|---|---|---|---|---|---|
| FewStepOPSD | 100 | 1.5692 / 1.5704 | 0.3520 / 0.3515 | 0.9205 / 0.9230 | 0.2968 / 0.2959 | 6.092 / 6.093 | 0.400 / 0.405 |
| 200 | 1.6052 / 1.6055 | 0.3643 / 0.3633 | 0.9338 / 0.9356 | 0.3071 / 0.3066 | 6.068 / 6.054 | 0.433 / 0.441 | |
| 300 | 1.6258 / 1.6243 | 0.3705 / 0.3698 | 0.9410 / 0.9416 | 0.3144 / 0.3129 | 5.984 / 5.983 | 0.474 / 0.486 | |
| 400 | 1.6406 / 1.6409 | 0.3733 / 0.3734 | 0.9464 / 0.9458 | 0.3208 / 0.3217 | 5.929 / 5.906 | 0.550 / 0.493 | |
| ReFL | 400 | 1.6166 / 1.6225 | 0.3651 / 0.3675 | 0.9355 / 0.9378 | 0.3160 / 0.3172 | 5.893 / 5.891 | 0.468 / 0.546 |
| s/iter | sampling (s) | update (s) |
| HPSv3 at | % change at against | |||||||||||
| trainer | it. | 4 | 8 | 16 | 32 | HPSv3 | ImgRwd | GenEval2 | Aesth. | CLIP | HPSv2 | Pick |
| FewStepGRPO | 400 | 11.58 | 11.15 | 10.49 | 10.09 | |||||||
| MeanFlowAWM | 100 | 12.30 | 12.36 | 12.17 | 11.85 | |||||||
| DMDAWM | 400 | 13.10 | 13.24 | 12.98 | 12.77 | |||||||
| MeanFlowNFT | 300 | 13.26 | 13.18 | 12.96 | 12.85 | |||||||
| DMDNFT | 400 | 13.49 | 13.49 | 13.34 | 13.13 | |||||||
| CLIPScore at iteration | |||||
| method | 0 | 100 | 200 | 300 | 400 |
| FewStepOPSD, | 0.2886 | 0.2959 | 0.3066 | 0.3129 | 0.3217 |
| ReFL, of | 0.2886 | 0.2967 | 0.3076 | 0.3127 | 0.3172 |
| ReFL, of HPSv2 | 0.2886 | 0.2845 | 0.2844 | 0.2863 | 0.2883 |
| MeanFlowNFT, weights | |||||
| (default), | 0.2886 | 0.2800 | 0.2815 | 0.2806 | 0.2814 |
Diffusion Model Fine-Tuning
Diffusion model fine-tuning adapts pretrained generative models to new concepts, tasks, domains, and preferences through targeted updates.
See all 188 papers