🤖 machine learning

Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models

This paper introduces Reinforce Adjoint Matching (RAM), a scalable RL post-training method for diffusion and flow-matching models that achieves superior alignment with rewards by deriving a simple consistency loss which corrects pretraining targets without requiring costly SDE rollouts, backward adjoint sweeps, or reward gradients.

Original authors: Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a master artist who has spent years learning to paint by copying thousands of photos. This artist is incredibly good at recreating what they see, but they don't necessarily know how to follow specific, tricky instructions like "paint a red zebra next to a blue truck" or "write the word 'HELLO' clearly on a shield."

To teach the artist these new skills, we usually use a method called Reinforcement Learning (RL). Think of this as a strict art critic who looks at the painting, gives it a score, and tells the artist, "Do better next time."

The Old Way: The Expensive, Slow Critic

Previously, training these AI artists with a critic was like trying to teach someone to swim by throwing them into the ocean and watching them flail.

The Process: The AI would generate a full image (the "swim"), the critic would score it, and then the system would try to figure out exactly which brushstroke caused the good or bad score.
The Problem: To do this, the AI had to simulate the entire painting process from start to finish, over and over again, just to get one piece of feedback. It was slow, computationally expensive, and often unstable (like trying to calculate the wind speed while the boat is sinking).

The New Way: RAM (Reinforce Adjoint Matching)

The authors of this paper, Andreas Bergmeister and his team, realized there's a much smarter way to do this. They found a "shortcut" that keeps the training fast and simple, like the original pre-training phase.

Here is how RAM works, using a simple analogy:

1. The "Clean Endpoint" Trick
Imagine the painting process is like a movie playing in reverse. The movie starts with a blank canvas (noise) and ends with a finished painting.

Old Method: The AI had to watch the whole movie, score the ending, and then rewind the movie frame-by-frame to see where it went wrong.
RAM Method: The authors realized that if you know the final painting (the "clean endpoint"), you don't need to watch the whole movie to understand the process. You can just take that finished painting, and mathematically "add noise" to it instantly to create a messy version, just like the original training did.

2. The "One-and-Done" Feedback
Instead of simulating the whole painting process to get feedback, RAM does this:

Generate: The AI paints a finished image (the endpoint).
Score: The critic gives it a score (e.g., "Great job on the text!").
Rewind Instantly: Instead of simulating the rewind, the math instantly creates a "noisy" version of that image.
Learn: The AI learns to turn that specific noisy version back into the high-scoring image.

The Magic Insight:
The paper proves that the rules for how to "add noise" to an image don't change, even when you are trying to make the image better. The only thing that changes is which images the AI decides to paint in the first place. Because the "noise rules" stay the same, the AI can use the same simple math it used during its original training, just with a new target.

Why This is a Big Deal

Speed: The paper claims RAM is 34 to 50 times faster than previous methods. It reaches the same level of skill in a fraction of the time.
Simplicity: It doesn't require complex, unstable calculations (like "SDE rollouts" or "backward adjoint sweeps"). It's just a simple regression task, similar to the original training.
Quality: When they tested it on Stable Diffusion 3.5M, the AI got much better at:
- Composability: Putting objects in the right places (e.g., a truck left of a refrigerator).
- Text Rendering: Writing words clearly in images.
- Human Preference: Making images that humans actually like looking at.

The Bottom Line

Think of RAM as giving the artist a "magic eraser" and a "magic scorecard." Instead of forcing the artist to repaint the whole canvas to see what went wrong, the magic eraser instantly shows them a messy version of their best work, and the scorecard tells them exactly how to fix that specific mess. This allows the AI to learn complex new skills without the heavy computational cost of previous methods.

Technical Summary: Reinforce Adjoint Matching (RAM)

Problem Statement

Diffusion and flow-matching models achieve high-fidelity generation because their pretraining relies on supervised regression: a clean sample is analytically noised, and the model regresses against a closed-form target. However, aligning these models with reward functions (e.g., for compositional correctness, text rendering, or human preference) via Reinforcement Learning (RL) post-training disrupts this efficient structure.

Existing RL post-training methods face significant computational and structural hurdles:

Policy Gradient Methods: Treat denoising as a Markov Decision Process, requiring stochastic SDE rollouts and optimizing variational lower bounds. This introduces high variance and requires many sampling steps.
Stochastic Optimal Control (Adjoint) Methods: Target the KL-regularized reward maximization objective exactly but require backward adjoint sweeps along stored SDE trajectories and reward gradients backpropagated through the sampling process. This is numerically unstable and computationally expensive, particularly for flow-matching models where noise schedules diverge near the noisy end.
Surrogate Losses: Recent approaches replace intractable likelihoods with surrogates (e.g., ELBOs or contrastive endpoints), sacrificing the clean regression structure that enables pretraining to scale.

The core challenge is to derive an RL post-training objective that retains the simplicity and scalability of the pretraining regression framework while effectively optimizing for rewards.

Methodology: Reinforce Adjoint Matching (RAM)

