Résumé technique : Échantillonnage à partir de modèles de langage par flot via des ponts conditionnés marginalement

Énoncé du problème

Les modèles de langage par flot (FLM) représentent une classe de modèles de langage non autoregressifs qui appliquent un flot continu ou un débruitage à des séquences de tokens encodés en one-hot. Contrairement aux modèles de diffusion continus génériques, les FLM possèdent une propriété structurelle unique : leur moyenne de débruitage, $m_t(x) = \mathbb{E}[X_0 | X_t = x]$, correspond exactement à la distribution marginale a posteriori sur les tokens propres à chaque position. Plus précisément, la coordonnée $(\ell, v)$ du vecteur de sortie représente $P(X_{0,\ell} = e_v | X_t = x)$.

Les procédures d'échantillonnage standard pour les FLM, héritées des conventions de diffusion continue (par exemple, les échantillonneurs de style DDPM), traitent cette moyenne de débruitage comme une estimation ponctuelle. Elles établissent un pont entre l'état bruité actuel et cette moyenne conditionnelle. Cependant, dans le contexte des séquences one-hot, la moyenne conditionnelle se situe à l'intérieur d'un produit de simplexes et ne constitue pas une séquence one-hot valide. Établir un pont vers cette « représentation d'incertitude » plutôt que d'échantillonner à partir de la distribution a posteriori résulte en une approximation appariée aux moments qui effondre la structure discrète inhérente du modèle de langage. Cette approche échoue à exploiter l'ensemble des informations prédictives a posteriori disponibles dans la sortie du FLM.

Méthodologie : Échantillonnage par pont conditionné marginalement (MCB)

Les auteurs proposent l'échantillonnage par pont conditionné marginalement (MCB), un échantillonneur prédictif a posteriori sans entraînement, conçu spécifiquement pour la structure des FLM.

Mécanisme central

Au lieu d'établir un pont vers la moyenne conditionnelle, l'échantillonneur MCB effectue une transition inverse en deux étapes à chaque pas de temps discret $k$ (passant du niveau de bruit $u_k$ à $u_{k+1}$) :

1. Échantillonnage d'une extrémité propre : L'échantillonneur évalue les marginales a posteriori token par token du FLM à l'état continu actuel $y_{t_k}$. Il échantillonne ensuite indépendamment un token valide one-hot $x^I_{0,\ell}$ pour chaque position $\ell$ à partir du posterior factorisé :
   $$q^I_{u_k}(x_0 | y) = \prod_{\ell=1}^L q_{u_k, \ell}(x_{0,\ell} | y)$$
   Cela produit une extrémité propre échantillonnée $x^I_0$ qui est une séquence one-hot valide.

2. Échantillonnage de l'état suivant via un pont analytique : L'échantillonneur tire ensuite le prochain état continu $y_{t_{k+1}}$ à partir du pont analytique d'Ornstein-Uhlenbeck (OU) conditionné à la fois par l'état actuel et par l'extrémité échantillonnée $x^I_0$. Le noyau de transition est :
   $$y_{t_{k+1}} | y_{t_k}, x^I_0 \sim \mathcal{N}\left( \frac{\sinh(\gamma_k)}{\sinh(u_k)}x^I_0 + \frac{\sinh(u_k - \gamma_k)}{\sinh(u_k)}y_{t_k}, \frac{2\sinh(\gamma_k)\sinh(u_k - \gamma_k)}{\sinh(u_k)}I_D \right)$$
   où $\gamma_k$ est la taille du pas de temps.

Intégration avec les contrôles de décodage

Puisque l'extrémité est échantillonnée à partir de distributions catégorielles explicites, le cadre MCB prend naturellement en charge les contrôles de décodage non autoregressifs standards sans réentraînement :

• Mise à l'échelle de température : Appliquée directement aux logits du posterior d'extrémité factorisé pour contrôler la netteté.
• Troncature du noyau (Nucleus Truncation) : Appliquée aux probabilités des tokens échantillonnés pour tronquer les queues de faible probabilité.

Analyse théorique et contributions

L'article fournit une justification théorique rigoureuse du MCB, séparant les erreurs d'approximation statistique des erreurs de procédure d'échantillonnage.

