Why few-step RL needs a framework

RL rollouts dominate the cost of post-training. Few-step generators make rollouts cheap, but they predict an average velocity, while every RL objective we want to use was written for an instantaneous one.

The few-step generator

What goes wrong with direct substitution

Paper Fig. 1: held-out GenEval2 against RL iteration (mixed reward), from the shared pre-RL few-step checkpoint. Dashed: original objective with the average velocity substituted. Solid: the few-step adaptation of this paper.

Three phases

Paper Fig. 2. (a) RL raises the relative mass of the high-reward mode; the objective is evaluated on the policy prediction \(\tilde x_\theta=X[\tilde x_t,v_\theta,0]\). (b) The gradient splits into a regression force toward \(\tilde x_{\mathrm{reg}}\) and a reward force along \(\tilde x_{\mathrm{rwd}}-\tilde x_{\mathrm{base}}\). (c) For a few-step model the prediction and anchors are computed at the end of the transport (u-form) or along its tangent (v-form).

Transformation to x-space

The objectives we consider are written in transition-density, velocity or reward space. Under the linear noise schedule a velocity instead identifies a point in x-space at any chosen level \(r\), so every objective can be expressed through one common prediction. §4.1

Two objectives in x-space

The decoupled loss decomposition

In x-space, the gradient of every objective considered here is the transposed Jacobian of the prediction applied to the sum of two vectors: a reward-free self-regression and a reward augmentation. §4.2

Original methods as anchors and weights

