Résumé Technique : Exploration et Exploitation de la Stabilité dans l'Appariement de Flux Latent

1. Énoncé du Problème

L'entraînement des modèles génératifs, en particulier les Modèles de Diffusion et leurs équivalents par Appariement de Flux (FM), reste coûteux en termes de calcul. Ces modèles nécessitent généralement des ensembles de données massifs, des ressources de calcul importantes, de longues durées d'entraînement et, dans les configurations conditionnelles, des annotations étendues. Bien que l'Appariement de Flux offre une formulation déterministe pour la génération d'échantillons via des Équations Différentielles Ordinaires (EDO) qui réduit souvent le nombre d'étapes d'échantillonnage par rapport aux Équations Différentielles Stochastiques (EDS), le coût de l'entraînement de ces modèles sur des ensembles de données à grande échelle demeure un goulot d'étranglement.

La question centrale abordée par ce travail est de savoir si la taille de l'ensemble de données et la capacité du modèle peuvent être considérablement réduites sans sacrifier la qualité de génération. Plus précisément, les auteurs examinent la stabilité des modèles d'Appariement de Flux Latent (LFM) face à des perturbations majeures, notamment le sous-échantillonnage de l'ensemble de données (élagage), les modifications architecturales et les configurations d'entraînement altérées. Ils postulent que si les trajectoires de transport LFM sont intrinsèquement stables, on pourrait exploiter cette propriété pour dériver des algorithmes d'entraînement et d'inférence plus efficaces.

2. Méthodologie

2.1. Fondement Théorique : Stabilité dans l'Appariement de Flux

Ce travail s'appuie sur l'observation que le Flux Rectifié admet une expression sous forme fermée pour le champ de vitesse. Pour un ensemble de données fini, le champ de vitesse optimal est une somme pondérée des échantillons d'entraînement, où les poids sont déterminés par une fonction softmax sur les distances au carré entre l'état actuel et les échantillons d'entraînement mis à l'échelle. Les auteurs notent que ce softmax devient "pointu" tôt dans la phase de transport, ce qui signifie qu'un seul échantillon d'entraînement domine souvent le champ de vitesse pour une trajectoire donnée.

Cette insight théorique suggère que la trajectoire de transport est largement déterminée par la présence d'échantillons spécifiques "dominants". Par conséquent, supprimer les échantillons non dominants (élagage) devrait laisser la majorité des trajectoires inchangées. Les auteurs formalisent cela en définissant une métrique d'affectation : pour une trajectoire commençant au bruit $x_0$, ils déterminent à quel échantillon d'entraînement $x_i$ elle aboutit. Ils émettent l'hypothèse que cette affectation reste invariante même lorsqu'une grande fraction de l'ensemble de données est supprimée, à condition que l'échantillon dominant pour cette trajectoire soit conservé.

2.2. Stratégies d'Élagage des Données

Pour exploiter cette stabilité, les auteurs proposent trois critères d'élagage de données adaptés des modèles discriminatifs mais reformulés pour l'objectif FM (qui utilise des chemins de bruit partagés et des pas de temps) :

1. Notation basée sur le Gradient (G) : Les échantillons sont classés par la magnitude de la norme du gradient de la perte.
2. Notation basée sur la Perte (L) : Les échantillons sont classés par la magnitude de la perte calculée le long des chemins de bruit partagés.
3. Notation basée sur le Clustering (C) : Les échantillons sont regroupés dans un espace d'embedding préentraîné (par exemple, CLIP). Les auteurs proposent deux variantes spécifiques :
   • Proportionnelle ($C_p$) : Sélectionne des échantillons proportionnellement à la taille du cluster.
   • Équilibrée ($C_b$) : Sélectionne un nombre égal d'échantillons de chaque cluster pour atténuer les déséquilibres distributionnels.
   • Basée sur le Noyau ($C_b^\kappa$) : Sélectionne des échantillons de manière gloutonne pour correspondre à la moyenne du noyau gaussien du cluster dans l'espace latent.

Les auteurs évaluent également l'"inverse" de ces critères (sélectionner les scores les plus bas) pour tester la robustesse de l'affirmation de stabilité.

2.3. Inférence du Grossier au Fin (C2F)

Exploitant la stabilité des trajectoires à travers les capacités des modèles, les auteurs proposent une stratégie d'inférence en deux étapes :

1. Modèle Grossier : Une architecture légère (par exemple, DiT-S/2) est entraînée sur un sous-ensemble élagué des données. Il gère la première phase de la trajectoire ($t \in [0, t_0)$).
2. Modèle Fin : Une architecture à haute capacité (par exemple, DiT-XL/2), préentraînée sur l'ensemble complet des données, gère la deuxième phase ($t \in [t_0, 1]$).

