Resumen Técnico: Alineación Distribucional como Criterio para el Diseño de Vectores de Tarea en el Aprendizaje en Contexto

Declaración del Problema

El Aprendizaje en Contexto (ICL) permite que los Modelos de Lenguaje Grandes (LLM) se adapten a nuevas tareas al anteponer demostraciones etiquetadas a la consulta de entrada. Aunque efectivo, el ICL incurre en costos de inferencia crecientes y sobrecargas de memoria a medida que aumenta la longitud del contexto, lo que obstaculiza su utilidad práctica bajo restricciones de recursos. Para abordar esto, investigaciones recientes han propuesto Vectores de Tarea (TV): representaciones compactas de estados ocultos extraídas de las activaciones internas de un LLM durante el ICL. Estos vectores encapsulan la información de la tarea inferida de las demostraciones, permitiendo que el modelo realice la adaptación de la tarea durante la inferencia sin demostraciones (zero-shot) sin volver a procesar las demostraciones.

Sin embargo, los métodos existentes para extraer vectores de tarea carecen de un criterio de diseño fundamentado. Actualmente, la calidad de un método de extracción de vectores de tarea se evalúa únicamente mediante la precisión en tareas posteriores. Esta métrica indirecta ofrece una comprensión limitada de por qué un método supera a otro o cómo diseñar mecanismos de extracción más efectivos. Existe la necesidad de un criterio directo y teóricamente fundamentado para evaluar y optimizar la calidad del vector de tarea.

Metodología

1. El Criterio Propuesto: Alineación Distribucional

Los autores postulan que el objetivo principal de un vector de tarea es condensar el efecto de las demostraciones de tal manera que la distribución predictiva resultante se alinee estrechamente con la del ICL estándar. Para cuantificar esta alineación, introducen dNTP (discrepancia en la Probabilidad del Siguiente Token).

• Definición: dNTP mide la divergencia de Kullback-Leibler (KL) entre la distribución predictiva del siguiente token bajo ICL ($P_{icl}$) y la distribución bajo inferencia basada en vector de tarea ($P_{tv}$).
  $$dNTP(f; Z) = \mathbb{E}_{x \sim D} \left[ D_{KL} \left( P_{icl}(\cdot | x, Z) \parallel P_{tv}(\cdot | x, f(Z)) \right) \right]$$
  donde $f$ es la función de extracción y $Z$ representa las demostraciones.
• Validación: El análisis empírico a través de múltiples modelos y puntos de referencia revela una fuerte correlación negativa entre dNTP y la precisión en tareas posteriores. Un dNTP más bajo predice consistentemente una mayor precisión, estableciendo a dNTP como un proxy válido para la calidad del vector de tarea que no requiere etiquetas del conjunto de prueba.

2. Método Propuesto: Vector de Tarea Lineal (LTV)

Motivados por el objetivo de minimizar dNTP, los autores proponen el Vector de Tarea Lineal (LTV). El método se deriva de un problema de optimización proxy.

• Fundamento: Minimizar dNTP directamente es intratable ya que carece de una solución de forma cerrada y requiere actualizaciones iterativas. Los autores demuestran que minimizar el Error Cuadrático Medio (MSE) entre los estados ocultos del ICL ($h_{icl}$) y el estado oculto sin demostraciones más el vector de tarea ($h_{zs} + v$) sirve como una cota superior para dNTP.
  $$LMSE(f) = \mathbb{E} \left[ \| h_{icl} - (h_{zs} + f(Z)) \|_2^2 \right]$$
• Mapeo Lineal: Para preservar la ventaja "libre de entrenamiento" del ICL (evitando actualizaciones de parámetros del modelo), LTV restringe la función de extracción $f$ a un mapeo lineal. El vector de tarea $v$ se estima como una transformación lineal del estado oculto sin demostraciones $h_{zs}$:
  $$v = W h_{zs}(x)$$
