Soft Head Selection for Injecting ICL-Derived Task Embeddings

1. El Problema

Las Grandes Modelos de Lenguaje (LLMs) se adaptan a tareas específicas principalmente mediante dos enfoques:

• Ajuste Fino Eficiente en Parámetros (PEFT): Como LoRA o (IA)³. Ofrecen alto rendimiento pero requieren entrenamiento y modifican los pesos del modelo.
• Aprendizaje en Contexto (ICL): Utiliza ejemplos en el prompt. Es libre de entrenamiento pero aumenta la longitud del prompt y la sobrecarga de inferencia.

Recientemente, se propuso la adaptación basada en embeddings impulsada por ICL, donde se extraen "vectores de tarea" de las activaciones intermedias del modelo (usando prompts de pocos ejemplos) e inyectan durante la inferencia. Sin embargo, los métodos existentes han demostrado ventajas empíricas limitadas frente al PEFT o al ICL tradicional. Un problema central es que estos métodos no identifican eficazmente dónde inyectar la información, ya que asumen que todas las cabezas de atención son igualmente relevantes o utilizan heurísticas fijas, ignorando que la importancia de las cabezas de atención varía significativamente entre tareas.

2. Metodología: SITE

Los autores proponen SITE (Soft head-selection for ICL-derived Task Embeddings), un método que identifica y utiliza las cabezas de atención relevantes para una tarea específica mediante una optimización continua basada en gradientes. El proceso consta de tres etapas:

1. Construcción de Embeddings de Tarea:

   • Se utilizan $M$ prompts de ICL (pocos ejemplos) para realizar una pasada hacia adelante.
   • Se extraen las activaciones del último token de todas las cabezas de atención en todas las capas.
   • Se promedian estas activaciones para obtener un "embedding de tarea" ($t^{(l,h)}$) para cada cabeza de atención $h$ en cada capa $l$.

2. Optimización de Selección de Cabezas Suave (Soft Head-Selection):

   • Se introduce una matriz aprendible $A$ con parámetros $\alpha^{(l,h)} \in [0, 1]$ para cada cabeza de atención.
   • Durante la optimización, se mantiene el LLM congelado y solo se actualiza $A$.
   • La inyección se realiza interpolando linealmente entre la activación original ($o^{(l,h)}$) y el embedding de tarea ($t^{(l,h)}$):
     $$o^{(l,h)} \leftarrow (1 - \alpha^{(l,h)}) \cdot o^{(l,h)} + \alpha^{(l,h)} \cdot t^{(l,h)}$$
   • Los parámetros $\alpha$ se optimizan mediante descenso de gradiente (Adam) minimizando la pérdida de entropía cruzada en una tarea de predicción de token siguiente (usando prompts de 0-shot). Esto permite que el modelo "aprenda" qué cabezas son críticas para la tarea.

3. Inferencia en 0-shot:

   • Una vez optimizados los parámetros $\alpha$, se aplican durante la inferencia real.
   • La inyección ocurre una sola vez en el último token del prompt de entrada inicial (asumiendo caché KV), inyectando la información de la tarea en la caché para guiar la generación posterior sin necesidad de ejemplos en el prompt.

3. Contribuciones Clave

• Selección de Cabezas Basada en Gradientes: A diferencia de métodos anteriores que buscan heurísticamente o usan espacios de búsqueda restringidos, SITE formula la selección de cabezas como un problema de optimización continua, logrando una identificación precisa de las cabezas relevantes.
• Rendimiento Superior con Menos Parámetros: SITE supera a los métodos de adaptación basados en embeddings anteriores y al ICL de pocos ejemplos, acercándose al rendimiento del PEFT, pero utilizando órdenes de magnitud menos parámetros entrenables (solo los escalares de selección de cabezas).
• Insights Mecanicistas (Interpretabilidad): Mediante análisis de "parcheo de activación" (activation patching), los autores demuestran que:
  • Las cabezas seleccionadas capturan información específica de la tarea.
  • Existe una fuerte especificidad de tarea: tareas similares comparten cabezas de atención importantes, mientras que tareas disímiles no. Esto extiende la comprensión de las funciones de las cabezas de atención desde una perspectiva agnóstica a una específica por tarea.

4. Resultados Experimentales

Los experimentos se realizaron en 12 LLMs (familias Llama-3.1, Mistral, Qwen3, Gemma-3) con tamaños de 4B a 70B parámetros, y en 4 benchmarks diversos (FV, ANLI, MMLU-Pro, Big-Bench Hard).

• Rendimiento General: SITE supera consistentemente al ICL de 10 ejemplos y a todos los métodos de adaptación basados en embeddings previos (como FV, TV, MTV, I2CL).
  • En el benchmark FV (57 tareas), SITE con $M=50$ supera al ICL de 10 ejemplos en un 13.26% en promedio.
  • En tareas de generación abierta y razonamiento, el rendimiento es comparable o superior al PEFT (LoRA, (IA)³).
• Eficiencia de Parámetros: Mientras que LoRA requiere ~3.4M de parámetros entrenables, SITE solo requiere ~1K (para Llama-3.1-8B), ya que solo optimiza los escalares de selección de cabezas.
• Generalización: El método funciona robustamente en modelos de diferentes tamaños y arquitecturas, mostrando mejoras de un 10-14% sobre el ICL de 10 ejemplos en todos los modelos probados.
• Análisis de Parcheo:
  • Intra-tarea: Parchear las cabezas de alto $\alpha$ (seleccionadas) mejora drásticamente el rendimiento, mientras que parchear las de bajo $\alpha$ lo degrada.
  • Inter-tarea: Usar parámetros de selección de una tarea similar mejora el rendimiento de la tarea objetivo, confirmando la relación funcional entre tareas similares.

5. Significado e Impacto

El trabajo de SITE representa un avance significativo en la adaptación de LLMs sin entrenamiento (training-free) y eficiente.

• Viabilidad Operativa: Al requerir solo la inyección de un embedding precalculado y parámetros de selección fijos, ofrece una alternativa ligera al PEFT para despliegues donde el entrenamiento es costoso o imposible.
• Comprensión del Modelo: Proporciona una nueva lente mecánica para entender cómo los LLMs procesan tareas, revelando que la funcionalidad de las cabezas de atención es altamente dependiente del contexto de la tarea, no universal.
• Eficiencia: Mantiene la eficiencia de la inferencia en 0-shot (sin longitud de prompt excesiva) mientras logra un rendimiento cercano al de modelos ajustados, resolviendo el compromiso tradicional entre eficiencia y rendimiento en la adaptación de modelos.

En resumen, SITE demuestra que la selección inteligente y aprendida de dónde inyectar información derivada de ICL es la clave para desbloquear el potencial de la adaptación basada en embeddings, superando las limitaciones de los enfoques anteriores.