Resumo Técnico: Poda de Nós Multigranular para Descoberta de Circuitos Causais

1. Definição do Problema

A descoberta de circuitos causais visa identificar subredes mínimas dentro de Grandes Modelos de Linguagem (LLMs) que impulsionam causalmente comportamentos específicos. Embora os métodos existentes tenham avançado na interpretabilidade mecanística, eles enfrentam duas limitações primárias:

1. Escalabilidade Computacional: As abordagens atuais dependem predominantemente da poda de arestas (remoção de conexões entre componentes). Como o número de arestas em um transformer cresce quadraticamente com o número de nós ($O(|N|^2)$), a poda ao nível de aresta torna-se computacionalmente intratável à medida que o tamanho do modelo aumenta. Além disso, esses métodos frequentemente exigem o armazenamento de representações do fluxo residual e ativações intermediárias destiladas, levando a um uso de memória proibitivo.
2. Limitações de Granularidade: Os métodos existentes operam tipicamente em uma granularidade grosseira, tratando cabeças de atenção inteiras ou blocos MLP como unidades atômicas. Isso ignora a possibilidade de que neurônios (nós) específicos dentro desses componentes sejam responsáveis por comportamentos alvo, resultando em circuitos recuperados que perdem estruturas finas e retêm componentes desnecessários.

2. Metodologia

Os autores propõem um framework de poda de nós multigranular que descobre circuitos mínimos através de múltiplos níveis de granularidade em um único passo de otimização unificado.

2.1 Design de Granularidade Hierárquica

Em vez de podar arestas entre componentes grosseiros, o framework aprende máscaras sobre nós em quatro níveis distintos:

• Nível de Bloco: Blocos de Atenção inteiros e Blocos MLP.
• Nível de Cabeça: Cabeças de Atenção individuais.
• Nível de Neurônio: Neurônios individuais dentro das camadas de Atenção e das duas camadas lineares dos blocos MLP (MLP1/Keys e MLP2/Values).

2.2 Configuração de Poda Causal de Dois Fluxos (Two-Stream)

Para identificar a importância causal sem introduzir mudanças de distribuição (um risco com a simples ablação zero), o método utiliza uma passagem direta de dois fluxos:

• Fluxo Limpo (Clean Stream): O modelo processa uma entrada projetada para elicitar um comportamento específico (ex: prever um ano > 50).
• Fluxo Corrompido (Corrupted Stream): O modelo processa uma versão minimamente perturbada da entrada (ex: inverter pronomes de gênero ou alterar o ano inicial) para elicitar uma previsão diferente.
• Interpolação: Para cada nó $n_i$, uma máscara aprendível $m_i \in [0, 1]$ interpola entre a ativação limpa ($h^{clean}$) e a ativação corrompida ($h^{corrupted}$):
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  A otimização busca encontrar máscaras que permitam ao modelo produzir a saída correta (correspondendo ao comportamento limpo) enquanto maximiza a esparsidade.

2.3 Otimização e Parametrização de Máscara

• Distribuição Hard Concrete: As máscaras são parametrizadas usando a distribuição Hard Concrete para permitir a otimização baseada em gradiente, enquanto induzem um comportamento aproximadamente binário (0 ou 1).
• Objetivo Unificado: O objetivo de treinamento combina três componentes:
  1. Perda de Fidelidade ($L_{faith}$): Uma divergência KL que garante que a distribuição de saída do circuito podado corresponda ao modelo base.
  2. Perda de Tarefa ($L_{task}$): Uma perda de ranking baseada em margem que garante que o circuito distinga corretamente entre opções de verdade fundamental (ground-truth) e distratores.
  3. Regularizador de Esparsidade ($L_0$): Incentiva as máscaras a serem zero, minimizando o tamanho do circuito.
• Consistência Hierárquica: Após o treinamento, as máscaras são binarizadas com uma restrição que garante que, se um nó pai (ex: um bloco MLP) for podado, todos os seus nós filhos (ex: neurônios dentro desse bloco) também sejam desativados.

3. Principais Contribuições

