Weighted Adjoint Matching: Finetuning Flow Matching Models with Time-Dependent Control Cost
Fine-tuning pretrained diffusion and flow matching models for reward maximization tends to collapse modes or shift mass between distant regions. This paper introduces Weighted Adjoint Matching, which modifies the stochastic optimal control formulation of Adjoint Matching by placing a time-dependent weight $w(t)$ on the quadratic control cost---expensive early to preserve coarse semantics, cheap late to refine details and increase reward. The authors prove equivalence to maximizing reward under pathwise KL regularization, and demonstrate improved reward-diversity tradeoffs on Gaussian mixtures and MNIST.
This paper makes a clean theoretical contribution: by adding a time-dependent weight to the control cost in the SOC formulation, the authors obtain an objective equivalent to KL-regularizing every sub-path $[0,t]$ of the generative process (Proposition 1). The resulting algorithm is a minimal modification of Adjoint Matching---dividing by $w(t)$ in the regression loss---and the Gaussian mixture experiment crisply demonstrates how standard tilting redistributes mass across modes regardless of spatial separation. However, the empirical evaluation stops at MNIST; the authors candidly acknowledge that scaling to large conditional models remains future work.
The theoretical grounding is the strongest asset. Proposition 1 shows that the weighted SOC problem is equivalent to an objective regularized by the KL divergence on every sub-path $[0,t]$, yielding a principled interpretation of why early-time regularization preserves diversity. The algorithmic change is minimal---"the extra time weight $w(t)$ that only appears at one place in the loss"---and the authors identify and fix a numerical instability caused by exploding $\gamma(t)$ factors via an auxiliary weight $\lambda(t)$. The Gaussian mixture example is pedagogically valuable: it analytically shows that tilting changes mode weights even when both modes are equidistant from the reward axis.
The primary weakness is limited experimental scope. The paper explicitly admits, "We have focused on the case of unconditional generative modeling and simple datasets for which balance could be evaluated precisely," and all comparisons are restricted to a low-dimensional Gaussian mixture and unconditional MNIST with a brightness reward. The choice of extremely large early weights---$w_{\text{steps}}(t)=10^5$ for $t
The evidence supports the core claim that time-dependent weighting improves the reward-diversity tradeoff relative to uniform Adjoint Matching. On the Gaussian mixture, adjoint matching empirically tracks the theoretical tilted-distribution curve, while weighted variants achieve higher reward at the same mode balance. The MNIST brightness task---despite its simplicity---shows that weighted adjoint matching preserves class ratios better: "adjoint matching generates 18.0% of zeros and only 3.4% of ones, while weighted adjoint matching ($t_s=0.45$) reaches a reward of 0.180 with 10.5% of zeros and 9.1% of ones." The comparison to related work is generally fair; the authors correctly distinguish methods that merely truncate early timesteps or use scalar RL signals from their principled path-KL formulation. One gap is the absence of any empirical comparison to the Wasserstein-regularized baseline discussed in Section 6.
The algorithm is presented with sufficient mathematical detail to reproduce, including pseudocode in Appendix D.5 and a full derivation of the functional derivative in Appendix D. The authors address numerical stability explicitly by introducing $\lambda(t)$ to cancel exploding $\gamma(t)$ factors near the endpoints. However, heavy AI assistance on the codebase is disclosed: while "the core implementation of Algorithm 1 was written by the authors, then reviewed and improved with AI assistance," the remainder of the code---including hyperparameter choices, checkpointing, and figure generation---was largely AI-written. Full architecture details and hyperparameters are referenced in Appendix E, though the provided text truncates before those details appear.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.