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

2026-05-07

AI Summary

This work addresses the training instability and computational inefficiency of flow-based generative models when aligning with human preferences by formulating the problem as an optimal control task over velocity fields. The authors propose a deterministic adjoint matching framework that directly regresses the control signal via value-gradient guidance and introduces a novel trajectory-end truncation strategy to reduce the computational cost of adjoint calculations. This approach enables flexible trade-offs between preference alignment and distribution preservation beyond conventional KL regularization. Experimental results on SiT-XL/2 and FLUX.2-Klein-4B demonstrate that the method significantly improves alignment performance while enhancing generation diversity and mode retention.

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.

Problem

• flow-based generative models
• human preference alignment
• optimal control
• distributional preservation
• diversity

Innovation

• adjoint matching
• flow-based generative models
• optimal control
• preference alignment
• truncated adjoint