Résumé Technique : Modélisation Générative en Une Étape via des Flots de Gradient de Wasserstein

Énoncé du Problème

Les modèles génératifs modernes, en particulier les modèles de diffusion et les méthodes basées sur les flots, ont atteint une qualité d'échantillonnage de pointe mais souffrent de coûts computationnels élevés et de latence en raison de leur dépendance à un échantillonnage itératif multi-étapes. Bien que des méthodes de distillation en une ou quelques étapes existent, elles héritent souvent de signaux d'entraînement issus de trajectoires initialement conçues pour une génération itérative ou reposent sur des dynamiques heuristiques (par exemple, les Modèles de Dérive) qui manquent de garanties de convergence fondées sur des principes et sont sujettes à l'effondrement de modes. Le défi central consiste à concevoir une dynamique d'entraînement permettant à un générateur d'apprendre une cartographie directe et en une seule étape d'une distribution de référence vers une distribution de données cible, tout en assurant une convergence stable et une couverture globale des modes.

Méthodologie : W-Flow

Les auteurs proposent W-Flow, un cadre qui découple la conception des dynamiques d'entraînement de la paramétrisation du générateur. L'idée centrale consiste à prescrire l'évolution de la distribution du modèle vers la distribution cible en utilisant des Flots de Gradient de Wasserstein (WGF), puis à entraîner un réseau de neurones statique pour compresser cette évolution multi-étapes en une seule étape d'inférence.

1. Dynamiques d'entraînement via le Flot de Gradient de Wasserstein

Au lieu d'apprendre directement une séquence d'étapes de débruitage ou des champs de vitesse, W-Flow définit la trajectoire d'entraînement comme une évolution en temps continu dans l'espace des mesures de probabilité.

• Fonctionnel d'Énergie : L'évolution est pilotée par la descente la plus raide d'un fonctionnel d'énergie $F(q) = D(q \| p)$, où $q$ est la distribution actuelle du modèle et $p$ la distribution de données cible.
• Champ de Vitesse : Le champ de vitesse en temps continu $V_t$ est dérivé de la première variation du fonctionnel d'énergie : $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$.
• Discrétisation : Le flot continu est discrétisé à l'aide d'un schéma d'Euler explicite. À chaque étape d'entraînement $k$, la distribution du modèle $q^{(k)}$ est mise à jour via une application de transport locale $g^{(k)}(x) = x + \eta V^{(k)}(x)$, où $\eta$ est la taille de l'étape.

2. Instantiation avec la Divergence de Sinkhorn

Pour garantir que le champ de vitesse soit calculable et géométriquement bien comporté en haute dimension, les auteurs instancient le fonctionnel d'énergie en utilisant la divergence de Sinkhorn ($S_\epsilon$), une métrique de transport optimal (OT) régularisée par entropie.

• Formulation du Champ de Vitesse : Le champ de vitesse induit est donné par la différence entre deux projections barycentriques :
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  Ici, $T^\epsilon_{q,p}$ transporte la masse de la distribution générée $q$ vers la cible $p$, tandis que $T^\epsilon_{q,q}$ représente un terme de self-transport qui impose la cohérence avec la distribution actuelle.
• Coordination Globale : Contrairement aux interactions locales par noyau (par exemple, dans MMD ou les Modèles de Dérive), le champ de vitesse de Sinkhorn est dérivé d'un plan de transport optimal global, garantissant que les mouvements des particules sont coordonnés à travers l'ensemble du lot pour prévenir l'effondrement de modes.
• Estimateur à Deux Lots : Pour éviter les artefacts d'auto-appariement où les particules s'apparient à elles-mêmes avec un coût nul, les auteurs emploient une stratégie à deux lots. Ils calculent le terme de self-transport $T^\epsilon_{q,q}$ en utilisant un deuxième lot indépendant d'échantillons issus de la distribution actuelle, plutôt que de masquer la diagonale de la matrice de coût d'un seul lot.

3. Entraînement du Générateur

Le générateur $f_\theta$ est entraîné pour imiter les mises à jour discrètes des particules définies par le WGF.

• Objectif : La fonction de perte minimise la distance au carré entre la sortie du générateur et la position de particule « cible » dérivée de la mise à jour de Sinkhorn :
  $$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$$
  où $\text{sg}(\cdot)$ désigne l'opérateur stop-gradient, traitant la position de particule mise à jour comme une cible de régression fixe.
