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

Resumen Técnico: Un vistazo cercano a la recuperación del modelo de mundo en planificadores de LLM con ajuste fino supervisado

Planteamiento del problema
Si bien el Ajuste Fino Supervisado (SFT, por sus siglas en inglés) ha mejorado las capacidades de planificación de extremo a extremo de los Modelos de Lenguaje Extensos (LLMs), sigue sin estar claro si estos modelos aprenden genuinamente a representar y razonar sobre los problemas de planificación que resuelven. Específicamente, ¿recuperan los LLMs ajustados mediante SFT un "modelo de mundo" —una comprensión interna de las transiciones de estado y la validez de las acciones— o simplemente aprenden a imitar secuencias de acciones válidas? Investigaciones previas sugieren que los LLMs tienen dificultades con la planificación en entornos de disparo cero (zero-shot) y pueden aprender soluciones de atajos en lugar de mecanismos de razonamiento robustos. Este trabajo investiga si el SFT sobre ejemplos de planes válidos induce la recuperación del modelo de mundo en dominios de planificación clásica.

Metodología
Los autores introducen un marco para evaluar la recuperación del modelo de mundo utilizando dos definiciones complementarias:

1. Representaciones internas: Si el LLM aprende representaciones lineales que codifican los valores de verdad de los predicados de estado y la validez de las acciones.
2. Capacidades generativas: Si el LLM asigna una alta probabilidad a las secuencias de acciones válidas y una baja probabilidad a las inválidas.

El estudio utiliza el modelo gemma2-9b-instruct, ajustado para generar planes de extremo a extremo (denominados "Plansformers") en dos entornos de planificación clásica: Blocksworld y Logistics. Los autores experimentan con tres distribuciones de datos de entrenamiento distintas para determinar su impacto en la recuperación:

• Planes óptimos: Generados por un solucionador heurístico (Fast Downward).
• Planes de caminata aleatoria (Random Walk): Generados mediante el muestreo de acciones válidas aleatorias hasta alcanzar un número máximo de pasos.
• *Planes de caminata aleatoria++ (Random Walk++):* Una caminata aleatoria modificada que invalida las acciones que conducen a estados previamente visitados para evitar ciclos y mejorar la cobertura del espacio de estados.

Además, los autores exploran la Cadena de Pensamiento de Estado (State-CoT), donde los estados intermedios se intercalan con las acciones en la secuencia de entrenamiento para mejorar la fundamentación.

Para medir la recuperación, los autores emplean:

• Sondas lineales (Linear Probes): Modelos de regresión logística entrenados sobre las activaciones de las capas ocultas para predecir la validez de las acciones y los valores de verdad de los predicados de estado.
• Clasificación generativa: Una métrica que evalúa si las probabilidades de los tokens del modelo clasifican las acciones válidas por encima de las inválidas al ser consultado con planes parciales.
• Pruebas fuera de la distribución (OOD): Evaluación en instancias de planificación con mayor complejidad (por ejemplo, más bloques) que las observadas durante el entrenamiento.

Resultos clave

1. Codificación lineal de la validez: El ajuste fino supervisado sobre secuencias de acciones válidas permite que los LLMs codifiquen linealmente la validez de las acciones en sus representaciones internas. Las puntuaciones F1 de las sondas alcanzan niveles cercanos a la perfección (~1.0) en las capas intermedias y superiores en todas las distribuciones de entrenamiento.
2. Representación parcial del estado: Los modelos aprenden a codificar linealmente los valores de verdad de algunos predicados de estado (por ejemplo, arm-empty, at ?pkg ?loc), pero el rendimiento varía significamente según la complejidad del predicado. Los predicados más simples con menos argumentos se codifican con mayor precisión que los complejos.
3. Divergencia entre métricas internas y generativas: Un hallazgo crítico es el desajuste entre las representaciones internas y los resultados generativos. Los modelos ajustados con planes óptimos poseen representaciones internas de alta calidad sobre la validez de las acciones (altas puntuaciones de sonda) pero fallan al utilizar las probabilidades de los tokens para distinguir entre acciones válidas e inválidas durante la generación. Por el contrario, los modelos ajustados con datos de caminata aleatoria muestran una clasificación generativa superior (clasificando las acciones válidas por encima de las inválidas) pero presentan dificultades con la generalización OOD.
4. Impacto de la cobertura del espacio de estados: Los datos de Random Walk++, que aseguran una cobertura más amplia del espacio de estados al eliminar ciclos, producen la recuperación más robusta del modelo de mundo. Este método logra un rendimiento superior tanto en la clasificación generativa dentro de la distribución como en las tareas de planificación OOD, en comparación con las caminatas aleatorias estándar o los planes óptimos.
5. Generalización OOD: Aunque las representaciones internas de la validez de las acciones se mantienen robustas ante cambios de distribución (datos OOD), el rendimiento generativo colapsa para los modelos entrenados con caminatas aleatorias estándar. Los modelos entrenados con Random Walk++ y planes óptimos muestran una mejor resiliencia a los cambios OOD, aunque el rendimiento sigue degradándose a medida que aumenta la complejidad del problema.

Significancia y afirmaciones
El artículo afirma proporcionar una "receta" para aplicar técnicas de interpretabilidad a los planificadores de LLM, ofreciendo una visión holística que combina el análisis de la representación interna con la evaluación generativa. Los autores concluyen que:

• El SFT sobre planes válidos es suficiente para que los LLM desarrollen representaciones internas que codifican linealmente la validez de las acciones y ciertos predicados de estado, lo que sugiere una forma de recuperación del modelo de mundo.
• Sin embargo, la presencia de representaciones internas no garantiza que el modelo las utilice eficazmente en sus probabilidades generativas; la distribución de entrenamiento (específicamente la cobertura del espacio de estados) influye fuertemente en esta alineación.
• Los hallazgos son "optimistas" respecto a la viabilidad del ajuste fino para mejorar la planificación, demostrando que los LLM pueden recuperar parcialmente modelos de mundo de planificación. No obstante, el trabajo también destaca limitaciones, señalando que el vínculo causal entre estas representaciones y el proceso de planificación del modelo sigue sin probarse, y que la generalización a un número significativamente mayor de parámetros o a la planificación en lenguaje natural sigue siendo una cuestión abierta.

Los autores declaran explícitamente que no pretenden haber resuelto la planificación OOD ni haber probado el razonamiento causal, sino que proporcionan evidencia empírica de que los LLM ajustados desarrollan representaciones internas estructuradas de la dinámica de planificación que son distintas, aunque relacionadas, con sus capacidades generativas.