Resumen Técnico: Head Pursuit: Probing Attention Specialization in Multimodal Transformers

Planteamiento del Problema

Los modelos generativos a gran escala, incluyendo tanto los modelos de lenguaje unimodales (LLMs) como los modelos de visión y lenguaje multimodales (VLMs), han alcanzado un rendimiento notable en diversas tareas. Sin embargo, sus mecanismos internos aún se comprenden solo parcialmente. Mientras que herramientas de interpretabilidad previas, como el Logit Lens, han revelado información semántica en los estados ocultos, estas suelen aplicarse de forma heurística a ejemplos individuales, lo que dificulta la generalización de los hallazgos o la cuantificación de la importancia de componentes específicos (como las cabezas de atención) en la configuración de la salida del modelo. Existe la necesidad de un método con principios sólidos para identificar qué cabezas de atención se especializan en atributos semánticos o visuales específicos y determinar si la manipulación de estas cabezas permite un control dirigido sobre el comportamiento del modelo sin necesidad de entrenamiento adicional.

Metodología

Los autores proponen un marco llamado Head Pursuit, que reinterpreta la práctica de sondear las activaciones intermedias a través del lente del procesamiento de señales dispersas (sparse signal processing). La metodología central involucra los siguientes pasos:

1. Descomposición Dispersa mediante SOMP: Los autores emplean la Búsqueda Ortogonal Simultánea por Emparejamiento (SOMP), un algoritmo codicioso para aproximar señales de alta dimensión con combinaciones lineales dispersas de elementos base.

   • Entrada: Para una cabeza de atención $h$ dada en la capa $l$, se analiza la matriz de activación $H_{h,l}$ (que representa $n$ muestras).
   • Diccionario: En lugar de aprender un diccionario a partir de las activaciones, el método utiliza la matriz de unembedding $D$ del modelo como un diccionario fijo. Cada fila de $D$ corresponde a un token del vocabulario, proporcionando un vínculo directo con conceptos semánticos interpretables por humanos.
   • Proceso: SOMP selecciona iterativamente átomos del diccionario (tokens) que se correlacionan máximamente con los residuales de la cabeza a través de todas las muestras del conjunto de datos. Esto produce una matriz de coeficientes dispersos $W^*$ tal que $H \approx W^*D$.
   • Generalización: A diferencia del Logit Lens, que proyecta un único vector residual, SOMP opera sobre múltiples muestras simultáneamente y selecciona múltiples direcciones semánticas, proporcionando una caracterización robusta y a nivel de conjunto de datos del rol funcional de una cabeza.

2. Selección y Clasificación de Cabezas:

   • Para identificar cabezas relevantes para un concepto objetivo específico (por ejemplo, "colores", "toxicidad", "números"), el diccionario de unembedding se restringe a los tokens asociados con ese concepto.
   • Se aplica SOMP a este diccionario específico del concepto.
   • Las cabezas se clasifican según la fracción de la varianza en sus activaciones explicada por la reconstrucción de SOMP. Una alta explicación de la varianza indica una fuerte especialización para ese concepto.

3. Intervención:

   • Una vez identificadas las top-$k$ cabezas especializadas, los autores intervienen reescalando su contribución al flujo residual.
   • Intervención Inhibitoria: Se invierte el signo de las representaciones de la cabeza ($\alpha = -1$) para suprimir el atributo objetivo.
   • Intervención de Mejora: Se amplifican las representaciones de la cabeza ($\alpha > 1$) para promover la generación del atributo objetivo.
   • Esta intervención se realiza durante la pasada hacia adelante (forward pass) sin actualizar los pesos del modelo.

Contribuciones Clave

