One-Sided Quantile Coupling for Flow Matching

Résumé Technique : Couplage par Quantiles Unilatéral pour le Flow Matching

Énoncé du Problème
Le Flow Matching (FM) entraîne des modèles génératifs en temps continu en effectuant une régression sur le champ de vitesse d'un chemin de probabilité reliant une distribution source simple (typiquement une gaussienne standard) à une distribution de données cible. La qualité du couplage — l'appariement des échantillons sources et cibles — influence de manière critique la géométrie des trajectoires d'entraînement. Le FM standard utilise un appariement aléatoire indépendant, ce qui entraîne souvent des chemins longs et conflictuels qui augmentent la variance du champ de vitesse cible et entravent l'efficacité de l'entraînement.

Bien que des couplages structurés inspirés du Transport Optimal (OT) puissent redresser ces chemins et améliorer la qualité de la génération, les méthodes existantes (par exemple, OT-CFM) reposent sur la résolution d'un problème d'assignation par mini-lot. Ces approches nécessitent la construction de matrices de coûts par paires et la résolution de problèmes d'assignation (via l'algorithme de Hongrois ou les itérations de Sinkhorn), ce qui entraîne des coûts computationnels qui évoluent au moins de manière quadratique (ou cubique pour les résolutions exactes) par rapport à la taille du lot. Cela limite la scalabilité et peut introduire de l'instabilité. L'article répond à la question de savoir si les bénéfices géométriques d'un couplage structuré peuvent être obtenus sans apparier deux lots pré-échantillonnés, évitant ainsi le coût élevé du transport par lot.

Méthodologie : Couplage par Quantiles pour le Flow Matching (QC-FM)
Les auteurs proposent le Quantile Coupling Flow Matching (QC-FM), un schéma de couplage unilatéral léger. Contra-rément aux méthodes traditionnelles qui échantillonnent à la fois un lot de bruit et un lot de données puis les ré-apparient, le QC-FM échantillonne uniquement le lot de données et construit directement les échantillons sources appariés.

1. Construction Unilatérale :

   • Étant donné un lot de données $\mathcal{B}_x$, la méthode projette les points de données sur $k$ directions orthogonales aléatoires (tranches) définies par une matrice $U$.
   • Pour chaque direction, les données sont classées au sein du lot. Ces rangs empiriques sont projetés sur des quantiles gausiens en utilisant la fonction de répartition inverse de la loi normale standard.
   • Le code latent source $e^{(i)}$ est construit pour satisfaire ces contraintes de quantiles projetés. Plus précisément, la projection de la source sur les directions de tranches correspond aux quantiles gausiens dérivés des rangs des données.
   • Les dimensions restantes (le complément orthogonal de l'espace de la tranche) sont complétées par l'échantillonnage d'une distribution gaussienne conditionnelle. Cela garantit que la source reste stochastique dans les directions non contraintes.

2. Extensions Hybrides Gaussiennes-Reste :
   Puisque le couplage QC complet (où tous les points sources sont construits via les rangs) s'écarte de la priorité gaussienne i.i.d. standard pour $k \ge 2$ en raison des dépendances entre les codes de tranches, les auteurs introduisent deux stratégies hybrides pour conserver les avantages du QC tout en préservant le prior :

   • QC-FM-Mixture : Applique la construction QC à un sous-ensemble « ancre » du lot (par exemple, 20 %) et apparie le reste des échantillons sources « restants » avec les échantillons de données restants en utilisant un appariement aléatoire indépendant (l'approche Baseline). Cela crée un mélange de couplages structurés et non structurés.
   • QC-FM-Adjacency : Applique le QC au sous-ensemble ancre mais apparie les échantillons sources restants en fonction de la géométrie locale induite par les ancres. Les points de données restants sont assignés à l'ancre la plus proche (dans l'espace structuré), et les échantillons sources restants sont appariés au sein de ces groupes. Cela propage le biais structuré sans dupliquer les cibles.

Contributions Clés

• Algorithme QC-FM : Un nouveau couplage local au lot qui impose un transport monotone unidimensionnel le long de tranches aléatoires. Il élimine le besoin de matrices de coûts par paires et de solveurs d'assignation, réduisant le coût computationnel dominant à $O(kB \log B)$ (tri) plutôt qu'à $O(B^2)$ ou $O(B^3)$.
• Analyse Théorique :
  • Réduction de la Variance : Les auteurs prouvent que pour un cadre fixe, le couplage QC élimine la variance de régression irréductible le long de chaque tranche sélectionnée, rendant le flux idéal exactement droit dans ces directions.
  • Préservation du Prior : Ils quantifient que pour $k=1$, la marge source est exactement gaussienne. Pour $k \ge 2$, l'écart par rapport au prior gaussien est confiné à la copule des codes de tranches, et le coût de transport de cet écart est borné.
• Schémas d'Entraînement Hybrides : L'introduction des variantes Mixture et Adjacency qui équilibrent l'injection de biais géométrique avec la préservation du signal gaussien explicite de la Baseline.

Résultats Expérimentaux
La méthode a été évaluée sur CIFAR-10, CelebA-64, FFHQ-64 et ImageNet-64 sous des budgets d'entraînement et des architectures équivalents.

• Performance : QC-FM-Mixture a systématiquement surpassé la Baseline de couplage indépendant et le solveur exact de l'OT-CFM sur les quatre jeux de données. Les améliorations par rapport à la Baseline vont de 6,3 % (CIFAR-10) à 12,9 % (FFHQ) en termes de FID-50K.
• Efficacité : En termes de coût computationnel, QC-FM est nettement plus rapide que les méthodes basées sur l'OT. À une taille de lot de 2048, QC-FM-Mixture était plus de 800 fois plus rapide que le solveur de Hongrois de l'OT et 35 fois plus rapide que le solveur de Sinkhorn.
• Diagnostics de Test (Toy) : Sur des jeux de données synthétiques (Checkerboard et 8-Gaussians), les variantes de QC-FM ont produit des trajectoires apprises presque droites (Ratio de Longueur de Chemin $\approx 1$) et une variance de la vitesse cible nettement plus faible que la Baseline, confirmant les bénéfices géométriques du couplage.

Signification et Revendications
L'article affirme que préserver la structure de rang projetée est un moyen simple, scalable et efficace d'injecter un biais géométrique utile dans les couplages de Flow Matching sans résoudre un problème de transport par lot.

• Scalabilité : QC-FM offre un juste milieu attractif entre la Baseline non structurée et les méthodes de transport par lot coûteuses. Il est computationnellement léger et contrôlable via le nombre de projections ($k$).
• Praticité : La méthode ne nécessite aucun composant appris pour l'inférence ; la génération commence toujours à partir de la gaussienne standard. Le couplage est construit à la volée pendant l'entraînement.
• Limitations : Les auteurs notent modestement que QC-FM est un substitut local au lot pour un couplage structuré, et non une approximation du transport optimal global. Les quantiles sont estimés au sein des mini-lots, et la théorie contrôle la variance par tranche plutôt que de garantir la droiture globale des trajectoires ou des améliorations spécifiques du FID. La sélection du nombre de tranches ($k$) et du ratio d'ancres ($p$) repose actuellement sur des heuristiques et un ajustement spécifique aux jeux de données.

En conclusion, ce travail démontre qu'un couplage par quantiles unilatéral peut améliorer la qualité des échantillons et l'efficacité de l'entraînement dans le Flow Matching en structurant l'appariement bruit-données par la préservation des rangs, offrant une alternative pratique aux méthodes de transport optimal computationnellement prohibitives.