Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis

Título: Localizando Reconhecimento de Tarefa e Aprendizado de Tarefa no Aprendizado em Contexto (ICL) via Análise de Cabeças de Atenção

1. O Problema

O Aprendizado em Contexto (In-Context Learning - ICL) é uma capacidade fundamental dos Grandes Modelos de Linguagem (LLMs) que lhes permite resolver novas tarefas a partir de demonstrações no prompt, sem atualização de pesos (fine-tuning). Embora o ICL seja amplamente utilizado, seu mecanismo interno permanece um "caixa preta".

Existem duas perspectivas dominantes na pesquisa atual que, até então, não foram reconciliadas:

1. Abordagem Introspectiva (Componente): Foca em componentes internos específicos, como as Cabeças de Indução (Induction Heads - IHs), identificando-as como cruciais para o ICL através de estudos de ablação. No entanto, essa abordagem muitas vezes não explica como essas cabeças realizam o ICL ou como se comportam sob diferentes entradas.
2. Abordagem Holística (Funcional): Decompõe o ICL em duas funções principais: Reconhecimento de Tarefa (Task Recognition - TR), que identifica o espaço de rótulos (ex: saber que as respostas são "positivo" ou "negativo"), e Aprendizado de Tarefa (Task Learning - TL), que aprende o mapeamento específico entre texto e rótulo. Esta abordagem não consegue, contudo, rastrear essas funções para componentes concretos do modelo (como cabeças de atenção específicas).

O problema central é a falta de um quadro unificado que conecte a precisão mecânica (quais cabeças fazem o quê) com a clareza funcional (como o modelo executa TR e TL).

2. Metodologia: TSLA e Análise Geométrica

Os autores propõem um novo framework baseado na Atribuição de Logits no Subespaço de Tarefa (Task Subspace Logit Attribution - TSLA) para identificar e localizar as cabeças de atenção especializadas em TR e TL.

• Definição de TR e TL via TSLA:
  • Reconhecimento de Tarefa (TR): O modelo deve reconhecer o conjunto de tokens de rótulos candidatos ($Y$). O método TSLA calcula a projeção da saída da cabeça de atenção no subespaço gerado pelos vetores de unembedding dos rótulos de demonstração. Cabeças com alta projeção neste subespaço são classificadas como cabeças TR.
  • Aprendizado de Tarefa (TL): O modelo deve aprender o mapeamento texto-rótulo específico. O TSLA mede a alinhação da saída da cabeça com a direção média entre o rótulo correto e os rótulos incorretos (dentro do subespaço de tarefa). Cabeças que aumentam o "gap" de logits entre o rótulo correto e os concorrentes são classificadas como cabeças TL.
• Análise Geométrica: Os autores utilizam análise geométrica dos estados ocultos (hidden states) para entender o papel causal:
  • As cabeças TR alinham os estados ocultos ao subespaço da tarefa (reduzindo o ângulo entre o estado e o subespaço dos rótulos).
  • As cabeças TL realizam uma rotação dentro do subespaço, orientando o estado oculto específico para o vetor de unembedding do rótulo correto.
• Validação Experimental:
  • Correlação e Sobreposição: Comparação das cabeças TR/TL identificadas com as Cabeças de Indução (IHs) e entre si.
  • Ablação de Canais: Remoção seletiva das cabeças TR e TL para observar o impacto na precisão e na "Taxa de Reconhecimento de Tarefa" (TR Ratio).
  • Perturbação de Entrada: Distorção dos textos de demonstração (quebra do mapeamento TL) ou alteração dos rótulos (quebra do reconhecimento TR) para testar a independência dos mecanismos.
  • Steering (Direcionamento): Injeção de vetores de tarefa (Task Vectors) derivados das cabeças TR e TL em prompts zero-shot para restaurar o desempenho.

3. Principais Contribuições

