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

Sintesi Tecnica: Un esame approfondito del recupero del modello del mondo nei pianificatori LLM sottoposti a Fine-Tuning Supervisionato

Problema
Sebbene il Fine-Tuning Supervisionato (SFT) abbia migliorato le capacità di pianificazione end-to-end dei Large Language Models (LLM), rimane poco chiaro se questi modelli apprendano realmente a rappresentare e ragionare sui problemi di pianificazione che risolvono. Nello specifico, i modelli LLM sottoposti a fine-tuning recuperano un "modello del mondo" — una comprensione interna delle transizioni di stato e della validità delle azioni — o si limitano ad apprendere sequenze di azioni valide per imitazione? La ricerca precedente suggerisce che gli LLM fatichino nella pianificazione in contesti zero-shot e possano apprendere soluzioni scorciatoia piuttosto che meccanismi di ragionamento robusti. Questo lavoro indaga se il SFT su esempi di piani validi induca il recupero del modello del mondo in domini di pianificazione classica.

Metodologia
Gli autori introducono un framework per valutare il recupero del modello del mondo utilizzando due definizioni complementari:

1. Rappresentazioni Interne: Se l'LLM apprenda rappresentazioni lineari che codificano i valori di verità dei predicati di stato e la validità delle azioni.
2. Capacità Generative: Se l'LLM assegni un'alta probabilità a sequenze di azioni valide e una bassa probabilità a quelle non valide.

Lo studio utilizza il modello gemma2-9b-instruct, sottoposto a fine-tuning per generare piani end-to-end (chiamati "Plansformers") in due ambienti di pianificazione classica: Blocksworld e Logistics. Gli autori sperimentano tre diverse distribuzioni di dati di addestramento per determinare il loro impatto sul recupero:

• Piani Ottimali: Generati da un risolutore euristico (Fast Downward).
• Piani Random Walk: Generati campionando azioni valide casuali fino a un numero massimo di passi.
• Piani Random Walk++: Una versione modificata del random walk che invalida le azioni che portano a stati precedentemente visitati per prevenire cicli e migliorare la copertura dello spazio degli stati.

Inoltre, gli autori esplorano lo State Chain-of-Thought (State-CoT), dove gli stati intermedi sono intercalati con le azioni nella sequenza di addestramento per migliorare la groundedness (radicamento).

Per misurare il recupero, gli autori impiegano:

• Sonde Lineari (Linear Probes): Modelli di regressione logistica addestrati sulle attivazioni degli strati nascosti per predire la validità delle azioni e i valori di verità dei predicati di stato.
• Classificazione Generativa: Una metrica che valuta se le probabilità dei token dell'LLM classificano le azioni valide con una probabilità superiore rispetto alle azioni non valide quando viene fornito un piano parziale.
• Test Out-of-Distribution (OOD): Valutazione su istanze di pianificazione con complessità aumentata (es. più blocchi) rispetto a quelle viste durante l'addestramento.

Risultati Chiave

1. Codifica Lineare della Validità: Il fine-tuning supervisionato su sequenze di azioni valide consente agli LLM di codificare linearmente la validità delle azioni nelle loro rappresentazioni interne. I punteggi F1 delle sonde raggiungono livelli quasi perfetti (~1.0) negli strati intermedi e superiori in tutte le distribuzioni di addestramento.
2. Rappresentazione Parziale dello Stato: I modelli apprendono a codificare linearmente i valori di verità di alcuni predicati di stato (es. arm-empty, at ?pkg ?loc), ma le prestazioni variano significativamente in base alla complessità del predicato. I predicati più semplici con meno argomenti sono codificati più accuratamente rispetto a quelli complessi.
3. Divergenza tra Metriche Interne e Generative: Un risultato critico è il disallineamento tra le rappresentazioni interne e gli output generativi. I modelli sottoposti a fine-tuning su piani ottimali possiedono rappresentazioni interne di alta qualità della validità delle azioni (alti punteggi delle sonde) ma falliscono nel utilizzare le probabilità dei token per distinguere le azioni valide da quelle non valide durante la generazione. Al contrario, i modelli sottoposti a fine-tuning su dati random walk mostrano una classificazione generativa superiore (classificando le azioni valide con priorità rispetto alle non valide) ma faticano con la generalizzazione OOD.
4. Impatto della Copertura dello Spazio degli Stati: I dati Random Walk++, che garantiscono una copertura più ampia dello spazio degli stati eliminando i cicli, producono il recupero più robusto del modello del mondo. Questo approccio ottiene prestazioni superiori sia nella classificazione generativa in-distribution che nei compiti di pianificazione OOD rispetto ai random walk standard o ai piani ottimali.
5. Generalizzazione OOD: Mentre le rappresentazioni interne della validità delle azioni rimangono robuste ai cambiamenti di distribuzione (dati OOD), le prestazioni generative crollano per i modelli addestrati su random walk standard. I modelli addestrati su Random Walk++ e Piani Ottimali mostrano una migliore resilienza agli shift OOD, sebbene le prestazioni degradino comunque all'aumentare della complessità del problema.

Significato e Rivendicazioni
Il documento sostiene di fornire una "ricetta" per applicare tecniche di interpretabilità ai pianificatori LLM, offrendo una visione olistica che combina l'analisi della rappresentazione interna con la valutazione generativa. Gli autori concludono che:

• Il SFT su piani validi è sufficiente affinché gli LLM apprendano rappresentazioni interne che codificano linearmente la validità delle azioni e certi predicati di stato, suggerendo una forma di recupero del modello del mondo.
• Tuttavia, la presenza di rappresentazioni interne non garantisce che il modello le utilizzi efficacemente nelle sue probabilità generative; la distribuzione di addestramento (specificamente la copertura dello spazio degli stati) influenza pesantemente questo allineamento.
• Le scoperte sono "ottimistiche" riguardo alla fattibilità del fine-tuning per migliorare la pianificazione, dimostrando che gli LLM possono parzialmente recuperare i modelli del mondo della pianificazione. Tuttavia, il lavoro evidenzia anche dei limiti, notando che il legame causale tra queste rappresentazioni e il processo di pianificazione del modello rimane non dimostrato, e che la generalizzazione a un numero significativamente maggiore di parametri o alla pianificazione in linguaggio naturale rimane una questione aperta.

Gli autori dichiarano esplicitamente di non pretendere di aver risolto la pianificazione OOD o di aver dimostrato il ragionamento causale, ma di fornire prove empiriche che i modelli LLM sottoposti a fine-tuning sviluppino rappresentazioni interne strutturate della dinamica di pianificazione che sono distinte, pur essendo correlate, alle loro capacità generative.