From Next Token Prediction to (STRIPS) World Models

Cette étude démontre que la prédiction du prochain token peut générer des modèles de monde capables de planification en comparant un Transformer symbolique (STRIPS) et un Transformer standard, révélant que ce dernier, notamment avec une attention « stick-breaking », surpasse le premier en optimisation et en généralisation tout en permettant l'extraction de modèles STRIPS fonctionnels.

L'essentiel

Voici une explication simple et imagée de cette recherche, comme si nous en discutions autour d'un café.

Le Grand Défi : Apprendre à un ordinateur à "comprendre" le monde

Imaginez que vous essayez d'enseigner à un enfant comment fonctionne un jeu de construction (comme des Lego). Vous ne lui donnez pas le manuel d'instructions. Vous lui montrez simplement des vidéos de quelqu'un qui construit et qui déconstruit des tours.

• Si l'enfant voit quelqu'un poser une brique rouge sur une brique bleue, il apprend que c'est possible.
• Si l'enfant voit quelqu'un essayer de poser une brique rouge sous une brique bleue (ce qui est impossible), il doit comprendre que c'est une erreur.

L'objectif de cette recherche est de savoir si les Intelligences Artificielles modernes (les grands modèles de langage comme ceux qui écrivent des poèmes ou répondent à vos questions) peuvent apprendre ces "règles du jeu" (ce qu'on appelle un modèle du monde) simplement en regardant des séquences d'actions, sans qu'on leur explique la physique des Lego.

Le problème, c'est que souvent, ces IA sont de superbes "parrots" (perroquets) : elles répètent ce qu'elles ont vu, mais elles ne comprennent pas pourquoi ça marche. Si on leur demande de planifier une nouvelle construction qu'elles n'ont jamais vue, elles échouent.

Les deux nouvelles méthodes proposées

Les chercheurs ont créé deux types d'architectures (deux façons de construire le cerveau de l'IA) pour résoudre ce problème dans un monde symbolique très précis (le langage STRIPS, utilisé en robotique).

1. Le "Transformer STRIPS" : L'architecte rigide

Imaginez un architecte très strict qui a un plan bleu imprimé dans sa tête.

• Comment ça marche : Il est programmé avec des règles très précises. Pour chaque pièce de Lego (chaque "atome"), il a une petite boîte dédiée qui dit : "Cette pièce peut être posée seulement si celle-ci est là".
• Le problème : C'est comme si on forçait l'enfant à apprendre par cœur des milliers de règles complexes. C'est très difficile à optimiser (à entraîner). Il faut beaucoup de vidéos (données) pour qu'il ne se trompe pas, et parfois, il reste bloqué. Il est très précis quand il comprend, mais il est lent et difficile à dresser.

2. Le "Transformer Stick-Breaking" : L'explorateur flexible

Imaginez maintenant un explorateur curieux qui n'a pas de plan préétabli, mais qui utilise une astuce magique appelée "attention stick-breaking" (l'attention à la cassure de bâton).

• L'analogie du bâton : Imaginez que vous avez un bâton de longueur 1. Vous devez le couper en plusieurs morceaux pour donner des parts à différentes actions passées.
  • Si une action récente a eu un impact fort (elle a cassé un mur), vous lui donnez la plus grosse part du bâton.
  • Si une action lointaine a eu un impact, elle ne reçoit qu'une miette.
  • L'astuce "stick-breaking" permet à l'IA de se concentrer uniquement sur l'action la plus récente et la plus importante qui a changé l'état des choses, en ignorant le bruit des actions anciennes.
• Le résultat : Cet explorateur apprend beaucoup plus vite, avec moins de données, et il est très doué pour généraliser. Il comprend la logique du jeu même si on lui présente des situations qu'il n'a jamais vues.

La Magie : De la prédiction à la planification

Le vrai tour de force de l'article, c'est ce qui se passe après l'entraînement.

1. L'entraînement : L'IA regarde des milliers de vidéos d'actions (positives et négatives) et apprend à dire "Oui, cette action est possible" ou "Non, c'est impossible".
2. L'extraction : Une fois entraînée, les chercheurs "déballent" la boîte noire de l'IA. Ils extraient les règles qu'elle a apprises pour créer un modèle symbolique (un vrai manuel d'instructions écrit en langage humain).
3. La Planification : Ils prennent ce manuel d'instructions et le donnent à un planificateur classique (un robot très fort en logique mais nul en apprentissage).
   • Résultat : L'IA a appris les règles, et le robot a utilisé ces règles pour résoudre des problèmes complexes, même dans des situations totalement nouvelles (des millions de combinaisons de départ différentes).

