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

Technische Samenvatting: Een nauwe blik op de herwinning van wereldmodellen in Supervised Fine-Tuned LLM-planners

Probleemstelling
Hoewel Supervised Fine-Tuning (SFT) de end-to-end planningscapaciteiten van Large Language Models (LLM's) heeft verbeterd, blijft het onduidelijk of deze modellen daadwerkelijk leren om de planningsproblemen die ze oplossen te representeren en erover te redeneren. In het bijzonder: leren gefinetunede LLM's een "wereldmodel" terug — een intern begrip van staatsovergangen en de geldigheid van acties — of leren ze slechts geldige actiesequenties na te bootsen? Voorlopig onderzoek suggereert dat LLM's moeite hebben met plannen in zero-shot settings en mogelijk shortcut-oplossingen leren in plaats van robuuste redeneermechanismen. Dit werk onderzoekt of SFT op voorbeelden van geldige plannen leidt tot de herwinning van een wereldmodel in klassieke planningsdomeinen.

Methodologie
De auteurs introduceren een framework om de herwinning van wereldmodellen te evalueren met behulp van twee complementaire definities:

1. Interne Representaties: Of de LLM lineaire representaties leert die de waarheidswaarden van staatspredicaten en de geldigheid van acties coderen.
2. Generatieve Capaciteiten: Of de LLM een hoge waarschijnlijkheid toekent aan geldige actiesequenties en een lage waarschijnlijkheid aan ongeldige sequenties.

De studie maakt gebruik van het gemma2-9b-instruct model, dat is gefinetuned om end-to-end plannen te genereren (genoemd "Plansformers") in twee klassieke planningsomgevingen: Blocksworld en Logistics. De auteurs experimenteren met drie verschillende trainingsdatadistributies om hun impact op herwinning te bepalen:

• Optimale Plannen: Gegenereerd door een heuristische solver (Fast Downward).
• Random Walk Plannen: Gegenereerd door willekeurige geldige acties te samplen tot een maximaal aantal stappen is bereikt.
• Random Walk++ Plannen: Een aangepaste random walk die acties die leiden naar eerder bezochte staten ongeldig maakt om cycli te voorkomen en de state space coverage te verbeteren.

Daarnaast verkennen de auteurs State Chain-of-Thought (State-CoT), waarbij tussenliggende staten worden geïntercaleerd met acties in de trainingssequentie om de verankering (groundedness) te verbeteren.

Om herwinning te meten, gebruiken de auteurs:

• Lineaire Probes: Logistische regressiemodellen getraind op verborgen laagactivaties om de geldigheid van acties en de waarheidswaarden van staatspredicaten te voorspellen.
• Generatieve Classificatie: Een metriek die evalueert of de tokenwaarschijnlijkheden van de modellen geldige acties hoger rangschikken dan ongeldige acties wanneer ze worden geprompt met gedeeltelijke plannen.
• Out-of-Distribution (OOD) Testing: Evaluatie op planningsinstanties met een verhoogde complexiteit (bijv. meer blokken) dan tijdens de training gezien.

Belangrijkste Resultaten

1. Lineaire Codering van Geldigheid: Supervised fine-tuning op geldige actiesequenties stelt LLM's in staat om de geldigheid van acties lineair te coderen in hun interne representaties. Probe F1-scores bereiken bijna perfecte niveaus (~1.0) in intermediaire en bovenste lagen over alle trainingsdistributies heen.
2. Partiële Staat Representatie: Modellen leren om de waarheidswaarden van sommige staatspredicaten (bijv. arm-empty, at ?pkg ?loc) lineair te coderen, maar de prestaties variëren aanzienlijk per predicaatcomplexiteit. Eenvoudigere predicaten met minder argumenten worden nauwkeuriger gecodeerd dan complexe predicaten.
3. Divergentie tussen Interne en Generatieve Metrieken: Een cruciaal bevinding is de mismatch tussen interne representaties en generatieve outputs. Modellen die zijn gefinetuned op optimale plannen bezitten hoogwaardige interne representaties van de geldigheid van acties (hoge probe scores), maar falen in het gebruiken van tokenwaarschijnlijkheden om geldige van ongeldige acties te onderscheiden tijdens generatie. Omgekeeld vertonen modellen die gefinetuned zijn op random walk data superieure generatieve classificatie (het hoger rangschikken van geldige acties dan ongeldige acties), maar worstelen ze met OOD-generalisatie.
4. Impact van State Space Coverage: Random Walk++ data, die een bredere state space coverage garandeert door cycli te elimineren, levert de meest robuuste wereldmodel herwinning op. Het bereikt superieure prestaties in zowel in-distribution generatieve classificatie als OOD plannings-taken vergeleken met standaard random walks of optimale plannen.
5. OOD Generalisatie: Hoewel interne representaties van de geldigheid van acties robuust blijven tegen distributieverschuivingen (OOD data), stort de generatieve prestatie in voor modellen die getraind zijn op standaard random walks. Modellen getraind op Random Walk++ en Optimale plannen vertonen een betere veerkracht tegen OOD-verschuivingen, hoewel de prestaties nog steeds afnemen naarmate de complexiteit van het probleem toeneemt.

Betekenis en Claims
Het artikel beweert een "recept" te bieden voor het toepassen van interpreteerbaarheidstechnieken op plannende LLM's, door een holistisch overzicht te bieden dat interne representatieanalyse combineert met generatieve evaluatie. De auteurs concluderen dat:

• SFT op geldige plannen voldoende is voor LLM's om interne representaties te leren die de geldigheid van acties en bepaalde staatspredicaten lineair coderen, wat wijst op een vorm van wereldmodel herwinning.
• De aanwezigheid van interne representaties echter niet garandeert dat het model deze effectief zal gebruiken in zijn generatieve waarschijnlijkheden; de trainingsdistributie (specifiek de state space coverage) beïnvloedt deze uitlijning zwaar.
• De bevindingen zijn "optimistisch" met betrekking tot de levensvatbaarheid van fine-tuning om plannen te verbeteren, en tonen aan dat LLM's gedeeltelijk planningswereldmodellen kunnen herwinnen. Echter, het werk benadrukt ook beperkingen, waarbij wordt opgemerkt dat de causale link tussen deze representaties en het planningsproces van het model nog niet bewezen is, en dat generalisatie naar aanzienlijk grotere parameteraantallen of natuurlijke taalplanning een open vraag blijft.

De auteurs verklaren expliciet dat zij niet beweren OOD-planning te hebben opgelost of causaal redeneren te hebben bewezen, maar eerder empirisch bewijs leveren dat gefinetunede LLM's gestructureerde interne representaties van planningsdynamiek ontwikkelen die verschillend zijn van, maar gerelateerd aan, hun generatieve capaciteiten.