Resumo Técnico: Head Pursuit: Probing Attention Specialization in Multimodal Transformers

Declaração do Problema

Modelos generativos de larga escala, incluindo tanto modelos de linguagem unimodais (LLMs) quanto modelos de visão-linguagem multimodais (VLMs), alcançaram um desempenho notável em diversas tarefas. No entanto, seus mecanismos internos permanecem apenas parcialmente compreendidos. Embora ferramentas de interpretabilidade anteriores, como o Logit Lens, tenham revelado informações semânticas nos estados ocultos, elas são frequentemente aplicadas de forma heurística a exemplos individuais, tornando difícil generalizar as descobertas ou quantificar a importância de componentes específicos (como cabeças de atenção) na moldagem da saída do modelo. Há uma necessidade de um método principiado para identificar quais cabeças de atenção se especializam em atributos semânticos ou visuais específicos e determinar se a manipulação dessas cabeças permite o controle direcionado sobre o comportamento do modelo sem treinamento adicional.

Metodologia

Os autores propõem um framework chamado Head Pursuit, que reinterpreta a prática de sondar ativações intermediárias através da lente do processamento de sinais esparsos. A metodologia central envolve as seguintes etapas:

1. Decomposição Esparsa via SOMP: Os autores empregam o Simultaneous Orthogonal Matching Pursuit (SOMP), um algoritmo guloso para aproximar sinais de alta dimensão com combinações lineares esparsas de elementos de base.

   • Entrada: Para uma determinada cabeça de atenção $h$ na camada $l$, a matriz de ativação $H_{h,l}$ (representando $n$ amostras) é analisada.
   • Dicionário: Em vez de aprender um dicionário a partir das ativações, o método utiliza a matriz de unembedding $D$ do modelo como um dicionário fixo. Cada linha de $D$ corresponde a um token no vocabulário, fornecendo um link direto para conceitos semânticos interpretáveis por humanos.
   • Processo: O SOMP seleciona iterativamente átomos do dicionário (tokens) que maximamente correlacionam com os resíduos da cabeça em todas as amostras do conjunto de dados. Isso gera uma matriz de coeficientes esparsa $W^*$ tal que $H \approx W^*D$.
   • Generalização: Diferente do Logit Lens, que projeta um único vetor residual, o SOMP opera em múltiplas amostras simultaneamente e seleciona múltiplas direções semânticas, proporcionando uma caracterização robusta e ao nível do conjunto de dados da função funcional de uma cabeça.

2. Seleção e Classificação de Cabeças:

   • Para identificar cabeças relevantes para um conceito alvo específico (ex: "cores", "toxicidade", "números"), o dicionário de unembedding é restrito aos tokens associados a esse conceito.
   • O SOMP é aplicado a este dicionário específico do conceito.
   • As cabeças são classificadas com base na fração de variância nas suas ativações explicada pela reconstrução do SOMP. Uma alta explicação de variância indica forte especialização para esse conceito.

3. Intervenção:

   • Uma vez identificadas as top-$k$ cabeças especializadas, os autores intervêm redimensionando a contribuição delas para o fluxo residual.
   • Intervenção Inibitória: O sinal das representações da cabeça é invertido ($\alpha = -1$) para suprimir o atributo alvo.
   • Intervenção de Reforço: As representações da cabeça são amplificadas ($\alpha > 1$) para promover a geração do atributo alvo.
   • Esta intervenção é realizada durante o passe direto (forward pass) sem atualizar os pesos do modelo.

Principais Contribuições

1. Enquadramento Teórico: O artigo enquadra o Matching Pursuit (MP) e sua variante multi-amostra, o SOMP, no contexto da interpretabilidade de modelos generativos, estabelecendo uma conexão formal com ferramentas padrão como o Logit Lens, enquanto o generaliza para lidar com múltiplas amostras e múltiplas direções semânticas simultaneamente.
2. Descoberta de Especialização: O método revela padrões consistentes de especialização de cabeças em transformers unimodais (Mistral-7B) e multimodais (LLaVA-NeXT, Gemma3, Qwen2.5-VL). Descobriu-se que cabeças específicas geram de forma confiável grupos de tokens semanticamente coerentes (ex: nomes de países, meses, números, sentimentos).
3. Controlabilidade: Os autores demonstram que intervir em uma fração muito pequena de cabeças (tão pouco 1% do total) pode controlar de forma confiável a supressão ou o reforço de conceitos alvo. Isso sugere que as camadas de atenção contêm uma estrutura linear altamente interpretável e manipulável.

Resultados

Os autores validaram sua abordagem em várias tarefas:

• Resposta a Perguntas (TriviaQA): Inverter os sinais de apenas 8 cabeças especializadas (0,8% do modelo) dedicadas a "nomes de países" degradou significativamente o desempenho em perguntas relacionadas a países, tendo um impacto mínimo em outras perguntas. A seleção de cabeças aleatórias ou a seleção baseada em Logit Lens não mostraram esse efeito direcionado.
• Mitigação de Toxicidade: Usando uma lista restrita de palavras-chave tóxicas para selecionar cabeças, os autores conseguiram reduzir com sucesso a frequência de gerações tóxicas nos conjuntos de dados RealToxicityPrompts (RTP) e Thoroughly Engineered Toxicity (TET). Inverter 32 cabeças especializadas reduziu as gerações tóxicas em aproximadamente 30-50% em comparação com as linhas de base, enquanto intervenções aleatórias frequentemente aumentavam a toxicidade.
• Classificação de Imagem (LLaVA): Inverter as top-32 cabeças selecionadas via SOMP para conjuntos de dados específicos (ex: MNIST, SVHN) interrompeu significativamente a precisão da classificação. O estudo descobriu que conjuntos de dados com semânticas relacionadas (ex: tarefas de reconhecimento de dígitos) compartilham conjuntos sobrepostos de cabeças especializadas, confirmando a sobreposição funcional entre domínios.
• Legenda de Imagens (Flickr30k):
  • Inibição: Inverter 16 cabeças especializadas para atributos como "cores", "sentimentos" ou "quantidade" removeu quase completamente essas palavras-chave das legendas, mantendo altos escores CIDEr (qualidade da legenda).
  • Reforço: Amplificar essas cabeças ($\alpha=5$) aumentou a presença de conceitos alvo em mais de 60% com apenas uma degradação marginal na qualidade da legenda.

Significância e Alegações

O artigo afirma que as camadas de atenção em grandes modelos generativos possuem um subespaço linear estruturado e alinhado a conceitos. Esta descoberta apoia a hipótese de que conceitos de alto nível são representados em subespaços residuais aproximadamente lineares.

A significância deste trabalho reside no fornecimento de ferramentas simples e livres de treinamento para entender e editar modelos de larga escala. Ao identificar que um pequeno número de unidades pode exercer um amplo controle semântico, os autores oferecem uma alternativa fundamentada às intervenções heurísticas baseadas em prompts ou ao fine-tuning. Esta abordagem permite:

• Interpretabilidade: Decompor o comportamento do modelo em um pequeno conjunto de elementos interpretáveis.
• Controle: Supressão ou reforço direcionado de atributos específicos (ex: toxicidade, conceitos visuais específicos) com impacto mínimo nas capacidades gerais do modelo.

Os autores mantêm-se modestos quanto às limitações, observando que o método assume linearidade (o que pode não capturar todas as estruturas não lineares), depende da qualidade do dicionário semântico e utiliza um redimensionamento global simples em vez de um direcionamento (steering) refinado. Eles sugerem que trabalhos futuros poderiam explorar intervenções específicas de posição ou aplicar estas técnicas a tarefas de geração de imagem.