Gradual Fine-Tuning for Flow Matching Models

TL;DR

Gradual Fine-Tuning (GFT) addresses the challenge of adapting flow matching models to new target distributions without sacrificing pretrained advantages or requiring memoryless noise schedules. By formulating a KL-based objective that linearly interpolates between the pretrained drift and the target drift through a temperature parameter, GFT yields a closed-form optimal drift as a convex combination of $v_q$ and $v_{ heta_0}$ and supports annealing to gradually shift from preservation to adaptation. The framework remains valid under arbitrary source-target couplings, including OT, and can be implemented with nonlinear fine-tuning methods like full fine-tuning or LoRA. Empirically, GFT improves convergence stability and shortens probability paths, enhancing inference efficiency while maintaining generation quality, across cross-domain and in-domain shifts on diverse datasets. Overall, GFT provides a theoretically principled, coupling-friendly, and practically effective pathway for scalable adaptation of flow-based generative models.

Abstract

Fine-tuning flow matching models is a central challenge in settings with limited data, evolving distributions, or strict efficiency demands, where unconstrained fine-tuning can erode the accuracy and efficiency gains learned during pretraining. Prior work has produced theoretical guarantees and empirical advances for reward-based fine-tuning formulations, but these methods often impose restrictions on permissible drift structure or training techniques. In this work, we propose Gradual Fine-Tuning (GFT), a principled framework for fine-tuning flow-based generative models when samples from the target distribution are available. For stochastic flows, GFT defines a temperature-controlled sequence of intermediate objectives that smoothly interpolate between the pretrained and target drifts, approaching the true target as the temperature approaches zero. We prove convergence results for both marginal and conditional GFT objectives, enabling the use of suitable (e.g., optimal transport) couplings during GFT while preserving correctness. Empirically, GFT improves convergence stability and shortens probability paths, resulting in faster inference, while maintaining generation quality comparable to standard fine-tuning. Our results position GFT as a theoretically grounded and practically effective alternative for scalable adaptation of flow matching models under distribution shift.

Figures (10)

• Figure 1: Comparison of fine-tuning methods for in-domain adaptation on Camelyon17. The shaded region of the path length graph (b) represents one standard deviation from the mean. Further results shown in Appendix \ref{['ap:small_results']}.
• Figure 2: Generated images of the Camelyon17 dataset with varying fine-tuning methods. (a) is trained from random initialization, (b, c) show cross-domain adaptation results, and (d, e) show in-domain adaptation results. Further generates images shown in Appendix \ref{['ap:images']}.
• Figure 3: Cross-domain adaptation results, as opposed to training from random initialization. Fine-tuning and pretraining are done on the Camelyon17 train dataset. The shaded regions of the path length graphs (b, c) represent one standard deviation from the mean.
• Figure 4: Comparison of fine-tuning methods for cross-domain adaptation on Camelyon17. The shaded region of the path length graph (b) represents one standard deviation from the mean. Further results shown in Appendix \ref{['ap:large_results']}.
• Figure 5: Cross-domain adaptation on RxRx1, showing the regularization effect of varying cooling schedules. Every experiment using the GFT objective has identical hyperparameters, apart from the minimum $\beta$ value reached by the cooling schedule. The shaded region of the path length graph (a) represents one standard deviation from the mean.

Theorems & Definitions (2)

• Theorem 4.1
• Proposition 4.2