Pour assurer une transition fluide au point de "couture" $t_0$, le modèle Grossier est affiné avec une fonction de perte de continuité. Cette perte minimise la différence entre les champs de vitesse des modèles Grossier et Fin à $t_0$. La trajectoire du modèle Fin est utilisée pour générer l'état cible $x_{t_0}$ par inversion d'EDO (intégration vers l'arrière), garantissant ainsi que le modèle Grossier apprend à s'aligner sur le chemin du modèle Fin.

3. Contributions Clés

1. Démonstration Empirique de la Stabilité LFM : Les auteurs montrent que des modèles LFM entraînés indépendamment convergent vers des trajectoires de transport hautement cohérentes sous le sous-échantillonnage de l'ensemble de données, les changements architecturaux et différentes initialisations aléatoires. Cette stabilité est liée aux propriétés théoriques de l'objectif FM.
2. Algorithmes d'Élagage Pratiques : Ils introduisent trois critères d'élagage (Gradient, Perte, Clustering) adaptés au FM. Ils démontrent qu'un élagage agressif (par exemple, supprimer 50-95 % des données) ne dégrade pas les performances et peut même améliorer la vitesse de convergence et la qualité de génération dans des budgets de calcul limités.
3. Accélération de l'Inférence par Approche du Grossier au Fin : Ils proposent une stratégie de génération à deux modèles qui divise la trajectoire. Cette approche permet une accélération significative de l'inférence (environ 2,15$\times$) tout en maintenant la qualité de génération, à condition que l'hypothèse de stabilité soit valable à travers les architectures de modèles.
4. Atténuation des Biais par Élagage Équilibré : Les auteurs montrent que l'élagage par clustering équilibré ($C_b$) peut atténuer les biais dans les données générées (par exemple, genre, teinte de peau) sans nécessiter d'équilibrage explicite sensible aux étiquettes, en égalisant la représentation des clusters sémantiques dans l'ensemble d'entraînement.

4. Résultats Expérimentaux

Les méthodes ont été évaluées sur CelebA-HQ, FFHQ et ImageNet-1K.

• Stabilité sous Élagage : Sur CelebA-HQ, des modèles entraînés sur des sous-ensembles disjoints (élagués à 50 %) ont produit des images visuellement similaires lorsqu'ils partaient du même germe de bruit $x_0$. La similarité cosinus ArcFace entre les modèles élagués et non élagués variait de 0,79 à 0,83, nettement supérieure aux paires aléatoires (0,37).
• Métriques de Performance :
  • CelebA-HQ : L'élagage de 50 % des données en utilisant le clustering équilibré ($C_b$) a légèrement amélioré la Distance Fréchet d'Inception (FID), passant de 24,24 (non élagué) à 22,80.
  • ImageNet : L'élagage de 75 % des données ($pr=0,75$) a entraîné une convergence plus rapide et a maintenu ou légèrement amélioré la FID par rapport à l'ensemble complet des données. Même avec un élagage de 95 %, le modèle a convergé plus rapidement initialement, bien que les performances se soient dégradées après un certain nombre d'itérations.
• Robustesse aux Perturbations :
  • Architecture : L'augmentation de la capacité du modèle (de DiT-S/2 à DiT-XL/2) a entraîné des changements de trajectoire minimes (similarité $s=0,81$). Le passage à une architecture U-Net a réduit la similarité ($s=0,55$) mais a préservé les attributs grossiers.
  • Espace Latent : La stabilité est sensible aux changements de paramétrisation de l'espace latent (par exemple, changer le germe du VAE ou appliquer des transformations inversibles), ce qui a rompu la stabilité.
  • Objectif : Le passage de l'Appariement de Flux à la diffusion basée sur le score a complètement altéré les trajectoires, indiquant que la stabilité est spécifique à l'objectif FM.
• Efficacité du Grossier au Fin : Sur CelebA-HQ, l'approche C2F avec $t_0=0,7$ a réduit le temps d'inférence de 93,95 ms/image (Fin uniquement) à 43,53 ms/image (accélération d'environ 2,15$\times$) tout en maintenant la FID.
• Équité : L'élagage équilibré ($C_b$) a réduit la divergence KL vers une distribution uniforme sur les attributs (genre, âge, teinte de peau) plus efficacement que les modèles non élagués, malgré son agnosticisme aux étiquettes.

5. Signification et Revendications

Les auteurs affirment que la stabilité observée est une propriété caractéristique du paradigme de l'Appariement de Flux Rectifié sous des représentations latentes fixes, plutôt qu'un artefact d'algorithmes d'entraînement spécifiques.

• Efficacité : La signification principale réside dans la capacité à réduire le temps d'entraînement et les besoins en données sans sacrifier la qualité. Les auteurs démontrent que pour de grands ensembles de données comme ImageNet, l'élagage peut accélérer la convergence, et pour des ensembles de données plus petits comme CelebA-HQ, il peut améliorer la qualité de génération.
• Accélération de l'Inférence : L'approche C2F offre une manière fondée de réduire les coûts d'inférence en déchargeant les premières étapes grossières de la génération vers des modèles légers, en s'appuyant sur la stabilité de la trajectoire pour garantir que le modèle à haute capacité n'a besoin que d'affiner les détails finaux.
• Contrôle des Biais : Ce travail suggère que l'élagage de l'ensemble de données, spécifiquement via le clustering équilibré, peut servir de mécanisme pour contrôler et atténuer les déséquilibres distributionnels dans les données générées, offrant une voie vers des modèles génératifs plus équitables sans rééquilibrage complexe sensible aux étiquettes.

L'article conclut que si la stabilité est robuste aux changements de données et de capacité du modèle, elle est sensible aux changements de la structure de l'espace latent et de l'objectif génératif lui-même. Cette insight permet aux praticiens de concevoir des pipelines plus économes en ressources tout en maintenant une génération de haute fidélité.