FewStepRL: Fine-tuning Few-step Diffusion Models with Reinforcement Learning

FewStepRL. Fine-tuning Few-step Diffusion Models with Reinforcement Learning. By Haitao Lin and six others, from Tencent Hunyuan, Westlake University, and one other institution. Published on the fifth of November.

Reinforcement learning has become a standard way to post-train diffusion models. Most of its objectives, however, roll out the current policy at every iteration, and with a sampler that takes many steps, those rollouts typically dominate the cost of training. Few-step models are the obvious way to cut that cost. The difficulty is a mismatch. A few-step model represents an average velocity over an interval, while the existing objectives were built for an instantaneous velocity and a many-step sampler. The objectives and the model are, in effect, speaking about different quantities.

FewStepRL adapts reinforcement learning objectives to few-step models systematically, rather than one at a time. The adaptation proceeds in three phases: a transformation, then a decomposition, then an adaptation. Each phase puts the objectives into a more common form, until a whole family of methods can be expressed in the same terms and carried over to few-step models together. The point of going in phases is that each one removes a different obstacle. The first removes the difference in what the objectives are stated over. The second removes the difference in how their gradients are built. The third removes the difference between the instantaneous velocity the objectives expect and the average velocity a few-step model provides.

The transformation moves everything into x-space. Each objective under consideration is expressed through a common policy prediction, by default at the clean end of the process, and at an explicitly marked level for objectives defined over transitions or flow maps. Objectives that were stated over densities, over velocities, or over rewards thereby become comparable.

The decomposition then writes the gradient of each objective as a sum of two parts: a self-regression toward a reference, and a reward augmentation. At the level of the gradient, a method reduces to a choice of anchors and weights. Methods that looked different on the surface now differ only in where they anchor the regression and how much weight they give each part. Finally, the adaptation computes each anchor from a few-step model, either at the end of its transport or along its tangent. The instantaneous velocity the original objectives need is recovered either from the network itself or from an auxiliary network.

Applying this yields few-step versions of FlowGRPO, AWM, DiffusionNFT, DiffusionOPSD and ReFL, all in one modular codebase. The experiments then compare them. When the reward is differentiable and its gradient is affordable, methods based on the reward gradient are the most effective. Otherwise, methods in the style of DiffusionNFT generally reach high rewards in relatively few iterations, despite a per-iteration cost that is sometimes higher.

A last finding concerns sampling at inference. Taking extra inference steps shows no systematic aggregate benefit in these experiments. Selected outputs, though, visibly improve when given two to four times as many sampling steps as were used in training.

Five separate methods, each originally built around an instantaneous velocity and a many-step sampler, end up as instances of one recipe: pick anchors, pick weights, and compute the anchors from a few-step model. That reduction is what makes the comparison possible at all, and the comparison gives a practical rule. Use the reward gradient when it is available and cheap, and reach for the DiffusionNFT style when it is not.