FewStepRL: Fine-tuning Few-step Diffusion Models with Reinforcement Learning
Haitao Lin, Minsi Ren, Yang Li, Xingtong Ge, Zheyuan Hu, Peiyan Hu, Tailin Wu
cs.LG, cs.CV
2026-11-06
Five diffusion RL objectives are recast for few-step models, and on mixed-reward GenEval2 the adapted losses beat a naive plug-in by up to 0.013.
RL post-training for a diffusion model resamples the current policy every iteration. A many-step sampler spends 32 or 64 network evaluations on that rollout, and the rollout dominates time and memory. Video makes the imbalance worse. Few-step generators cut sampling to one to four evaluations by learning an average velocity over an interval. MeanFlow fits that average directly. DMD and AnyFlow distill the teacher transport into a flow map that jumps a long interval in one call. The average velocity is a secant of the sampling trajectory. Where the path bends, the secant and the instantaneous velocity come apart.
FlowGRPO, AWM, DiffusionNFT, DiffusionOPSD and ReFL are written for an instantaneous velocity or for many-step transitions. Paste an average velocity into those formulas and three choices stay open: which evaluation receives the gradient, where the reference policy is scored, and which pairs (t, r) are trained. Figure 1 starts both the unmodified objective and the adapted one from the same few-step checkpoint. Held-out GenEval2 is lower for the unmodified objective in all three families.
Under a linear noise schedule a velocity is an affine transport. Subtract (t − r) times the velocity from a state at noise level t and the result sits at level r. A squared error on velocities is then a squared distance between those points, times a time factor. Objectives written on densities, velocities or rewards become comparable on one clean sample: which point the prediction should approach, and which it should leave. The default level is the clean end. Transition and flow-map objectives mark r explicitly.
The gradient is the Jacobian of that prediction times two forces. Self-regression pulls toward an anchor, usually the old policy's prediction at the same state, so the reward cannot move the sample arbitrarily far. Reward augmentation moves along the line from a baseline anchor to the rewarded sample. The sign follows the group advantage: above the group mean attracts, below it repels. A high raw score is still repelled when the sample loses to its group. A new algorithm is then a new set of anchors and weights. AWM sets the self-regression weight to zero and uses a separate velocity penalty to stay near the old policy. DiffusionNFT places both anchors on the old policy's prediction and clips the reward coefficient with the advantage.
A few-step network emits an average velocity u(x, t, r). The u-form queries only r = 0 and jumps to the clean sample, so the other r values the sampler uses get no gradient. The v-form first recovers an instantaneous velocity. MeanFlow estimates how the average velocity changes with t by a finite difference on the frozen old policy. The correction is detached and costs two extra forwards per query. DMD fits an auxiliary velocity vφ by flow matching on the old policy's rollouts, five steps per outer iteration in the reported setup, and uses vφ only as the baseline anchor. The prediction that carries the gradient still has to come from the student. Otherwise the reward never reaches the parameters being updated.
The outer loop is fixed. The old policy rolls out and is scored, vφ is fitted when the method needs it, inner queries accumulate the loss, one optimizer step is taken, and an EMA refreshes the old policy. Swapping anchors yields few-step versions of FlowGRPO, AWM, DiffusionNFT, DiffusionOPSD and ReFL. OPSD needs a reward that is differentiable in the clean sample. It walks normalized reward-gradient steps to an ascent point and a descent point, then fits the policy to those points. ReFL keeps a single ascent step, turns self-regression off, fixes the mixing weight at 1, and evaluates the reward on the old policy's prediction.
Image runs start from one AnyFlow-distilled SD3.5-Medium checkpoint. Figure 1 plots held-out GenEval2 (Soft-TIFA) under the mixed reward, out to 400 iterations. Dashed curves paste the average velocity into the original objective. Solid curves are the adaptation. The annotated gap is adaptation minus the naive objective.
| Method | Metric | Result |
| FewStepGRPO vs FlowGRPO | held-out GenEval2, best and step 400 | +0.005 |
| MeanFlowAWM vs AWM | same | +0.003 |
| MeanFlowNFT vs DiffusionNFT | best checkpoint | +0.005 |
| MeanFlowNFT vs DiffusionNFT | step 400 | +0.013 |
GRPO and AWM stay at a few thousandths. NFT opens to 0.013 by step 400. The adapted curve is higher in every panel.
When the reward is differentiable and its gradient is affordable, reward-gradient methods are the most effective, and ReFL alone is already a strong baseline. Otherwise the setup uses group advantages. The image experiments also include a non-differentiable OCR reward. DiffusionNFT-style training generally reaches high rewards in relatively few iterations, sometimes at a higher cost per iteration. DMD variants are faster per iteration on images and slower on video, and they more often come with larger held-out degradation. The video backbone is AnyFlow-distilled Wan2.1-T2V-1.3B. Sampling at two to four times the training step count does not systematically raise aggregate scores. Selected outputs do look better.
Few-step models are there to make sampling cheap. RL that still rolls out dozens of steps gives that saving back during training. The five objectives share rollout, scoring and evaluation, so a new variant is mostly new anchors. Use reward gradients, including plain ReFL, when the reward is differentiable and the backward pass is cheap. Use an NFT-style group advantage when the score is not differentiable, as with OCR, and check whether a DMD variant is buying speed with held-out regression. A finer sampler at test time should not be assumed to raise the aggregate number.
A gap of 0.013 GenEval2 shows that the naive substitution is lossy. It also shows that the adaptation is a repair. The split is useful when a few-step student is already in hand and rollout cost is what blocks post-training.
Figure 1 is one mixed reward and one distilled checkpoint. Gaps from +0.003 to +0.013 show that the naive substitution loses. They do not support switching every few-step post-training run onto this framework. Solid curves are the v-form and dashed curves are the u-form, so the gap mixes training every (t, r) with the reward rewrite. No ablation separates the two.
ReFL inside the framework is a detached proxy. The reward derivative is taken at the old policy's prediction and then shifted onto the anchor. It lines up with backpropagating through the current prediction only when the EMA decay is 0, and the magnitudes match only when the scale factor matches too.
Self-regression is offered as the guard against reward hacking. The paper does not run a separate hacking stress test. MeanFlow spends two extra forwards per query, and DMD keeps a second network chasing the rollouts. How much of the few-step saving those costs consume is reported only as a direction: faster per iteration on images, slower on video. There is no single wall-clock comparison. The students are AnyFlow distillations of SD3.5-Medium and Wan2.1-T2V-1.3B. Consistency models and adversarial few-step students are outside the experiments.