RL for Diffusion modelsPart 1

Chapter 1

Why a reward if we have DPO

Pretraining teaches a model to connect text and images: ask for a cat and it draws a cat. But among the images that match a prompt, some look better and some look worse: flat lighting here, a plastic texture there. To make the model produce the good variants more often, we need the next training stage: preference optimization.

One of its methods is DPO. We collect winner and loser pairs and train the model to move toward the first and away from the second. Data is easy to collect: images from any model will do.

But DPO has two weak spots. The model easily drifts away from the pretraining distribution and has to be held back, with SFT on winner images or methods like ThermoDPO (more on this in the previous post on preference optimization). And more importantly, a pair only says "this image is better than that one." What exactly to fix, the pair doesn't say.

For example, we ask for "four cups" and the model draws three. We could collect a pair "four cups are better than three," but from it the model won't understand that its mistake was in the counting. And such pairs first have to be found somewhere.

The error, on the other hand, is easy to count: a detector finds the cups in the image and compares their number with the prompt. That is a reward (a number that says how well the image solves a specific task). In the same way, OCR checks the text on signs, a separate model scores aesthetics, another one checks hand anatomy.

If a problem can be scored with a number, it can be optimized. That is what RL does. Let's assume we already have a good reward model (a scorer: it looks at the finished image and returns that number).

How do we update the generator so that it gets high reward more often? That is what this post is about, the first of two on RL for diffusion. How to build and balance the reward models themselves is the topic of the second one.

There are two paths, and the choice depends on whether you can take a gradient of the reward. If it is a differentiable network, its gradient can be pushed back through generation straight into the model weights. If the reward is just a number (like the cup counter), what is left is to work with probabilities: make good generations more likely and bad ones rarer.

The first methods took the first path. Today that path is more of a historical reference. But it is exactly their limitations that show where current methods came from and why they are built the way they are. So that is where we start.

Chapter 2

One step at the end: ReFL

Diffusion draws an image over dozens of steps: an image gradually emerges from noise . The steps do different work. In the early ones the model decides what will be in the frame and where: the composition, object placement, lighting. How many cups there will be is decided here. In the late ones details and texture are filled in.

The reward model only looks at the finished image. For its score to affect the composition, the gradient has to travel back through the whole chain, all the way to the early steps.

But pushing a gradient through dozens of steps is expensive. The first to bring a reward gradient into the generator was ReFL, from the same paper that introduced the ImageReward reward model. The solution is blunt: if we can't get through the whole trajectory, let's take one step.

We pick a random late step , run generation up to it without gradients, predict the clean image from in one step, compute the reward on it, and update the model along its gradient [ReFL, Sec. 3]:

Which step to pick

Why a late one? At early steps the prediction is very blurry. The reward model can barely tell anything apart on such an image, so there is nothing to teach the model. At the last steps the problem is the opposite: the image is already formed, and any update is a little bit of texture.

What is left is the late but not the very last stretch, where the image is already readable but not yet finished [ReFL, Fig. 4]. For example, with 50 generation steps, roughly steps 30 through 40.

Hence the main limitation. Whatever differentiable reward we take, by this point the composition is already chosen: if the model drew three cups, a fourth one won't appear. ReFL can only make those three neater.

It improves what can still be changed at late steps, detail and texture, but it can't fix the overall composition. If you optimize HPS or PickScore, the model usually adds a sharp burst of detail and texture: the scene stays the same, it just looks more "rendered."

What if we shift the window earlier?

An idea suggests itself. As training goes on, the model draws sharper and sharper images already at early steps. Could we gradually move the training window toward the start of the trajectory, and reach the composition that way?

And this really works: if you fine-tune the model on such trajectories, ReFL starts working at earlier steps too. But the earlier the step, the further the one-step prediction is from a real image, and the reward model was trained on real images.

For its score to be trustworthy, we need a real : generation has to be honestly run to the end and the gradient pushed back along it. And that means through many steps at once.

Chapter 3

To the end, but only the last steps: DRaFT

