Résumé technique : Les Transformers de prédiction de latent suivant apprennent des modèles du monde compacts

Énoncé du problème

Les architectures Transformer standard remplacent la récurrence par une auto-attention et un mécanisme de mémoire qui s'étend avec la longueur de la séquence. Bien que cela permette une parallélisation exceptionnelle et une puissance prédictive accrue, cela élimine le biais inductif inhérent à la compression de l'historique en états latents compacts dotés de règles de transition cohérentes. Par conséquent, les Transformers apprennent souvent des raccourcis complexes et spécifiques à la tâche qui s'ajustent bien aux données d'entraînement mais généralisent mal, échouant à former des « modèles du monde » ou des états de croyance cohérents nécessaires à la planification et au raisonnement à long terme. L'objectif standard de prédiction du token suivant ne garantit pas que les états internes cachés du modèle évoluent comme des statistiques suffisantes pour l'historique passé, ni qu'ils puissent prédire récursivement les observations futures.

Méthodologie : Prédiction de latent suivant (NextLat)

Les auteurs introduisent la Prédiction de latent suivant (NextLat), un cadre qui enrichit l'entraînement standard par token suivant avec des prédictions auto-supervisées dans l'espace latent. L'idée centrale consiste à entraîner le Transformer à apprendre des représentations latentes (spécifiquement, les états cachés de la dernière couche $h_t$) qui sont prédictifs du prochain état latent étant donné le prochain token.

Fondement théorique

L'article établit que si un modèle optimise à la fois la cohérence du token suivant (prédire $X_{t+1}$ à partir de $h_t$) et la cohérence de transition (prédire $h_{t+1}$ à partir de $h_t$ et $X_{t+1}$), les états cachés $h_t$ sont théoriquement garantis de converger vers des états de croyance. Un état de croyance est défini comme une statistique suffisante de l'historique $X_{1:t}$ pour prédire le futur $X_{t+1:T}$. Contrairement à l'entraînement autorégressif standard, qui permet des recherches ad hoc de tokens passés sans compresser l'information, NextLat impose un biais inductif récursif où l'état latent doit contenir toutes les informations nécessaires pour transitionner vers l'état suivant et prédire le token suivant.

Implémentation

NextLat fonctionne en ajoutant un objectif auxiliaire léger à la boucle d'entraînement standard du Transformer :

