Reinforce Adjoint Matching: Scaling Diffusion RL

This work introduces Reinforce Adjoint Matching (RAM), a novel method for post-training diffusion and flow-matching models using reinforcement learning, designed to overcome the scalability and complexity limitations of existing RL approaches. While diffusion model pre-training is inherently scalable due to its supervised regression nature (noising a clean sample and regressing against a closed-form target), RL post-training has traditionally relied on costly SDE rollouts, reward gradients, or surrogate losses, sacrificing this regression structure.

Theoretical Foundations

The core idea of RAM is to extend the scalable regression structure of pre-training to RL post-training. The authors leverage the principle that, under KL-regularized reward maximization, the optimal generative process exponentially tilts the clean-endpoint distribution towards samples with higher reward, while leaving the noising law unchanged. This means the problem can be framed as finding a new velocity field that corresponds to this tilted distribution.

1. Optimal Distribution and Velocity Fields: The optimal distribution is an exponential tilting of the reference (pre-trained) distribution. The paper establishes a key relationship: the difference between the velocity field of the optimal process and the velocity field of the reference process is equivalent to the gradient of a value function, V(xt). This value function represents the expected reward minus a ‘path cost’, which measures the L2 norm difference between the optimal and reference models along a sample path.
2. Adjoint Matching Challenge: The goal is to estimate this gradient of the value function (the “adjoint”) and use it as a regression target. Traditional adjoint matching methods require the reward function to be differentiable and rely on stochastic SDE rollouts with expensive and numerically unstable backpropagation.
3. RAM’s Solution - Approximating the Adjoint: RAM tackles this by splitting the value function into two terms: a reward term (non-differentiable) and a path cost term (differentiable).
   • Reward Term Gradient: The gradient of the expected reward term is estimated using the REINFORCE identity: ∇E[R] = E[R * score]. The crucial component here is the “bridge score” (∇log p(x0|xt)), which is the score of the conditional density of a clean sample x0 given a noisy sample xt. This is computed using Bayes’ theorem, relating it to the analytical noising kernel and the model’s velocity field. A biased but practically low-variance estimator is preferred over an unbiased but high-variance one.
   • Path Cost Term: The path cost can be differentiated, but in high-dimensional settings (like image generation), it proved to be computationally expensive, numerically unstable, and introduced high variance. The path cost could easily dominate the reward signal.
4. Key Approximation: RAM’s critical innovation is to approximate the value function by dropping the path cost term. This simplifies the adjoint to just the gradient of the expected reward. This approximation leads to a consistency loss that corrects the pre-training target with the reward.

The RAM Training Objective and Procedure

The resulting RAM objective is a simple consistency loss that modifies the pre-training target. For a noisy sample xt, the model is trained to regress against a target which is the reference model’s prediction plus a term proportional to the computed reward R(x0) and a residual (x0 - xt). This residual essentially steers the model in a more rewarding direction.

Training Procedure:

1. Sample Endpoint x0: A clean sample x0 is generated from the current fine-tuned model, typically using a deterministic ODE sampler (e.g., probability flow ODE).
2. Evaluate Reward R(x0): The reward for x0 is computed.
3. Noise x0 to xt: As in pre-training, x0 is noised by linearly interpolating it with a Gaussian noise sample x1 at a random timestep t to get xt.
4. Regress: The current model predicts the velocity field for xt, and this prediction is regressed against the modified target. This process is highly efficient: once x0 and R(x0) are obtained (the most expensive part), multiple xt samples and regression targets can be cheaply generated, providing a rich, less correlated learning signal.

Empirical Results and Advantages

RAM was evaluated on Stable Diffusion 3.5M across several benchmarks: complex composability (Geneval), visual text rendering, and human preference. It demonstrated significant performance gains:

• Efficiency: RAM achieved the highest reward with up to 50x fewer training steps (GPU hours) compared to Flow-GRPO and other policy gradient baselines. This speedup stems from RAM’s ability to use deterministic ODE sampling for x0 and then generate independent training targets, unlike policy gradient methods that require stochastic SDE rollouts and yield highly correlated gradients.
• Quality: RAM also performed better on image quality metrics, indicating less reward hacking and more realistic generated images compared to other methods.

Discussion on KL-Regularization

Crucially, dropping the path cost in the adjoint approximation does not imply abandoning KL-regularization. The underlying optimality condition (relating velocity fields to the gradient of the value function) is a consequence of KL-regularized reward maximization. The reference_model term in the RAM loss serves as the anchor, ensuring the fine-tuned model stays within the support of the pre-trained distribution, thus implicitly maintaining KL-regularization. This design prevents severe reward hacking seen in some other methods.

• Tags:
• Reinforcement Learning
• Generative Models
• Feature
• Scalability
• Post-Training
• Image Generation
• Diffusion Models
• Flow Matching
• Adjoint Matching
• REINFORCE