Résumé technique : Flux conservateurs

Énoncé du problème

La modélisation générative moderne est principalement définie par le paradigme du transport : apprendre une application d'une distribution de référence simple (généralement un bruit gaussien) vers une distribution de données complexe. Bien que les modèles de diffusion basés sur les scores et les flux aient atteint des résultats de pointe en inversant un processus de bruitage progressif, cette approche nécessite intrinsèquement d'apprendre une application de transport globale sur une famille de distributions intermédiaires.

Les auteurs identifient une limitation fondamentale dans la discrétisation des dynamiques à temps continu utilisées dans ces modèles. Plus précisément, bien que les dynamiques de Langevin à temps continu préservent théoriquement la distribution de données cible $p$ comme mesure invariante, leurs approximations explicites à temps discret (par exemple, l'algorithme de Langevin non ajusté, ULA) échouent à le faire exactement. L'erreur de discrétisation provoque un dérive des échantillons générés par rapport à la variété des données réelles au fil du temps, rompant l'exigence selon laquelle le processus de génération doit rester au sein de la distribution de données à chaque étape.

En outre, les auteurs soutiennent que dans de nombreux contextes pratiques — tels que les systèmes de récupération augmentée ou les pipelines génératifs latents — la génération ne commence pas nécessairement à partir d'un état de bruit « dépourvu d'information ». Au contraire, des états initiaux valides existent souvent sous la forme de points de données récupérés, de voisins les plus proches ou de codes latents qui se situent déjà près de ou sur la distribution cible. L'article postule que pour de tels scénarios, un changement de paradigme du transport (bruit vers données) vers l'invariance (données vers données) est plus naturel et efficace.

Méthodologie

L'article propose les Flux conservateurs, un cadre où la génération est effectuée par une dynamique stochastique discrète qui laisse la distribution de données invariante. Le système est initialisé à partir d'états supportés par les données ($x_0 \sim p$) et évolue tout en préservant la loi de probabilité à chaque étape ($x_j \sim p$ pour tout $j \ge 1$).

Le cadre introduit deux mécanismes d'échantillonnage distincts préservant la probabilité, qui peuvent fonctionner directement au-dessus de modèles d'appariement de flux préentraînés existants sans réentraînement :

1. Dynamique de Langevin corrigée (dMALA)

Cette méthode adapte l'algorithme de Langevin ajusté par Metropolis (MALA) à la modélisation générative où la densité explicite $p(x)$ n'est pas disponible.

• Mécanisme : Elle utilise un champ de score appris $s_\theta$ ou, plus pratiquement, un débruiteur $D_\sigma$ entraîné pour prédire des échantillons propres à partir d'observations bruitées. Le débruiteur est dérivé directement du pont utilisé lors de l'entraînement par appariement de flux.
• Correction : Pour restaurer l'invariance, la méthode emploie une correction de Metropolis-Hastings. Au lieu d'exiger la différence de log-densité $\log p(y) - \log p(x)$, qui est inconnue, la méthode utilise l'intégration linéaire du champ de score conservateur (ou du résidu du débruiteur) le long du chemin entre l'état actuel $x$ et la proposition $y$.
• Cible d'invariance : Ce mécanisme préserve une distribution lissée par un filtre gaussien $p_\sigma = p * \mathcal{N}(0, \sigma^2 I)$, où $\sigma$ est déterminé par le niveau de bruit du débruiteur.
• Efficacité : En approximant l'intégrale linéaire par la règle du trapèze, le taux d'acceptation se simplifie en une fonction des résidus au carré du débruiteur, ne nécessitant que deux évaluations du débruiteur par étape.

2. Flux Prédicteur-Correcteur

Cette méthode cible directement la distribution de données originale $p$, plutôt qu'une version lissée.

• Mécanisme : Elle fonctionne comme un processus en deux étapes :
  1. Prédicteur : Une quantité contrôlée de bruit est ajoutée à l'échantillon actuel $x \sim p$, le déplaçant vers une distribution intermédiaire de « pont » $p_\tau$ (où $\tau < 1$).
  2. Correcteur : Une application de flux apprise (ou une application de solution directe) est appliquée pour transporter l'échantillon bruité de retour vers la distribution de données $p$ (temps $t=1$).
• Invariance : Sous un modèle de flux exact, ce noyau est strictement préservateur de probabilité. La méthode repose sur la capacité du modèle de flux à mapper avec précision les échantillons du pont de retour vers la variété des données.
• Flexibilité : Contrairement à la nature locale d'acceptation/rejet du dMALA, cette approche permet des mouvements plus grands et non locaux en sortant temporairement de la variété pour revenir via le flux appris.

Contributions clés

1. Changement de paradigme : L'article reformule la modélisation générative du « transport depuis le bruit » vers une « sonde stationnaire de la distribution de données ». L'objet central est une dynamique stochastique qui préserve la loi pertinente, initialisée à partir d'états supportés par les données.
2. Deux mécanismes préservant la probabilité : Les auteurs développent et formalisent deux échantillonneurs spécifiques (dMALA et Prédicteur-Correcteur) qui corrigent les erreurs de discrétisation des dynamiques de Langevin standard et de la génération basée sur les flux.
3. Réutilisation de modèles préentraînés : Une contribution pratique cruciale est qu'aucun des deux mécanismes ne nécessite l'entraînement d'un nouveau modèle génératif à partir de zéro. Tous deux peuvent extraire le score, le débruiteur ou l'application de flux nécessaire directement à partir de points de contrôle d'appariement de flux préentraînés existants (par exemple, SoFlow, SiT).
4. Correction théorique : L'article aborde explicitement l'échec des dynamiques de Langevin discrètes à préserver la loi cible et fournit la dérivation mathématique pour corriger cela via l'intégration linéaire et les flux de pont.

Résultats expérimentaux

Le cadre a été validé sur une cible synthétique en forme de rouleau suisse, ImageNet-256 et Oxford Flowers-102.

• Cible synthétique (rouleau suisse) : Les expériences ont démontré que tandis que l'ULA standard dérive loin de la distribution cible, les dynamiques corrigées (dMALA) et le flux Prédicteur-Correcteur restent sur la variété correcte, préservant la distribution sur de longues chaînes.
• ImageNet-256 :
  • Les deux méthodes ont systématiquement amélioré les échantillonneurs de base des modèles de fond préentraînés (SoFlow-XL/2-cond et SiT-REPA-XL).
  • Le dMALA a obtenu les scores de distance Fréchet Inception (FID) les plus bas, en particulier avec des niveaux de lissage plus faibles ($\sigma=0.3$), indiquant une haute fidélité à la distribution de données.
  • Le Prédicteur-Correcteur a obtenu des scores Inception (IS) et une précision plus élevés, suggérant qu'il génère des échantillons plus diversifiés et distincts, bien qu'avec un FID légèrement plus élevé que le dMALA dans certaines configurations.
  • Comportement à long terme : Les deux méthodes ont montré une dégradation gracieuse avec l'augmentation de la longueur de la chaîne ($K$), cohérente avec l'accumulation d'erreurs d'approximation dans le modèle de fond préentraîné, mais ont maintenu des performances nettement supérieures aux bases de référence.
• Efficacité computationnelle : La méthode Prédicteur-Correcteur s'est révélée compétitive par rapport à l'intégration directe d'équations différentielles ordinaires (ODE). Par exemple, 200 évaluations de fonctions de réseau (NFE) de la méthode Prédicteur-Correcteur ont obtenu un FID plus faible (4,91) que l'intégration Euler avant sur 800 étapes (6,51) du même modèle de fond.

Importance et revendications

Les auteurs positionnent ce travail comme un nouveau paradigme pour la modélisation générative. Ils affirment que, en déplaçant l'objectif de la transformation du bruit en données vers le maintien de la distribution de données sous des dynamiques apprises, le cadre permet :

• Variation fidèle des échantillons : En maintenant le processus sur les régions de haute probabilité de la variété des données.
• Amélioration de la fidélité des données : Grâce à la correction des erreurs de discrétisation qui affectent les dynamiques de Langevin standard.
• Réutilisation efficace : La capacité d'améliorer les modèles préentraînés existants sans réentraînement.

L'article reconnaît modestement des limites, notant que l'efficacité des échantillonneurs est liée à la qualité du modèle de flux ou basé sur les scores préentraîné sous-jacent. Les erreurs d'approximation dans le modèle de score peuvent s'accumuler au fil des applications répétées du noyau. De plus, la validation actuelle est restreinte aux domaines continus (images), et l'extension du paradigme aux domaines discrets comme la génération de texte est identifiée comme un travail futur. Les auteurs soulignent que leur contribution est l'unification des principes de conservation dans un cadre génératif organisé autour de la sonde stationnaire plutôt que du bruitage inverse.