Résumé Technique : Modélisation du Langage avec des Flows Hypersphériques (S-FLM)

Énoncé du Problème

Les modèles autoregressifs (AR) dominent la modélisation du langage en raison de leur évolutivité et de leur évaluation rapide de la vraisemblance, mais ils souffrent d'une génération séquentielle et d'un manque de contexte bidirectionnel lors de l'inférence. Les modèles de diffusion discrets offrent une génération parallèle et un contexte bidirectionnel, mais sont limités par des contraintes de traçabilité qui les forcent à échantillonner à partir de marginales factorisées, réduisant ainsi leur expressivité par rapport aux modèles AR.

Les récents Modèles de Flows pour le Langage (FLM) tentent de combler ce fossé en appliquant des flows continus au langage, transportant le bruit vers les données via une Équation Différentielle Ordinaire (ODE) déterministe. Cependant, les FLM existants présentent deux lacunes critiques :

1. Coût Computationsnel : Ils représentent les tokens comme des vecteurs one-hot de dimension $|V|$ (taille du vocabulaire). Pour de grands vocabulaires (100k–200k), cela nécessite de stocker et de manipuler des vecteurs de haute dimension, rendant l'entraînement significativement plus lent que pour les modèles de diffusion discrets ou les modèles AR.
2. Ambiguïté Sémantique : L'ajout de bruit gaussien à des vecteurs one-hot manque d'interprétation sémantique claire. Contrairement aux images, où le bruit dégrade progressivement la structure, l'ajout de bruit gaussien à des embeddings one-hot équidistants n'a pas de sens géométrique intuitif.

Méthodologie

Les auteurs proposent le Modèle de Flow de Langage Hypersphérique (S-FLM), un modèle de flow riemannien défini sur l'hypersphère unité $S^{d-1}$ plutôt que sur un espace euclidien ou des espaces d'états discrets.

Mécanismes de Base

• Espace Latent Hypersphérique : Au lieu de vecteurs one-hot, les tokens sont mappés vers des embeddings de norme unitaire $\hat{e}_v \in S^{d-1}$. Le processus forward transporte les embeddings propres vers une a priori uniforme sur la sphère via une Interpolation Linéaire Sphérique (SLERP), plutôt qu'une interpolation linéaire avec du bruit gaussien.
• Matching de Flow Riemannien : Le modèle apprend un champ de vitesse $u_t$ sur la variété $S^{d-1}$ qui définit une ODE. La vitesse conditionnelle est dérivée en utilisant la carte logarithmique sur la sphère :
  $$u_{t|1}(z_t | z_1) = \frac{\dot{\alpha}_t}{1-\alpha_t} \log_{z_t}(z_1)$$
  où $z_1$ est l'embedding propre et $z_t$ est le latent bruité.
• Objectif d'Entraînement : Le débruiteur $p_\theta^{1|t}$ est entraîné en utilisant l'Entropie Croisée (CE) pour prédire le token propre $x_\ell$ étant donné le latent bruité $z_t$. Cela évite le minimum trivial où tous les embeddings s'effondrent en un point, un risque lors de la régression directe des vitesses contre des embeddings apprenables.
• Échantillonnage : Lors de l'inférence, le modèle intègre le champ de vitesse appris depuis un bruit uniforme sur la sphère jusqu'à la variété des données. Le champ de vitesse est calculé en marginalisant sur les vecteurs tangents pointant vers chaque embedding propre, pondérés par l'a posteriori $p_\theta^{1|t}$.
  • Vitesse Exacte : Somme sur l'ensemble du vocabulaire (coûteuse en calcul).
  • Vitesse Top-k : Restreint la somme aux $k$ tokens les plus probables. Le décodage Top-1 correspond au décodage glouton.
  • Décodage Stochastique : Échantillonne un seul token depuis l'a posteriori pour définir la direction de la vitesse.

Innovations Architecturales et d'Entraînement

