A retrieval conditioned rebinding circuit for dynamic entity tracking in large language models

Resumen Técnico: Un Circuito de Revinculación Condicionado por la Recuperación para el Seguimiento Dinámico de Entidades en Modelos de Lenguaje de Gran Escala

Planteamiento del Problema
Los Modelos de Lenguaje de Gran Escala (LLM) deben interpretar correctamente el contexto y recuperar información relevante vinculando las entidades con sus atributos y actualizando estos vínculos a medida que los estados cambian. Si bien trabajos previos han demostrado que los LLM pueden realizar el seguimiento de entidades, su comportamiento sigue siendo inconsistente en diversos entornos, y los mecanismos computacionales subyacentes para manejar transiciones de estado explícitas en contextos naturalistas no se comprenden completamente. Específicamente, no está claro si los modelos mantienen un "modelo de situación" actualizado globalmente después de cada cambio de estado (por ejemplo, una operación de intercambio o swap) o si emplean un mecanismo más dirigido para actualizar los vínculos solo cuando se requiere la recuperación.

Metodología
Los autores investigan este problema utilizando una tarea de seguimiento de estado dinámico que involucra cajas y objetos, adaptada de Kim y Schuster (2023). El estudio emplea intervenciones causales en las activaciones de cuatro modelos con ajuste de instrucciones: Gemma2-9B-it, Gemma3-12B-it, Llama3.2-3B-it y Llama3.1-8B-it. El análisis procede en cuatro etapas:

1. Análisis de Mediación Causal: Utilizando intervenciones de intercambio (interchange interventions), los autores rastrean el flujo de información a nivel de token. Construyen entradas contrafácticas para determinar cómo la información asociada con las cajas de referencia, las cajas objetivo del intercambio y los objetos objetivo influye en la predicción final.
2. Parcheo de Rutas (Path Patching): Para localizar el mecanismo a nivel de cabeza de atención, los autores aplican el parcheo de rutas (V-patching y Q-patching) para identificar un circuito compacto. Seleccionan iterativamente las rutas que impactan más negativamente en el rendimiento cuando se corrompen, aislando un subconjunto de cabezas de atención responsables de la tarea.
3. Intervenciones de Intercambio Específicas por Rol: Los autores prueban las funciones de los grupos de cabezas identificados mediante el parcheo de patrones de atención entre ejecuciones originales y contrafácticas. Esto distingue entre las cabezas que median la selección de la dirección (routing) y aquellas que transfieren el contenido del objeto.
4. Intervenciones de ID de Vinculación: Adaptando el marco de Feng y Steinhardt (2024), los autores realizan intervenciones causales en los subespacios de Consulta (Q) y Clave (K) de grupos de cabezas específicos. Estiman las direcciones de los ID de vinculación y los vectores de desplazamiento para probar si la recuperación de la dirección depende de la coincidencia Q/K de los identificadores de vinculación.

Contribuciones Clave y Resultados

• Evidencia de la Revinculación Condicionada por la Recuperación: El análisis de mediación causal revela que la información sobre las cajas de referencia y las cajas objetivo del intercambio no fluye directamente desde sus menciones iniciales, sino que emerge después de la operación de intercambio y se propaga hacia la predicción. Por el contrario, la información del objeto se enruta de manera más directa. Esta disociación respalda la hipótesis de la "Revinculación Condicionada por la Recuperación": el modelo no recodifica globalmente todo el estado post-intercambio. En su lugar, preserva los vínculos originales y aplica un remapeo sobre los IDs de vinculación solo durante la recuperación de la respuesta.
• Identificación de un Circuito Compacto: El parcheo de rutas identifica un circuito de cabezas de atención específico compuesto por el 3–10% de los componentes totales del modelo (por ejemplo, 38 cabezas en Gemma-9B) que recupera un rendimiento cercano al del modelo completo (0.89 de precisión frente a 1.00 del modelo completo). Este circuito consta de cinco grupos funcionales:
  • Grupo E (Ancla de Vinculación): Escribe la información de identidad de la caja.
  • Grupo D (Transmisor de Posición de Intercambio): Propaga la información del intercambio.
  • Grupo C (Actualizador de Posición): Moldea la atención hacia las posiciones de vinculación actualizadas.
  • Grupo B (Desreferenciador): Actúa como un puntero, mediando la selección de la dirección.
  • Grupo A (Recuperador de Respuesta): Recupera el token del objeto objetivo.
• Mecanismo de Enrutamiento de Direcciones: Las intervenciones de intercambio confirman que el Grupo B funciona como un puntero, sensible a la selección de la dirección. El parcheo de sus patrones de atención desplaza el enfoque del modelo hacia la posición de origen correcta sin transferir el contenido contrafáctico del objeto.
• Firmas de los Subespacios Q/K: Las intervenciones de ID de vinculación revelan que el mecanismo de puntero depende de la coincidencia Q/K, pero la implementación varía según la familia de modelos:
  • Modelos Gemma: Exhiben una firma clara de coincidencia Q/K donde la información de vinculación se expresa tanto en los subespacios de consulta como de clave. Desplazar el vector de consulta redirige efectivamente la atención.
  • Modelos Llama: Dependen principalmente de los vectores de clave para la información de vinculación. Las intervenciones en el lado de la clave producen los mayores desplazamientos de atención, mientras que los efectos en el lado de la consulta son más débiles.
• Generalizabilidad: El mecanismo de revinculación condicionada por la recuperación se observa en las familias Gemma y Llama, lo que sugiere que es una estrategia computacional compartida para el seguimiento de estado dinámico, a pesar de las diferencias en las firmas específicas de los circuitos.

Significancia y Reivindicaciones
El artículo afirma proporcionar una explicación interpretable, a nivel de circuito, de cómo los LLM realizan el seguimiento de estado dependiente del contexto. Los hallazgos sugieren que los LLM no necesariamente mantienen una simulación mental completa y globalmente actualizada del mundo después de cada cambio de estado. En su lugar, utilizan una estrategia de "revinculación" donde la información relacionada con el intercambio se preserva y se aplica localmente en el momento de la recuperación para redirigir la atención hacia el objeto correcto a través de IDs de vinculación actualizados.

Los autores señalan que, si bien la estrategia computacional (revinculación condicionada por la recuperación) parece consistente entre las familias de modelos, la implementación a nivel de circuito (específicamente la dependencia de los subespacios Q frente a K) difiere. Esto implica que el mecanismo puede ser una solución compartida al problema del seguimiento dinámico, pero su realización depende de arquitecturas de modelos o distribuciones de preentrenamiento específicas.

Limitaciones
Los autores reconocen que su configuración de tarea está altamente controlada (un solo intercambio, objetos de un solo token, cajas nombradas alfabéticamente) y puede no capturar la gama completa de mecanismos en el discurso natural. Además, las direcciones de los ID de vinculación utilizadas en las intervenciones se estiman bajo un supuesto de representación lineal aproximada, lo que podría subestimar los mecanismos si la información de vinculación se representa de forma no lineal o se distribuye en múltiples subespacios.