Résumé Technique : Apprentissage sur la Variété avec le Flow Matching Riemannien

1. Énoncé du Problème

Les avancées récentes dans la modélisation générative se sont tournées vers l'exploitation d'encodeurs de représentations sémantiques de haute dimension gelés (ex: DiNOv2, SigL2) pour améliorer la fidélité de la synthèse et la vitesse de convergence. Cependant, les cadres standards de Diffusion Transformers (DiT) et de Flow Matching (FM) ne parviennent pas à converger efficacement lorsqu'ils sont entraînés directement sur ces représentations.

Des travaux antérieurs, spécifiquement les Autoencodeurs de Représentation (RAE), ont attribué cet échec à un goulot d'étranglement de capacité, proposant que les modèles nécessitent une « mise à l'échelle de la largeur » (augmentation de la largeur du transformer pour correspondre à la dimensionnalité du latent) pour résoudre le problème. Cet article conteste cette hypothèse, arguant que la cause profonde n'est pas une capacité de modèle insuffisante, mais un désalignement géométrique fondamental, nommé Interférence Géométrique.

Le problème central provient du fait que les encodeurs de représentation projettent les caractéristiques sur une variété hypersphérique ($S^{d-1}$) présentant une variance radiale quasi nulle (une géométrie de « coquille dure »). En revanche, le Flow Matching euclidien standard suppose un chemin de probabilité linéaire (une corde) entre les données et le bruit gaussien. Cette interpolation linéaire coupe à travers l'intérieur à faible densité de l'hypersphère, forçant le modèle à apprendre des champs de vitesse dans des régions indéfinies de l'espace des caractéristiques. Par conséquent, le modèle gaspille sa capacité à minimiser les erreurs radiales (qui sont structurellement nulles sur la variété) plutôt qu'à apprendre les transitions angulaires sémantiques.

2. Méthodologie : Flow Matching Riemannien avec Régularisation de Jacobi (RJF)

Pour résoudre l'Interférence Géométrique sans modifications architecturales comme la mise à l'échelle de la largeur, les auteurs proposent le RJF, un cadre qui aligne le processus génératif avec la géométrie intrinsèque de l'espace latent.

2.1 Flow Matching Riemannien (Correction de Trajectoire)

Au lieu d'une interpolation linéaire dans l'espace euclidien, le RJF définit le chemin de probabilité directement sur la variété hypersphérique en utilisant l'Interpolation Linéaire Sphérique (SLERP).

• Trajectoires Géodésiques : La trajectoire entre un point de donnée $x$ et le bruit $\epsilon$ (tous deux projetés sur la sphère unité) suit la géodésique (arc de grand cercle) plutôt qu'une corde.
• Vélocité de l'Espace Tangent : Le champ de vitesse cible est calculé en dérivant le chemin SLERP. Cela garantit que le vecteur de vitesse prédit appartient strictement à l'espace tangent de la variété ($v_t \cdot x_t = 0$), éliminant ainsi le besoin d'apprendre des composantes radiales.
• Échantillonnage : Lors de l'inférence, le modèle utilise un intégrateur de Géodésique (Expontential Map) pour mettre à jour l'état latent, garantissant que la trajectoire reste sur la surface de la variété tout au long du processus de génération.

2.2 Régularisation de Jacobi (Correction de Courbure)

Les auteurs observent que la simple correction de la trajectoire est insuffisante car l'erreur quadratique moyenne (MSE) standard traite les erreurs de vélocité de manière uniforme. Sur une hypersphère à courbure positive, les géodésiques convergent (focusing), ce qui signifie que les erreurs de vélocité se propagent de manière non linéaire.

• Champs de Jacobi : Pour tenir compte de cela, les auteurs introduisent un facteur de pondération dérivé des champs de Jacobi, qui quantifient la séparation entre les géodésiques causée par les perturbations de vélocité.
• Repondération de la Perte : La fonction de perte est multipliée par un facteur de pondération géométrique $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$, où $\Omega$ est la distance géodésique. Cela réduit le poids des erreurs près des données ($t=0$) là où le focus atténue les perturbations, et donne la priorité à la précision près du bruit ($t=1$), effectuant ainsi un recuit du signal d'apprentissage pour correspondre à la courbure de la variété.