• Solución de Forma Cerrada: El mapeo lineal óptimo $W^*$ se calcula mediante regresión ridge sobre un conjunto de consultas de entrenamiento sin etiquetas. La regresión tiene como objetivo la diferencia entre el estado oculto del ICL y el estado oculto sin demostraciones ($h_{icl} - h_{zs}$).
  $$W^* = \arg\min_W \left( \| Y - W H \|_F^2 + \lambda \| W \|_F^2 \right)$$
  donde $H$ contiene estados ocultos sin demostraciones y $Y$ contiene las diferencias objetivo. Esto produce una solución de forma cerrada, garantizando eficiencia.

Contribuciones Clave

1. Métrica dNTP: Introducción de una métrica que cuantifica la discrepancia entre las distribuciones de inferencia basadas en TV y las basadas en ICL. El documento demuestra que dNTP es un proxy de rendimiento fuerte, correlacionándose negativamente con la precisión en tareas posteriores.
2. Vector de Tarea Lineal (LTV): Desarrollo de un método de extracción de vectores de tarea que minimiza dNTP mediante un mapeo lineal de forma cerrada. Este enfoque evita la optimización iterativa y el ajuste fino del modelo.
3. Superioridad Empírica: Experimentos extensivos que muestran que LTV supera a las líneas base existentes libres de entrenamiento (Function Vector, Task Vector, State Vector, I2CL) en ocho puntos de referencia de clasificación y cinco LLM.
4. Transferencia Inter-modelo: Demostración de que los vectores de tarea extraídos de un modelo más grande (por ejemplo, 72B) pueden transferirse a un modelo más pequeño (por ejemplo, 7B) para mejorar el rendimiento, abordando escenarios donde el modelo de inferencia tiene capacidad o longitud de contexto limitadas.

Resultados Experimentales

• Rendimiento en Clasificación: A través de cinco LLM (incluyendo variantes de LLaMA-2/3 y Qwen) y ocho puntos de referencia (por ejemplo, SST-2, AGNews, DBPedia), LTV logró la mayor precisión promedio. Específicamente, en LLaMA-2-13B, LTV mejoró la precisión promedio en un 9.2% sobre la mejor línea base (State Vector), alcanzando el 79.46%.
• Eficiencia: LTV mantiene una latencia de inferencia comparable a la inferencia sin demostraciones (aprox. 0.0226s por muestra), significativamente más rápida que el ICL. Aunque la extracción requiere un costo de pre-cálculo único, este se amortiza sobre el conjunto de prueba.
• Tareas de Regresión: LTV también fue evaluado en tareas de regresión lineal y ReLU, logrando el Error Cuadrático Medio (MSE) más bajo entre todas las líneas base de vectores de tarea.
• Robustez: El método permanece robusto ante variaciones en hiperparámetros (número de consultas de entrenamiento $N$ y regularización $\lambda$) y el número de demostraciones ($k$).
• Transferibilidad: Transferir LTV de un modelo Qwen-2.5-72B a un modelo Qwen-2.5-7B mejoró la precisión promedio de clasificación del modelo más pequeño en un 6.43%, igualando incluso el rendimiento del propio LTV del modelo más grande en algunos casos.

Significado y Afirmaciones

El documento afirma que la alineación distribucional es un criterio fundamental y útil para diseñar métodos de extracción de vectores de tarea. Al desplazar el enfoque de evaluación de la precisión indirecta en tareas posteriores a la discrepancia distribucional directa (dNTP), los autores proporcionan un marco fundamentado para comprender y mejorar los vectores de tarea.

El significado de LTV radica en su capacidad para:

1. Cerrar la brecha entre el objetivo teórico del ICL (inferencia bayesiana implícita) y la implementación práctica (vectores compactos) sin requerir actualizaciones del modelo.
2. Ofrecer una alternativa computacionalmente eficiente al ICL, reduciendo la latencia de inferencia mientras se mantiene una alta precisión.
3. Habilitar la transferencia inter-modelo, sugiriendo que las representaciones de tarea extraídas de modelos más grandes y capaces pueden servir como un recurso para modelos más pequeños con contexto o capacidad limitados.

Los autores concluyen que LTV ofrece un marco más eficiente y efectivo para la compresión de tareas en LLM, validando la hipótesis de que minimizar la discrepancia entre las distribuciones de TV e ICL conduce a una adaptación de tareas superior.