Résumé Technique : Améliorer le Flux Rectifié avec des Conditions aux Limites

Énoncé du Problème

Le Flux Rectifié (Rectified Flow - RF), également connu sous le nom de Flow Matching, est apparu comme un cadre de modélisation générative puissant qui apprend un champ de vitesse pour transformer une distribution de bruit en une distribution de données cibles via une équation différentielle ordinaire (ODE). Malgré son succès dans des tâches telles que la génération de texte-vers-image et de vidéo, les auteurs identifient une limitation critique dans les implémentations standards du RF : le champ de vitesse appris ne satisfait souvent pas les conditions aux limites théoriquement dérivées.

Plus précisément, sous une interpolation linéaire entre le bruit ($t=0$) et les données ($t=1$), le champ de vitesse optimal doit satisfaire :

1. Limite Gauche ($t=0$) : $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Limite Droite ($t=1$) : $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

Les paramétrages de champs de vitesse par réseaux de neurones standards n'imposent pas intrinsèquement ces contraintes. Comme illustré dans l'article, la vitesse prédite à $t=1$ dévie souvent de l'application identité attendue. Cette violation est particulièrement préjudiciable lors de l'échantillonnage stochastique (par exemple, via la dynamique de Langevin). Puisque l'estimation de la fonction de score (via la formule de Tweedie) implique une division par $(1-t)$, une vitesse inexacte près de $t=1$ provoque la divergence de l'estimation du score. Cela conduit à un échantillonnage instable, résultant en des images sur-lissées ou des artefacts de type "dessin animé", surtout lors de l'utilisation de pas de temps plus grands ou d'échelles de bruit plus importantes.

Méthodologie

Pour remédier à cela, les auteurs proposent le Modèle de Flux Rectifié à Limites Imposées (Boundary-enforced Rectified Flow Model). Cette approche incorpore explicitement les contraintes de bord dans la paramétrisation du champ de vitesse à travers deux variantes distinctes :

1. Modèle de Flux Rectifié à Limites par Masquage (Mask-based)

Cette variante impose les deux conditions aux limites ($t=0$ et $t=1$) en structurant le champ de vitesse comme une combinaison linéaire de termes de bord et d'une sortie de réseau de neurones apprenable :
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Où :

• $m_\theta(\mathbf{x}, t)$ est un réseau de neurones standard (par exemple, U-Net ou DiT).
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ sont des fonctions scalaires conçues pour s'annuler ou s'activer à des limites spécifiques (par exemple, $g(1)=f(0)=h(0)=h(1)=0$ et $g(0)=f(1)=1$).
• Les auteurs suggèrent un choix "Standard Cosinus" : $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, et $h(t) = \sin(\pi t)$.

Cette construction garantit que les conditions aux limites sont satisfaites par conception, quel que soit le résultat du réseau de neurones.

2. Modèle de Flux Rectifié à Limites par Soustraction (Subtraction-based)

Inspirée par l'observation selon laquelle l'imposition de la seule condition $t=1$ apporte des bénéfices empiriques significatifs, les auteurs proposent une variante plus simple, à limite unique :
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
Cette formulation satisfait intrinsèquement $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ car le terme $m_\theta(\mathbf{x}, 1)$ s'annule à $t=1$. Bien que cette approche offre simplicité de conception et flexibilité, elle induit un coût computationnel, nécessitant une passe avant supplémentaire pour calculer $m_\theta(\mathbf{x}, 1)$ pour chaque évaluation.

Impact sur l'Échantillonnage Stochastique

En imposant $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$, le Modèle de Flux Rectifié à Limites garantit que le numérateur de la formule de Tweedie ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) tend vers zéro lorsque $t \to 1$. Cela empêche la fonction de score $\nabla \log \rho_t$ de diverger, stabilisant ainsi les échantillonneurs stochastiques (échantillonneurs SDE) et permettant une génération robuste même avec des échelles de bruit ou des pas de temps plus importants.

Contributions Clés

1. Identification d'une Limitation : L'article souligne que les champs de vitesse non contraints violent les conditions aux limites théoriques du Flux Rectifié, entraînant des trajectoires d'ODE inexactes et un échantillonnage stochastique instable.
2. Paramétrage Novateur : L'introduction du Modèle de Flux Rectifié à Limites, qui offre deux variantes pratiques (par Masquage et par Soustraction) pour imposer ces conditions avec une modification minimale du code.
3. Stabilisation des Fonctions de Score : Démontrer que l'imposition des limites stabilise directement la fonction de score près de $t=1$, atténuant les problèmes de divergence qui affectent l'échantillonnage stochastique du RF classique.
4. Validation Empirique : Des expériences exhaustives montrant des améliorations constantes de la qualité de génération et de la stabilité de l'échantillonnage sur différents jeux de données et échelles de modèles.

Résultats Expérimentaux

Les auteurs ont évalué leurs modèles sur CIFAR-10 et ImageNet (256$\times$256 et 512$\times$512) en utilisant des architectures de Diffusion Transformer (DiT).

• Améliorations Quantitatives :

  • Sur ImageNet 256$\times$256, le Modèle de Flux Rectifié à Limites par Soustraction a atteint un FID de 6,32 (échantillonnage ODE) et 6,38 (échantillonnage SDE), contre 6,87 et 6,79 respectivement pour la base de référence RF classique.
  • Le modèle par Masquage a atteint un FID de 6,63 (ODE) et 6,18 (SDE).
  • Des améliorations ont également été observées pour l'Inception Score (IS), la Précision et le Rappel sur les deux jeux de données.
  • La variante par Soustraction a systématiquement obtenu les scores IS les plus élevés.

• Robustesse de l'Échantillonnage :

  • Dans des expériences utilisant l'échantillonneur "Overshooting" avec des échelles de bruit ($c$) variables, les modèles de Flux Rectifié à Limites ont démontré une performance nettement plus stable et une dégradation moindre de la qualité des échantillons par rapport au RF classique à mesure que la stochasticité augmentait.
  • Les résultats qualitatifs montrent que les modèles de Flux Rectifié à Limites produisent des images plus nettes avec moins d'artefacts et une meilleure fidélité des objets, particulièrement dans les contextes d'échantillonnage stochastique où le RF classique a tendance à trop lisser.

• Scalabilité :

  • Les bénéfices du Modèle de Flux Rectifié à Limites sont restés constants lors du passage à des modèles plus larges (DiT-L, DiT-XL) et à des résolutions plus hautes (512$\times$512).

Signification et Revendications

L'article affirme que le Modèle de Flux Rectifié à Limites offre une solution simple et facilement implémentable à une limitation fondamentale du Flux Rectifié. En imposant explicitement les conditions aux limites, la méthode :

1. Améliore la Qualité de Génération : Surpasse systématiquement le RF classique en termes de FID et d'autres métriques.
2. Permet un Échantillonnage Stochastique Stable : Résout la divergence de la fonction de score près du temps terminal, permettant des stratégies d'échantillonnage plus robustes.
3. Est Facile à Adapter : Peut être intégré dans les modèles RF existants (utilisant U-Net ou DiT) avec des modifications architecturales minimales.

Les auteurs présentent ce travail comme un raffinement nécessaire pour le déploiement pratique du Flux Rectifié, particulièrement dans les scénarios nécessitant un échantillonnage stochastique de haute qualité. Ils suggèrent que bien que leur attention se porte sur le RF, le concept central d'imposition des conditions aux limites pourrait potentiellement être étendu aux modèles de diffusion avec des adaptations mineures. Des travaux futurs sont proposés pour explorer l'imposition de limites adaptatives et l'application à la génération de texte-vers-vidéo.