1. Localização Mecanística: Derivação do método TSLA para identificar cabeças de atenção específicas responsáveis pelos componentes TR e TL, mapeando estudos de perturbação de entrada para componentes concretos.
2. Unificação de Paradigmas: Demonstra que as Cabeças de Indução (IHs) são, na verdade, um subconjunto das cabeças de Reconhecimento de Tarefa (TR). Isso reconcilia a importância das IHs com a decomposição funcional TR/TL.
3. Mecanismos Geométricos Distintos: Revela que TR e TL operam através de transformações geométricas diferentes nos estados ocultos: TR alinha o estado ao subespaço da tarefa, enquanto TL rotaciona o estado dentro desse subespaço em direção ao rótulo correto.
4. Dependência de Tarefa: Mostra que a eficácia dos vetores de tarefa varia conforme o tipo de tarefa (classificação vs. geração), com TR sendo crucial para reconhecimento de rótulos fixos e TL sendo dominante em tarefas de geração aberta.

4. Resultados Chave

• Especialização das Cabeças:
  • As cabeças TR aparecem significativamente em camadas mais profundas do modelo do que as cabeças TL e IHs.
  • Existe uma forte sobreposição e correlação entre cabeças TR e IHs, enquanto as cabeças TL são distintas e variam mais entre tarefas.
  • Cabeças TR focam sua atenção nos rótulos de demonstração (para definir o espaço de busca), enquanto cabeças TL focam no texto da consulta (query) para extrair semântica e mapear para o rótulo.
• Ablação e Independência:
  • Remover cabeças TR causa um colapso tanto na precisão quanto na Taxa de Reconhecimento de Tarefa (TR Ratio) (o modelo começa a prever rótulos fora do conjunto de demonstração, como "random guess").
  • Remover cabeças TL reduz a precisão, mas mantém a TR Ratio alta (o modelo sabe quais são os rótulos possíveis, mas falha em escolher o correto).
  • Isso confirma que TR e TL são componentes separáveis e causalmente distintos.
• Steering e Zero-Shot:
  • A baixa performance zero-shot é causada principalmente pela falha no Reconhecimento de Tarefa. Injetar vetores derivados de cabeças TR restaura o desempenho para níveis de ICL.
  • Em tarefas de geração de texto (onde não há um conjunto fixo de rótulos), vetores derivados de cabeças TL são mais eficazes do que os de TR, pois capturam o mapeamento texto-rótulo complexo.
• Geometria:
  • A adição de saídas de cabeças TR aumenta drasticamente a alinhamento com o subespaço da tarefa.
  • A adição de saídas de cabeças TL aumenta o Logit Difference (discriminação entre rótulos) sem necessariamente alterar o alinhamento global com o subespaço.

5. Significado e Impacto

Este trabalho oferece uma explicação unificada e interpretável para o funcionamento do ICL em LLMs. Ao decompor o processo em Reconhecimento de Tarefa e Aprendizado de Tarefa e localizá-los em cabeças de atenção específicas, o artigo:

• Resolve contradições anteriores: Explica por que as Cabeças de Indução são tão importantes (elas são os motores do reconhecimento de espaço de rótulos) e por que às vezes falham (se o reconhecimento falhar, o aprendizado não importa).
• Fornece ferramentas de controle: Permite que pesquisadores e engenheiros intervenham seletivamente no modelo (via ablação ou steering) para melhorar o reconhecimento de tarefas ou o mapeamento específico, dependendo da necessidade.
• Avança a Interpretabilidade: Move a análise de ICL de observações macroscópicas (perturbação de entrada) para uma compreensão microscópica e geométrica (dinâmica de estados ocultos e contribuições de cabeças), estabelecendo as bases para o desenvolvimento de modelos mais robustos e controláveis.

Em resumo, o paper demonstra que o ICL não é um fenômeno monolítico, mas sim uma orqueração de mecanismos especializados onde cabeças profundas garantem que o modelo "saiba qual tarefa está fazendo" (TR) e cabeças subsequentes garantem que ele "faça a tarefa corretamente" (TL).