Resumen Técnico: Mecanismos de Mezcla en la Recuperación de Entidades en Contexto

Enunciado del Problema

Los modelos de lenguaje (LM) destacan en el razonamiento en contexto, una capacidad que depende fundamentalmente de la habilidad para vincular entidades relacionadas dentro de un texto (por ejemplo, vincular "Ana" con "tarta" en "Ana ama la tarta") y posteriormente recuperarlas basándose en una consulta (por ejemplo, responder "¿Quién ama la tarta?"). Aunque investigaciones anteriores establecieron que los LM utilizan un mecanismo posicional—recuperando entidades basándose en su índice en la secuencia de contexto—esta visión se derivó en gran medida de estudios que involucraban listas cortas de grupos de entidades (por ejemplo, $n=2$ o $3$).

El problema central abordado en este trabajo es la fiabilidad de este mecanismo posicional en entornos más complejos y realistas. A medida que aumenta el número de grupos de entidades vinculadas, los LM tienen dificultades para razonar sobre contextos largos, exhibiendo fenómenos como el efecto "perdido en el medio". Los autores investigan si el mecanismo posicional se generaliza a secuencias más largas o si los LM emplean estrategias alternativas para mantener la precisión en la recuperación.

Metodología

Los autores emplean técnicas de interpretabilidad mecánica, específicamente intervenciones de intercambio, para sondear la estructura causal interna de los LM.

1. Configuración de la Tarea: El estudio utiliza tareas de razonamiento en contexto con plantillas que involucran "Grupos de Entidades" (tuplas de entidades que llenan roles específicos, por ejemplo, Nombre + Comida). Se consulta al modelo para recuperar una entidad objetivo basándose en una entidad de consulta dentro de un grupo específico.
2. Diseño Contrafactual: Para distinguir entre hipótesis competitivas, los autores construyen conjuntos de datos emparejados de entradas originales y contrafactuales. Estos pares están diseñados de tal manera que tres mecanismos hipotéticos—Posicional, Léxico y Reflejo—predirían salidas distintas si se realizara una intervención de intercambio sobre sus respectivas variables intermedias.
   • Mecanismo Posicional ($P$): Recupera la entidad en el índice del grupo de la consulta contrafactual.
   • Mecanismo Léxico ($L$): Recupera la entidad vinculada a la entidad de consulta contrafactual (por ejemplo, si la consulta cambia de "Tim" a "Ana", recupera lo que "Ana" ama).
   • Mecanismo Reflejo ($R$): Recupera la entidad objetivo mediante un puntero directo que se origina en la propia objetivo, desreferenciado a través de la consulta.
3. Protocolo de Intervención: Los autores realizan intervenciones de intercambio en el flujo residual del último token en capas específicas ($\ell$) identificadas como el punto justo antes de que ocurra la recuperación. Parchean el vector del flujo residual de la ejecución contrafactual en la ejecución original y observan la distribución del siguiente token resultante.
4. Modelado Causal: Se desarrolla un modelo causal ($M$) para simular el comportamiento del LM. Este modelo combina los tres mecanismos como una mezcla ponderada:
   • El término posicional se modela como una distribución gaussiana centrada en el índice de la consulta, con una varianza dependiente del índice.
   • Los términos léxico y reflejo se modelan como distribuciones one-hot.
   • El modelo se entrena para predecir la distribución de logits del LM utilizando la Divergencia de Jensen–Shannon (JSD) como métrica de pérdida.
5. Alcance: Los experimentos cubren nueve modelos de tres familias (Llama, Gemma, Qwen) que van desde 2B hasta 72B parámetros, evaluados en diez tareas de vinculación distintas.

Contribuciones Clave

1. Identificación de Tres Mecanismos Distintos

El artículo desafía la visión predominante de que la recuperación es únicamente posicional. Demuestra que los LM emplean una mezcla de tres mecanismos:

