Résumé Technique : Fine-Tuning Graduel pour les Modèles de Flow Matching

Énoncé du Problème

Le fine-tuning de modèles de flow matching pré-entraînés présente un défi important dans des scénarios caractérisés par des données limitées, des distributions évolutives ou des contraintes computationnelles. Bien que les avancées récentes dans le fine-tuning basé sur des récompenses (par exemple, l'Adjoint Matching, les approches par apprentissage par renforcement) aient montré des promesses, elles souffrent souvent de limitations théoriques et d'instabilités pratiques. Plus précisément, les méthodes existantes échouent fréquemment à garantir simultanément la correction théorique, la stabilité de l'entraînement, l'efficacité de l'inférence et la préservation de la diversité.

Un contexte critique, bien qu'encore peu exploré, est le fine-tuning lorsque des échantillons directs de la distribution cible sont disponibles (par exemple, dans l'adaptation de domaine ou le dérive de données), éliminant ainsi le besoin de fonctions de récompense externes. Le fine-tuning standard dans ce contexte traite souvent le modèle pré-entraîné comme une simple initialisation, ignorant ses dynamiques apprises (telles que les chemins de probabilité courts), ce qui conduit à une instabilité élevée et à un potentiel effondrement de mode (mode collapse). Il existe un besoin pour un cadre principled qui équilibre l'adaptation à une nouvelle distribution cible tout en préservant les propriétés bénéfiques du modèle pré-entraîné.

Méthodologie : Fine-Tuning Graduel (GFT)

Les auteurs proposent le Fine-Tuning Graduel (GFT), un cadre qui utilise une régularisation par recuit pour créer une séquence d'objectifs intermédiaires contrôlés par la température. Cette approche interpole de manière fluide entre les dynamiques pré-entraînées et les dynamiques cibles.

Formulation Théorique

Le GFT est dérivé d'un problème d'optimisation sur l'espace des chemins induit par les flux stochastiques. L'objectif minimise la divergence de Kullback-Leibler (KL) entre le processus fine-tuné $P_\theta$ et le processus cible $P_q$, régularisée par la divergence KL entre $P_\theta$ et le processus pré-entraîné $P_{\theta_0}$ :

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

où $\beta \in \mathbb{R}^+$ est un paramètre de température.

Le Théorème 4.1 établit que la minimisation de cet objectif produit un champ de dérive optimal en forme fermée $v^*_\theta$ qui est une combinaison convexe du champ de vecteurs cible $v_q$ et du champ de vecteurs pré-entraîné $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)$$

Régularisation par Recuit (Annealed Regularization) : Pour atteindre la convergence, le GFT emploie un programme de refroidissement dépendant du temps où $\beta$ est progressivement recuit d'une valeur initiale élevée vers $0^+$.

• $\beta$ élevé : Pénalise fortement l'écart par rapport au modèle pré-entraîné, empêchant l'effondrement prématuré et assurant la stabilité.
• $\beta$ faible ($\to 0^+$) : Relâche la contrainte, permettant au modèle de récupérer asymptotiquement le champ de vecteurs cible $v_q$.

Reformulation Traitable

Puisque le champ de vecteurs cible $v_q$ est inconnu et que l'objectif est défini sur les mesures de chemin, les auteurs appliquent le Théorème de Girsanov pour reformuler l'objectif en termes de termes de dérive. Cela aboutit à une fonction de perte traitable impliquant la distance $L_2$ au carré entre les dérives.

De plus, pour permettre un entraînement et une inférence efficaces, le GFT utilise les principes du Conditional Flow Matching (CFM). En se conditionnant sur des couplages source-cible arbitraires (par exemple, des plans de Transport Optimal), la perte devient :

$$\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]$$

