Résumé technique : Alignement de la géométrie latente pour l'appariement de flux sphérique en génération d'images

Énoncé du problème

L'appariement de flux latent (et les modèles de diffusion) pour la génération d'images opère typiquement en transportant un bruit gaussien vers des latents de Variational Autoencoder (VAE) le long de trajectoires linéaires dans un espace euclidien. Cependant, cette approche néglige deux propriétés structurelles critiques des latents VAE de haute dimension :

1. Concentration de la mesure : Tant le prior gaussien que les latents de données encodés se concentrent dans des coques sphériques minces plutôt que de remplir le volume de l'espace. Une ligne droite euclidienne (corde) reliant deux points sur ces coques traverse des rayons intérieurs qu'aucune des deux distributions n'occupe, contraignant le modèle à apprendre des transitions à travers des régions de l'espace latent rarement peuplées.
2. Dominance directionnelle : La sortie du décodeur est principalement sensible à la direction (composante angulaire) d'un token latent plutôt qu'à sa longueur (composante radiale). Des expériences consistant à échanger le rayon et la direction de tokens provenant de voisins de même classe démontrent que préserver la direction tout en modifiant le rayon entraîne un changement perceptuel minime, tandis que modifier la direction provoque des décalages sémantiques significatifs.

L'appariement de flux linéaire standard ignore ces faits, allouant une part significative de la cible de vitesse d'entraînement (jusqu'à 50-90 % selon le tokenizer) au mouvement radial. Puisque le décodeur est moins sensible au rayon, cette supervision est inefficace, conduisant à une convergence plus lente et à une qualité d'image sous-optimale (FID) par rapport aux méthodes respectant la géométrie sous-jacente.

Méthodologie

Les auteurs proposent l'Appariement de Flux Sphérique, une méthode qui aligne la géométrie latente avec la structure sphérique des données sans modifier l'architecture de diffusion ni nécessiter d'encodeurs auxiliaires. L'approche se compose de quatre composants clés :

1. Projection sphérique token par token :
   La méthode projette chaque token latent $z$ d'un VAE préentraîné sur une hypersphère de rayon fixe $S^{d-1}(\sqrt{d})$. Cela est réalisé via une étape de normalisation $L_2$ insérée entre l'encodeur figé et le décodeur :
   $$z_{i,j} \leftarrow \sqrt{d} \cdot \frac{z_{i,j}}{\|z_{i,j}\|}$$
   Cela garantit que tous les tokens se situent sur une sphère dont le rayon correspond au rayon de concentration du prior gaussien standard.

2. Prior sphérique :
   Au lieu d'utiliser un bruit gaussien brut, le point d'extrémité du bruit $z_0$ est construit en échantillonnant un bruit gaussien isotrope $\epsilon \sim \mathcal{N}(0, I)$ et en le projetant sur la sphère :
   $$z_0 = \sqrt{d} \cdot \frac{\epsilon}{\|\epsilon\|_2}$$
   Cela résulte en une distribution uniforme sur la sphère, correspondant à la géométrie des données projetées.

3. Affinage du décodeur :
   Pour s'adapter à la contrainte géométrique, le décodeur est affiné pendant quelques époques (5 époques dans les expériences) tandis que l'encodeur et la couche de projection restent figés. L'objectif de reconstruction conserve les pertes de pixels standard (L1, LPIPS, adversaire) mais supprime le terme KL, car le latent projeté est déterministe.

4. Transport Slerp et vitesse tangente :
   Le chemin de transport entre les points d'extrémité du bruit et des données est remplacé par l'Interpolation Linéaire Sphérique (slerp), qui suit l'arc géodésique sur la sphère.

   • Chemin géodésique : Le chemin $z_t$ reste sur la sphère pour tout $t \in [0, 1]$.
   • Cible de vitesse : La vitesse conditionnelle $u_t$ est la dérivée temporelle de la géodésique, projetée sur l'espace tangent $T_{z_t}S^{d-1}$. Cela garantit que la vitesse est purement angulaire, avec une composante radiale nulle par construction.
   • Entraînement et échantillonnage : Le modèle est entraîné pour prédire cette vitesse tangente. Lors de l'inférence, l'échantillonneur utilise l'application exponentielle pour intégrer la vitesse, garantissant que la trajectoire reste sur la sphère sans nécessiter de corrections de projection radiale à chaque étape.

Contributions clés

L'article identifie et quantifie un « décalage rayon-coque » dans l'appariement de flux latent standard et introduit une intervention géométrique pour le résoudre :

1. Identification du décalage : Les auteurs démontrent que les trajectoires linéaires forcent le modèle à prédire un mouvement radial auquel le décodeur est insensible, gaspillant la capacité et ralentissant la convergence.
2. Projection token par token : Ils introduisent une méthode pour contraindre les latents VAE préentraînés à une sphère de rayon fixe en n'utilisant que l'affinage du décodeur, évitant ainsi la nécessité de réentraîner l'encodeur ou d'utiliser des encodeurs de représentation auxiliaires.
3. Appariement de flux basé sur Slerp : Ils proposent d'entraîner l'appariement de flux le long de géodésiques slerp avec des cibles de vitesse tangente, garantissant que l'ensemble de la trajectoire respecte la variété sphérique.
4. Validation empirique : La méthode améliore les performances sur plusieurs familles de tokenizers (FLUX.2, VA-VAE, REPA-E FLUX.1) et échelles de modèles (SiT-B à SiT-XL) sans modifications architecturales.

Résultats

Les expériences sur la génération conditionnelle par classe d'ImageNet-256 démontrent des améliorations constantes :

• Améliorations FID : La méthode sphérique-slerp obtient des scores FID inférieurs à ceux de l'appariement de flux linéaire standard sur toutes les configurations testées. Par exemple, sur SiT-B/2 avec FLUX.2, le FID passe de 26,35 (Linéaire) à 20,55 (Sphérique-slerp).
• Efficacité de l'entraînement : La méthode atteint un FID cible de 30 en environ 2,2 fois moins d'étapes d'entraînement par rapport à la base linéaire standard.
• Généralisation : Des gains sont observés sur différentes échelles de guidage (CFG) et tailles de modèles (B et XL), indiquant que l'avantage est inhérent à la géométrie latente plutôt qu'à un réglage spécifique des hyperparamètres.
• Études d'ablation :
  • L'échange de décodeurs entre les pipelines standard et sphérique dégrade les performances dans les deux sens, confirmant que le flux appris est étroitement couplé à la géométrie latente spécifique.
  • L'utilisation d'un chemin « décomposé en coque » (interpolant rayon et direction séparément) sur des latents standards ne produit pas les mêmes gains, confirmant que l'élimination complète du degré de liberté radial est cruciale.
  • La qualité de reconstruction (rFID) des décodeurs sphériques affinés reste compétitive, montrant que la contrainte géométrique ne compromet pas sévèrement la capacité du VAE à reconstruire des images.

Importance

L'article établit la géométrie du support latent comme un axe de conception critique pour la modélisation générative. Il soutient que la forme de la distribution latente définit l'espace que le générateur doit traverser, et que l'ignorance de cette géométrie (en supposant une linéarité euclidienne) conduit à l'inefficacité.

L'importance de ce travail réside dans sa capacité à améliorer la qualité de génération d'images et la vitesse d'entraînement sans :

• Modifier l'architecture de diffusion (par exemple, le backbone SiT).
• Ajouter des encodeurs auxiliaires (contrairement à l'appariement de flux riemannien sur des caractéristiques DINOv2 figées).
• Réentraîner l'encodeur VAE à partir de zéro.

En projetant simplement les latents préentraînés existants sur une sphère et en alignant le chemin de flux sur les géodésiques résultantes, la méthode atteint des résultats de pointe avec une surcharge computationnelle minimale. Les auteurs suggèrent que les travaux futurs devraient explorer ces contraintes géométriques sous des régimes de compression plus forts et concevoir des tokenizers qui correspondent nativement à cette structure sphérique.