FewStepRL Brings Reinforcement Learning to Few-Step Diffusion Models

Reinforcement learning has become an important tool for post-training diffusion models. In a typical setup, the current policy is repeatedly rolled out, its outputs are scored, and the resulting signal is used to update the model. When generation relies on a long denoising trajectory, these rollouts can dominate the cost of training. Few-step diffusion models appear to offer a natural solution by shortening the trajectory, but they introduce a modeling mismatch that cannot be solved simply by reducing the number of sampling steps.

A many-step model is commonly described through an instantaneous velocity at a particular point in the transport process. A few-step model instead learns an average velocity across a finite interval. Existing reinforcement-learning objectives were largely designed around the former interpretation. Applying them directly to the latter can therefore make the training signal inconsistent with the behavior represented by the model.

FewStepRL addresses this problem as a general adaptation framework rather than as a single new objective. Its procedure has three main parts:

• Transformation into x-space. The objectives are rewritten through a shared policy prediction. The default representation is at the clean-sample endpoint, while transition and flow-map objectives retain an explicitly identified level in the process. This creates a common language for objectives originally defined over densities, velocities, or rewards.
• Gradient decomposition. The gradient of each objective is expressed as the sum of a self-regression term toward a reference and a reward-augmentation term. At the gradient level, the key design choices become the selection of anchors and their weights.
• Anchor adaptation. Anchors are computed from the few-step model at the end of its transport or along its tangent. When an instantaneous velocity is required, it can be recovered from the model itself or supplied by an auxiliary network.

Using this formulation, the authors implement few-step variants of FlowGRPO, AWM, DiffusionNFT, DiffusionOPSD, and ReFL within one modular codebase. The practical advantage is that researchers can compare different reinforcement-learning objectives without rebuilding a separate adaptation pipeline for every method. They can instead vary the reference anchors, weighting strategy, and source of velocity estimates.

The reported results point to a conditional rather than universal winner. When the reward is differentiable and computing its gradient is affordable, reward-gradient-based methods are the most effective. If the reward is not differentiable, or its gradient is too costly, DiffusionNFT-style methods generally achieve high rewards in relatively few iterations, although their per-iteration cost can sometimes be higher.

The study also examines sampling beyond the step count used during training. Increasing the number of inference steps does not produce a systematic aggregate improvement in the experiments. Nevertheless, selected outputs can visibly improve when inference uses two to four times as many steps as training. This suggests that extra steps may help particular samples without guaranteeing a broad performance gain.

The broader contribution of FewStepRL is a reusable way to reason about reinforcement learning for fast diffusion models. It shows that few-step training is not merely a matter of shortening the sampler. The definitions of velocity, the location of the policy prediction, the reward signal, and the reference used for regression must all remain consistent. That perspective can help practitioners balance reward optimization against the computational advantages that make few-step diffusion models attractive in the first place.

Source: Hugging Face Daily Papers