Résumé technique : Couplages de suppression de divergence pour le flux rectifié

Énoncé du problème

Le flux rectifié (RectFlow) vise à accélérer la modélisation générative en apprenant des champs de vitesse à trajectoires rectilignes, permettant une génération d'échantillons de haute qualité avec peu, voire un seul, pas d'intégration numérique (NFE). La méthode repose sur une procédure itérative de « reflu » : un modèle de base entraîné avec un couplage indépendant génère de nouvelles paires source-cible, qui sont ensuite utilisées pour réentraîner un modèle doté d'un champ de vitesse plus rectiligne.

Cependant, l'article identifie un goulot d'étranglement fondamental : la distorsion du couplage. En pratique, les trajectoires générées par le modèle de base se courbent et s'entremêlent souvent. Cette distorsion est fréquemment associée à des régions de divergence non nulle dans le champ de vitesse appris. Dans ces régions, l'expansion ou la contraction locale provoque l'écartement ou la compression des trajectoires, détournant les particules de leurs points d'arrivée idéaux. Par conséquent, le modèle réentraîné hérite de ces erreurs, et des rounds supplémentaires de rectification produisent des rendements décroissants car le modèle effectue essentiellement un « reflu » d'un transport déjà courbé. La qualité de l'intégrateur vers l'avant utilisé pour générer les paires de couplage, plutôt que la fonction de perte d'entraînement ou la capacité du modèle, limite l'efficacité de RectFlow.

Méthodologie : DS-RectFlow

Les auteurs proposent DS-RectFlow (Flux Rectifié de Suppression de Divergence), une méthode de correction hors ligne conçue pour générer des paires de couplage plus propres durant le processus de reflu, sans modifier l'objectif d'entraînement, l'architecture du modèle ou la procédure d'inférence.

Mécanisme central

La méthode est fondée sur l'observation que la courbure des trajectoires (distorsion du couplage) est pilotée par la composante compressible du champ de vitesse. En utilisant la décomposition de Helmholtz, le champ de vitesse $v$ est divisé en une partie de transport sans divergence $u$ et une partie dipolaire $\nabla \phi$ qui porte toute la compressibilité du champ. L'article soutient que la suppression de l'amplitude de la divergence $|\nabla \cdot v|$ maintient les trajectoires rectilignes.

L'intégrateur de suppression de divergence

Lors de la génération hors ligne des paires de couplage (étape de « reflu »), DS-RectFlow remplace l'intégrateur d'Euler standard par un intégrateur d'Euler de suppression de divergence :

1. Correction d'état : À chaque pas d'intégration, au lieu de progresser directement depuis l'état courant $x_t$, l'algorithme recherche un état voisin $x^*_t$ où la divergence estimée $|\hat{\nabla} \cdot v_\theta(x^*_t, t)|$ est minimisée.
2. Estimation de Hutchinson : Pour éviter le coût élevé du calcul exact de la divergence en hautes dimensions, la méthode utilise l'estimateur de trace de Hutchinson pour approximer $\nabla \cdot v$. Elle génère $m$ perturbations gaussiennes autour de l'état courant et sélectionne le candidat minimisant l'amplitude estimée de la divergence.
3. Intégration : L'étape d'Euler est ensuite calculée à partir de cet état corrigé $x^*_t$ en utilisant le champ de vitesse original, non modifié : $x_{t+\Delta t} = x^*_t + v_\theta(x^*_t, t) \cdot \Delta t$.

Détails de mise en œuvre

• Hors ligne uniquement : La correction est appliquée uniquement lors de la génération de l'ensemble de données de couplage pour le réentraînement.
• Inférence : Le modèle final entraîné est déployé en utilisant une intégration d'Euler standard sans surcharge computationnelle supplémentaire. La logique de suppression de divergence n'est pas présente au moment de l'inférence.
• Contrôle du budget : La correction est appliquée uniquement jusqu'à un temps $t_{stop}$ (généralement 0,5) et est régie par un paramètre de rayon de recherche $\delta$.

Contributions clés

