🤖 machine learning

Efficient Adjoint Matching for Fine-tuning Diffusion Models

This paper introduces Efficient Adjoint Matching (EAM), a method that significantly accelerates reward fine-tuning for diffusion models by reformulating the underlying stochastic optimal control problem with a linear base drift and modified terminal cost, thereby eliminating costly backward adjoint simulations and enabling faster convergence while maintaining or improving performance on standard text-to-image benchmarks.

Original authors: Jeongwoo Shin, Dongsoo Shin, Joonseok Lee, Jaewoong Choi, Jaemoo Choi

Published 2026-05-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jeongwoo Shin, Dongsoo Shin, Joonseok Lee, Jaewoong Choi, Jaemoo Choi

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 chef (the pretrained AI model) who is incredibly good at cooking generic meals. They can make a perfect steak or a beautiful cake just by following a standard recipe. However, if you ask them to cook a "spicy vegan burger that looks like a sunset," they might struggle to match your specific, fancy taste.

To fix this, you hire a food critic (the reward model) to taste the dishes and tell the chef how to improve. The goal is to teach the chef to cook exactly what you want.

The Old Way: "Adjoint Matching" (AM)

The previous method, called Adjoint Matching, was like a very strict, slow, and exhausting training regimen for the chef.

The Forward Journey (The Cooking): To teach the chef, the system had to simulate the entire cooking process from raw ingredients to the final dish, step-by-step, using a very slow, random method (like stirring the pot with a spoon that occasionally drops the food). This took a long time and required a lot of computer power.
The Backward Journey (The Critique): Once the dish was done, the system had to replay the entire cooking process in reverse, minute by minute, to figure out exactly which step went wrong and how to fix it. This "rewind and analyze" step was also incredibly slow and computationally heavy.

The Problem: This two-step process (slow forward cooking + slow backward analysis) made training very expensive and slow. It was like trying to learn a new recipe by cooking it 100 times and then rewinding the video 100 times to find the mistakes.

The New Way: "Efficient Adjoint Matching" (EAM)

The authors of this paper realized the slowness came from the specific "rules of the kitchen" (the base drift) the chef was forced to follow. They decided to redesign the kitchen rules to make training faster without changing the final goal.

Here is how EAM works, using simple analogies:

1. The "Shortcut" Cooking (Forward Simulation)
Instead of forcing the chef to cook the dish from scratch using the slow, random method, EAM lets the chef use a fast, deterministic oven to cook the final dish in just a few steps.

The Magic Trick: Once the final dish (the image) is ready, the system doesn't need to simulate the messy cooking process to get the intermediate steps. Instead, it simply "sprinkles some noise" onto the finished dish to create the intermediate states. It's like saying, "Here is the perfect cake; now imagine what it looked like halfway through baking by adding a little flour back in." This skips the slow, random cooking simulation entirely.

2. The "Instant" Critique (Backward Adjoint)
In the old method, the critic had to rewind the video to give feedback. In EAM, because the kitchen rules were changed to be simpler (linear), the critic can give feedback instantly.

The Magic Trick: The system now has a "cheat sheet" (a closed-form formula). Instead of replaying the video to find the error, the critic just looks at the final dish and the cheat sheet, and instantly knows exactly how to adjust the recipe. No rewinding, no slow analysis.

The Results

By making these two changes, the authors found that:

Speed: The new method trains 4 times faster than the old one. It's like going from driving a car in heavy traffic to taking a direct highway.
Quality: Despite being faster, the resulting AI models are just as good, or even better, at following human preferences. They produce images that are more beautiful, follow prompts more accurately, and look more natural.
Details: Because the new method preserves the connection between the final image and the "noisy" steps better, the AI keeps fine details (like the texture of fur or the reflection in water) that the old method sometimes smoothed out.

Summary

The paper introduces Efficient Adjoint Matching (EAM). It solves the problem of slow AI training by changing the underlying rules of how the AI "thinks" during training.

Old Way: Slow, random cooking + Slow, backward video analysis.
New Way: Fast, direct cooking + Instant, formula-based feedback.

The result is a smarter, faster way to teach AI models to create art that humans actually like, without needing supercomputers to run for days.

Technical Summary: Efficient Adjoint Matching for Fine-tuning Diffusion Models

Problem Statement

Reward-based fine-tuning has become the standard approach for aligning pretrained diffusion and flow models with human preferences in text-to-image (T2I) generation. Among these, Adjoint Matching (AM) offers a principled formulation by casting reward fine-tuning as a Stochastic Optimal Control (SOC) problem. AM aims to transform a pretrained model into a target distribution $p^\star_1(x) \propto e^{\beta r(x)}p_{data}(x)$ , where $r(x)$ is a reward function.

However, AM suffers from significant computational inefficiencies that hinder its practical application:

