Two ways to learn
There are two ways to acquire a skill. The first is learning from examples: collect samples of others doing a task and train a model to match the distribution of those samples. The second is learning from outcomes: try something, score the result, and adjust based on the score.
Learning from examples requires a dataset of good examples. Learning from outcomes only requires a way to score candidate outputs, which is often much easier to obtain. But scores only help if the learner already produces reasonable outputs. If the first attempts are random, almost all of them are bad and their scores carry little information. Learning from outcomes therefore needs a strong prior that already produces plausible candidates.
Humans learn most skills this way. We first imitate others and then improve through practice and feedback. Reasoning LLMs are trained the same way: they are pretrained on the web and then post-trained with RL against verifier and preference rewards.
Continuous generative modeling
We expect the same approach to work for continuous generative modeling, which covers even more applications than the discrete case. Besides images, video and 3D, continuous generative models are now used for protein structure prediction and design, trajectory generation in robotics, and world models for simulation and planning.
Diffusion and flow-matching models reduce continuous generative modeling to regression, the simplest supervised learning problem, much as autoregressive models reduce discrete sequence modeling to classification. Their pretraining is well understood and scales well. A clean sample is corrupted with Gaussian noise, and the model regresses against a closed-form target.
Post-training has so far had no comparably simple algorithm. Existing methods either give up the simple regression of pretraining and rely on costly SDE rollouts or reward gradients, or they keep a regression but optimize a surrogate loss.
We show that KL-regularized reward maximization leads to a regression objective with a closed-form target, as in pretraining.
Setup
We start from a pretrained diffusion or flow-matching model, given by a velocity field v^{\mathrm{ref}} (the now standard parameterization, although any other would work as well). Integrating v^{\mathrm{ref}} from Gaussian noise produces samples X_0 \sim p^{\mathrm{ref}}. We also have a scalar reward r that scores samples. The standard RL post-training objective maximizes the expected reward while staying close to the pretrained model in KL divergence,
Its maximizer p^\ast(x) \propto p^{\mathrm{ref}}(x)\exp(r(x)) is the pretrained distribution tilted toward high-reward samples. We want the velocity field v^\ast that generates p^\ast. With samples from p^\ast, we could fit v^\ast with the pretraining loss, but we have none. All we can do is draw samples from a model, the pretrained one or the one being trained, and evaluate their rewards.
In principle, samples from the pretrained model would suffice. Weighting the flow-matching loss on X_0 \sim p^{\mathrm{ref}} by \exp(r(X_0)) gives, up to a constant factor, the flow-matching loss for p^\ast, whose minimizer is v^\ast. But the high-reward samples that dominate p^\ast are often rare under p^{\mathrm{ref}}. A few samples with large weights then dominate the loss, and its estimates have high variance. We therefore train on-policy, on samples from the current model. As training progresses, the model moves toward p^\ast, and its samples land where p^\ast has mass. Training alternates between generating samples with the current model, scoring them with the reward, and updating the model on them.
The RAM regression target
Using stochastic optimal control, as in Adjoint Matching (Domingo-Enrich et al., 2025), we derive a fixed-point equation for the optimal velocity field:
The expectation is over sampled paths that pass through X_t = x. Estimating the gradient of the path cost requires SDE rollouts with backward adjoint sweeps, or it has high variance. At image scale this is too expensive or too noisy, so RAM drops the path cost. This is exact at initialization, where the model equals the reference and the path cost is zero. In general, the resulting fixed point agrees with the KL-regularized optimum to first order in the reward scale. The regularization toward v^{\mathrm{ref}} now comes from the instantaneous deviation on the left.
To compute the remaining gradient of the expected reward without differentiating the reward function, we use the REINFORCE identity,
Under Gaussian noising, the score of p_{0|t} has a closed form in terms of the velocity field. Let \epsilon = (x - (1-t)X_0)/t be the noise that takes X_0 to x. By Bayes' rule,
Tweedie's formula expresses the marginal score through the expected noise, and the velocity field determines the expected noise,
Substituting both into the score and writing x = (1-t)X_0 + t\epsilon gives
The prefactor cancels the \tfrac{t}{1-t} in the fixed-point equation, which without the path cost becomes
The right side is the covariance of the reward and the pretraining target \epsilon - X_0 given X_t = x. RAM solves this equation by regression, with the current model v^\theta in place of v^\ast and its samples in place of p^\ast. This gives the RAM loss
where \mathrm{sg} is the stop-gradient. A high reward pulls the model toward this sample's pretraining target \epsilon - X_0, and a low reward pushes it away.
The target depends only on the final sample X_0 and the noisy state X_t, not on the path between them. So we do not need to sample X_0 given X_t. Instead, we sample X_0 from the current model and noise it as in pretraining, X_t = (1-t)\,X_0 + t\,\epsilon with \epsilon \sim \mathcal N(0, I) and t \in (0, 1). This is valid because, under the KL-optimal process, the conditional distribution of X_t given X_0 is the same as in pretraining. Only the distribution of X_0 changes. Drawing X_0 takes dozens of network evaluations and a potentially expensive reward evaluation, while noising it is cheap. We therefore reuse each sample for many draws of (\epsilon, t). A training step looks as follows.
x0 = ode_sample(model) # on-policy endpoint
r = reward_fn(x0) # scalar (no grad)
for _ in range(K): # in parallel
t = sample_t() # random t in (0, 1)
eps = randn_like(x0) # gaussian noise
xt = (1 - t) * x0 + t * eps # analytic noising
v_hat = model_ref(xt, t) + r * ((eps - x0) - model(xt, t))
loss = (model(xt, t) - v_hat.detach()) ** 2
loss.mean().backward()In practice, the model generates samples conditioned on a prompt. For each prompt, we draw a group of samples (24 images in our text-to-image experiments) and subtract the group's mean reward from each reward. A sample with the average reward then has reward zero and leaves the target at v^{\mathrm{ref}}. This centering is necessary for training to work, and it is now the standard approach in RL for generative models.
Post-training Stable Diffusion 3.5M
We post-train Stable Diffusion 3.5M on three tasks: compositional generation (GenEval), visual text rendering (OCR) and human-preference alignment (PickScore). RAM reaches the highest reward on all three, tied with Advantage Weighted Matching on OCR. On GenEval and PickScore, its image quality is comparable to or better than that of the baselines. On OCR, the reward pushes models to trade some image quality for legibility. RAM reaches Flow-GRPO's peak reward in 50×, 48× and 34× fewer training steps, and each step costs slightly less compute.
Compositional image generation
Generate images that satisfy complex prompts: the right objects, with the right attributes, in the right counts and spatial relations. Scored by the GenEval verification pipeline.
a photo of three fire hydrants
a photo of a truck left of a refrigerator
Visual text rendering
Render text into the scene that is legible and matches the prompt. The reward is how accurately an OCR model reads the requested string from the generated image.
zoo sign reading "Lion Habitat Zone"
medieval shield reading "Defender Of The Realm"
Human preference alignment
Produce images that humans prefer for a given prompt. Scored by PickScore, a preference model trained on a large dataset of human pairwise comparisons.
a Ferrari car made out of wood
cyberpunk Casablanca, futuristic street
A simpler derivation of the target
There is another way to derive the RAM target, without stochastic optimal control. It leads to an exact equation that suggests a different way of solving the optimization problem, which Tilt Matching (Potaptchik et al., 2025) explores. With one approximation, the equation becomes the fixed point of the RAM loss. Personally, I find this derivation simpler, because it stays within flow matching and needs no stochastic process. It may look longer than the derivation above, but only because that one takes the fixed-point equation from stochastic optimal control as given, while this one derives each step from scratch. In the paper, we use it to characterize the approximation error of the RAM fixed point. Here, I present it as a standalone derivation of the RAM objective.
We tilt the pretrained distribution with strength \lambda \in [0, 1], p^\lambda(x) \propto p^{\mathrm{ref}}(x)\exp(\lambda\, r(x)). This interpolates between the pretrained distribution (\lambda = 0) and the KL optimum (\lambda = 1). Let v^\lambda be the flow-matching velocity field of p^\lambda,
At \lambda = 0 this is v^{\mathrm{ref}}, and at \lambda = 1 it is v^\ast. If we know how v^\lambda changes with \lambda, we can compute the difference between the two as
Our goal is therefore to compute the derivative \partial_\lambda v^\lambda_t(x).
The derivative of the velocity field
We start by computing how the density of X_0 given X_t = x changes with \lambda. By Bayes' rule and writing out the tilt,
where Z^\lambda_t(x) is the integral of the numerator over x_0. In the logarithm of this density, only \lambda\, r(x_0) and \log Z^\lambda_t(x) depend on \lambda. Differentiating the logarithm and then applying the chain rule gives
Differentiating Z^\lambda_t(x) under the integral brings down a factor r(x_0), so
Now we write the velocity field as an integral over x_0, v^\lambda_t(x) = \int (\epsilon - x_0)\, p^\lambda_{0|t}(x_0 \mid x)\, dx_0 with \epsilon = (x - (1-t)\,x_0)/t. Only the density depends on \lambda, so
The second line uses the integral for v^\lambda_t(x), and the third the integral for \partial_\lambda \log Z^\lambda_t(x). Written as an expectation,
Approximating the integral
Substituting this derivative into the integral gives an exact equation for v^\ast - v^{\mathrm{ref}}. Tilt Matching derives the same equation and reaches the optimum by annealing through a sequence of intermediate models. If we instead approximate the integral with the right-endpoint rule, which evaluates the integrand once, at \lambda = 1, we obtain the RAM fixed point,
This is the equation that the RAM loss solves by regression, with the current model in place of the optimum. The right-endpoint rule is exact when v^\lambda is affine in \lambda, for example for a Gaussian reference and a linear reward. In general, it is correct to first order in the reward scale.
Thanks to Carles Domingo-Enrich for feedback on this post.