A Close Look At World Model Recovery In Supervised Fine-Tuned LLM Planners

Résumé Technique : Un regard approfondi sur la récupération de modèles de monde dans les planificateurs LLM par ajustement fin supervisé

Énoncé du problème
Bien que l'ajustement fin supervisé (SFT - Supervised Fine-Tuning) ait amélioré les capacités de planification de bout en bout des grands modèles de langage (LLM), il reste incertain si ces modèles apprennent réellement à représenter et à raisonner sur les problèmes de planification qu'ils résolvent. Plus précisément, les LLM ajustés par SFT récupèrent-ils un « modèle de monde » — une compréhension interne des transitions d'état et de la validité des actions — ou se contentent-ils d'apprendre à imiter des séquences d'actions valides ? Des recherches antérieures suggèrent que les LLM éprouvent des difficultés avec la planification en mode zero-shot et peuvent apprendre des solutions de raccourci plutôt que des mécanismes de raisonnement robustes. Ce travail étudie si le SFT sur des exemples de plans valides induit la récupération d'un modèle de monde dans les domaines de la planification classique.

Méthodologie
Les auteurs introduisent un cadre pour évaluer la récupération du modèle de monde en utilisant deux définitions complémentaires :

1. Représentations internes : Si le LLM apprend des représentations linéaires qui encodent les valeurs de vérité des prédicats d'état et la validité des actions.
2. Capacités génératives : Si le LLM assigne une probabilité élevée aux séquences d'actions valides et une probabilité faible aux séquences invalides.

L'étude utilise le modèle gemma2-9b-instruct, ajusté pour générer des plans de bout en bout (appelés « Plansformers ») dans deux environnements de planification classique : Blocksworld et Logistics. Les auteurs expérimentent trois distributions de données d'entraînement distinctes pour déterminer leur impact sur la récupération :

• Plans optimaux : Générés par un solveur heuristique (Fast Downward).
• Plans de marche aléatoire (Random Walk) : Générés en échantillonnant des actions valides aléatoires jusqu'à ce qu'un nombre maximal d'étapes soit atteint.
• Plans de marche aléatoire++ (Random Walk++) : Une marche aléatoire modifiée qui invalide les actions menant à des états déjà visités afin d'éviter les cycles et d'améliorer la couverture de l'espace d'états.

De plus, les auteurs explorent la Chaîne de Pensée d'État (State-Chain-of-Thought ou State-CoT), où les états intermédiaires sont entrelacés avec les actions dans la séquence d'entraînement pour améliorer l'ancrage.

Pour mesurer la récupération, les auteurs emploient :

• Sondes linéaires (Linear Probes) : Modèles de régression logistique entraînés sur les activations des couches cachées pour prédire la validité des actions et les valeurs de vérité des prédicats d'état.
• Classification générative : Une métrique évaluant si les probabilités de jetons (tokens) du modèle classent les actions valides plus haut que les actions invalides lors de la génération.
• Tests hors distribution (OOD - Out-of-Distribution) : Évaluation sur des instances de planification avec une complexité accrue (par exemple, plus de blocs) par rapport à celles vues durant l'entraînement.

Résultats clés

1. Encodage linéaire de la validité : L'ajustement fin supervisé sur des séquences d'actions valides permet aux LLM d'encoder linéairement la validité des actions dans leurs représentations internes. Les scores F1 des sondes atteignent des niveaux proches de la perfection (~1,0) dans les couches intermédiaires et supérieures à travers toutes les distributions d'entraînement.
2. Représentation partielle de l'état : Les modèles apprennent à encoder linéairement les valeurs de vérité de certains prédicats d'état (ex: arm-empty, at ?pkg ?loc), mais les performances varient considérablement selon la complexité du prédicat. Les prédicats simples avec moins d'arguments sont encodés plus précisément que les plus complexes.
3. Divergence entre les métriques internes et génératives : Une découverte critique est le décalage entre les représentations internes et les sorties génératives. Les modèles ajustés sur des plans optimaux possèdent des représentations internes de haute qualité de la validité des actions (scores de sonde élevés) mais échouent à utiliser les probabilités de jetons pour distinguer les actions valides des actions invalides lors de la génération. À l'inverse, les modèles ajustés sur des données de marche aléatoire montrent une classification générative supérieure (classant les actions valides plus haut que les invalides) mais peinent avec la généralisation OOD.
4. Impact de la couverture de l'espace d'états : Les données Random Walk++, qui assurent une couverture plus large de l'espace d'états en éliminant les cycles, produisent la récupération la plus robuste du modèle de monde. Elles obtiennent une performance supérieure tant en classification générative en distribution qu'en tâches de planification OOD par rapport aux marches aléatoires standards ou aux plans optimaux.
5. Généralisation OOD : Bien que les représentations internes de la validité des actions restent robustes aux changements de distribution (données OOD), la performance générative s'effondre pour les modèles entraînés sur des marches aléatoires standards. Les modèles entraînés sur Random Walk++ et les plans optimaux montrent une meilleure résilience aux changements OOD, bien que les performances se dégradent toujours à mesure que la complexité du problème augmente.

Signification et affirmations
L'article affirme fournir une « recette » pour appliquer les techniques d'interprétabilité aux planificateurs LLM, offrant une vue holistique qui combine l'analyse de la représentation interne et l'évaluation générative. Les auteurs concluent que :

• Le SFT sur des plans valides est suffisant pour que les LLM apprennent des représentations internes qui encodent linéairement la validité des actions et certains prédicats d'état, suggérant une forme de récupération de modèle de monde.
• Cependant, la présence de représentations internes ne garantit pas que le modèle les utilisera efficacement dans ses probabilités génératives ; la distribution d'entraînement (spécifiquement la couverture de l'espace d'états) influence lourdement cet alignement.
• Les résultats sont « optimistes » quant à la viabilité de l'ajustement fin pour améliorer la planification, démontrant que les LLM peuvent partiellement récupérer des modèles de monde de planification. Toutefois, le travail souligne également des limites, notant que le lien causal entre ces représentations et le processus de planification du modèle reste non prouvé, et que la généralisation à des nombres de paramètres nettement plus élevés ou à la planification en langage naturel reste une question ouverte.

Les auteurs déclarent explicitement qu'ils ne prétendent pas avoir résolu la planification OOD ni prouvé le raisonnement causal, mais qu'ils fournissent des preuves empiriques que les LLM ajustés développent des représentations internes structurées de la dynamique de planification qui sont distinctes de, bien que liées à, leurs capacités génératives.