We introduce Weighted Adjoint Matching, a method to finetune a pretrained diffusion or flow matching model given a differentiable reward. We derive it from Adjoint Matching, which reformulates finetuning as a stochastic optimal control problem with a quadratic control cost. Building on the observation that these models generate coarse features before fine details, we place a time-dependent weight on the control cost. A high cost early in the generative process preserves balance and diversity in the coarse features, which carry most of the semantics, while a low cost at the end lets the model increase the reward by refining details. We prove that our method is equivalent to maximizing the reward under time-varying Kullback--Leibler regularization on the generative paths. Experiments on Gaussian mixtures and on MNIST show a better tradeoff between reward and mode balance than Adjoint Matching.
Figures & tables
Figure 1: Trajectories of a flow matching model, in black, first sample the coarse features, illustrated here as the class of the digit, and only later its fine details, here differences in brightness. Suppose we want to finetune the model to produce brighter digits. Adjoint matching [ 35 ] , in red , penalizes the control uniformly in time. Early in the process, when the coarse features are not yet settled, a small push is enough to move a trajectory towards the class that is brighter on average, here the eights. Weighted adjoint matching, in green , makes the control expensive at the beginning of the process, so that the finetuned trajectories follow the base process until the coarse features are settled, then it refines only the details, leading to a more balanced distribution.
Figure 2: Red component proportional to probability that trajectory going through the point ends up in upper mode, blue to probability for lower mode, and transparency proportional to density of marginal distribution.
Figure 3: Average reward with respect to weight of positive mode for adjoint matching and weighted adjoint matching, using different functions t↦w(t) .
Figure 4: Average brightness reward on the MNIST dataset with respect to class count ratio nmaxnmin .
Algorithm 1 Weighted Adjoint Matching to finetune a flow matching model
Figure 5: Plot of E[r(X^)] with respect to E[1x1>0(X^)] for adjoint matching and weighted adjoint matching with a linear weighting function.
Figure 6: Average reward with respect to weight of positive mode for adjoint matching and weighted adjoint matching, using different functions t↦w(t) .
Figure 7: Sample of base distribution (blue) and finetuned models (orange), for adjoint matching and weighted adjoint matching for different values of a and ts .
Figure 8: Average reward with respect to ratio nmaxnmin .
Predicted class (%)
Method
a
R^
0
1
2
3
4
5
6
7
8
9
Base
–
0.140
10.2
9.9
9.9
10.8
9.4
9.1
9.8
10.0
10.1
10.9
Adjoint
5.7
0.159
14.8
5.8
12.0
12.0
7.4
8.5
10.0
8.0
12.1
9.4
Weighted, ts=0.25
21
0.161
10.6
6.6
11.4
11.8
8.9
8.1
10.2
9.1
13.0
10.4
Weighted, ts=0.3
34
0.162
10.5
7.4
10.6
10.9
9.4
8.5
9.2
9.2
12.9
11.4
Weighted, ts=0.35
52
0.161
10.6
8.8
10.2
11.0
9.2
7.7
9.8
9.0
12.3
11.2
Appendix
Table 1: Predicted class of the generated samples, in percent, for the run of each method whose average luminosity is closest to 0.16.
Predicted class (%)
Method
a
R^
0
1
2
3
4
5
6
7
8
9
Base
–
0.140
10.2
9.9
9.9
10.8
9.4
9.1
9.8
10.0
10.1
10.9
Adjoint
10
0.176
18.0
3.4
12.5
12.3
7.0
7.9
10.7
6.5
13.1
8.6
Weighted, ts=0.25
37
0.181
12.0
3.4
13.7
11.3
8.2
7.8
8.0
7.7
17.4
10.6
Weighted, ts=0.3
57
0.180
11.9
6.1
11.2
11.6
8.1
6.4
8.9
8.4
16.0
11.2
Weighted, ts=0.35
92
0.180
10.6
7.8
11.0
10.7
8.5
7.6
9.7
8.6
14.2
11.4
Appendix
Table 2: Predicted class of the generated samples, in percent, for the run of each method whose average luminosity is closest to 0.18.
Predicted class (%)
Method
a
R^
0
1
2
3
4
5
6
7
8
9
Base
–
0.857
10.4
10.1
9.4
10.8
9.5
9.0
9.9
9.8
10.3
10.7
Adjoint
8.6
0.880
5.3
19.6
7.8
9.9
9.6
9.6
8.9
10.8
7.8
10.6
Weighted, ts=0.25
25
0.880
8.8
14.8
8.1
10.4
12.0
9.0
7.3
10.4
7.6
11.6
Weighted, ts=0.3
43
0.880
9.3
12.2
8.8
11.1
10.2
10.0
9.6
10.7
8.1
10.1
Weighted, ts=0.35
63
0.880
9.8
11.2
9.0
10.5
10.6
10.5
9.4
10.4
8.4
10.2
Appendix
Table 3: Predicted class of the generated samples, in percent, for the run of each method whose average darkness is closest to 0.88.
Figure 9: Generated samples grouped by predicted class, for the run of each method whose average luminosity is closest to 0.16. Percentages are computed on the 1000 displayed samples, and therefore differ from those of Table 1 , computed on 50000 samples.
Figure 10: Generated samples grouped by predicted class, for the run of each method whose average luminosity is closest to 0.18. Percentages are computed on the 1000 displayed samples, and therefore differ from those of Table 2 , computed on 50000 samples.
Figure 11: Generated samples grouped by predicted class, for the run of each method whose average darkness is closest to 0.88. Percentages are computed on the 1000 displayed samples, and therefore differ from those of Table 3 , computed on 50000 samples.
We propose a deterministic adjoint matching framework that formulates human preference alignment for flow-based generative models as an optimal control problem over velocity fields. One can directly regress the control toward a value-gradient-induced target under the current policy, leading to a simple and stable training objective. Building on this perspective, we introduce a truncated adjoint scheme that focuses computation on the terminal portion of the trajectory, where reward-relevant signals concentrate, which yields substantial computational savings while preserving alignment quality. We further generalize the framework beyond standard KL-based regularization, allowing more flexible trade-offs between alignment strength and distributional preservation. Experiments on SiT-XL/2 and FLUX.2-Klein-4B demonstrate consistent gains across multiple alignment metrics, along with substantially improved diversity and mode preservation.
Reward fine-tuning has become a common approach for aligning pretrained diffusion and flow models with human preferences in text-to-image generation. Among reward-gradient-based methods, Adjoint Matching (AM) provides a principled formulation by casting reward fine-tuning as a stochastic optimal control (SOC) problem. However, AM inevitably requires a substantial computational cost: it requires (i) stochastic simulation of full generative trajectories under memoryless dynamics, resulting in a large number of function evaluations, and (ii) backward ODE simulation of the adjoint state along each sampled trajectory. In this work, we observe that both bottlenecks are closely tied to the non-trivial base drift inherited from the pretrained model. Motivated by this observation, we propose Efficient Adjoint Matching (EAM), which substantially improves training efficiency by reformulating the SOC problem with a linear base drift and a correspondingly modified terminal cost. This reformulation removes both sources of inefficiency; it enables training-time sampling with a few-step deterministic ODE solver and yields a closed-form adjoint solution that eliminates backward adjoint simulation. On standard text-to-image reward fine-tuning benchmarks, EAM converges up to 4x faster than AM and matches or surpasses it across various metrics including PickScore, ImageReward, HPSv2.1, CLIPScore and Aesthetics.
Jeongwoo Shin, Dongsoo Shin, Yuchen Zhu +5
Seoul National University · Georgia Institute of Technology · Sungkyunkwan University
We study the problem of training diffusion and flow generative models to sample from target distributions defined by an exponential tilting of a base density; a formulation that subsumes both sampling from unnormalized densities and reward fine-tuning of pre-trained models. This problem can be approached from a stochastic optimal control (SOC) perspective, using adjoint-based or score matching methods, or from a non-equilibrium thermodynamics perspective. We provide a unified framework encompassing these approaches and make three main contributions: (i) bias-variance decompositions revealing that Adjoint Matching/Sampling and Novel Score Matching have finite gradient variance, while Target and Conditional Score Matching do not; (ii) norm bounds on the lean adjoint ODE that theoretically support the effectiveness of adjoint-based methods; and (iii) adaptations of the CMCD and NETS loss functions, along with novel Crooks and Jarzynski identities, to the exponential tilting setting. We validate our analysis with reward fine-tuning experiments on Stable Diffusion 1.5 and 3.
Carles Domingo-Enrich, Yuanqi Du, Michael S. Albergo
Microsoft Research New England · Harvard University, Kempner Institute