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

Resumo Técnico: Um Olhar Atento sobre a Recuperação de Modelos de Mundo em Planejadores de LLM com Ajuste Fino Supervisionado

Definição do Problema
Embora o Ajuste Fino Supervisionado (SFT) tenha melhorado as capacidades de planejamento end-to-end de Grandes Modelos de Linguagem (LLMs), ainda não está claro se esses modelos realmente aprendem a representar e raciocinar sobre os problemas de planejamento que resolvem. Especificamente, será que os LLMs ajustados via SFT recuperam um "modelo de mundo" — uma compreensão interna de transições de estado e validade de ações — ou eles apenas aprendem a mimetizar sequências de ações válidas? Pesquisas anteriores sugerem que os LLMs têm dificuldade com planejamento em configurações zero-shot e podem aprender soluções de atalho em vez de mecanismos de raciocínio robustos. Este trabalho investiga se o SFT em exemplos de planos válidos induz a recuperação de modelos de mundo em domínios de planejamento clássico.

Metodologia
Os autores introduzem um framework para avaliar a recuperação de modelos de mundo usando duas definições complementares:

1. Representações Internas: Se o LLM aprende representações lineares que codificam os valores de verdade dos predicados de estado e a validade das ações.
2. Capacidades Generativas: Se o LLM atribui alta probabilidade a sequências de ações válidas e baixa probabilidade a sequências inválidas.

O estudo utiliza o modelo gemma2-9b-instruct, ajustado para gerar planos end-to-end (chamados de "Plansformers") em dois ambientes de planejamento clássico: Blocksworld e Logistics. Os autores experimentam com três distribuições distintas de dados de treinamento para determinar seu impacto na recuperação:

• Planos Ótimos: Gerados por um solver heurístico (Fast Downward).
• Planos de Caminhada Aleatória (Random Walk): Gerados amostrando ações válidas aleatórias até atingir uma contagem máxima de passos.
• Planos de Caminhada Aleatória++ (Random Walk++): Uma caminhada aleatória modificada que invalida ações que levam a estados previamente visitados para evitar ciclos e melhorar a cobertura do espaço de estados.

Adicionalmente, os autores exploram o Cadeia de Pensamento de Estado (State Chain-of-Thought ou State-CoT), onde estados intermediários são intercalados com ações na sequência de treinamento para aumentar o embasamento (groundedness).

Para medir a recuperação, os autores utilizam:

• Sondas Lineares (Linear Probes): Modelos de regressão logística treinados nas ativações das camadas ocultas para prever a validade da ação e os valores de verdade dos predicados de estado.
• Classificação Generativa: Uma métrica que avalia se as probabilidades de tokens do modelo classificam ações válidas com maior prioridade do que ações inválidas ao serem estimuladas com planos parciais.
• Teste Fora da Distribuição (Out-of-Distribution - OOD): Avaliação em instâncias de planejamento com aumento de complexidade (ex: mais blocos) em relação ao que foi visto durante o treinamento.

Resultados Principais

1. Codificação Linear de Validade: O ajuste fino supervisionado em sequências de ações válidas permite que os LLMs codifiquem linearmente a validade da ação em suas representações internas. As pontuações F1 das sondas alcançam níveis próximos da perfeição (~1.0) em camadas intermediárias e superiores em todas as distribuições de treinamento.
2. Representação Parcial de Estado: Os modelos aprendem a codificar linearmente os valores de verdade de alguns predicados de estado (ex: arm-empty, at ?pkg ?loc), mas o desempenho varia significativamente conforme a complexidade do predicado. Predicados mais simples com menos argumentos são codificados com maior precisção do que os complexos.
3. Divergência entre Métricas Internas e Generativas: Uma descoberta crítica é o descompasso entre as representações internas e as saídas generativas. Modelos ajustados com planos ótimos possuem representações internas de alta qualidade da validade da ação (altas pontuações de sonda), mas falham em usar as probabilidades de tokens para distinguir ações válidas de inválidas durante a geração. Por outro lado, modelos ajustados com dados de caminhada aleatória mostram uma classificação generativa superior (ranqueando ações válidas acima de inválidas), mas enfrentam dificuldades com a generalização OOD.
4. Impacto da Cobertura do Espaço de Estados: Dados de Random Walk++, que garantem uma cobertura mais ampla do espaço de estados ao eliminar ciclos, resultam na recuperação mais robusta do modelo de mundo. Este método alcança desempenho superior tanto na classificação generativa in-distribution quanto em tarefas de planejamento OOD em comparação com caminhadas aleatórias padrão ou planos ótimos.
5. Generalização OOD: Embora as representações internas da validade da ação permaneçam robustas a mudanças de distribuição (dados OOD), o desempenho generativo colapsa para modelos treinados com caminhadas aleatórias padrão. Modelos treinados com Random Walk++ e Planos Ótimos mostram maior resiliência a mudanças OOD, embora o desempenho ainda degrade conforme a complexidade do problema aumenta.

Significância e Alegações
O artigo afirma fornecer uma "receita" para aplicar técnicas de interpretabilidade em planejadores de LLM, oferecendo uma visão holística que combina a análise de representação interna com a avaliação generativa. Os autores concluem que:

• O SFT em planos válidos é suficiente para que os LLMs aprendam representações internas que codificam linearmente a valididade da ação e certos predicados de estado, sugerindo uma forma de recuperação de modelo de mundo.
• No entanto, a presença de representações internas não garante que o modelo as utilizará efetivamente em suas probabilidades generativas; a distribuição de treinamento (especificamente a cobertura do espaço de estados) influencia fortemente esse alinhamento.
• As descobertas são "otimistas" quanto à viabilidade do ajuste fino para melhorar o planejamento, demonstrando que os LLMs podem recuperar parcialmente modelos de mundo de planejamento. Contudo, o trabalho também destaca limitações, observando que o elo causal entre essas representações e o processo de planejamento do modelo permanece não comprovado, e que a generalização para quantidades significativamente maiores de parâmetros ou planejamento em linguagem natural continua sendo uma questão aberta.

Os autores declaram explicitamente que não pretendem ter resolvido o planejamento OOD ou provado o raciocínio causal, mas sim fornecer evidências empíricas de que os LLMs ajustados via SFT desenvolvem representações internas estruturadas da dinâmica de planejamento que são distintas, porém relacionadas, às suas capacidades generativas.