Resumen Técnico: ¿Rastrean los Modelos de Lenguaje las Entidades a Través de Cambios de Estado?

Enunciado del Problema

El rastreo de entidades (RE) es una habilidad fundamental requerida para el razonamiento complejo, que implica la capacidad de mantener y actualizar los estados de los objetos dentro de un contexto a medida que ocurren operaciones (por ejemplo, mover, añadir, eliminar). Si bien investigaciones anteriores han examinado cómo los modelos de lenguaje (ML) basados en transformadores manejan el enlace estático de entidades (asociar un objeto con una ubicación sin cambios de estado) o han resuelto problemas de RE utilizando lenguajes sintéticos simplificados y modelos "de juguete", existe un conocimiento limitado sobre cómo los ML preentrenados no de juguete manejan tareas de RE realistas expresadas en lenguaje natural con múltiples operaciones que cambian el estado. Específicamente, sigue sin estar claro si estos modelos construyen incrementalmente un estado del mundo a medida que se despliega el contexto o si emplean una estrategia diferente para resolver consultas.

Metodología

Los autores investigan los mecanismos subyacentes al RE en modelos no de juguete (Gemma-2-2B, CodeLlama-13B y Llama-3.1-70B) utilizando un conjunto de datos derivado del "benchmark" de "cajas" (Kim & Schuster, 2023). El conjunto de datos consiste en una descripción inicial de objetos en cajas, seguida de una secuencia de operaciones que cambian el estado (PUT, REMOVE, MOVE), y una consulta que pregunta por el contenido de una caja específica.

El estudio emplea un enfoque dual que combina el análisis conductual y la interpretabilidad mecánica:

1. Sondeo Conductual:

   • Sondas Globales vs. Locales: Se entrenan sondas lineales en la corriente residual del token final para determinar si el modelo codifica un estado global (ubicación de todos los objetos) o solo estados locales (contenido de la caja consultada).
   • Sondeo de Estados Previos: Se entrenan sondas para detectar si los estados previos (estados antes de la operación final) son decodificables secuencialmente a través de las capas, probando si el modelo procesa los cambios de estado de forma incremental o en paralelo.

2. Análisis Mecánico:

   • Rastreo de Circuitos (Parcheo de Rutas): Se utiliza para identificar las cabezas de atención específicas responsables de la operación PUT, comparándolas con el circuito para el enlace de entidades estático (DESCRIPCIÓN).
   • Análisis de Subespacios: Se utilizan el enmascaramiento contrafactual basado en desiderata (DCM) y el parcheo de subespacios para determinar si los circuitos PUT y DESCRIPCIÓN comparten los mismos subespacios de comunicación para la información posicional.
   • Investigación del Mecanismo REMOVE: Debido a la dificultad de analizar la "ausencia" de señales, los autores emplean análisis de logits/rank, sondas ternarias (clasificando estados como {existe, eliminado, no existe}) e intervenciones causales (anulando los pesos de la sonda) para caracterizar cómo se implementan las operaciones REMOVE.

Hallazgos Clave

1. Estrategia de Rastreo No Incremental

Contrario a la hipótesis de que los ML construyen incrementalmente estados del mundo a medida que procesan tokens o capas, los resultados apoyan una estrategia no incremental:

• A través de Tokens: Las sondas globales fallan al decodificar el estado completo del mundo en la corriente residual del token final, mientras que las sondas locales (dirigidas a la caja consultada) logran una alta precisión. Esto sugiere que los modelos no mantienen un estado global, sino que consolidan dinámicamente la información relevante solo cuando se presenta la consulta.
• A través de Capas: El sondeo de estados previos revela que los modelos no procesan los cambios de estado secuencialmente a través de las capas. En cambio, los estados previos se decodifican con una precisión similar a los estados finales en las mismas capas, lo que indica un procesamiento paralelo de todas las operaciones relevantes en el token final.

2. Mecanismos de las Operaciones