1. Framework Multigranular: A introdução de um framework de otimização unificado que aprende máscaras simultaneamente através de blocos de transformer, cabeças de atenção, blocos MLP e neurônios individuais.
2. Escalabilidade e Eficiência: O método evita a escala quadrática da poda de arestas e a necessidade de armazenar ativações intermediárias para todas as arestas. Ele reduz o uso de memória da GPU em 3× a 11× em comparação com métodos anteriores, permitindo a descoberta de circuitos no Llama 3.1-8B usando uma única GPU de 40GB.
3. Descoberta de Grão Fino: A abordagem identifica circuitos significativamente mais compactos do que os baselines, alcançando especificamente alta esparsidade nos níveis de bloco MLP e neurônio, algo que métodos anteriores falharam em realizar.

4. Resultados Experimentais

O método foi avaliado em três tarefas padrão de descoberta de circuitos (Identificação de Objeto Indireto - IOI, Pronomes de Gênero - GP, e Maior Que - GT) usando modelos GPT-2, GPT-2-XL e Llama.

4.1 Esparsidade e Compactação

• Poda de MLP e Neurônio: O método proposto supera significativamente os baselines de poda de aresta (EAP, EP) e a poda de peso não estruturada (DiscoGP). No cenário menos favorável, ele podou 33,34% mais blocos MLP e 59,8% mais neurônios do que os métodos anteriores.
• Estruturas Dependentes de Tarefa: Os circuitos descobertos exibem topologias distintas baseadas na tarefa:
  • IOI: Dominado por computação MLP distribuída; a atenção é amplamente podada, exceto nas camadas finais.
  • Pronomes de Gênero (GP): Altamente localizado e descontínuo; a computação MLP é restrita a algumas camadas, enquanto a atenção é esparsa.
  • Maior Que (GT): Exibe esparsidade extrema com blocos contíguos de camadas inativas e padrões de execução do tipo "skip-like".

4.2 Fidelidade e Consistência

• Desempenho: Os circuitos podados mantêm alta precisão na tarefa (frequentemente >95%) e baixa divergência KL em relação ao modelo completo.
• Consistência: Experimentos através de cinco sementes aleatórias mostram que os circuitos descobertos (métricas de fidelidade e perfis de esparsidade) são altamente estáveis.
• Análise Qualitativa: No Llama 3.2-1B, as cabeças de atenção sobreviventes na tarefa IOI foram encontradas realizando papéis funcionais específicos (IO-Movers, S-Inhibitors), espelhando circuitos descobertos manualmente no GPT-2, validando a capacidade do método de recuperar insights mecanísticos.

4.3 Eficiência de Computação

• Memória: O método requer pelo menos 3× menos memória do que os baselines de poda de aresta (ex: 6,27 GB vs. 72,79 GB para EAP no GPT-2).
• Tempo: O treinamento completa em menos de 6 minutos, significativamente mais rápido que os métodos de poda de aresta (aprox. 50 minutos), embora ligeiramente mais lento que o método EAP de passagem única (que troca tempo por um uso massivo de memória).

5. Significância e Alegações

O artigo afirma que sua principal significância reside em superar os gargalos de escalabilidade e granularidade dos atuais métodos de descoberta de circuitos. Ao mudar da poda de arestas para a poda de nós multigranular, os autores demonstram que:

• Muitos neurônios retidos por métodos de granularidade grosseira são desnecessários e podem ser removidos com impacto negligenciável no desempenho.
• Modelos maiores (ex: Llama 8B) suportam circuitos específicos de tarefa tão fortes ou mais fortes quanto modelos menores, sugerindo que o aumento de escala introduz redundância em vez de exigir circuitos mais densos.
• O método fornece um caminho computacionalmente viável para descobrir circuitos causais de grão fino em modelos de grande escala, revelando que diferentes comportamentos são codificados por padrões distintos de atividade de atenção, MLP e neurônios.

Os autores reconhecem uma limitação: o framework identifica quais nós são necessários, mas não recupera explicitamente a estrutura de interação (arestas) entre os nós sobreviventes. Eles sugerem como uma direção futura a combinação da seleção de nós com a recuperação de arestas para inferir estruturas de interação.