Continuous Adversarial Flow Models

1. Problématique et Contexte

Les modèles de flux (Flow Matching - FM) ont connu un succès récent dans la génération d'images et de vidéos en apprenant un champ de vitesse pour transporter une distribution de bruit vers une distribution de données via une équation différentielle ordinaire (ODE). Cependant, une limitation critique subsiste, particulièrement dans les modalités visuelles :

• Génération hors distribution (OOD) : Sans guidance (conditionnement), les modèles FM ont tendance à générer des échantillons qui ne respectent pas fidèlement la distribution sous-jacente des données réelles.
• Cause racine : L'objectif d'entraînement standard des modèles FM repose sur une perte Moyenne des Erreurs Quadratiques (MSE), qui mesure une distance euclidienne isotrope. Théoriquement, cette perte converge vers le flux réel avec une capacité infinie, mais avec une capacité finie (réseaux de neurones), elle induit une généralisation incorrecte car elle ignore la structure de la variété (manifold) des données.
• Limites des solutions existantes :
  • Les Autoencodeurs de représentation nécessitent un espace latent.
  • Le Flow Matching Riemannien exige une définition manuelle de la variété.
  • Les pertes perceptuelles fixes peuvent être "hackées" par le générateur, créant des artefacts.
  • Les Adversarial Flow Models (AFM) existants fonctionnent en temps discret, ce qui rend l'entraînement instable lorsque l'intervalle de temps tend vers zéro et ne résout pas complètement le problème de la disparition du gradient.

L'objectif est donc de concevoir un modèle continu capable d'apprendre une généralisation plus fidèle à la distribution des données sans dépendre d'une guidance externe forte.

2. Méthodologie : Continuous Adversarial Flow Models (CAFM)

Les auteurs proposent les CAFM, une extension des modèles de flux adversariaux au domaine du temps continu.

A. Architecture et Concepts Clés

Le modèle CAFM est un Flux Normalisant Continu (CNF) qui intègre une ODE du bruit vers les données. Il se compose de deux réseaux :

1. Le Générateur (G) : Identique aux modèles FM, il prédit le champ de vitesse $v_t$ à l'état $x_t$.
2. Le Discriminateur (D) : Contrairement aux AFM discrets qui discriminent les échantillons $x_t$, le discriminateur CAFM opère dans l'espace des dérivées (espace des vitesses).

B. Discrimination via Produit Jacobien-Vecteur (JVP)

L'innovation centrale réside dans la manière dont le discriminateur évalue la "réalité" d'un flux. Au lieu de classifier $x_t$, le discriminateur $D(x_t, t)$ apprend un potentiel scalaire. La discrimination se fait via le Produit Jacobien-Vecteur (JVP) :

$$D_{\text{jvp}}(x_t, t, \dot{x}_t, \dot{t}) = \frac{\partial D(x_t,t)}{\partial x_t}\dot{x}_t + \frac{\partial D(x_t,t)}{\partial t}\dot{t}$$

• Entraînement : Le JVP est évalué en utilisant la vitesse conditionnelle réelle $\bar{v}_t$ comme vecteur tangent pour les données réelles, et la vitesse prédite $G(x_t, t)$ pour les données générées.
• Objectif Adversarial : Le discriminateur apprend à attribuer une valeur de potentiel plus élevée aux directions de flux réalistes. Le générateur est optimisé pour maximiser ce potentiel.
• Avantage mathématique : Cette formulation en espace dérivé permet d'éviter le problème de la disparition du gradient (vanishing gradient) même lorsque les distributions de support ne se chevauchent pas parfaitement, car le gradient par rapport au générateur dépend linéairement du Jacobien du discriminateur, qui reste non nul tant que $G \neq v_{\text{true}}$.

C. Objectifs de Perte

Les objectifs finaux pour le discriminateur ($D$) et le générateur ($G$) sont :

• $D$ : Minimise une fonction de contraste bornée (type Least Squares) entre les JVP des données réelles et générées, plus une pénalité de centrage pour éviter la dérive des logits.
• $G$ : Minimise la même fonction de contraste (pour tromper $D$) plus un terme de régularisation de transport optimal ($L_{ot}$) qui encourage la solution de norme minimale.
• Note importante : Pour le post-entraînement, le terme de transport optimal ($\lambda_{ot}$) est mis à zéro pour éliminer tout biais vers la norme euclidienne, ne laissant que l'objectif adversarial pour guider la généralisation.

