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

TL;DR

This paper introduces a deterministic adjoint matching framework for flow-based generative models, improving human preference alignment through a control perspective, computational efficiency, and flexible regularization.

Contribution

It presents a novel control-based training method with a truncated adjoint scheme and generalized regularization, enhancing alignment and diversity in flow models.

Findings

Consistent improvements in alignment metrics across experiments.

Substantial computational savings with the truncated adjoint scheme.

Enhanced diversity and mode preservation in generated outputs.

Abstract

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.