1. S-arch (Squelette Hypersphérique) : Une variante du Transformer où les activations intermédiaires sont contraintes à $S^{d-1}$. Elle remplace les résidus additifs par des interpolations normalisées et paramètre les couches d'attention/MLP comme des rotations. Cela aligne l'architecture avec la variété d'entrée, améliorant la qualité des échantillons.
2. Troncature du Programme de Bruit : Les auteurs dérivent une borne théorique $\alpha^\star(\delta)$ basée sur la géométrie de haute dimension. Ils montrent qu'en haute dimension, l'a posteriori s'effondre en une distribution one-hot dès que le niveau de bruit descend en dessous d'un certain seuil. L'entraînement est tronqué aux niveaux de bruit au-dessus de cette borne pour éviter d'apprendre à partir de tâches de débruitage triviales.
3. Programme de Bruit Adaptatif : Inspiré par InfoNoise, le programme d'entraînement est adapté pour allouer plus d'échantillons aux niveaux de bruit où le gradient de perte est le plus raide, améliorant l'efficacité de l'entraînement.

Contributions Clés

1. Cadre S-FLM : Un modèle de langage basé sur des flows opérant sur des embeddings de dimension $d$ sur une hypersphère, éliminant le besoin de matérialiser des vecteurs one-hot de dimension $|V|$. Cela réduit les coûts d'entraînement pour les rendre comparables aux modèles de diffusion discrets.
2. Squelette S-arch : Une architecture novatrice qui maintient les activations sur l'hypersphère, produisant de meilleures performances sur des tâches de raisonnement (GSM8K) et de modélisation du langage (OpenWebText) par rapport aux squelettes DiT standards.
3. Troncature Théorique du Bruit : Une heuristique principée pour tronquer le programme de bruit basée sur la taille du vocabulaire et la dimension de l'embedding, ce qui améliore significativement les performances sur les tâches de raisonnement.
4. Gains de Performance : S-FLM améliore substantiellement les modèles de flows continus précédents. Sur GSM8K, il atteint ~18 % de précision avec un décodage top-1 (contre <1 % pour les FLM antérieurs), réduisant l'écart avec les modèles de diffusion masqués (MDLM, Duo) à température standard ($T=1$).

Résultats Expérimentaux

• Raisonnement Sudoku : S-FLM atteint des performances comparables aux FLM précédents et aux modèles de diffusion discrets (Duo), les meilleurs résultats de diffusion continue atteignant 94,8 % de précision sur des puzzles faciles et 45,0 % sur des puzzles difficiles lors de l'utilisation de troncature et de programmes adaptatifs.
• GSM8K (Raisonnement Mathématique) :
  • Les modèles continus antérieurs (FLM, CANDI) ont atteint <1 % de précision.
  • S-FLM avec décodage top-1 atteint 18,0 % de précision, égalant la performance de MDLM et Duo à $T=1$.
  • Un écart de performance subsiste à basse température ($T=0,1$), où les modèles discrets (Duo, MDLM) atteignent 33–36 %, tandis que S-FLM plafonne autour de 18 %.
• OpenWebText (Modélisation du Langage) : S-FLM correspond à la frontière de Perplexité Générative (Gen. PPL) / Entropie des FLM précédents à haut Nombre d'Évaluations de Fonction (NFE) et les surpasse à faible NFE, où les modèles continus précédents présentaient des frontières instables.
• Efficacité d'Entraînement : S-FLM s'entraîne plus rapidement que les FLM (qui utilisent des vecteurs one-hot) et est comparable en vitesse aux modèles de diffusion discrets.

Importance et Revendications

L'article revendique que S-FLM représente une avancée significative pour les modèles de langage à flows continus en abordant les problèmes d'évolutivité et d'interprétation sémantique des approches précédentes. En opérant sur l'hypersphère, le modèle évite le goulot d'étranglement computationnel des vecteurs one-hot de haute dimension et fournit un mécanisme d'injection de bruit géométriquement cohérent.

Les auteurs notent que si S-FLM réduit l'écart avec les modèles de diffusion discrets et les modèles AR en termes de perplexité générative et de précision de raisonnement à des températures standard, un écart substantiel subsiste lors du décodage à basse température. Le travail est présenté comme une avancée méthodologique rendant les modèles de flows continus plus viables pour la modélisation du langage à grand vocabulaire, bien que les auteurs restent modestes concernant l'écart de performance restant par rapport aux modèles autoregressifs de l'état de l'art dans les domaines vérifiables.