Resumen Técnico: La puntuación de contribución de logit identifica cabezales de recuperación no literales

1. Planteamiento del problema

Los Modelos de Lenguaje de Gran Escala (LLMs) frecuentemente sintetizan respuestas a partir del significado semántico de un fragmento de contexto relevante en lugar de copiar literalmente los tokens. Aunque el trabajo actual de interpretabilidad mecánica ha identificado "cabezales de recuperación" (retrieval heads) responsables del acceso al contexto, los métodos de detección actuales están fundamentalmente limitados a la recuperación literal.

Los detectores existentes (por ejemplo, Wu et al., 2025; Fu et al., 2025) dependen de observables de pesos de atención. Estos premian a los cabezales donde el token de origen atendido coincide con el token de la respuesta generada. Este criterio de "coincidencia de tokens" identifica con éxito cabezales que copian y pegan información (recuperación literal), pero falla al detectar cabezales que realizan una recuperación no literal. En escenarios no literales, un modelo puede atender a una frase de origen (por ejemplo, "Torre Eiffel") pero sintetizar un token de respuesta diferente (por ejemplo, "Yuki") a través de su circuito de valor de salida (OV). La puntuación basada en la atención captura dónde lee un cabezal, pero ignora qué escribe, perdiendo el mecanismo crítico mediante el cual ocurre la síntesis no literal.

2. Metodología: Puntuación de Contribución de Logit (LOCOS)

Los autores presentan la Puntuación de Contribución de Logit (LOCOS), un detector "consciente de la escritura" (write-aware) que evalúa los cabezales de atención basándose en su contribución al logit de la respuesta final, en lugar de su distribución de atención.

Mecanismo Central

LOCOS descompone la operación de un cabezal de atención en dos circuitos:

1. Circuito QK: Determina dónde lee el cabezal (pesos de atención $\alpha$).
2. Circuito OV: Determina qué escribe el cabezal en el flujo residual (proyección de salida $W_O v$).

El método puntúa un cabezal $(l, h)$ proyectando la salida de su circuito OV sobre la dirección de desembedding del token de la respuesta correcta ($u_{y_t}$).

El procedimiento de tres pasos

1. Contribución de logit por posición:
   Para un paso de decodificación $t$ que genera el token $y_t$, la contribución de una posición de origen $j$ a través del cabezal $(l, h)$ se calcula como:
   $$\phi^{(l,h)}_{t,j} = u_{y_t}^\top o^{(l,h)}_{t,j} = \alpha^{(l,h)}_{t,j} \cdot u_{y_t}^\top W^{(l,h)}_O v^{(l,h)}_{t,j}$$
   Este escalar $\phi$ captura el proceso completo: un cabezal puede tener una alta atención ($\alpha$) pero una contribución cercana a cero si su salida OV es ortogonal a la dirección de la respuesta. Por el contrario, un cabezal con una atención moderada puede tener una alta contribución si su circuito OV transforma la representación atendida en un vector alineado con la respuesta.

2. Contraste espacial:
   Para aislar las señales específicas de recuperación, LOCOS compara las contribuciones de las posiciones de la aguja (el fragmento de contexto relevante $[s_\tau, e_\tau)$) contra las posiciones fuera de la aguja (off-needle).
   $$\Phi^{(l,h),+}_t = \sum_{j \in \text{needle}} \phi^{(l,h)}_{t,j}$$
   $$\Phi^{(l,h),-}_t = \frac{e_\tau - s_\tau}{N_t - (e_\tau - s_\tau)} \sum_{j \notin \text{needle}} \phi^{(l,h)}_{t,j}$$
   El factor de reescalado asegura que la suma fuera de la aguja sea comparable a la suma de la aguja por longitud. La puntuación para un solo paso es la diferencia: $\Phi^{(l,h),+}_t - \Phi^{(l,h),-}_t$. Este contraste espacial (aguja frente a fuera de la aguja dentro de un mismo paso) reemplaza el contraste temporal (paso de respuesta frente a no respuesta) utilizado en trabajos previos.

3. Agregación:
   La puntuación final $S_{l,h}$ es la media de estos contrastes espaciales en todos los pasos de decodificación donde el modelo genera una respuesta correcta, agregada sobre todos los ensayos exitosos en un benchmark. Las puntuaciones negativas indican cabezales cuyas contribuciones al logit de la respuesta se originan predominantemente en posiciones fuera de la aguja.

3. Contribuciones clave

