RAM: First-Principles Approach to 50x Faster RL Post-Training for Diffusion Models

The Pre-Training and Post-Training Dilemma for Diffusion Models

Diffusion models have become the standard paradigm for generating continuous representations — outside of language modeling, they're used to generate continuous data across virtually every domain. Diffusion models are a class of Markov chain-based generative models whose core idea is to capture data distributions by progressively adding noise (the forward process) and then learning to reverse this denoising (the backward process). Since DDPM was introduced in 2020, diffusion models have become the de facto standard for continuous data across nearly all domains, including images, audio, video, and molecular design. Their pre-training scalability stems from the simplicity of their objective: minimizing the mean squared error between the model's predicted noise and the ground-truth noise (the score matching objective), which scales well with both model size and data size.

Andreas Bergmeister, a PhD student at TU Munich, identifies a core contradiction in his latest research: diffusion model pre-training is extremely simple and scalable, yet reinforcement learning-based post-training is complex, tedious, and difficult to scale.

This work, called Reinforce Adjoint Matching (RAM), attempts to return to first principles — examining exactly what makes diffusion model pre-training so scalable, and then deriving a training objective that makes RL post-training equally simple and scalable.

Why does post-training matter so much for diffusion models? The authors offer an intuitive framing: generative models have two learning paradigms — learning from samples (matching distributions on a dataset, i.e., pre-training) and learning from outcomes (generating outputs, scoring them, and adjusting the model accordingly, i.e., post-training). In the LLM space, RLHF (Reinforcement Learning from Human Feedback) has been validated as the key technique for aligning models with human preferences, and systems like GPT-4 and Claude rely on this paradigm. Transferring a similar approach to diffusion models presents unique challenges: the generation process is fundamentally the solving of a stochastic differential equation (SDE) or ordinary differential equation (ODE) across tens to hundreds of denoising steps, and standard policy gradient methods require backpropagation or Monte Carlo estimation over the entire generation trajectory — an enormous computational cost. In domains like robotics, where diffusion models generate execution policies and collecting imitation learning datasets is expensive, "learning from outcomes" is especially critical.

From Optimal Distributions to Value Function Gradients

The post-training objective can be formalized as: given a pre-trained diffusion model and a reward function, find a distribution that maximizes expected reward while staying close to the pre-trained model (minimizing KL divergence). KL divergence (Kullback-Leibler Divergence) is an information-theoretic measure of the difference between two probability distributions. In the post-training optimization objective, this KL constraint is crucial — it ensures the model stays within the support of the reference distribution, prevents drift toward regions that generate unreasonable samples, and is the primary mechanism against reward hacking.

The optimal solution to this optimization problem is a classical result: exponential tilting of the reference distribution, where the optimal distribution's density is proportional to the reference model density multiplied by the exponentiated reward: π*(x) ∝ π_ref(x) · exp(r(x)/β), where β is a temperature parameter controlling regularization strength. This result corresponds to the Gibbs distribution in statistical mechanics, a special case of the maximum entropy principle in information theory, and is the theoretical cornerstone connecting generative models with reinforcement learning.

The authors follow the Flow Matching framework — a modern variant of diffusion models that learns a deterministic velocity field to continuously map a noise distribution to a data distribution. The latest text-to-image models like Stable Diffusion 3 and Flux both adopt this framework. By defining forward noising processes for both the reference and optimal distributions, and using the score-velocity relationship (the score function, i.e., the gradient of the log probability density ∇_x log p(x), has a definite algebraic relationship with the velocity field), subtracting the velocity fields of the optimal and reference processes yields, after algebraic derivation, exactly the gradient of the value function (the adjoint). This means: if you can estimate the gradient of the value function, you can train the model via a regression objective.

The original Adjoint Matching method converts this into an on-policy regression objective, estimating the value function from the current policy and then regressing its gradient. The challenge is: how to cheaply and stably estimate this adjoint when the reward is non-differentiable.

Core Innovation: Discarding the Path Cost

The value function consists of two components: a non-differentiable reward term and a differentiable path cost term. The authors use the additivity of expectations to split and estimate these two terms independently.

For the reward term, the authors use the REINFORCE / log-derivative trick — a classic technique that allows gradient estimation in expectation form even when the reward function is non-differentiable (e.g., human ratings, rule-based judgments, black-box renderers): ∇θ E{xp_θ}[f(x)] = E_{xp_θ}[f(x) · ∇_θ log p_θ(x)]. This converts the gradient of the expected reward into an expectation of the reward multiplied by the bridge score. Computing the score of clean samples conditioned on noisy samples via Bayes' theorem — reversing the conditioning direction — allows the model's velocity field to directly express the result.

The key decision that makes RAM truly scalable is: directly discarding the path cost term. The authors candidly note that the project took a long time precisely because the team tried various exact and approximate methods to estimate the path cost, but none of them could truly scale.

Why is the path cost so intractable? In stochastic optimal control theory, the path cost measures the trajectory deviation cost from the reference process to the current process — mathematically equal to the integral of the L2 norm of the difference between the two velocity fields along the entire trajectory. In high-dimensional settings like text-to-image generation, Stable Diffusion 3.5's output has approximately 65,000 dimensions (latent space 16×64×64 ≈ 65,536 dimensions). Even if the model deviates only slightly, integrating over 20–50 steps causes the path cost to accumulate to thousands of times the magnitude of the reward signal, producing catastrophic variance explosion. The path cost from the difference between the reference and current model velocity fields is enormous, and when this high-variance scalar is subtracted from the reward, the reward signal is nearly drowned out — like receiving a reward of 100 for one image and -200 for the next, leaving the learning signal completely directionless. This has deep connections to the "credit assignment problem" and "curse of dimensionality" in reinforcement learning.