1. Encuadre Teórico: El artículo enmarca el Matching Pursuit (MP) y su variante multiesmple SOMP dentro del contexto de la interpretabilidad de modelos generativos, estableciendo una conexión formal con herramientas estándar como el Logit Lens mientras lo generaliza para manejar múltiples muestras y múltiples direcciones semánticas simultáneamente.
2. Descubrimiento de Especialización: El método revela patrones consistentes de especialización de cabezas tanto en transformadores unimodales (Mistral-7B) como multimodales (LLaVA-NeXT, Gemma3, Qwen2.5-VL). Se encuentra que cabezas específicas generan de manera fiable grupos de tokens semánticamente coherentes (por ejemplo, nombres de países, meses, números, sentimientos).
3. Controlabilidad: Los autores demuestran que intervenir en una fracción muy pequeña de cabezas (tan solo el 1% del total) puede suprimir o mejorar de manera fiable conceptos específicos. Esto sugiere que las capas de atención contienen una estructura lineal altamente interpretable y manipulable.

Resultados

Los autores validaron su enfoque en varias tareas:

• Preguntas y Respuestas (TriviaQA): Invertir los signos de solo 8 cabezas especializadas (0.8% del modelo) dedicadas a "nombres de países" degradó significativamente el rendimiento en preguntas relacionadas con países, teniendo un impacto mínimo en otras preguntas. La selección de cabezas aleatorias o la selección basada en Logit Lens no mostraron este efecto dirigido.
• Mitigación de Toxicidad: Utilizando una lista restringida de palabras clave tóxicas para seleccionar cabezas, los autores lograron reducir con éxito la frecuencia de generaciones tóxicas en los conjuntos de datos RealToxicityPrompts (RTP) y Thoroughly Engineered Toxicity (TET). Invertir 32 cabezas especializadas redujo las generaciones tóxicas aproximadamente entre un 30% y un 50% en comparación con las líneas base, mientras que las intervenciones aleatorias a menudo aumentaron la toxicidad.
• Clasificación de Imágenes (LLaVA): Invertir las 32 mejores cabezas seleccionadas mediante SOMP para conjuntos de datos específicos (por ejemplo, MNIST, SVHN) interrumpió significativamente la precisión de la clasificación. El estudio encontró que los conjuntos de datos con semánticas relacionadas (por ejemplo, tareas de reconocimiento de dígitos) compartían conjuntos superpuestos de cabezas especializadas, confirmando la superposición funcional entre dominios.
• Generación de Descripciones de Imágenes (Flickr30k):
  • Inhibición: Invertir 16 cabezas especializadas para atributos como "colores", "sentimientos" o "cantidad" eliminó casi por completo estas palabras clave de las descripciones, manteniendo puntuaciones CIDEr (calidad de la descripción) elevadas.
  • Mejora: Amplificar estas cabezas ($\alpha=5$) aumentó la presencia de los conceptos objetivo en más de un 60% con solo una degradación marginal en la calidad de la descripción.

Significado y Reivindicaciones

El artículo afirma que las capas de atención en los grandes modelos generativos poseen un subespacio lineal estructurado y alineado con conceptos. Este hallazgo respalda la hipótesis de que los conceptos de alto nivel se representan en subespacios residuales aproximadamente lineales.

La importancia de este trabajo radica en proporcionar herramientas simples y sin necesidad de entrenamiento para comprender y editar modelos a gran escala. Al identificar que un pequeño número de unidades puede ejercer un control semántico amplio, los autores ofrecen una alternativa con principios frente a las intervenciones heurísticas basadas en prompts o el ajuste fino (fine-tuning). Este enfoque permite:

• Interpretabilidad: Descomponer el comportamiento del modelo en un pequeño conjunto de elementos interpretables.
• Control: Supresión o mejora dirigida de atributos específicos (por ejemplo, toxicidad, conceptos visuales específicos) con un impacto mínimo en las capacidades generales del modelo.

Los autores mantienen la modestia respecto a las limitaciones, señalando que el método asume linealidad (lo que puede no capturar todas las estructuras no lineales), depende de la calidad del diccionario semántico y utiliza un reescalado global simple en lugar de un direccionamiento (steering) de grano fino. Sugieren que el trabajo futuro podría explorar intervenciones específicas de posición o aplicar estas técnicas a tareas de generación de imágenes.