• Detección consciente de la escritura: El artículo establece que identificar la recuperación no literal requiere medir la salida del circuito OV, no solo los pesos de atención.
• Algoritmo LOCOS: Un método de una sola pasada hacia adelante (single-forward-pass) que puntúa cabezales mediante su contribución proyectada al logit de la respuesta, contrastando fuentes dentro y fuera de la aguja.
• Validación causal: A través de experimentos de ablación de la media, los autores demuestran que LOCOS identifica cabezales que son causalmente críticos para la recuperación no literal, mientras que las líneas base basadas en la atención fallan en hacerlo.

4. Resultados experimentales

El método fue evaluado en el benchmark NoLiMa (una tarea de aguja en un pajar no literal) a través de tres familias de modelos: Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B) y OLMo-3.1 (32B).

Rendimiento de la ablación (NoLiMa)

• Degradación superior: La ablación de la media de los top-$k$ cabezales LOCOS causa un declive más pronunciado en ROUGE-L en comparación con todas las líneas base basadas en la atención (Wu/NIAH-scored, Wu/NoLiMa-scored y cabezales aleatorios).
• Caso de estudio Qwen3-8B: Ablar los 50 mejores cabezales LOCOS colapsa el ROUGE-L de 0.401 a 0.000. En contraste, la línea base de atención más fuerte (Wu/NIAH) retiene un ROUE-L de 0.292 a la misma profundidad.
• Especificidad: La ablación de los cabezales LOCOS degrada severamente la recuperación contextual, pero deja intactos el recuerdo paramétrico (por ejemplo, asociaciones ciudad-país) y el razonamiento aritmético. Esto confirma que los cabezales identificados son específicos de la recuperación, no genéricamente importantes para la salida del modelo.

Especificidad Literal vs. No Literal

• La ablación de los top-$k$ cabezales LOCOS degrada tanto la recuperación no literal (NoLiMa) como la literal (NIAH).
• Sin embargo, la caída es significativamente más pronunciada en NoLiMa. Esto indica que, aunque LOCOS captura cabezales relevantes para ambos regímenes, identifica de manera única el subconjunto de cabezales responsables de la síntesis no literal que los detectores de coincidencia de tokens pasan por alto.

Transferencia a tareas posteriores (Downstream)

• En los benchmarks de contexto largo posteriores (MuSiQue y BABILong), la ablación de los top-50 cabezales LOCOS degrada significativamente el rendimiento, particularmente en la familia Qwen3 (por ejemplo, la precisión de MuSique en Qwen3-8B cae de 0.55 a 0.08).
• El ranking de LOCOS frente a las líneas base basadas en la atención varía según el benchmark, lo que sugiere que las tareas de QA de múltiples saltos reclutan tanto circuitos de recuperación literal como no literal.

Distribución de capas

• Las puntuaciones de LOCOS están fuertemente concentradas en las capas tardías para la familia Qwen3 y Gemma-3-27B.
• Los autores señalan que esta concentración puede estar influenciada por la "asunción de camino directo" (los cabezales más cercanos a la salida tienen un impacto más directo en los logits), pero los experimentos de parcheo causal sugieren que el patrón de capas tardías no es meramente un artefacto.

5. Significado y afirmaciones

El artículo afirma que LOCOS proporciona la primera identificación mecánica de los cabezales de recuperación no literal.

• Limitación del trabajo previo: Los detectores existentes son "ciegos" a la recuperación no literal porque premian los cabezales donde el token atendido coincide con el token generado. Esto omite los cabezales que atienden a una fuente y sintetizan una respuesta diferente.
• Perspectiva mecánica: Al medir la proyección del circuito OV sobre la dirección de la respuesta, LOCOS captura la operación de "escritura" esencial para la síntesis.
• Utilidad práctica: Los cabezales identificados son causalmente críticos para el razonamiento de contexto largo. Este hallazgo tiene implicaciones para:
  • Interpretabilidad: Proporcionar un mapa más preciso de cómo los LLMs acceden y sintetizan el contexto.
  • Eficiencia: Permitir una compresión de la caché KV más dirigida, preservando solo los cabezales específicos responsables de la recuperación, en lugar de heurísticas uniformes o basadas en la atención.
  • Mitigación de alucinaciones: Ofrecer un mecanismo para identificar e intervenir en los cabezales que fallan al recuperar correctamente el contexto relevante.

Los autores mantienen la modestia respecto a la generalización, señalando que sus hallazgos son específicos para las arquitecturas probadas (transformers estándar de solo decodificador) y el benchmark NoLiMa. No afirman que LOCOS identifique todos los mecanismos de recuperación, sino que logra aislar el subconjunto omitido por los criterios de coincidencia de tokens.