Cette formulation préserve les garanties théoriques de l'objectif non conditionnel tout en permettant l'utilisation de couplages efficaces (comme l'OT) pour accélérer la génération.

Extension à l'Inférence à Peu d'Étapes

Le cadre s'étend naturellement aux modèles de Mean Flow (MF), qui permettent une inférence en peu d'étapes. Les auteurs prouvent que le champ de vitesse moyenne optimal pour les modèles MF suit également une forme de combinaison convexe, permettant d'appliquer le GFT à la génération en une ou quelques étapes sans sacrifier la correction théorique.

Principales Contributions

Le papier expose quatre contributions principales :

1. Convergence Théorique : Preuve que le GFT converge vers la distribution cible sous des hypothèses standards, avec un taux de convergence explicite dépendant de la température $\beta$.
2. Équivalence de Couplage : Établissement que le GFT produit des résultats équivalents lorsqu'on se conditionne sur des couplages source-cible arbitraires, permettant l'utilisation de techniques d'entraînement (comme l'OT) qui améliorent l'efficacité de l'inférence.
3. Extension à Peu d'Étapes : Démonstration de l'extension naturelle du GFT au fine-tuning pour les méthodes d'inférence à peu d'étapes (par exemple, Mean Flow) sans compromettre les garanties théoriques.
4. Supériorité Empirique : Validation expérimentale complète montrant que le GFT atteint une qualité de génération comparable ou supérieure aux autres méthodes tout en offrant une meilleure stabilité d'entraînement, une plus grande diversité de génération et des chemins de probabilité systématiquement plus courts.

Résultats Expérimentaux

Les auteurs ont évalué le GFT à travers plusieurs configurations, incluant l'adaptation cross-domaine et in-domaine sur les jeux de données de référence WILDS (Camelyon17, RxRx1, FMoW) et des données simulées en 2D.

• Stabilité et Convergence : Le GFT a démontré une stabilité significativement plus élevée par rapport au fine-tuning CFM standard et aux méthodes basées sur la récompense (ex: Adjoint Matching). Des métriques telles que la variance instantanée et la corrélation de Spearman entre les époques et le FID indiquent une convergence monotone et stable.
• Qualité de Génération : Dans l'adaptation cross-domaine, le GFT a obtenu des scores FID comparables au fine-tuning standard mais avec des longueurs de chemin moyennes plus courtes, indiquant une inférence plus efficace. Dans l'adaptation in-domaine, le GFT a égalé la précision du fine-tuning complet tout en maintenant des longueurs de chemin plus faibles.
• Diversité : Sur des cibles gaussiennes bimodales en 2D, le GFT et le CFM standard ont réussi à couvrir les deux modes avec une grande diversité. En revanche, les méthodes basées sur la récompense comme l'Adjoint Matching ont souffert d'un effondrement de mode, générant des échantillons étroitement regroupés autour d'un seul mode.
• Efficacité : Le GFT a conservé l'efficacité d'entraînement sans simulation du CFM. Alors que l'Adjoint Matching nécessitait des simulations d'ODE coûteuses pour chaque étape (résultant en des temps d'entraînement un ordre de grandeur plus élevés), le GFT est resté efficace.
• Inférence à Peu d'Étapes : Appliqué aux modèles Mean Flow, le GFT a obtenu un FID comparable aux modèles itératifs mais avec une réduction de 50 à 100 fois du nombre d'évaluations de modèle (NFE) lors de l'inférence.
• Contrôle de la Régularisation : Les expériences faisant varier la température minimale $\beta$ ont révélé un compromis contrôlable entre le FID final et la longueur de chemin moyenne, permettant aux utilisateurs de régler explicitement l'équilibre entre l'exactitude de la génération et la vitesse d'inférence.

Signification et Revendications

Le papier positionne le GFT comme une alternative simple, mais théoriquement fondée et pratiquement efficace, pour l'adaptation scalable des modèles de flow matching sous dérive de distribution.

Les auteurs affirment que le GFT comble les lacunes critiques de la littérature existante en satisfaisant tous les desiderata clés d'une méthode de fine-tuning idéale :

• Théoriquement Correct : Il converge de manière prouvée vers la distribution cible avec un taux contrôlable.
• Stable : Il présente une convergence hautement monotone, évitant les oscillations et l'instabilité communes au fine-tuning standard.
• Efficace : Il maintient l'entraînement sans simulation du CFM et supporte une inférence efficace via les couplages OT et les méthodes à peu d'étapes.
• Diversifié : Il préserve la diversité du modèle pré-entraîné, évitant l'effondrement de mode observé dans les approches basées sur la récompense.

En comblant le fossé entre le pré-entraînement de haute qualité et le fine-tuning, le GFT offre une solution robuste pour adapter les modèles génératifs à des domaines complexes pilotés par les données, où un contrôle nuancé sur le résultat du fine-tuning est requis. Ce travail établit une fondation pour l'adaptation scalable des modèles génératifs basés sur les flux.