Task Vectors, Learned Not Extracted: Performance Gains and Mechanistic Insight

Resumen Técnico: Vectores de Tarea Aprendidos, no Extraídos

1. Planteamiento del Problema

Los Grandes Modelos de Lenguaje (LLMs) poseen la capacidad de realizar nuevas tareas a partir de demostraciones en contexto (In-Context Learning o ICL) sin entrenamiento explícito. Investigaciones recientes sugieren que el modelo comprime estas demostraciones en Vectores de Tarea (TVs), representaciones vectoriales compactas que guían las predicciones.

Sin embargo, los métodos existentes para obtener estos vectores presentan limitaciones críticas:

• Extracción opaca e ineficiente: Los métodos actuales (como Vanilla TV o Function Vectors) extraen los vectores de estados ocultos o salidas de cabezas de atención mediante procesos de filtrado u optimización complejos y dependientes de la calidad de las representaciones del modelo.
• Falta de explicación mecánica: Aunque se sabe que inyectar estos vectores mejora el rendimiento, no se comprende cómo interactúan con los componentes internos del modelo (bajo nivel) ni cómo se propagan a través de las capas para influir en la salida final (alto nivel).
• Rigidez: Los vectores extraídos suelen ser sensibles a la capa de inyección y a la posición del token, limitando su flexibilidad.

2. Metodología

Los autores proponen un cambio de paradigma: en lugar de extraer vectores de las representaciones del modelo, entrenar directamente Vectores de Tarea Aprendidos (LTVs).

• Entrenamiento Directo (LTV): Se define un vector $\theta$ que se inyecta en los estados ocultos de una capa específica (o múltiples capas) y posiciones. Este vector se optimiza mediante descenso de gradiente para minimizar la pérdida de entropía cruzada en la predicción de la etiqueta correcta para una consulta de zero-shot:
  $$\text{Minimizar } -\log p(y_q | x_q, \theta, L, P)$$
  Donde $L$ son las capas y $P$ las posiciones de inyección. A diferencia de los métodos anteriores, esto no depende de la calidad de las demostraciones en contexto (ICL) para la extracción.
• Análisis Mecanístico:
  • Nivel Bajo (Interacciones): Se investiga cómo los LTVs interactúan con los circuitos de las cabezas de atención. Se utiliza la descomposición de los circuitos OV (Value-Output) de las cabezas de atención para reconstruir el efecto del vector.
  • Nivel Alto (Propagación): Se analiza la dinámica de los estados ocultos a medida que el vector se propaga por las capas restantes. Se propone y verifica la hipótesis de que la propagación es predominantemente lineal, descomponiéndose en dos operaciones: rotación (alinear el vector con el subespacio de la tarea) y estiramiento (ajustar la magnitud).

3. Contribuciones Clave

1. Superioridad del LTV: Demostración de que los vectores entrenados directamente superan consistentemente a los vectores extraídos (Vanilla TV y Function Vectors) en precisión, tanto en tareas de clasificación como de generación.
2. Flexibilidad y Escalabilidad: Los LTVs funcionan eficazmente en capas arbitrarias (incluso tardías), en posiciones de tokens distintas al último token, y pueden inyectarse en prompts de ICL para mejorar aún más el rendimiento.
3. Descubrimiento de Mecanismos de Bajo Nivel:
   • Se identifica que los circuitos OV de las cabezas de atención son el canal principal a través del cual los TVs influyen en el flujo residual.
   • Se define un subconjunto de "cabezas clave" (aproximadamente el 10% superior) que son decisivas para aprovechar el vector inyectado. Estas cabezas muestran una menor tendencia al "sumidero de atención" (attention sink) y se enfocan más en las posiciones finales.
4. Descubrimiento de Mecanismos de Alto Nivel:
   • A pesar de las no linealidades de los Transformers, la propagación de los TVs es casi lineal.
   • Dinámica Rotación-Estiramiento: Los vectores inyectados en capas tempranas sufren una fuerte rotación para alinearse con los subespacios relevantes de la tarea (direcciones de desincrustación de etiquetas). A medida que se avanza hacia capas más profundas, el efecto de rotación disminuye y el estiramiento (escalado de magnitud) se convierte en el factor dominante para modificar las logits de salida.

4. Resultados Experimentales

• Rendimiento: En modelos como Llama3.1-8B, los LTVs alcanzan o superan el rendimiento del ICL completo en diversas tareas (SST-2, TREC, Capital, Antónimos, etc.). En tareas de generación compleja (conjunto Myopic), los LTVs superan tanto a las líneas base de TV extraídos como al ICL estándar.
• Eficiencia: El método LTV actúa como una técnica de Fine-Tuning eficiente en parámetros (PEFT), requiriendo optimizar solo $d$ parámetros (la dimensión del vector), logrando una latencia y uso de memoria comparables o mejores que Prefix Tuning y LoRA.
• Generalización y Composicionalidad:
  • Los LTVs muestran una fuerte separación entre tareas diferentes y agrupamiento dentro de la misma tarea.
  • Se observa que la generalización entre tareas depende del solapamiento en el espacio de etiquetas (ej. SNLI y RTE comparten etiquetas binarias y sus vectores son similares).
  • Se demuestra composicionalidad: La suma de vectores de tareas relacionadas (ej. Inglés $\to$ Francés + Masculino $\to$ Femenino) produce un vector compuesto que funciona en una tarea de composición (Inglés Masculino $\to$ Francés Femenino).
• Validación de Linealidad: La reconstrucción de los efectos de los TVs utilizando operadores lineales (matrices de rotación y estiramiento) recupera la mayor parte del rendimiento original, confirmando la hipótesis de linealidad en la propagación.

5. Significado e Impacto

Este trabajo ofrece una doble contribución fundamental para la comprensión y aplicación de los LLMs:

1. Práctica: Proporciona un método eficiente y superior para obtener vectores de tarea que permiten controlar y adaptar el comportamiento de los modelos sin reentrenamiento masivo, superando las limitaciones de los métodos de extracción basados en ICL.
2. Teórica/Mecanística: Ofrece una lente principista para entender el In-Context Learning. Al revelar que los TVs operan a través de circuitos específicos de atención (OV) y que su propagación sigue una dinámica lineal de rotación y estiramiento, el trabajo desmitifica la "caja negra" de cómo los modelos utilizan las demostraciones en contexto para realizar inferencias correctas. Esto conecta la geometría de los estados ocultos con la mecánica de las cabezas de atención, proporcionando una base sólida para futuras investigaciones en interpretabilidad y control de modelos.