The decomposition applies to the instantaneous-velocity originals, and also to the reward-free distillation losses (with \(\omega\equiv1\) and the teacher's prediction \(\tilde x_{\mathrm{real}}=X[\tilde x_t,v_{\mathrm{real}},0]\) as the rewarded anchor).

From the worked examples of §4.2 (AWM, DiffusionNFT, MeanFlow, DMD) and §4.3 (DiffusionOPSD, ReFL). In DiffusionNFT both anchors coincide at the old prediction, so it regularises toward the point from which it measures the reward; its reward coefficient saturates at \(1/\beta\). AWM has no self-regression in its matching loss; its velocity-space proximal penalty plays that role.

An intuitive example

Adaptation to few-step models

The framework needs instantaneous velocities, but a few-step policy outputs an average velocity \(u_\theta(\tilde x_t,t,r)\). The rollout sample needs no reconstruction and the teacher stays instantaneous, so only \(\tilde x_\theta\) and the old-policy anchors must be rebuilt. §4.3

The three few-step forms

Paper Tab. 1, with \(J=J[u_{\mathrm{old}},t,r]\) and \(v_{\mathrm{ind}}=u_{\mathrm{old}}+J\). A few-step method does not introduce a new loss: it specifies how the four points are computed and which entries \(u_\theta(\cdot,t,r)\) receive a gradient.

One training step for every trainer

class FewStepRLTrainer:
    estimator = VInd()                 # Stage 3(ii)-(iii): v-ind or v-aux
    n_it = n_ts                        # default; FewStepGRPO uses M_step, FewStepOPSD uses 1

    def step(self, prompts):                         # ONE outer iteration
        batch = self.rollout(prompts)                # Stage 1, detached
        R = reward(decode(batch.x_hat), batch.c)     # Stage 2
        batch.A, batch.r_hat = self.advantage(R, batch.c)
        self.estimator.fit(batch)                    # v-aux only
        for q in self.queries(batch):                # Stage 3(i): n_it queries
            p = self.anchors(q)                      # Stage 3(ii)-(iii): x_θ, x_reg, x_base, x_rwd
            W_reg, W_rwd, omega = self.weights(q, p) # Stage 3(iv)
            loss = self.loss(p, W_reg, W_rwd, omega) # Stage 4: the two squares
            accumulate(mean_batch(loss + lam_kl * self.kl(q)) / self.n_it)
        self.update()                                # one step on θ
        self.refresh_old()                           # u_old ← EMA of u_θ

Abridged from the paper's Code 1. Each method overrides only the components in which it differs; see the components table.

Seven trainers, one training step

Each original objective becomes a few-step trainer by choosing its anchors, weights and instantaneous-velocity estimator. Five trainers are group-advantage-based (the reward enters only through scalar advantages of a prompt group); two are reward-gradient-based (the update direction comes from \(\nabla R\)). Tab. 5

Family map

The v-aux variant of FewStepGRPO is not used because \(v_\phi\) is fitted on endpoint-renoised states rather than replayed sampler states.

Components of each trainer Tab. 2

Click a trainer to open its page; hover to highlight its column.

\(\tilde x_{\mathrm{ind}}|_{J=0}=X[\tilde x_t,u_{\mathrm{old}}(\tilde x_t,t,r),0]\), used with the J-free \((t,r)\) query. \(a_+=\hat r/\gamma_+\), \(a_-=(1-\hat r)/\gamma_-\) with adaptive normalisers \(\gamma_\pm\); \(c_{\mathrm{adv}}\) undoes the compression of the advantage in \(\hat r\). In FewStepOPSD and ReFL the sign of \(\omega\) is fixed by construction, not by \(A\).

Network evaluations per inner iteration Tab. 4

† J dropped because it cancels, which leaves the loss and its gradient unchanged (App. C). Cost of one outer iteration: \(C=C_{\mathrm{rollout}}+c_{\mathrm{extra}}+n_{\mathrm{it}}(c_\theta+c_{\mathrm{old}}+c_J+c_\phi+c_{\mathrm{ref}})\).

Images: a differentiable mixed reward

SD3.5-Medium distilled with AnyFlow, fine-tuned on CLIPScore + HPSv2 + PickScore. The reward is differentiable, so all seven trainers apply. §6.1, §6.3

Results at iteration 400 Tab. 6(b)

Click a column to sort. Bold as in the paper: best among the trained methods per comparable column. "h": wall-clock hours to iteration 400 on 8 GPUs; GenEval2 best with its iteration; GenEval2 has a resolution of about 0.03.

What the numbers say §6.3

1. Adaptation is necessary. FewStepGRPO raises the training reward from 1.450 to 1.461 and HPSv3 from 10.30 to 11.58, whereas FlowGRPO with the average velocity substituted leaves the reward at 1.438 and lowers HPSv3 to 9.62. MeanFlowNFT reaches 1.509 against 1.440 for DiffusionNFT, and is higher on HPSv3 (13.23 vs. 11.79), ImageReward (1.385 vs. 1.317) and best GenEval2 (0.411 vs. 0.399, within resolution).
2. Reward gradient wins where it exists. FewStepOPSD and ReFL raise the training reward to 1.64 and 1.62 against 1.46–1.52 for the five group-advantage methods, and also lead on ImageReward and GenEval2. DMDNFT and MeanFlowNFT come next (1.51–1.52, GenEval2 0.41 at iteration 400).
3. CLIPScore moves only with \(\nabla R\). None of the five group-advantage runs increases CLIPScore (0.289 → 0.279–0.285), whereas FewStepOPSD and ReFL raise it to 0.317–0.322.

Cost of one iteration Tab. 7

Same prompt, same noise, nine models

Released checkpoints, 512², shared prompt and seed per row (prompts curated, seeds not picked per method); held-out ImageReward under each image. DMDAWM samples carry a faint 16-px latent-patch lattice; its column is shown after a Gaussian blur (σ = 2 px) and its ImageReward is recomputed on the blurred images. SD3.5-M uses its default 40 steps + CFG for reference.

Images: a non-differentiable OCR reward

The reward is the OCR accuracy of the text rendered in the image. It has no gradient with respect to the clean sample, so FewStepOPSD and ReFL do not apply and the five group-advantage trainers compete. §6.1, §6.3

OCR accuracy within 200 iterations Tab. 6(a)

Left: s/iter, OCR accuracy (online) at four iterations, and the first evaluated iteration with OCR ≥ 0.85 with the wall-clock hours to reach it. Right: OCR and held-out evaluators at the first saved checkpoint whose OCR accuracy exceeds 0.82; FewStepGRPO and MeanFlowAWM do not cross it within 200 iterations. Bold as in the paper; at iteration 200 MeanFlowNFT and DMDAWM are indistinguishable at the curve's resolution (about ±0.01) and both are bold.

Fast reward gains come with held-out degradation

Before / after: rendering “FewStepRL”

AnyFlow init

Drag across the image. Each trainer is shown at an early checkpoint of its released run with training-time OCR accuracy 0.82–0.88 (FewStepGRPO, which peaks at 0.78, at its released checkpoint), before the late-stage reward hacking. OCR = PP-OCRv5 accuracy on the quoted word for this image. Same prompt and seed for all.

Video: Wan2.1-T2V-1.3B

The same seven trainers on a 4-step AnyFlow-distilled Wan2.1-T2V-1.3B at 480×832 and 81 frames. Only part of the video reward is differentiable, which limits the reward-gradient methods. §6.2, §6.4

Results at iteration 250 Tab. 8

GPU-h = GPUs × s/iter × 250 / 3600 is the comparable cost; s/iter is compared only among the 32-GPU rows. The four reward components are raw scorer outputs; the sum uses the logged z-scores for MQ and TA. Bold as in the paper.

Same prompt, same noise: eight models in sync

Held-out VBench, dimension by dimension Tab. S2

What the numbers say §6.4

1. FewStepGRPO is the least effective: the only group-advantage method that ends below the shared initialisation (15.89 vs. 16.54), with TA dropping from 1.85 to 1.23.
2. DMDNFT is the most effective on the training reward (17.93, +1.39, leading three of four terms), and its gain coincides with larger held-out degradation: VBench overall 83.38 (last of seven) and dynamic degree 53.33 (lowest), against 83.70 and 58.89 for MeanFlowAWM, the best on VBench.
3. FewStepOPSD and ReFL gain the least reward after FewStepGRPO (16.26, 16.46): the video reward gradient needs several numerical simplifications (HPSv3-only gradient, fp32 reward model, 21 of 81 frames decoded one latent frame at a time). TA rises above the initialisation (1.91, 1.89); ReFL is second on VBench overall (83.69).
4. The cost ordering flips. On video the DMD-based methods are slower than their MeanFlow counterparts (821 vs. 734–737 s/iter): fitting \(v_\phi\) costs about 144 s per iteration, more than the ~60 s saved on \(J\). The reward-gradient methods are the two most expensive in GPU-hours.

4 steps vs. 50 steps (illustration)

Left: original Wan2.1-T2V-1.3B with 50 steps and CFG 5.0 (100 network evaluations); right: MeanFlowAWM, MeanFlowNFT and DMDNFT with 4 steps, ordered by mean pixel difference to the AnyFlow clip over 240 clips (3.4, 11.6, 24.7). Clips selected for a large difference of MeanFlowNFT and DMDNFT from AnyFlow; on average it is smaller. The paper does not evaluate the 50-step model.

Does a 4-step model improve with more steps?

Each fine-tuned checkpoint is sampled at \(M_{\mathrm{step}}\in\{4,8,16,32\}\) on the same shifted grid, 256 DrawBench prompts, seed 12345, no retraining. References: the AnyFlow init and SD3.5-Medium with CFG 4.5. §6.5

Four conclusions §6.5

1. Different training stages scale differently. SD3.5-Medium rises sharply from 4 to 8 steps (HPSv3 −4.49 → 9.96 by 32 steps) yet even at 32 steps (64 evaluations) stays below every fine-tuned 4-step model on HPSv3, ImageReward, aesthetic, HPSv2 and PickScore. The AnyFlow init degrades with more steps (ImageReward 1.11 → 1.02, GenEval2 0.363 → 0.320).
2. No systematic gain on the training reward. From 4 to 32 steps HPSv3 changes by +1.0% to −6.8% for six of seven trainers, against −17% for the AnyFlow init, which suggests that the decline is inherited from the initialisation rather than introduced by RL.
3. OCR-sensitive prompts expose a sharper failure. Selected DMDNFT and MeanFlowNFT samples collapse at 32 steps to little more than foreground text (HPSv3 12.18 → 0.32 and 11.36 → 0.32) while CLIPScore rises.
4. Visible sample-level scaling exists. Many samples look cleaner with 2–4× the training step count, although the aggregate rewards barely change.

Paper Tab. S3: HPSv3 at each step count and the % change of every dimension from 4 to 32 steps. "it.": the checkpoint used.

Which trainer should I use?

Guidance distilled from the experiments, each from a single run. Answer two questions to get the paper's recommendation. §6.6

Open problems §7

1. DMD-based variants need extra forward and backward passes and memory for \(v_\phi\); how to reduce this overhead is open.
2. Reward-gradient methods need a differentiable reward and a gradient that fits in memory; on video, MQ and TA give no gradient and per-frame decoding makes the gradient deviate from a joint decode.
3. With a scalar group advantage, whether every reward summand takes effect depends on summand weighting, the global-batch standard deviation and small within-group spreads.
4. A model fine-tuned at one step count does not improve systematically with more steps; how to obtain test-time scaling cheaply remains open.

4 steps, no CFG

Images (SD3.5-Medium, 512²) and videos (Wan2.1-T2V-1.3B, 81 frames) from the released FewStepRL checkpoints. Hover a strip to stop it and reveal the true colours; click a frame for its trainer, prompt and seed.

Prompts were curated; per prompt the seed was chosen among 8 with the paper's reward models and visual inspection. Image frames come from the four trainers whose samples are free of the faint 16-px latent-patch lattice (DMDNFT, MeanFlowNFT, MeanFlowAWM, FewStepGRPO); video frames cover all seven trainers. The text in the last strip — “FewStepRL” and “Hunyuan” — is rendered by the models themselves.

Code, checkpoints and citation

bash scripts/train/meanflownft_mix.sh 8
python inference.py scripts/eval/configs/sample_sd35m_fewstep.yaml \
  --prompt "..." --override num_stages=3 \
  stage3_lora_path=checkpoints/meanflownft_mix/checkpoint-400/generator_ema.pt

python inference.py scripts/eval/configs/sample_wan.yaml \
  --prompt "..." --override mode=meanflow_nft \
  meanflow_nft_path=checkpoints/dmdnft/checkpoint-250/generator_ema.pt

Released checkpoints

Each entry is checkpoints/<trial>/checkpoint-<step>/generator_ema.pt (EMA Stage-3 LoRA; SHA-256 in checkpoints/SHA256SUMS) in the Hugging Face repository of the column: Delcher/fewsteprl-sd35 (images) or Delcher/fewsteprl-wan21 (videos). Image LoRAs (≈ 43 MB) apply on top of SD3.5-Medium + the AnyFlow Stage-1/2 LoRAs; the DMD auxiliary velocity networks are in checkpoints/aux_vel. Video LoRAs (≈ 93 MB) apply on top of AnyFlow-Wan2.1-T2V-1.3B (non-commercial licence) and are all at iteration 250, the iteration of Tab. 8. The paper's image tables report fixed iterations (400 for mixed reward, ≤ 200 for OCR). Every sample on this page comes from these checkpoints.

BibTeX

@article{lin2026fewsteprl,
  title   = {FewStepRL: Fine-tuning Few-step Diffusion Models with Reinforcement Learning},
  author  = {Lin, Haitao and Ren, Minsi and Li, Yang and Ge, Xingtong and
             Hu, Zheyuan and Hu, Peiyan and Wu, Tailin},
  journal = {arXiv preprint arXiv:2611.08240},
  year    = {2026}
}

Acknowledgements

Built on AnyFlow, MeanFlow, DiffusionNFT, VideoAlign, HPSv3, VBench, diffusers and peft. AnyFlow-Wan and the reward models are under their own licenses.