Résumé Technique : Modélisation Générative par Appariement de Flux

Énoncé du Problème

La modélisation générative de l'état de l'art actuelle, en particulier dans les domaines de haute dimension comme la synthèse d'images, la conception de protéines et la robotique, repose fortement sur les modèles basés sur le score. Ces modèles apprennent la fonction de score de Stein, $\nabla \log p_{\text{data}}(x)$, typiquement via l'appariement de scores. Une fois appris, ce score est utilisé pour piloter des processus de diffusion (par exemple, la dynamique de Langevin) afin d'échantillonner à partir de la distribution cible $p_{\text{data}}$.

Cependant, la fonction de score n'est pas le champ vectoriel unique capable de générer $p_{\text{data}}$. Selon l'équation de Fokker–Planck (EFP), tout champ vectoriel $f_\theta$ satisfaisant la condition $\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$ aura $p_{\text{data}}$ comme distribution stationnaire. Cela implique l'existence d'une famille infinie de champs vectoriels génératifs (champs non conservatifs) capables de générer la même distribution cible mais induisant des dynamiques différentes (par exemple, des taux de mélange, des motifs de circulation ou des dépendances dirigées distincts).

Les objectifs d'appariement de scores existants forcent le modèle à converger vers un point unique dans cet espace ($\nabla \log p_{\text{data}}$), pénalisant toute déviation même si cette déviation préserve la distribution cible. Cela restreint la capacité d'encoder directement dans le champ génératif des biais inductifs utiles, des priors structurels ou des propriétés dynamiques (comme un mélange plus rapide ou des dépendances causales).

Méthodologie : Appariement de Flux

Les auteurs proposent l'Appariement de Flux, un nouveau paradigme qui généralise l'appariement de scores en apprenant n'importe quel champ vectoriel au sein de la famille de champs générant la distribution cible, plutôt que de forcer le modèle à correspondre au score point par point.

Objectif Central

Au lieu de minimiser la divergence de Fisher entre le modèle $f_\theta$ et le score $\nabla \log p_{\text{data}}$, l'Appariement de Flux minimise la divergence du flux de probabilité. La condition pour qu'un champ soit génératif est :
$$\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$$
Cette condition est plus faible que l'égalité point par point, admettant un "noyau" d'infinité de champs vectoriels valides.

Déduction Théorique

Pour optimiser cette condition, les auteurs définissent une Divergence de Fisher Projetée :
$$\tilde{J}(\theta) := \mathbb{E}_{x \sim p_{\text{data}}} \left[ \| \Pi_{\text{flux}} f_\theta(x) - \nabla \log p_{\text{data}}(x) \|^2 \right]$$
où $\Pi_{\text{flux}}$ projette un champ vectoriel sur l'espace des champs de gradient partageant la même divergence de flux. Bien que $\tilde{J}(\theta)$ soit théoriquement solide, il est intraitable en haute dimension.

Les auteurs dérivent une Perte d'Appariement de Flux ($L_{\text{flux}}$) traitable en :

1. Définissant le champ résiduel $u_\theta = f_\theta - \nabla \log p_{\text{data}}$ et l'opérateur de Langevin-Stein $r_\theta = \nabla \cdot u_\theta + u_\theta \cdot \nabla \log p_{\text{data}}$.
2. Utilisant des simulations de diffusion $d x_t = \nabla \log p_{\text{data}}(x_t) dt + \sqrt{2} dW_t$ pour propager $r_\theta$ à travers la distribution des données.
3. Appliquant l'intégration par parties et un astuce de stop-gradient pour convertir l'objectif en une forme évitant la rétropropagation à travers des dérivées de trajectoire coûteuses.

La perte finale est :
$$L_{\text{flux}}(\theta) := -\mathbb{E}_{t \sim q, x_0 \sim p_{\text{data}}, x_t | x_0} \left[ \frac{1}{q(t)} u_\theta(x_0)^\top \text{sg}\left( \frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t) \right) \right]$$
où $\text{sg}$ désigne l'opérateur stop-gradient. Les auteurs prouvent que $\nabla_\theta \tilde{J}(\theta) = 2 \nabla_\theta L_{\text{flux}}(\theta)$.

Recuit de Bruit et Implémentation Pratique

La méthode s'étend aux settings à recuit de bruit (similaires aux modèles de diffusion) en apprenant un continuum de champs pour des niveaux de bruit $\sigma$.

• Approximation du Score : Comme $\nabla \log p_{\text{data}}$ est inconnu, les auteurs utilisent une approximation non paramétrique par Estimation de Densité de Noyau (KDE) sur les mini-lots pour estimer le score.
• Réduction de Variance : Pour estimer le terme de trajectoire $\frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t)$, les auteurs emploient un estimateur de chaîne croisée utilisant tous les échantillons d'un mini-lot, pondérés par les probabilités de transition gaussiennes, plutôt que de s'appuyer sur une seule chaîne.