• Operación PUT: El mecanismo PUT se asemeja al circuito de enlace de entidades encontrado en trabajos anteriores (Prakash et al., 2024). Utiliza un mecanismo de "mirada hacia atrás" donde grupos específicos de atención (A, B, C, D) vinculan nuevos objetos a los IDs de caja utilizando información posicional (IDs de orden). Crucialmente, los circuitos PUT y DESCRIPCIÓN comparten subespacios funcionales para transmitir información posicional.
• Operación REMOVE: La operación REMOVE se implementa mediante una "etiqueta de supresión global".
  • El análisis de logits muestra que cuando un objeto es eliminado, su logit es suprimido, pero esta supresión no es estrictamente local; afecta al objeto globalmente independientemente del contexto de la caja.
  • Las sondas ternarias identifican una señal de "etiqueta de eliminación" ubicada principalmente en el token del objeto (no en el token del ID de caja).
  • Las intervenciones causales confirman que anular la etiqueta de eliminación en el token del objeto impide que el modelo prediga el objeto eliminado, mientras que las intervenciones en el token del ID de caja no tienen efecto. Esto implica que el modelo utiliza una heurística global frágil para suprimir objetos en lugar de un mecanismo de desvinculación preciso y específico de la caja.

3. Modos de Fallo Predichos y Observados

La visión mecánica sobre la "etiqueta de eliminación global" predice modos de fallo específicos que no son evidentes en las pruebas conductuales estándar:

• Eliminación sin Operación (No-op Remove): Eliminar un objeto de una caja que nunca fue mencionada en el contexto (o eliminar un objeto que no existe) conduce a la supresión global de ese tipo de objeto, causando errores en otras cajas.
• Objetos con Etiquetas Compartidas: Si el mismo tipo de objeto existe en múltiples cajas, eliminarlo de una caja hace que el modelo lo suprima globalmente, lo que lleva a predicciones incorrectas para la otra caja.
• Reintroducción: Eliminar un objeto y luego volver a colocarlo en una caja diferente falla porque la etiqueta de supresión global persiste.
  Las pruebas conductuales en estos escenarios específicos confirmaron altas tasas de degeneración, validando la hipótesis mecánica.

Contribuciones

1. Evidencia de Rastreo No Incremental: El artículo proporciona evidencia empírica de que los ML resuelven tareas de RE recuperando tokens objetivo en paralelo en el paso final de la consulta, en lugar de mediante la construcción acumulativa del mundo a través de tokens o capas.
2. Caracterización Mecánica de las Operaciones: El estudio caracteriza PUT como un circuito de enlace similar al rastreo de entidades estático e identifica REMOVE como un mecanismo de supresión global impulsado por etiquetas a nivel de objeto.
3. Puente entre Análisis Conductual y Mecánico: El trabajo demuestra un bucle de retroalimentación donde las ideas mecánicas (la etiqueta de eliminación global) predicen modos de fallo novedosos, que luego son validados mediante nuevas pruebas conductuales. A la inversa, los resultados conductuales informan la formulación de hipótesis mecánicas.
4. Soluciones Intervencionistas: Los autores muestran que anular la etiqueta de eliminación global mediante intervención puede corregir parcialmente estos modos de fallo, ofreciendo una prueba de concepto para mejorar la robustez del modelo a través de la comprensión mecánica.

Significado

El artículo destaca una limitación arquitectónica fundamental en los ML actuales basados en transformadores: resuelven tareas fundamentalmente secuenciales (rastrear cambios de estado a lo largo del tiempo) utilizando estrategias no secuenciales y paralelas. Esta dependencia de la agregación del contexto relevante en el último token, en lugar de mantener un estado del mundo persistente, conduce a comportamientos frágiles cuando las operaciones son ambiguas o cuando los objetos se comparten entre contextos. El trabajo ilustra cómo combinar la evaluación conductual con la interpretabilidad mecánica puede revelar modos de fallo ocultos y guiar el desarrollo de capacidades de razonamiento más robustas en los modelos de lenguaje.