Gradual Fine-Tuning for Flow Matching Models

TL;DR

This paper introduces Gradual Fine-Tuning (GFT), a theoretically grounded framework for effectively adapting flow matching models to new distributions, improving convergence stability and inference speed without sacrificing quality.

Contribution

GFT provides a novel, temperature-controlled interpolation method for fine-tuning flow models, with proven convergence and practical benefits over existing methods.

Findings

GFT enhances convergence stability during fine-tuning.

GFT reduces probability path length, speeding up inference.

Generation quality remains comparable to standard fine-tuning.

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.