The authors propose Reinforce Adjoint Matching (RAM), a method that extends the pretraining regression structure to RL post-training without requiring SDE rollouts, backward adjoint sweeps, or reward gradients.

Theoretical Foundation

The method is grounded in KL-regularized reward maximization formulated as a stochastic optimal control problem.

Optimal Process Structure: Under KL regularization, the optimal generative process tilts the distribution of clean endpoints ( $X_0$ ) toward high-reward samples ( $p^*_0 \propto p_{ref}^0 \exp(r)$ ) but leaves the conditional law of noisy states given a clean endpoint unchanged. Specifically, the trajectory from a fixed $X_0$ to a noisy state $X_t$ remains the same analytic Gaussian kernel used in pretraining.
Adjoint Matching Condition: The optimal control $u_t$ satisfies a fixed-point condition where the control equals the negative scaled gradient of the value function (the adjoint): $u_t = -\sigma_t \nabla_x V_t(x)$ .
REINFORCE Decomposition: To avoid computing $\nabla r$ (which may be non-differentiable), the authors apply the log-derivative (REINFORCE) identity to the value function. This decomposes the adjoint into a reward term weighted by the backward bridge score and a path-cost correction term.

The RAM Objective

RAM approximates the adjoint by dropping the path-cost correction (justified by low variance and stability at scale) and utilizing the structural insight that the noising kernel remains analytic.

Endpoint Sampling: Instead of simulating a controlled SDE, RAM draws a clean endpoint $X_0$ from the current model's distribution using any off-the-shelf ODE sampler.
Analytic Noising: The sampled $X_0$ is noised analytically to generate training states $X_t = (1-t)X_0 + t\epsilon$ , exactly as in pretraining.
Closed-Form Target: The regression target for the velocity field $v_\theta$ is derived by substituting the reward term and the Bayes bridge score (which has a closed-form expression for linear noising) into the adjoint condition.

The resulting loss function is:
$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$
where $X_0 \sim p_\theta^0$ , $\epsilon \sim \mathcal{N}(0, I)$ , and $\text{sg}(\cdot)$ denotes the stop-gradient operator.

Key Efficiency Mechanisms

No SDE Rollouts: Training states are generated by noising a single endpoint $K$ times with independent noise vectors. This yields $K$ conditionally independent training samples per model forward pass, providing more gradient signal per step than methods relying on single SDE trajectories.
No Reward Gradients: The method works with black-box rewards.
No Backward Sweeps: The analytic nature of the noising kernel eliminates the need for backward adjoint integration.

Key Contributions

Structural Extension: The paper demonstrates that the regression structure of pretraining extends to RL post-training under KL regularization. The optimal process only changes the endpoint distribution, not the conditional noising law.
RAM Algorithm: A novel consistency loss that corrects the pretraining target with reward information using a closed-form Bayes bridge score, avoiding the computational bottlenecks of existing methods.
Theoretical Justification: A derivation showing that RAM's fixed point corresponds to the KL-regularized optimum, with exactness guaranteed for Gaussian references under linear rewards and first-order agreement generally.
Variance Analysis: An analysis of path-cost estimators showing that dropping the path-cost correction (as RAM does) significantly reduces variance compared to full-horizon or local adjoint estimators, making it the only estimator viable for image-scale post-training.

Experimental Results

The authors evaluated RAM on Stable Diffusion 3.5 Medium (SD3.5M) across three tasks:

Compositional Image Generation (GenEval): Measuring object, attribute, and spatial relation correctness.
Visual Text Rendering (OCR): Measuring legibility of text in generated images.
Human Preference Alignment (PickScore): Aligning with human preferences.

Performance:

Reward Maximization: RAM achieved the highest training rewards on all three tasks, outperforming baselines including Flow-GRPO, DiffusionNFT, and AWM.
Training Efficiency: RAM reached the peak reward of Flow-GRPO in 50× fewer steps for GenEval, 48× fewer for OCR, and 34× fewer for PickScore.
Image Quality: Unlike some baselines that suffered from "reward hacking" (degrading aesthetic quality for task metrics), RAM maintained or improved image quality metrics (Aesthetic, DeQA, ImageReward, HPSv2) while maximizing task rewards.
CFG Compatibility: RAM remains compatible with Classifier-Free Guidance (CFG) after post-training, whereas baselines like DiffusionNFT and AWM saw their task rewards drop significantly when CFG was applied.

Significance and Claims

The paper claims that RAM represents a paradigm shift in RL post-training for diffusion and flow-matching models. By leveraging the specific structural properties of the optimal control problem, RAM restores the "simple regression" paradigm that made pretraining scalable.

The authors assert that:

RAM eliminates the need for costly SDE rollouts and backward adjoint sweeps, which have been the primary bottlenecks for scaling RL to high-resolution image generation.
The method is "simple and scales," achieving state-of-the-art results with significantly fewer training steps and comparable per-step compute costs to existing methods.
The approach successfully aligns models with complex rewards (composability, text, preferences) without visible quality degradation or reward hacking, demonstrating that the pretraining structure is robust enough to handle post-training alignment when the correct theoretical conditions are met.