Technical Summary: Gradual Fine-Tuning for Flow Matching Models

Problem Statement

Fine-tuning pretrained flow matching models presents a significant challenge in scenarios characterized by limited data, evolving distributions, or computational constraints. While recent advances in reward-based fine-tuning (e.g., Adjoint Matching, reinforcement learning approaches) have shown promise, they often suffer from theoretical limitations and practical instabilities. Specifically, existing methods frequently fail to simultaneously guarantee theoretical correctness, training stability, inference efficiency, and diversity preservation.

A critical, yet underexplored setting is fine-tuning when direct samples from the target distribution are available (e.g., in domain adaptation or dataset drift), removing the need for external reward functions. Standard fine-tuning in this context often treats the pretrained model merely as an initialization, ignoring its learned dynamics (such as short probability paths), which leads to high instability and potential mode collapse. There is a need for a principled framework that balances adapting to a new target distribution while preserving the beneficial properties of the pretrained model.

Methodology: Gradual Fine-Tuning (GFT)

The authors propose Gradual Fine-Tuning (GFT), a framework that utilizes annealed regularization to create a temperature-controlled sequence of intermediate objectives. This approach smoothly interpolates between the pretrained dynamics and the target dynamics.

Theoretical Formulation

GFT is derived from an optimization problem over the path space induced by stochastic flows. The objective minimizes the Kullback-Leibler (KL) divergence between the fine-tuned process $P_\theta$ and the target process $P_q$, regularized by the KL divergence between $P_\theta$ and the pretrained process $P_{\theta_0}$:

$$\min_{\theta} \text{KL}(P_\theta \| P_q) + \beta \text{KL}(P_\theta \| P_{\theta_0})$$

where $\beta \in \mathbb{R}^+$ is a temperature parameter.

Theorem 4.1 establishes that minimizing this objective yields a closed-form optimal drift $v^*_\theta$ which is a convex combination of the target vector field $v_q$ and the pretrained vector field $v_{\theta_0}$:

$$v^*_\theta(X_t, t) = \left( \frac{1}{1 + \beta} \right) v_q(X_t, t) + \left( \frac{\beta}{1 + \beta} \right) v_{\theta_0}(X_t, t)$$

Annealed Regularization: To achieve convergence, GFT employs a time-dependent cooling schedule where $\beta$ is gradually annealed from a high initial value to $0^+$.

• High $\beta$: Strongly penalizes deviation from the pretrained model, preventing premature collapse and ensuring stability.
• Low $\beta$ ($\to 0^+$): Relaxes the constraint, allowing the model to asymptotically recover the target vector field $v_q$.

Tractable Reformulation

Since the target vector field $v_q$ is unknown and the objective is defined over path measures, the authors apply Girsanov's Theorem to reformulate the objective in terms of drift terms. This results in a tractable loss function involving the squared $L_2$ distance between drifts.

Furthermore, to enable efficient training and inference, GFT utilizes Conditional Flow Matching (CFM) principles. By conditioning on arbitrary source-target couplings (e.g., Optimal Transport plans), the loss becomes:

$$\mathcal{L}_\pi(\theta) = \mathbb{E}_{Z \sim \pi, X_t \sim P(\cdot|Z), t \sim U[0,1]} \left[ \frac{1}{2} \left( \|v_\theta - v_q(\cdot|Z)\|^2 + \beta \|v_\theta - v_{\theta_0}\|^2 \right) \right]$$

This formulation preserves the theoretical guarantees of the unconditional objective while allowing for the use of efficient couplings (like OT) to accelerate generation.

Extension to Few-Step Inference

The framework naturally extends to Mean Flow (MF) models, which enable few-step inference. The authors prove that the optimal average velocity field for MF models also follows a convex combination form, allowing GFT to be applied to one-step or few-step generation without sacrificing theoretical correctness.

Key Contributions

The paper outlines four primary contributions:

1. Theoretical Convergence: Proof that GFT converges to the target distribution under standard assumptions, with an explicit convergence rate dependent on the temperature $\beta$.
2. Coupling Equivalence: Establishment that GFT yields equivalent results when conditioning on arbitrary source-target couplings, enabling the use of training techniques (like OT) that enhance inference efficiency.
3. Few-Step Extension: Demonstration of the natural extension of GFT to fine-tuning for few-step inference methods (e.g., Mean Flow) without compromising theoretical guarantees.
4. Empirical Superiority: Comprehensive experimental validation showing GFT achieves comparable or superior generation quality to other methods while providing improved training stability, generation diversity, and consistently shorter probability paths.

Experimental Results

The authors evaluated GFT across multiple settings, including cross-domain and in-domain adaptation on WILDS benchmark datasets (Camelyon17, RxRx1, FMoW) and 2D simulated data.

• Stability and Convergence: GFT demonstrated significantly higher stability compared to standard CFM fine-tuning and reward-based methods (e.g., Adjoint Matching). Metrics such as instantaneous variance and Spearman correlation between epochs and FID indicated monotonic, stable convergence.
• Generation Quality: In cross-domain adaptation, GFT achieved FID scores comparable to standard fine-tuning but with lower average path lengths, indicating more efficient inference. In in-domain adaptation, GFT matched the accuracy of full fine-tuning while maintaining lower path lengths.
• Diversity: On 2D bimodal Gaussian targets, GFT and standard CFM successfully covered both modes with high diversity. In contrast, reward-based methods like Adjoint Matching suffered from mode collapse, generating samples tightly clustered within a single mode.
• Efficiency: GFT retained the simulation-free training efficiency of CFM. While Adjoint Matching required expensive ODE simulations for each step (resulting in training times an order of magnitude higher), GFT remained efficient.
• Few-Step Inference: When applied to Mean Flow models, GFT achieved comparable FID to iterative models but with a 50-100x reduction in model evaluations (NFEs) during inference.
• Regularization Control: Experiments varying the minimum temperature $\beta$ revealed a controllable trade-off between final FID and average path length, allowing users to explicitly tune the balance between generation accuracy and inference speed.

Significance and Claims

The paper positions GFT as a simple yet theoretically grounded and practically effective alternative for the scalable adaptation of flow matching models under distribution shift.

The authors claim that GFT addresses critical gaps in existing literature by satisfying all key desiderata for an ideal fine-tuning method:

• Theoretically Correct: It provably converges to the correct target distribution with a controllable rate.
• Stable: It exhibits highly monotonic convergence, avoiding the oscillations and instability common in standard fine-tuning.
• Efficient: It maintains the simulation-free training of CFM and supports efficient inference via OT couplings and few-step methods.
• Diverse: It preserves the diversity of the pretrained model, avoiding the mode collapse seen in reward-based approaches.

By bridging the gap between high-quality pretraining and fine-tuning, GFT offers a robust solution for adapting generative models to complex, data-driven domains where nuanced control over the fine-tuning outcome is required. The work establishes a foundation for the scalable adaptation of flow-based generative models.