Les résultats clés (en images)

• Les modèles classiques (sans l'astuce "stick-breaking") : Ils sont excellents pour répéter ce qu'ils ont vu (comme un perroquet). Mais si on leur demande de jouer avec une séquence d'actions plus longue que dans l'entraînement, ils paniquent et oublient tout. C'est comme un enfant qui sait compter jusqu'à 10, mais qui s'arrête net à 11.
• Le modèle "Stick-Breaking" : Il comprend la logique profonde. Même si on lui donne une séquence d'actions très longue, il continue de bien jouer. Il a appris le principe, pas juste la séquence.
• Le modèle "STRIPS" (l'architecte rigide) : Il finit par comprendre, mais il faut beaucoup plus de temps et de données pour y arriver. De plus, il est parfois capricieux (il ne converge pas toujours vers la solution parfaite).

En résumé

Cette recherche prouve que l'on peut transformer un simple "prédicteur de texte" (qui devine le mot suivant) en un véritable modèle du monde capable de planifier.

C'est comme si on apprenait à un ordinateur à jouer aux échecs non pas en lui donnant les règles, mais en lui faisant regarder des milliers de parties. Grâce à la méthode "Stick-Breaking", l'IA ne se contente pas de mémoriser les coups ; elle comprend la dynamique du jeu et peut ensuite inventer de nouvelles stratégies pour gagner, même contre des adversaires qu'elle n'a jamais vus.

C'est une étape importante pour créer des IA qui ne se contentent pas de parler, mais qui comprennent et agissent dans le monde réel.

Résumé technique

Voici un résumé technique détaillé de l'article "From Next Token Prediction to (STRIPS) World Models" de Carlos Núñez-Molina, Vicenç Gómez et Hector Geffner.

1. Problématique et Contexte

La question centrale de ce travail est de savoir si la prédiction du prochain token (next-token prediction), telle qu'elle est effectuée par les grands modèles de langage (LLM) et les architectures Transformer, permet d'acquérir de véritables modèles du monde capables de soutenir la planification (planning).

Bien que des études antérieures aient montré que les Transformers pouvaient apprendre des représentations d'états latents dans des environnements contrôlés (comme les échecs), ces représentations n'ont souvent pas été jugées suffisantes pour une planification robuste. L'objectif est de déterminer si l'apprentissage par prédiction de séquences d'actions peut conduire à l'extraction de modèles symboliques exacts (modèles STRIPS) qui permettent de résoudre des problèmes de planification sur un nombre exponentiel d'états initiaux et d'objectifs non vus lors de l'entraînement.

Le cadre d'étude est celui des modèles d'actions STRIPS propositionnels, où les états sont définis par des ensembles d'atomes booléens et les actions par des listes de préconditions, d'effets ajoutés et d'effets supprimés.

2. Méthodologie et Architectures Proposées

Les auteurs introduisent deux architectures Transformer spécialisées pour apprendre des modèles STRIPS à partir de traces d'actions (séquences d'actions applicables et non applicables) sans observation directe des états (sauf via des actions de configuration spéciales).

A. Le Transformer STRIPS (Modèle Symbolique Aligné)

Cette architecture est conçue pour être symboliquement alignée avec la structure théorique des domaines STRIPS, s'appuyant sur des résultats liant les Transformers à attention dure (hard-attention) et le langage formel B-RASP (qui capture les langages sans étoile).

• Mécanisme : Chaque tête d'attention correspond à un atome du domaine. Le modèle utilise une attention dure masquée pour identifier, pour chaque atome et chaque position de la trace, l'action précédente la plus récente qui affecte cet atome.
• Inductive Bias : Elle intègre explicitement la logique des préconditions et des effets. Les paramètres du modèle encodent directement les relations précondition/effet.
• Extraction : Le modèle STRIPS est extrait directement en binarisant les paramètres appris ($\theta$).

B. Le Transformer Stick-Breaking (SB) (Modèle Standard)

C'est une architecture Transformer standard de type "decoder", mais avec deux modifications clés pour imiter l'attention dure tout en restant différentiable :

• Attention Stick-Breaking : Remplace l'attention Softmax par un mécanisme d'attention basé sur le "stick-breaking" (Tan et al., 2025). Ce mécanisme approxime l'attention dure en se concentrant sur le prédécesseur le plus récent ayant un score élevé, préservant ainsi le biais de récence nécessaire pour suivre l'état du monde.
• Absence de structure symbolique explicite : Contrairement au Transformer STRIPS, il n'y a pas d'alignement direct entre les têtes d'attention et les atomes. Le modèle apprend des représentations latentes.
• Extraction du modèle : Pour extraire un modèle STRIPS, une phase de "sondage d'état" (state probing) est nécessaire. Le modèle prédit l'applicabilité d'actions de test (test-p) pour reconstruire l'état à chaque étape de la trace, permettant ensuite de déduire les préconditions et effets par consensus majoritaire.

C. Apprentissage et Configuration

• Tâche : Prédire si une action donnée est applicable (0) ou non (1) étant donné la séquence d'actions précédentes.
• Données : Des traces positives (séquences valides) et négatives (séquences contenant une action inapplicable).
• Encodage des états : Pour intégrer l'information sur l'état initial, les auteurs utilisent des actions de configuration (init-p, test-p) qui ne sont pas prédites mais servent à encoder l'état initial et à vérifier l'état final.
• Fonction de perte : Utilisation de la Focal Loss pour gérer le déséquilibre entre actions applicables et non applicables.

3. Résultats Clés

Les expériences ont été menées sur cinq domaines classiques de planification (Blocksworld, Ferry, Npuzzle, Maze, Logistics) avec des tailles de problèmes variables.

1. Performance de Généralisation (Longueur de séquence) :

   • Les Transformers standards (avec attention Softmax et encodages positionnels) obtiennent une haute précision sur les traces d'entraînement mais échouent à généraliser aux traces plus longues (problème de généralisation compositionnelle).
   • Le Transformer SB (avec attention Stick-Breaking) atteint une précision quasi-parfaite sur les traces d'entraînement et généralise exceptionnellement bien aux traces de test beaucoup plus longues.
   • Le Transformer STRIPS généralise également bien, mais nécessite des ensembles de données plus importants pour converger et est plus difficile à optimiser.

2. Extraction et Planification :

   • Les modèles STRIPS extraits des deux architectures permettent d'utiliser des planificateurs STRIPS hors de la boîte (off-the-shelf, comme Mimir avec l'heuristique FF) pour résoudre des problèmes.
   • Généralisation Combinatoire : Les modèles réussissent à planifier sur un nombre exponentiel d'états initiaux et d'objectifs jamais vus pendant l'entraînement.
   • Précision de Planification : Le Transformer SB extrait des modèles qui obtiennent souvent 100% de précision de planification, surpassant le Transformer STRIPS qui montre une plus grande variabilité et des taux d'échec légèrement plus élevés dans certains cas complexes.

