Soft Head Selection for Injecting ICL-Derived Task Embeddings

Título: Seleção Suave de Cabeças para Injeção de Embeddings de Tarefa Derivados de ICL (SITE)

1. O Problema

As Grandes Modelos de Linguagem (LLMs) são tradicionalmente adaptados a tarefas específicas por meio de duas abordagens principais:

1. Fine-Tuning Eficiente em Parâmetros (PEFT): Como LoRA ou Adapters. Oferece alto desempenho, mas requer treinamento e atualização de pesos (ou módulos adicionais), o que pode ser custoso em termos de armazenamento e computação.
2. Aprendizado em Contexto (ICL): Usa exemplos no prompt para guiar o modelo sem treinamento. É flexível, mas aumenta o custo de inferência (overhead) e o comprimento do prompt, além de ser menos eficiente em tarefas complexas.

Recentemente, surgiu uma terceira abordagem: Adaptação Baseada em Embeddings Acionada por ICL. Esta técnica extrai "vetores de tarefa" (embeddings) das ativações intermediárias do modelo durante o ICL (few-shot) e os injeta na inferência zero-shot.

• Limitação Atual: Métodos existentes para essa abordagem (como FV, TV, MTV) frequentemente falham em superar consistentemente o ICL few-shot ou o PEFT. Eles muitas vezes dependem de heurísticas para decidir onde injetar esses embeddings (em quais camadas ou cabeças de atenção), sem uma seleção precisa, resultando em desempenho inconsistente e ineficiência.

2. Metodologia: SITE (Soft Head Selection for ICL-derived Task Embeddings)

O SITE propõe uma nova metodologia baseada em gradiente para identificar e injetar embeddings de tarefa de forma otimizada. O processo ocorre em três etapas:

• Etapa 1: Construção do Embedding de Tarefa

  • O modelo executa prompts de ICL few-shot (ex: 50 prompts com 10 pares entrada-saída).
  • As ativações do último token de cada cabeça de atenção em cada camada são extraídas.
  • Essas ativações são médias para criar um embedding de tarefa específico para cada cabeça de atenção ($t^{(l,h)}$).

• Etapa 2: Otimização da Seleção Suave de Cabeças (Soft Head Selection)

  • Em vez de selecionar cabeças fixamente (binário), o SITE introduz uma matriz de parâmetros aprendíveis $A$, onde cada entrada $\alpha^{(l,h)}$ é um parâmetro contínuo entre 0 e 1.
  • Este parâmetro controla o grau de interpolação linear entre a ativação original da cabeça ($o^{(l,h)}$) e o embedding de tarefa ($t^{(l,h)}$):
    $$o^{(l,h)} \leftarrow (1 - \alpha^{(l,h)}) \cdot o^{(l,h)} + \alpha^{(l,h)} \cdot t^{(l,h)}$$
  • O modelo base (LLM) permanece congelado. Apenas a matriz $A$ é otimizada via descida de gradiente (Adam) para minimizar a perda de previsão do próximo token em prompts zero-shot.
  • Isso permite que o algoritmo aprenda quais cabeças de atenção são mais relevantes para uma tarefa específica e quanto da informação da tarefa deve ser injetada nelas.

• Etapa 3: Inferência Zero-Shot

  • Durante a inferência real, o modelo usa apenas o prompt de entrada (sem exemplos).
  • O embedding de tarefa pré-computado é injetado nas cabeças selecionadas (com pesos $\alpha$) apenas uma vez, no último token do prompt inicial.
  • A informação da tarefa é então propagada através do cache KV (Key-Value) durante a geração autoregressiva, sem necessidade de mais intervenções.

3. Contribuições Principais

1. Método SITE: Uma técnica de adaptação baseada em embeddings que supera o ICL few-shot e se aproxima do desempenho do PEFT, utilizando drasticamente menos parâmetros treináveis (apenas os escalares de seleção de cabeça, ~1K parâmetros vs. milhões no LoRA).
2. Seleção Baseada em Gradiente: Substitui heurísticas de busca por uma otimização contínua suave, permitindo uma identificação precisa das cabeças de atenção relevantes para cada tarefa.
3. Insights Mecanísticos: Através de análises de "activation patching" (correção de ativação), o paper demonstra que:
   • A importância das cabeças de atenção varia significativamente entre tarefas (especificidade da tarefa).
   • Tarefas semanticamente similares compartilham cabeças de atenção importantes, enquanto tarefas dissimilares não.
4. Generalidade: O método foi validado em 12 LLMs diferentes (famílias Llama, Mistral, Qwen, Gemma) com tamanhos de 4B a 70B parâmetros.

4. Resultados Experimentais

Os experimentos foram conduzidos em 4 benchmarks (FV, ANLI, MMLU-Pro, Big-Bench Hard) e 12 modelos.

• Desempenho:
  • O SITE supera consistentemente todos os métodos baseados em embeddings anteriores e o ICL few-shot (10-shot).
  • No benchmark FV (57 tarefas), o SITE com $M=50$ prompts atingiu 90.02% de precisão, superando o ICL 10-shot (76.76%) em mais de 13 pontos percentuais.
  • O desempenho é comparável ao LoRA (PEFT) em tarefas de geração e compreensão, embora em tarefas de raciocínio complexo (MMLU-Pro, BBH) o LoRA ainda tenha uma leve vantagem.
• Eficiência:
  • Parâmetros: O SITE otimiza apenas ~1.02K parâmetros (vs. ~3.4M do LoRA).
  • Tempo: O tempo de inferência é quase idêntico ao do zero-shot (sem exemplos), sendo significativamente mais rápido que o ICL few-shot e métodos que exigem injeção em múltiplos tokens.
• Análise de Patching:
  • Intra-tarefa: Injetar embeddings apenas nas cabeças de alta $\alpha$ (selecionadas) melhora drasticamente o desempenho. Injetar em cabeças de baixa $\alpha$ degrada o desempenho.
  • Inter-tarefa: Usar cabeças selecionadas de uma tarefa similar (ex: tradução Inglês-Francês) em outra tarefa similar (Inglês-Alemão) mantém o alto desempenho, enquanto usar cabeças de tarefas não relacionadas falha.

5. Significado e Impacto

O trabalho do SITE é significativo por várias razões:

• Ponte entre ICL e PEFT: Demonstra que é possível obter o melhor dos dois mundos: a flexibilidade e o baixo custo de parâmetros do ICL/Embeddings, com a eficácia de adaptação do PEFT.
• Interpretabilidade Mecanística: Oferece uma nova perspectiva sobre como os LLMs processam tarefas, mostrando que a funcionalidade das cabeças de atenção não é estática, mas altamente dependente do contexto da tarefa.
• Eficiência Operacional: Para aplicações em produção, o SITE permite adaptar modelos a novas tarefas sem retrainar o modelo inteiro ou carregar grandes vetores de adaptação, mantendo a latência de inferência próxima ao zero-shot.
• Limitações e Futuro: O método ainda requer alguns dados rotulados para construir os embeddings e depende de acesso às ativações internas do modelo (limitando-o a modelos open-source). O trabalho sugere futuras investigações em tarefas de raciocínio complexo e geração de longo contexto.

Em resumo, o SITE representa um avanço na adaptação de LLMs, provando que a seleção inteligente e suave de componentes internos do modelo (cabeças de atenção) pode ser mais eficaz do que a injeção bruta de informações ou o ajuste massivo de pesos.