Sélection de Champs Spécifiques

L'Appariement de Flux apprend une famille de champs. Pour sélectionner un champ spécifique avec des propriétés désirées, les auteurs proposent deux stratégies :

1. Perte Spécifique à l'Application : Ajout de pertes auxiliaires (par exemple, pour la vitesse de mélange ou des contraintes structurelles) à l'objectif d'Appariement de Flux.
2. Paramétrisation du Modèle : Contraintes architecturales (par exemple, masques d'attention causaux) qui imposent des structures spécifiques directement dans $f_\theta$.

Contributions Clés

1. Paradigme d'Appariement de Flux : Introduction d'un cadre de modélisation générative apprenant des champs vectoriels au-delà du score en faisant correspondre la divergence du flux de probabilité.
2. Dérivation Efficace de Perte : Dérivation d'une perte d'Appariement de Flux évolutive préservant la géométrie de la divergence de Fisher tout en permettant des champs vectoriels non conservatifs.
3. Extension à Recuit de Bruit : Extension de la méthode aux distributions à recuit de bruit, permettant l'application à des ensembles de données d'images complexes et de haute dimension.
4. Flexibilité Démontrée : Preuves empiriques montrant que l'Appariement de Flux permet des applications impossibles sous l'appariement de scores standard, incluant :
   • Mélange Plus Rapide : Apprentissage de champs vectoriels accélérant la convergence d'échantillonnage.
   • Interprétabilité : Ajustement de modèles structurés et biologiquement significatifs (par exemple, vitesse de l'ARN) sans approximations neuronales boîte noire.
   • Dynamiques Structurées : Intégration de dépendances dirigées (par exemple, autorégression temporelle) directement dans le champ génératif.

Résultats

L'article évalue l'Appariement de Flux à travers cinq configurations :

1. Champs Contrôlables (Données Jouet) : Sur un mélange gaussien 2D, l'Appariement de Flux attribue une perte nulle aux champs préservant la distribution cible mais ayant des dynamiques différentes (par exemple, rotation, circulation), tandis que l'appariement de scores pénalise ces déviations.
2. Champs Interprétables (Vitesse de l'ARN) : Appliqué à la génomique de cellule unique, l'Appariement de Flux ajuste les paramètres du modèle dynamique scVelo par descente de gradient. Il surpasse l'ajustement scVelo standard basé sur EM en termes de cohérence et de correction aux frontières croisées sur cinq ensembles de données, démontrant que les gains proviennent de la procédure d'optimisation plutôt que d'une expressivité accrue du modèle.
3. Génération Non Restreinte (Images) : Sur CIFAR-10 et CelebA, l'Appariement de Flux sert d'objectif autonome. Bien qu'il atteigne des scores FID compétitifs, il est actuellement 3–4 fois plus lent et utilise 2–3 fois plus de mémoire pendant l'entraînement par rapport à l'Appariement de Scores par Débruitage (DSM). Cependant, il s'étend aux distributions de haute dimension.
4. Mélange Rapide : En optimisant un proxy de mélange, l'Appariement de Flux apprend des champs atteignant des scores FID raisonnables avec significativement moins d'étapes d'échantillonnage que l'Appariement de Flux standard ou le DSM.
5. Intégration de Structure : Dans une tâche de génération de trajectoires (ressorts non linéaires), l'Appariement de Flux intègre avec succès un masque d'attention causale pour imposer une autorégression temporelle. En revanche, appliquer le même masque au DSM dégrade les performances car la Jacobienne symétrique de la fonction de score entre en conflit avec les dépendances dirigées.

Importance et Revendications

Les auteurs affirment que l'Appariement de Flux ouvre une "nouvelle dimension" en modélisation générative en traitant le champ vectoriel lui-même comme un choix de conception plutôt que comme une cible fixe.

• Au-delà de l'Échantillonnage : La valeur principale n'est pas nécessairement de remplacer le DSM pour la génération d'images standard (où le DSM est hautement optimisé), mais de permettre des dynamiques aux propriétés utiles.
• Flexibilité : Il permet aux praticiens d'imposer des biais inductifs, des priors structurels et des contraintes mécanistes directement dans le processus génératif.
• Applications : L'article met en évidence des cas d'usage spécifiques débloqués par cette flexibilité : échantillonnage plus rapide via des dynamiques non réversibles, modèles scientifiques interprétables (comme la vitesse de l'ARN) et modèles génératifs avec des structures causales ou temporelles prescrites.

L'article conclut que bien que l'Appariement de Flux entraîne actuellement des coûts computationnels plus élevés, sa capacité à apprendre des champs génératifs structurés et non conservatifs offre un outil puissant pour les applications où les dynamiques de génération sont aussi importantes que la distribution de la sortie.