It's worth emphasizing that discarding the path cost does not mean discarding the entire KL regularization. RAM's fixed point is itself a corollary of the KL-regularized optimal solution, holding only for exponentially tilted distributions. By contrast, baselines like Diffusion NFT lack anchoring to the reference model, which explains why reward hacking is far less severe in RAM than in other baselines — reward hacking refers to models finding opportunistic ways to maximize a proxy reward without genuinely completing the task, such as generating images that appear to contain text but are actually unreadable to fool an OCR reward model.

Training Pipeline and 50x Speedup

RAM's final loss depends on only three points: a clean sample x0 drawn from the current model, a Gaussian noise sample x1, and a noised sample xd — no complete stochastic rollout is required.

The specific pipeline: use ODE (rather than stochastic sampling) to sample x0 from the current model, evaluate the reward, then construct noised samples via linear interpolation just as in pre-training, and finally perform regression. ODE sampling (deterministic solvers) produces a fully deterministic trajectory given an initial noise input — it's the default choice for inference in models like Flux and SD3, and also sidesteps the numerical instability of Flow Matching stochastic schedulers where the diffusion coefficient diverges at the noise end. Once the expensive endpoint samples and rewards are obtained, multiple training targets can be cheaply constructed — for each endpoint sample, multiple (the authors use 8) mutually independent noised samples can be generated. Multiple linear interpolations of the same endpoint with independent Gaussian noise vectors are, from an information-theoretic perspective, equivalent to multiple independent re-samples of the same observation, providing high-quality decorrelated learning signals.

This is the fundamental source of the speedup. The authors compare RAM against Flow-GRPO, a representative policy gradient method:

• Policy gradient methods must perform stochastic sampling, but flow matching models' stochastic schedulers diverge at the noise end with coefficients approaching infinity, making simulation difficult and endpoints blurry. More critically, training steps can only be computed at fixed points along the sampled trajectory, and these points are highly correlated (one point is just a low-noise version of another), yielding very weak learning signals. This connects to the inherent high-variance problem of REINFORCE — DeepSeek-R1 and GRPO dramatically reduce variance for language models via within-group relative rewards as baselines, but the decorrelation problem in diffusion models' continuous action spaces is even more challenging.
• RAM first samples endpoints with ODE, then jumps backward to construct any number of training steps. Since independent Gaussian noise vectors are used during noising, samples are mutually independent with low correlation, providing richer learning signals.

The experimental results are impressive: on benchmarks including GenEval (complex compositional prompts like "red zebra"), visual text rendering, and human preference models, RAM achieves equal or higher rewards using only 1/50th of the GPU hours, with each training step being slightly cheaper than Flow-GRPO. RAM also performs better on aesthetic and human preference metrics, effectively avoiding the problem of "text rendered correctly but image distorted."

Engineering Implementation Details

In Q&A, the authors clarified implementation details. The algorithm actually requires three models:

1. Current model (the object being optimized)
2. Reference model (the pre-trained model, providing KL anchoring)
3. Exponential Moving Average (EMA) model (used for sampling, improving training stability)

The EMA (Exponential Moving Average) model is a common training stabilization technique that uses an exponentially weighted average of historical weights rather than instantaneous current weights for inference — similar to the momentum mechanism in optimization — and effectively suppresses parameter oscillations during training. It has been widely validated for diffusion model training. A fourth EMA model is used for evaluation but is not required.

Since post-training uses LoRA fine-tuning, maintaining multiple model copies is inexpensive. LoRA (Low-Rank Adaptation) dramatically reduces trainable parameters by adding two low-rank matrices in parallel with the original weight matrices (W = W_0 + BA, where rank r is much smaller than the original dimensions) — for diffusion models with billions of parameters, full fine-tuning requires hundreds of gigabytes of VRAM, while LoRA requires only tens of megabytes of additional parameters, making it feasible to maintain multiple sets of model parameters on a single machine: just one base model plus different LoRA weights. The authors also note that although sampling is off-policy, importance sampling ratios are not used, as these ratios are difficult to reliably estimate for diffusion models.

Sampling settings: classifier-free guidance, 20 sampling steps, 24 rollouts per prompt, with mean reward as a baseline subtracted to reduce variance.

Regarding the future direction for path cost, the authors believe it's better to approach from other perspectives on RAM's fixed point rather than struggling to estimate it more stably. They mention that using an EMA or older version of the current model as a substitute for the reference model while reducing regularization strength aligns with recent "tilt matching" work on scaling rewards and tilting in multiple steps to reduce cumulative approximation error — an approach reminiscent of trust region methods' design philosophy. Trust region methods (like TRPO and PPO) ensure learning stability by constraining the magnitude of policy changes at each update step, which is fundamentally analogous to RAM's KL regularization framework.

Summary

RAM's value lies in returning to first principles: identifying the path cost in high-dimensional spaces as the fundamental reason RL for diffusion models is difficult to scale, and boldly discarding it while preserving the essence of KL regularization. This seemingly radical simplification, combined with ODE sampling and the construction of decorrelated training targets, ultimately achieves order-of-magnitude efficiency gains. For domains like robot control and text-to-image generation that rely on "learning from outcomes," RAM provides a simpler, more scalable path for diffusion model alignment.

Key Takeaways