Forward Simulation Bottleneck: AM requires simulating full generative trajectories using a stochastic differential equation (SDE) solver under "memoryless" dynamics. This necessitates a large number of function evaluations (NFEs) and requires storing the entire trajectory in memory.
Backward Simulation Bottleneck: To compute the control signal, AM requires a backward Ordinary Differential Equation (ODE) simulation of the adjoint state along the sampled trajectory. This involves computing Jacobian-Vector Products (JVPs) through the pretrained network at every integration step, which is computationally expensive.

The authors observe that these bottlenecks are intrinsically tied to the non-trivial base drift inherited from the pretrained model in standard AM formulations.

Methodology: Efficient Adjoint Matching (EAM)

The paper proposes Efficient Adjoint Matching (EAM), which reformulates the SOC problem to eliminate both computational bottlenecks while preserving the theoretical guarantee of reaching the reward-tilted target distribution.

1. Redesigning Base Dynamics

EAM introduces two structural conditions for the base dynamics:

Linearity: The base drift $b(x,t)$ is redesigned to be linear in $x$ (i.e., $b(x,t) = D(t)x$ ).
Memorylessness: The base dynamic remains memoryless, allowing intermediate states to be sampled directly from the endpoint.

By fixing the diffusion coefficient $\sigma(t)$ to match the standard setup, the authors derive a specific family of admissible linear drifts (Eq. 15) that satisfy these conditions. This linear drift enables a closed-form solution for the adjoint state, eliminating the need for backward ODE simulation and JVP evaluations.

2. Terminal Cost Correction

Changing the base drift alters the terminal distribution of the uncontrolled system. To ensure the optimal controlled distribution still matches the desired reward-tilted distribution $p^\star_1$ , EAM derives a corrected terminal cost $g(x)$ . The gradient of this cost, $\nabla g(x)$ , includes terms for the pretrained score, the new base terminal score, and the reward gradient (Eq. 16).

3. Training Algorithm

The EAM training loop (Algorithm 1) proceeds as follows:

Endpoint Generation: Instead of a full SDE rollout, the endpoint image $X_1$ is generated using a few-step deterministic ODE solver (e.g., 10 NFEs).
Intermediate State Construction: Intermediate states $X_t$ are sampled directly from the original noising kernel $q_t(\cdot | X_1)$ conditioned on the generated endpoint, bypassing the need for forward SDE simulation.
Closed-Form Adjoint: The adjoint state is computed via a closed-form scaling of the endpoint gradient $\nabla g(X_1)$ , removing the backward simulation entirely.
Control Parameterization: The control $u(x,t)$ is parameterized via a trainable velocity model (e.g., LoRA) initialized from the pretrained model, allowing fine-tuning to start from the pretrained dynamics.

Key Contributions

Algorithm Proposal: The introduction of EAM, a reward-gradient-based fine-tuning algorithm that achieves the same theoretical objectives as AM but with significantly reduced computational cost.
Theoretical Characterization: A characterization of linear base drifts that satisfy the requirements for efficient trajectory construction (memorylessness) and closed-form adjoint computation.
Empirical Validation: Demonstration that EAM converges up to 4 $\times$ faster than AM on standard T2I benchmarks while matching or surpassing its performance across multiple metrics.

Results

The authors evaluated EAM on Stable Diffusion 3.5-Medium (SD3.5-M) using the Pick-a-Pic dataset, comparing it against AM and the baseline SD3.5-M.

Efficiency: EAM reduces per-iteration training cost by up to 4 $\times$ . This is achieved by replacing 40-step SDE simulations with 10-step ODE endpoint sampling and eliminating backward adjoint simulations.
Performance: On benchmarks including PickScore, ImageReward, HPSv2.1, CLIPScore, and Aesthetics, EAM consistently matches or outperforms AM.
- In single-reward (PickScore) settings, EAM achieved a PickScore of 23.16 compared to AM's 23.00.
- In multi-reward settings (PickScore + HPSv2.1 + Aesthetics), EAM achieved a PickScore of 22.93 and an Aesthetics score of 6.56, outperforming AM (6.21).
Qualitative Improvements: Visual comparisons indicate that EAM better preserves fine-grained details and complex compositions compared to AM, which sometimes smooths out details. This is attributed to EAM's consistency between training (endpoint-conditioned noising) and inference trajectories.

Significance and Claims

The paper claims that the primary limitations of Adjoint Matching are not inherent to the SOC formulation itself but are artifacts of the specific choice of base drift. By redesigning the base drift to be linear and memoryless, and correspondingly adjusting the terminal cost, EAM removes the two major computational bottlenecks (forward SDE simulation and backward adjoint simulation) without sacrificing the theoretical alignment guarantees.

The authors position EAM as a practical, scalable solution for reward fine-tuning that retains the sample efficiency of gradient-based methods while drastically reducing the computational overhead, making high-quality alignment of diffusion models more accessible. The work does not claim to solve all optimization challenges (e.g., buffer replay or adaptive schedules are noted as future work) but establishes a new baseline for efficiency in this domain.