3. Principales Contributions

1. Analyse Géométrique de l'Échec de Convergence : L'article identifie l'Interférence Géométrique comme le principal goulot d'étranglement pour la diffusion standard sur les encodeurs de représentation. Il démontre que l'échec provient d'objectifs euclidiens forçant les modèles à apprendre à travers l'intérieur hors-variété de l'hypersphère, plutôt que d'un manque de capacité du modèle.
2. Cadre RJF : La proposition d'un cadre d'entraînement géométriquement cohérent combinant le Flow Matching Riemannien (pour l'alignement de la trajectoire) et la Régularisation de Jacobi (pour la propagation des erreurs sensible à la courbure).
3. Efficacité sans Mise à l'Échelle de la Largeur : La démonstration que les architectures DiT standards (ex: DiT-B avec 131M paramètres) peuvent converger efficacement sur des latents sémantiques de haute dimension sans la mise à l'échelle de la largeur coûteuse en calcul proposée par les travaux précédents.

4. Résultats Expérimentaux

Les expériences ont été menées sur ImageNet-1K à une résolution de 256×256 en utilisant diverses échelles de modèles (DiT-B, DiT-L, DiT-XL) et des encodeurs de représentation (DiNOv2-B, SigLIP, MAE).

• Convergence et Performance FID :
  • DiT-B (131M) : Le Flow Matching euclidien standard sur les caractéristiques DiNOv2-B a échoué à converger (FID ~24.21). Le RJF a atteint un FID de 6.77 en 80 époques et 3.37 (avec guidage) en 200 époques.
  • DiT-XL (677M) : Le RJF a atteint un FID de 3.62 en 80 époques sans guidage, surpassant la ligne de base euclidienne (FID 4.28) et le LightningDiT basé sur VAE (FID 4.29).
  • Mise à l'échelle : La méthode a montré des améliorations constantes à travers les tailles de modèles, atteignant des scores d'Inception (IS 186.2) et de Précision (0.82) de pointe.
• Études d'Ablation :
  • Radial vs Angulaire : Le masquage de la perte radiale n'a apporté que des gains marginaux (FID 21.99), confirmant que le chemin linéaire lui-même (la corde) est le problème principal.
  • Trajectoire : Passer au Flow Matching Riemannien (SLERP) seul a considérablement amélioré la performance (FID 7.06), validant l'hypothèse du désalignement de trajectoire.
  • Régularisation : L'ajout de la Régularisation de Jacobi a permis d'affiner davantage les résultats (FID 6.77).
• Généralisation : La méthode a convergé avec succès sur d'autres encodeurs avec des contraintes hypersphériques, notamment SigLIP (FID 10.39) et MAE (FID 19.82), là où les méthodes standards échouaient complètement.

5. Signification et Revendications

L'article affirme que la difficulté de l'entraînement des modèles de diffusion sur les encodeurs de représentation est fondamentalement géométrique, et non architecturale. En respectant la topologie hypersphérique de l'espace latent, les Diffusion Transformers standards peuvent apprendre des sémantiques de haute fidélité efficacement.

Les auteurs affirment que le RJF débloque le potentiel des architectures DiT standards, leur permettant de égaler ou de surpasser les performances de modèles beaucoup plus larges ou de pipelines complexes basés sur des VAE. Ce travail établit que la modélisation générative efficace sur les variétés de représentation nécessite un alignement géométrique (trajectoires correctes et pertes sensibles à la courbure) plutôt qu'une mise à l'échelle brute de la capacité. Cette approche permet une synthèse de haute fidélité avec moins de paramètres et une convergence plus rapide, comme en témoigne le modèle DiT-B atteignant un FID de 3.37 là où les méthodes précédentes n'arrivaient pas à converger.