• Posicional: Efectivo para entidades al principio y al final del contexto, pero se vuelve ruidoso y difuso para posiciones intermedias a medida que aumenta el número de grupos ($n$).
• Léxico: Recupera el contraparte vinculado de la entidad de consulta (por ejemplo, "tarta" $\to$ "Ana"). Este mecanismo proporciona una señal nítida, one-hot.
• Reflejo: Utiliza un puntero directo desde la entidad objetivo hacia sí misma, que se recupera a través de la consulta. Esto es particularmente crucial cuando la objetivo aparece antes que la consulta en una cláusula (por ejemplo, "Tim ama el té" $\to$ "¿Quién ama el té?"), un escenario donde la copia léxica es arquitectónicamente difícil para los modelos autoregresivos.

2. El Patrón de Fiabilidad en "U"

A través de extensos experimentos de intervención, los autores descubren un patrón consistente:

• Primera y Última Posición: La recuperación está dominada por el mecanismo posicional, que es altamente concentrado y fiable.
• Posiciones Intermedias: La señal posicional se degrada significativamente. Para compensar, los LM dependen en gran medida de los mecanismos léxico y reflejo, que proporcionan señales más nítidas para desambiguar la entidad correcta.
• Interacción: Los mecanismos exhiben una "sinergia competitiva". Cuando los índices coinciden, las señales se amplifican; cuando divergen, se suprimen mutuamente.

3. Modelo Causal de Predicción del Siguiente Token

Los autores construyen un modelo causal que combina estos tres mecanismos. Este modelo logra un 95% de acuerdo (medido por la Similitud de Jensen–Shannon) con las distribuciones reales del siguiente token de los LM en diversas tareas y tamaños de modelo. En contraste, los modelos que dependen únicamente del mecanismo posicional (la visión predominante) alcanzan puntuaciones significativamente más bajas (~45%), desempeñándose peor que una línea base uniforme en algunas configuraciones.

4. Generalización a Entradas Naturalistas

El estudio valida que estos hallazgos se mantienen en entornos más naturalistas donde los grupos de entidades se intercalan con texto "relleno" (oraciones abiertas). Aunque la adición de tokens de relleno aumenta el ruido y debilita ligeramente el mecanismo léxico, la dependencia fundamental de la mezcla de tres mecanismos persiste, ofreciendo una explicación mecánica para el efecto "perdido en el medio".

Resultados

• Prevalencia del Mecanismo: En posiciones intermedias de contextos largos, el mecanismo posicional representa solo ~20% del comportamiento del modelo, mientras que los mecanismos léxico y reflejo se convierten en los impulsores principales de la recuperación correcta.
• Robustez del Modelo: El patrón de mezcla es consistente entre familias de modelos (Llama, Gemma, Qwen) y escalas (de 2B a 72B parámetros).
• Validación Refleja: A través de diseños contrafactuales específicos donde la respuesta contrafactual no existe en el contexto original, los autores prueban la existencia de un mecanismo de puntero reflejo distinto del propio token de respuesta.
• Estudios de Ablación: Eliminar cualquier mecanismo individual del modelo causal degrada significativamente el rendimiento, confirmando que los tres son necesarios para una abstracción precisa del comportamiento del LM.

Significado y Afirmaciones

El artículo afirma establecer una imagen más completa de cómo los LM vinculan y recuperan entidades en contexto. Al ir más allá de la visión posicional simplista, los autores proporcionan una explicación mecánica tanto de las fortalezas (recuperación robusta mediante múltiples señales) como de las fragilidades (ruido en posiciones intermedias que conduce a errores de "perdido en el medio") de los LLM en entornos de contexto largo.

El trabajo sugiere que el fenómeno "perdido en el medio" no es meramente un fallo de atención, sino una degradación específica de la señal posicional, que solo es parcialmente compensada por los mecanismos léxico y reflejo. Los autores publican su código y datos para facilitar investigaciones adicionales sobre los mecanismos causales del razonamiento en contexto.