Going through dozens of steps at once is expensive. But between "one step" and "all steps" there is a middle ground: run generation to the end, but let the gradient through only the last sampling steps. That is what DRaFT does, in its DRaFT-K variant. The reward is computed on the real : we follow the actual path to the image instead of guessing it in one step.

But there is still no fourth cup. We still update only the end of the trajectory, and the number of cups in the frame was decided long before it. The model draws texture or fingers a bit better, and that is it.11 DRaFT-LV: average the reward. A simple modification of DRaFT that is worth knowing for any RL. After getting without gradients, noise it times with different noise and denoise again in one step. Average the reward over all predictions. This estimate has lower variance, and the update is cleaner: DRaFT-LV learns roughly twice as fast as ReFL [DRaFT, Fig. 4].

Why not just take more steps

It seems that could just be increased: four, five, ten steps, until we reach the step where the number of cups is chosen. But the gradient from to an earlier is a product of the Jacobians of all steps in between:

If these matrices aren't close to identity, the gradients grow sharply after a dozen multiplications and the updates stop working [DRaFT, Fig. 8]. Add the memory needed to store all intermediate activations. In the DRaFT paper itself, small , down to , turn out to be the best trade-off [DRaFT, Fig. 9].

The early steps, where the number of cups is decided, are separated from the reward by a dozen such multiplications. And each of them adds to the explosion.

Chapter 4

Early steps without the explosion: DRTune, ShortFT

The explosion lives in the product of Jacobians, not in the length of the path as such. So there are two ways around it: break the product or shorten the chain itself. DRTune does the first, ShortFT the second. Both reach the early steps, where it is decided how many cups will be in the frame.

DRTune: stop-gradient at the network input

DRTune stops the gradient at the denoiser input [DRTune, Sec. 3.2]. Then the gradient doesn't pass through the network at every step and doesn't multiply Jacobians. It travels along the trajectory only through the linear part of the sampling step, and the norms stop growing [DRTune, Fig. 7].

Not all steps are trained, only a few spread across the trajectory, including early ones. Generation can be stopped early and predicted in one step.

This way we update the model on a reasonable, not too noisy image and reach the early steps, which means the composition. The authors show that training the early steps is what matters for low-level rewards, where methods with shallow supervision fail [DRTune, Table 1 and Fig. 4].

ShortFT: a short trajectory through a shortcut model

ShortFT attacks the problem from the other side. Instead of a fifty-step trajectory, take a few-step model with 4–10 steps that preserves the trajectory of the original model. Through such a short chain the gradient can be pushed all the way, from the reward to the noise, and to the step where the number of cups is chosen.

How do we transfer the result to the full model? The shortcut can be built with LoRA distillation, for example Hyper-SD. On top of it we jointly train another LoRA, this one on the reward. After training, the distillation LoRA is removed and the trained one is kept, and it transfers to the full model.

How easily this works depends on the distillation method. The authors separately close the gap between training and inference with a progressive scheme [ShortFT, Sec. 3.3 and Fig. 6].

Why not distill before RL?

A natural question: why not take a distilled model right away and run RL on it? The problem is diversity. A distilled model at four steps gives almost no diversity. Without it exploration doesn't work: the model has nothing to compare and nothing to choose from.

That is why distillation is usually done after RL. Besides, a model after RL has already partly collapsed and is easier to distill.

What flow matching changed

Flow matching trajectories are straighter. In diffusion, the predicted image at early steps can be very far from the final one. With flow models, an almost finished composition is already visible early in the trajectory: fine details aren't there yet, but the scene is clear. Rewards can be computed on such predictions.

That is why for flow models full rollouts (generation runs from noise to the finished image) became feasible even without a shortcut LoRA. The trajectories are good and straight to begin with, so you can generate with ten steps. The images will be a bit blurry, but the model learns to make them sharper on its own over time.

Where this approach is still useful

There is also a practical takeaway. If the task is about specific details, full RL may not be needed: per-pixel gradients of a differentiable reward are very dense and work well on small things. For composition you need something smarter.