1. Modèle de dynamique latente ($p_\psi$) : Un simple Perceptron Multicouche (MLP) est entraîné pour prédire le prochain état caché $\hat{h}_{t+1}$ étant donné l'état caché courant $h_t$ et le prochain token $X_{t+1}$.
2. Fonction de perte : La perte totale combine :
   • Perte de token suivant ($L_{next-token}$) : Entropie croisée standard pour la prédiction de token.
   • Perte d'état caché suivant ($L_{next-h}$) : Une perte Smooth L1 régressant l'état prédit $\hat{h}_{t+1}$ contre le vrai état suivant $h_{t+1}$ (avec un stop-gradient appliqué à la cible pour éviter l'effondrement représentationnel).
   • Perte de divergence KL ($L_{KL}$) : Un objectif complémentaire assurant que la distribution de tokens prédite à partir de l'état latent prédit correspond à la distribution issue de l'état latent vrai, agissant comme un signal de supervision douce.
3. Supervision multi-étapes : Le modèle peut être entraîné à dérouler récursivement la dynamique latente sur $d$ étapes, fournissant des signaux de gradient plus riches pour les dépendances à long terme.

Crucialement, NextLat laisse l'architecture du Transformer, sa procédure d'inférence et son efficacité d'entraînement parallèle inchangées. Le modèle de dynamique latente est utilisé uniquement pendant l'entraînement pour façonner les représentations ; lors de l'inférence, le modèle décode de manière autorégressive comme d'habitude.

Contributions clés

1. Garantie théorique : L'article prouve que l'optimisation pour la cohérence du token suivant et du latent suivant façonne de manière prouvée les représentations des transformers en états de croyance, une propriété non garantie par la prédiction de token seule.
2. Cadre pratique : NextLat introduit un objectif auxiliaire computationnellement efficace qui injecte un biais inductif récursif sans altérer la structure d'entraînement parallèle ou d'inférence du Transformer.
3. Décodage auto-spéculatif de longueur variable : En apprenant une dynamique latente cohérente, NextLat permet une prédiction récursive dans l'espace latent. Cela permet un décodage auto-spéculatif de longueur variable, où le modèle peut ébaucher des tokens au-delà de l'horizon d'entraînement fixe, accélérant considérablement l'inférence.

Résultats expérimentaux

Les auteurs évaluent NextLat sur quatre domaines : modélisation du monde, raisonnement, planification et modélisation du langage.

• Modélisation du monde (Taxis de Manhattan) : NextLat apprend des cartes internes nettement plus cohérentes avec le vrai modèle du monde (topologie des rues) par rapport aux bases de référence standard GPT, MTP et JTP. Il atteint une compression de séquence supérieure (0,71 contre 0,65 pour GPT) et un rang latent effectif beaucoup plus faible (52,7 contre 160,1 pour GPT), indiquant des représentations plus compactes et efficaces.
• Raisonnement (Countdown) : Sur la tâche de raisonnement mathématique, NextLat surpasse largement les bases de référence (54,8 % de précision à $d=1$ contre 33,1 % pour GPT). Il démontre une planification anticipée supérieure, réduisant les « compromis regrettables » où les modèles forcent des équations finales invalides pour correspondre à une cible.
• Planification (Graphes Path-Star) : NextLat maintient des taux de résolution proches de 100 % sur diverses topologies de graphes, alors que des bases de référence comme BST et JTP échouent sur des graphes plus grands ($G_{7,7}$). Cela suggère que NextLat évite l'apprentissage de raccourcis « Hans le Malin » courant dans la prédiction de l'espace des tokens.
• Modélisation du langage (FineWeb-Edu) :
  • Précision en aval : NextLat préserve mieux la perplexité du token suivant que MTP et JTP tout en réalisant des gains modestes en précision zero-shot en aval (59,21 % en moyenne contre 58,82 % pour GPT).
  • Vitesse d'inférence : En tirant parti du décodage auto-spéculatif de longueur variable, NextLat atteint une inférence jusqu'à 3,3 fois plus rapide par rapport à l'échantillonnage autorégressif standard. Cette accélération persiste même lors de l'ébauche au-delà de l'horizon d'entraînement, une capacité dont MTP et JTP sont dépourvus en raison de leurs horizons fixes dans l'espace des tokens.

Importance et revendications

L'article présente NextLat comme un paradigme simple mais efficace pour apprendre des représentations compactes, prédictives et généralisables dans les Transformers. Son importance principale réside dans l'unification du biais inductif récursif des modèles du monde compacts avec l'évolutivité et le parallélisme des Transformers.

Les auteurs affirment que NextLat offre une solution à la nature « myope » de la prédiction de token suivant, encourageant les modèles à former des états internes structurés et prédictifs plutôt que de s'appuyer sur des corrélations de tokens superficielles. En permettant un décodage spéculatif de longueur variable sans modifications architecturales, NextLat fournit une voie pratique vers une inférence plus rapide et un meilleur raisonnement à long terme. L'article conclut que NextLat est une étape vers des objectifs d'entraînement qui dotent les modèles autorégressifs de représentations plus simples, plus compactes et donc plus généralisables des systèmes complexes.

Les auteurs restent modestes concernant les orientations futures, notant qu'ils n'ont pas exploré des architectures complexes de dynamique latente (en utilisant uniquement des MLP simples) et que la conception optimale de l'objectif (par exemple, la nécessité de la perte KL à plus grande échelle) nécessite une étude systématique supplémentaire. Ils reconnaissent également que, bien que NextLat améliore les représentations, la complexité circuitale sous-jacente du Transformer à profondeur fixe reste bornée par la classe TC0, bien que le modèle de dynamique latente co-entraîné (RNN) semble capable de généraliser au-delà de ces limites.