1. Optimalité prédictive a posteriori : Les auteurs prouvent que le MCB est l'échantillonneur prédictif a posteriori exact induit par le posterior d'extrémité factorisé représenté par le FLM. Bien qu'il partage la même moyenne conditionnelle que le pont DDPM standard, le MCB préserve la variance de l'extrémité, alors que le DDPM l'effondre.
2. Quantification de l'erreur de factorisation : Le coût du remplacement du vrai posterior joint propre par le produit des marginales de tokens est montré comme étant exactement l'information mutuelle conditionnelle entre les positions de tokens. L'erreur dans le noyau de pont à un pas est bornée par cette information mutuelle.
3. Amélioration de la divergence KL dans l'espace des trajectoires : En utilisant une comparaison d'espace de trajectoires de Girsanov, les auteurs prouvent que, sous des marginales a posteriori exactes, l'échantillonneur MCB engendre une divergence KL dans l'espace des trajectoires par rapport au processus inverse vrai qui est pas plus grande que celle de l'échantillonneur standard à moyenne conditionnelle (DDPM).
   • L'amélioration est stricte chaque fois que les observations intermédiaires du pont coordonnée par coordonnée révèlent une information supplémentaire sur le token propre au-delà de l'estimation ponctuelle de grille figée.
   • Cela découle d'un argument de projection $L_2$ : conditionner sur l'état intermédiaire $X_{t,\ell}$ (via l'extrémité échantillonnée) fournit un prédicteur strictement meilleur pour $X_{0,\ell}$ que de conditionner uniquement sur le point de grille plus grossier $X_{t_k}$.

Résultats expérimentaux

Les auteurs ont évalué le MCB sur des checkpoints de FLM préentraînés entraînés sur les ensembles de données LM1B et OWT, en maintenant les modèles fixes pour isoler les effets de la procédure d'échantillonnage.

• Compromis Qualité-Diversité : Sur LM1B, le MCB a systématiquement surpassé l'échantillonneur ODE standard (qui utilise la moyenne conditionnelle).
  • Perplexité générative (PPL) : Le MCB a atteint une perplexité générative plus faible, en particulier lorsque le nombre d'étapes de discrétisation augmentait. Des températures plus basses ($\tau=0.7$) ont produit la PPL la plus faible parmi les variantes MCB.
  • Entropie : Alors que l'échantillonneur ODE standard souffrait de générations dégénérées à faible entropie à de faibles nombres d'étapes (s'effondrant vers un texte répétitif), le MCB a maintenu une entropie stable plus proche de la distribution des données.
  • Compromis : Le MCB a fourni une courbe de compromis lisse où la mise à l'échelle de température contrôlait l'équilibre entre la qualité (PPL) et la diversité (entropie), évitant l'effondrement observé dans la ligne de base.
• Efficacité : La méthode nécessite le même nombre d'évaluations de modèle que l'échantillonnage standard et ne nécessite pas d'entraînement.

Signification et affirmations

L'article affirme que la signification principale de ce travail réside dans l'interprétation correcte de la sortie des modèles de langage par flot. Les auteurs soutiennent que les sorties des FLM ne doivent pas être traitées uniquement comme des moyennes conditionnelles à valeurs dans un simplexe pour l'établissement de ponts, mais plutôt comme des distributions directes sur les extrémités de pont one-hot propres.

En exploitant cette structure discrète, le MCB :

1. Améliore la qualité de l'échantillonnage : Il atteint un meilleur compromis qualité-diversité que les échantillonneurs standards sans réentraînement.
2. Permet un contrôle flexible : Il fournit une interface principée et non autoregressive pour les contrôles de décodage au niveau du token (température, échantillonnage par noyau) qui sont natifs au cadre prédictif a posteriori.
3. Garantie théorique : Il offre une réduction prouvable de l'erreur KL dans l'espace des trajectoires par rapport à l'approche à moyenne conditionnelle, fondée sur l'avantage informationnel du conditionnement sur les états intermédiaires.

Les auteurs reconnaissent une limitation : l'injection de stochasticité discrète à chaque étape crée un processus de branchement non lisse, ce qui peut poser des défis pour les techniques d'accélération de l'inférence telles que la distillation de cartes de flot qui reposent sur des trajectoires lisses et déterministes.