A reliable differentiable reward for complex global properties is hard to get; today mostly large VLMs manage it. Rewards can be combined: with a full rollout, some scorers are computed only on the finished image, at the end of the trajectory.

But our cup counter isn't a network. It returns "three" or "four," and it has no gradient.

Chapter 5

When you can't take a gradient of the reward: Flow-GRPO

We don't need the reward gradient if we have something else: a number for the finished sample and the probabilities of the steps that led to it. Then we can make high-reward trajectories more likely and low-reward ones rarer. This is policy gradient, and there are three reasons it pushed out direct backpropagation.

First, many oracles are binary: four cups or not, the texture is good or not. Such things don't have a differentiable description. Second, even for a trained reward model, its gradients are hard to trust. The model easily finds directions in which the reward goes up while the image doesn't get better: reward hacking.

Third, we want online control: generate rollouts with the current model and fix its own problems.

So we formulate denoising as a Markov decision process. The state is the current noisy sample, the action is the next step, the policy (the rule for choosing an action; here it is the model itself), and the reward comes at the end:

The gradient of the expected reward is written with the REINFORCE trick:

how much better the sample is than the average for the prompt ( is the baseline)make every step of this trajectory more likely

There isn't a single derivative of here: the reward enters as a number.

Flow-GRPO

Flow-GRPO is now one of the main RL methods for flow models. We do a full rollout with a small number of steps, compute the reward, and update the whole model through policy gradient. On top of that comes GRPO: the advantage (how much better a sample is than the group mean) is computed within a group of rollouts for one prompt.

The first problem is sampling. Flow matching generates with an ODE: one noise gives one image. To get different trajectories we would have to take different noises, but that is closer to searching for a lucky seed than to improving the model.

We use a well-known trick: every ODE has an equivalent SDE with the same marginal distributions [Flow-GRPO, Sec. 4]. Then one noise can produce different trajectories, and the steps become Gaussian, with a computable .

The noise level is set as:

The parameter controls how diverse the trajectories are. If it is too small, roughly below 0.5, rollouts are almost identical. If it is too large, samples start to break [Flow-GRPO, Fig. 7b]. In practice I use of around 0.7.

The second problem is the cost of rollouts. The same observation as in ShortFT helps here: flow matching trajectories are straight. In Flow-GRPO, training rollouts use 10 steps, while inference keeps the original 40 [Flow-GRPO, Sec. 4 and Fig. 7a]. Coarse rollouts are enough to rank samples within a group.

Then comes ordinary GRPO. For the prompt "four cups" we generate a group of rollouts: the detector gives 1 where there are four cups and 0 where there are three. We normalize the rewards within the group and get the advantage : positive for four cups, negative for three. Every step of the trajectory gets the same advantage, including the early ones where the number of cups was decided.

Mechanically it looks like this. Do a full rollout and store the steps. Compute the reward on the final image. Then walk along the trajectory, compute at each step, accumulate gradients, and take one optimizer step. The gradient isn't pushed through the chain of steps, so there is no explosion.

The result is noticeable. For SD3.5-M, Flow-GRPO raises GenEval accuracy from 63% to 95%, and text rendering from 59% to 92% [Flow-GRPO, Table 2]. The authors also add a KL term to the reference model (a frozen copy of the original model): it holds back reward hacking [Flow-GRPO, Fig. 6].

Where Flow-GRPO is inconvenient

The method has two problems. First, we optimize the reverse process rather than the forward one, although the model was pretrained on the forward process. This puts strong constraints on the sampler: only the Euler discretization of the SDE works, otherwise the math doesn't add up [Flow-GRPO, Sec. 4].

Second, rollouts are done with CFG. So training involves both the conditional and the unconditional branches. The step probability is computed from the velocity after CFG, and both branches have to be stored and updated. That means twice the memory and twice the compute.

In short: we need the probability of every step and a special sampler that provides it. But the reward is a number on the finished image. Can we get by with just that?

Chapter 6

Without probabilities or a sampler: DiffusionNFT, AWM