1. Diagnostic des limites de RectFlow : L'article identifie que la qualité du couplage généré par le modèle de base, spécifiquement la présence d'une divergence non nulle causant une compression/expansion des trajectoires, est le principal goulot d'étranglement empêchant RectFlow d'atteindre son potentiel théorique de génération en un seul pas.
2. Stratégie de correction hors ligne : Il introduit une nouvelle correction hors ligne qui atténue la composante divergente du champ de vitesse durant la génération du couplage. Contrairement aux raffinements au moment de l'inférence (par exemple, solveurs d'ordre supérieur ou guidage), cette méthode déplace le coût de la correction entièrement vers la phase de génération des données d'entraînement.
3. Surcharge d'inférence nulle : La méthode produit un modèle qui s'exécute à un coût horaire identique à celui de RectFlow standard lors du déploiement, car le chemin d'inférence reste un simple pas d'Euler.
4. Restauration des effets cumulatifs : En fournissant des paires de couplage plus propres, DS-RectFlow restaure le bénéfice « cumulatif » de la rectification, où chaque round de reflu améliore significativement la rectilignité du champ, un phénomène souvent perdu dans RectFlow standard en raison de la distorsion accumulée.

Résultats expérimentaux

Les auteurs évaluent DS-RectFlow sur des benchmarks synthétiques 2D et des tâches de génération d'images (CIFAR-10, CelebA-64).

• Benchmarks 2D (Damier & GMM) :

  • Sur le benchmark Damier, DS-RectFlow réduit significativement la distance de Wasserstein tranchée (SWD) par rapport à RectFlow standard. À $NFE=1$, il atteint une SWD de 0,102 (contre 0,165 pour la version standard), égalant efficacement les performances d'une base de référence de Matching de Flux à 20 pas, mais en un seul pas.
  • Sur le benchmark de croisement GMM, DS-RectFlow montre initialement une erreur plus élevée à faible NFE (en raison de la capture plus fidèle de décisions complexes d'attribution de modes) mais converge rapidement vers des performances supérieures. Dès le deuxième round de reflu ($k=2$), il réalise une amélioration de 3x de la SWD à $NFE=1$ par rapport à RectFlow standard.
  • Les visualisations confirment que DS-RectFlow produit des modes plus nets et mieux séparés, avec moins d'échantillons « fuites » dans des régions interdites.

• Génération d'images (CIFAR-10 & CelebA-64) :

  • CIFAR-10 : DS-RectFlow atteint systématiquement une distance Fréchet Inception (FID) plus faible sur tous les budgets NFE. Avec un solveur hors ligne d'Euler-20, la méthode réduit le FID de 17,41 à 12,03 à $NFE=1$ (une amélioration relative de 31 %). Notamment, le modèle DS-RectFlow à $NFE=1$ surpasse le modèle standard à tous les pas NFE de l'exploration.
  • CelebA-64 : Des tendances similaires sont observées. À $NFE=2$, DS-RectFlow atteint un FID de 4,90, surpassant la meilleure performance du modèle standard sur tous les NFE (qui est de 5,64 à $NFE=20$).
  • Les résultats démontrent que la correction hors ligne se transfère directement dans un modèle capable de génération de haute qualité en un seul pas.

Importance et revendications

L'article revendique que la qualité du couplage, et non la capacité d'entraînement ou le budget d'inférence, est le principal goulot d'étranglement de RectFlow. En traitant la distorsion induite par la divergence hors ligne, DS-RectFlow permet les améliorations cumulatives que la rectification promet mais délivre rarement en pratique.

Les auteurs soulignent que cette approche offre un avantage pratique : elle apporte les bénéfices d'un intégrateur corrigé de la divergence (trajectoires plus rectilignes, meilleur alignement des modes) tout en maintenant le profil d'inférence à latence nulle de RectFlow standard. La correction est un coût unique payé lors de la génération des données, amorti sur l'ensemble de la session d'entraînement. L'article suggère que ce mécanisme pourrait être applicable à d'autres méthodes basées sur le flux, telles que les interpolants stochastiques et le matching de flux de transport optimal, partout où des composantes de dérive divergentes dégradent la qualité des échantillons.