Weighted Adjoint Matching: Finetuning Flow Matching Models with Time-Dependent Control Cost

Eliot Beyler; Francis Bach

stat.MLOct 14, 2026

Weighted Adjoint Matching: Finetuning Flow Matching Models with Time-Dependent Control Cost

Authors: Eliot Beyler, Francis Bach

Organizations: Inria, Ecole Normale Supérieure, PSL Research University

Abstract

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.

Panel of Figure 2 — 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.

Panel of Figure 3 — Figure 3: Average reward with respect to weight of positive mode for adjoint matching and weighted adjoint matching, using different functions $t\mapsto w(t)$ .

Figure 4: Average brightness reward on the MNIST dataset with respect to class count ratio

\frac{n_{\min}}{n_{\max}}

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

\left\{\begin{array}[]{ll}X^{(m)}_{t_{0}}&=X^{(m)}_{0}\sim\mathcal{N}(0,I),\\ X^{(m)}_{t_{k+1}}&=X^{(m)}_{t_{k}}+h\left(2v^{\textnormal{finetune}}_{\theta}(X^{(m)}_{t_{k}},t_{k})-\frac{\dot{\alpha}_{t_{k}}}{\alpha_{t_{k}}}X^{(m)}_{t_{k}}\right)+\sqrt{h}\gamma(t_{k})\mathcal{N}(0,I).\end{array}\right.

(34)

Appendix

Algorithm 1 Weighted Adjoint Matching to finetune a flow matching model

Figure 5: Plot of

\mathbb{E}[r(\hat{X})]

with respect to

\mathbb{E}[\mathbb{1}_{x_{1}>0}(\hat{X})]

for adjoint matching and weighted adjoint matching with a linear weighting function.

Panel of Figure 6 — Figure 6: Average reward with respect to weight of positive mode for adjoint matching and weighted adjoint matching, using different functions $t\mapsto 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 $t_{s}$ .

Figure 8: Average reward with respect to ratio

\frac{n_{\min}}{n_{\max}}

			Predicted class (%)
Method	$a$	$\hat{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, $t_{s}=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, $t_{s}=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, $t_{s}=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$	$\hat{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, $t_{s}=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, $t_{s}=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, $t_{s}=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$	$\hat{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, $t_{s}=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, $t_{s}=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, $t_{s}=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.

Panel of Figure 9 — 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 $1\,000$ displayed samples, and therefore differ from those of Table 1 , computed on $50\,000$ samples.

Panel of Figure 10 — 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 $1\,000$ displayed samples, and therefore differ from those of Table 2 , computed on $50\,000$ samples.

Panel of Figure 11 — 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 $1\,000$ displayed samples, and therefore differ from those of Table 3 , computed on $50\,000$ samples.

Diffusion Model Fine-Tuning

Diffusion model fine-tuning adapts pretrained generative models to new concepts, tasks, domains, and preferences through targeted updates.

See all 181 papers

Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipeline
May 7, 2026Zhengyi Guo, Jiayuan Sheng, David D. Yao +1Fine-Tuning Flow-Based Generative Modeling
Efficient Adjoint Matching for Fine-tuning Diffusion Models
May 12, 2026Jeongwoo Shin, Dongsoo Shin, Yuchen Zhu +5Diffusion Model Fine-Tuning RL for Diffusion Models
Seoul National University · Georgia Institute of Technology · Sungkyunkwan University
A unified perspective on fine-tuning and sampling with diffusion and flow models
Apr 30, 2026Carles Domingo-Enrich, Yuanqi Du, Michael S. AlbergoStochastic Optimal Control Diffusion Sampling
Microsoft Research New England · Harvard University, Kempner Institute

May 7, 2026J/K move · Enter open · S save

Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipeline

Zhengyi Guo, Jiayuan Sheng, David D. Yao, Wenpin Tang

Fine-Tuning Flow-Based Generative Modeling

Open paper AskPDF arXiv

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.

Abstract

Diffusion Model Fine-Tuning

Explore similar work

Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipeline

Efficient Adjoint Matching for Fine-tuning Diffusion Models

A unified perspective on fine-tuning and sampling with diffusion and flow models