We can, if we learn on the forward process rather than the reverse one. Store only clean samples and their rewards, noise them, and train the model with the same loss as in pretraining, only weighted by the reward. That is what DiffusionNFT and AWM do, and in the end diffusion RL turns out to be weighted flow matching.

DiffusionNFT: learning on the forward process

DiffusionNFT drops the likelihood of the reverse process. Rollouts are generated with any sampler, even a high-order ODE. We store only the clean samples and their rewards. Then we noise these samples and train the model with the ordinary flow matching loss (the pretraining loss: how far the predicted velocity is from the true one).

How do we separate good samples from bad ones? First, rewards in a group are normalized and turned into an optimality probability . We could simply fine-tune on the good examples. But rejection fine-tuning without negatives quickly collapses the model [DiffusionNFT, Sec. 4.4]. We need to push away from three cups just as much as we pull toward four.

For this, the trained velocity is written as the old model's velocity plus an increment . The increment with a plus sign defines an implicit positive policy, with a minus sign a negative one:

The positive branch is trained on samples with weight , the negative one with weight :

Let's see where the increment moves at the start of training:

For a good sample , and the model is pulled toward its flow matching target (what the loss pulls the prediction toward). For a bad one it is pushed away. The old model stays as an anchor: training moves away from it rather than from nowhere. That is why it is more stable than DPO on the same samples or fine-tuning on positives only.

The old model that generates rollouts is updated softly, with an EMA [DiffusionNFT, Sec. 3.3]. If you copy the current weights into it right away, the reward grows fast at first, but then training falls apart [DiffusionNFT, Fig. 8].

The benefits are immediate. No constraints on the sampler. No issues with CFG: the whole training runs without it, and only the conditional branch is trained. In effect, CFG gets distilled along the way.

Rollout diversity comes from the model itself rather than from artificial SDE noise. The timesteps for the loss are the same as in sampling. DiffusionNFT is up to 25 times more efficient than Flow-GRPO and raises GenEval from 0.24 to 0.98 within a thousand steps without CFG [DiffusionNFT, Figs. 1 and 6].

Advantage Weighted Matching: just weighted flow matching

The next step makes the connection explicit. Advantage Weighted Matching shows that DDPO-like methods, including Flow-GRPO, are essentially doing flow matching with noisy targets [AWM, Sec. 2]. Hence high variance and slow convergence.

The AWM recipe: in the policy gradient, replace the intractable with its diffusion ELBO estimate. That estimate is exactly the pretraining loss. The result is ordinary flow matching, with each sample weighted by its advantage.

With the model is pulled toward this sample's target, with it is pushed away. As in Flow-GRPO, a KL term to the reference model is added to the loss. The sampler can be anything: ODE, SDE, few-step. The weight is the same trick used to reweight timesteps in the ordinary diffusion loss.

On SD3.5-M and FLUX this is up to 34 times faster than Flow-GRPO at the same generation quality [AWM, Table 2 and Fig. 5].

If this is just SFT, why the complexity?

If high-reward samples amount to SFT on positives, can we just do that directly? Yes: AWM shows exactly that. Write out the ELBO estimate of the policy gradient and this is what remains: move away from the examples we don't like and toward the ones we do, guided only by the reward.

The positive and negative branches of DiffusionNFT are more of an intuitively found way to stabilize this. And the main advantage of online policy gradient stays: the reward can be non-differentiable; it only sets the direction of movement.

What if samples are generated in advance?

How often should reward-labeled samples be refreshed? Can we prepare them in advance, as in DPO? The most reliable option is always online. You can try a one-step lag, but it is hard to make stable: even that delay can collapse the model.

So, if you can, generate rollouts with the current model, or rather its slow EMA copy, as in DiffusionNFT. A KL term to the reference model protects against collapse on the model's own synthetic data: the model doesn't drift too far from it. This regularization works much better here than in DPO.

Delivering the reward ended up as a weight on the pretraining loss. So you have to watch the same things as in pretraining: the data and what the reward teaches.