D. Efficacité et Mise en Œuvre

• Post-entraînement : La méthode est principalement conçue pour affiner (fine-tune) des modèles FM pré-entraînés. Cela permet de bénéficier de la convergence rapide du FM tout en corrigeant la généralisation via l'adversarialité.
• Calcul : Le JVP est calculé efficacement via la différenciation automatique en mode direct (forward-mode), permettant de calculer la sortie du réseau et son JVP en un seul passage avant.
• Stabilité : L'utilisation de RMSNorm au lieu de LayerNorm dans le discriminateur s'est révélée cruciale pour la stabilité de l'entraînement.

3. Contributions Clés

1. Première application de l'entraînement adversarial en temps continu : Extension des AFM (discrets) vers un cadre continu, résolvant les problèmes d'instabilité liés à la limite $|t-s| \to 0$.
2. Discrimination dans l'espace des vitesses : Introduction du JVP comme mécanisme de discrimination, alignant l'objectif sur la dynamique du flux plutôt que sur la position statique.
3. Amélioration de la généralisation : Démonstration que le remplacement de la perte MSE par un objectif adversarial dynamique permet de générer des échantillons beaucoup plus fidèles à la distribution des données, sans nécessiter de guidance forte.
4. Efficacité pratique : Une méthode de post-entraînement légère (quelques époques) capable d'améliorer significativement les modèles existants (SiT, JiT).

4. Résultats Expérimentaux

Les auteurs ont évalué CAFM sur plusieurs tâches de génération d'images.

A. Génération sur ImageNet (256px)

• Sans guidance (Guidance-free) :
  • SiT (Latent Space) : Le FID passe de 8.26 (modèle FM de base) à 3.63 après 10 époques de post-entraînement CAFM.
  • JiT (Pixel Space) : Le FID passe de 7.17 à 3.57.
  • Comparaison : Ces résultats surpassent l'état de l'art (SOTA) dans les deux espaces (latent et pixel) sans guidance.
• Avec guidance (CFG) :
  • Amélioration du FID de 2.06 à 1.53 pour SiT et de 1.86 à 1.80 pour JiT.
  • CAFM permet d'atteindre de meilleurs scores avec des niveaux de guidance (CFG) plus faibles que les modèles de base.

B. Génération Texte-vers-Image

Sur le modèle Z-Image (512px) :

• GenEval : Score global amélioré de 0.81 à 0.85 (sans guidance) et de 0.81 à 0.85 (avec guidance).
• DPG-Bench : Score global amélioré de 83.7 à 85.2.
• Les résultats qualitatifs montrent une meilleure cohérence des objets, des attributs et des relations spatiales.

C. Entraînement à partir de zéro

Bien que moins efficace que le post-entraînement, CAFM peut être utilisé pour l'entraînement initial. Cependant, il converge plus lentement que le FM standard, confirmant l'hypothèse que le FM est optimal pour l'apprentissage de base, tandis que CAFM excelle dans le raffinement de la généralisation.

5. Signification et Impact

Ce travail représente une avancée significative dans la théorie et la pratique des modèles génératifs continus :

• Dépassement de la limite Euclidienne : Il démontre que la perte MSE, bien que théoriquement correcte pour la convergence, est sous-optimale pour la généralisation en pratique. L'introduction d'un discriminateur appris dynamiquement permet de capturer la structure de la variété des données mieux qu'une métrique fixe.
• Alternative aux GANs et Guidance : CAFM offre une voie pour obtenir une haute fidélité d'échantillonnage sans dépendre exclusivement de la guidance (qui modifie la distribution cible) ou des GANs (souvent instables et limités en temps discret).
• Efficacité pour l'industrie : La capacité d'améliorer massivement les performances de modèles existants (comme SiT et JiT) avec seulement 10 époques de fine-tuning en fait une méthode très attractive pour les applications de production, permettant d'obtenir des résultats de qualité SOTA à moindre coût computationnel.

En résumé, les Continuous Adversarial Flow Models combinent la stabilité et la flexibilité des modèles de flux continus avec la sensibilité aux détails perceptuels des réseaux adversariaux, résolvant le problème de la génération hors distribution et établissant de nouveaux standards de qualité pour la génération d'images.