3. Comparaison des Architectures :

   • Bien que le Transformer STRIPS intègre un biais inductif symbolique fort, le Transformer SB s'avère plus facile à entraîner, plus robuste et plus performant en termes de généralisation.
   • L'attention Stick-Breaking s'avère être le composant critique pour permettre la généralisation à long terme, là où l'attention Softmax échoue.

4. Contributions et Signification

• Preuve de concept pour les modèles du monde : L'article démontre que la simple prédiction de token, lorsqu'elle est couplée à une architecture appropriée (attention Stick-Breaking), suffit à apprendre des modèles du monde symboliques exacts capables de soutenir la planification.
• Supériorité de l'attention Stick-Breaking : Il établit que pour les tâches nécessitant un suivi d'état à long terme (comme la planification), l'attention dure (ou son approximation différentiable via le stick-breaking) est supérieure à l'attention Softmax standard, qui souffre de problèmes de généralisation de longueur.
• Pont entre Deep Learning et Planification Symbolique : La méthode permet d'extraire des modèles STRIPS interprétables et utilisables par des planificateurs classiques à partir de données brutes d'actions, comblant le fossé entre les approches neuronales et symboliques.
• Généralisation Compositionnelle : Les modèles appris ne mémorisent pas simplement des séquences, mais apprennent la dynamique sous-jacente du domaine, permettant de résoudre des problèmes combinatoires bien au-delà de la distribution d'entraînement.

En conclusion, cette étude suggère que les Transformers, lorsqu'ils sont correctement conçus (notamment via l'attention Stick-Breaking), peuvent acquérir une compréhension profonde des dynamiques d'environnements discrets, transformant la prédiction de token en un véritable moteur de modélisation du monde pour la planification.