• Guidage sans Classifieur (CFG) : Le cadre prend en charge le CFG en modifiant directement le champ de vitesse. Au lieu de mélanger linéairement les distributions, les auteurs proposent un guidage de vitesse, qui ajoute un terme proportionnel à la différence entre les champs de vitesse conditionnels et inconditionnels. Cette approche s'avère récupérer la distribution cible exponentiellement inclinée standard lors de l'utilisation de la divergence KL et fournit une extension géométriquement cohérente pour la divergence de Sinkhorn.

Contributions Clés

1. Cadre Principiel en Une Étape : W-Flow établit un nouveau paradigme pour la modélisation générative en une étape en ancrant les dynamiques d'entraînement dans les flots de gradient de Wasserstein, offrant une interprétation variationnelle claire absente des approches heuristiques comme les Modèles de Dérive.
2. Dynamiques Basées sur Sinkhorn : L'instanciation du WGF avec la divergence de Sinkhorn fournit une règle de mise à jour efficace basée sur le transport optimal, qui capture les écarts distributionnels globaux et améliore la couverture des modes par rapport aux méthodes basées sur des noyaux ou des interactions locales.
3. Convergence Théorique : Les auteurs prouvent que les dynamiques de particules à échantillon fini de W-Flow convergent vers les dynamiques distributionnelles en temps continu sous des hypothèses de régularité appropriées, fournissant une justification théorique pour le processus d'apprentissage.
4. Estimateur à Deux Lots : Une nouvelle stratégie à deux lots pour estimer les termes de self-transport est introduite, éliminant le biais d'auto-appariement inhérent aux estimateurs à lot unique et aux heuristiques de masquage de diagonale.

Résultats Expérimentaux

Les auteurs évaluent W-Flow sur ImageNet 256×256 et FFHQ.

• Performance de Pointe : W-Flow atteint un nouveau FID de pointe de 1,29 (échelle XL) et 1,35 (échelle L) pour la génération conditionnelle par classe en une seule étape sur ImageNet, surpassant nettement les méthodes en une étape précédentes (par exemple, Modèles de Dérive, MeanFlow) et dépassant de nombreux modèles de diffusion multi-étapes (par exemple, LightningDiT, SiT) malgré l'utilisation d'une seule évaluation de fonction (NFE).
• Efficacité d'Échantillonnage : Comparé aux modèles de diffusion multi-étapes ayant des scores FID similaires, W-Flow produit un échantillonnage environ 100 fois plus rapide (par exemple, 77,88 images/seconde contre 0,93 image/seconde pour SiT-XL/2 sur un seul GPU H100).
• Couverture des Modes : Dans des expériences avec des distributions déséquilibrées (par exemple, 95 % de visages de personnes âgées, 5 % de visages d'enfants), W-Flow capture avec succès les modes minoritaires, alors que les Modèles de Dérive présentent un effondrement de modes, ignorant entièrement la classe minoritaire.
• Vitesse de Convergence : W-Flow converge substantiellement plus vite pendant l'entraînement, atteignant la performance d'un Modèle de Dérive entraîné pendant 1280 époques en seulement 384 époques.
• Transfert de Domaine : La méthode démontre une forte capacité dans les tâches de transfert de domaine (par exemple, la traduction d'âge sur FFHQ), produisant des trajectoires de transport plus courtes et plus cohérentes par rapport aux bases de référence heuristiques.

Importance et Revendications

L'article revendique que les flots de gradient de Wasserstein fournissent une fondation principielle et efficace pour une modélisation générative rapide et de haute fidélité. En séparant la conception des dynamiques d'entraînement de la paramétrisation du générateur, W-Flow démontre qu'un générateur en une seule étape peut être entraîné pour atteindre des performances comparables à des processus itératifs multi-étapes complexes. Les auteurs soulignent que la coordination globale fournie par les dynamiques de transport optimal est la clé pour stabiliser l'entraînement et assurer une couverture complète des modes, répondant ainsi aux limitations critiques des approches en une étape existantes. Les résultats suggèrent que W-Flow peut servir d'alternative évolutive pour des applications nécessitant une génération de haute qualité à faible latence, telles que les outils créatifs en temps réel et la simulation.