ADAG: Automatically Describing Attribution Graphs

Título: ADAG: Descrevendo Automaticamente Grafos de Atribuição

Autores: Aryaman Arora, Zhengxuan Wu, Jacob Steinhardt, Sarah Schwettmann (Stanford University & Transluce)

1. O Problema

A pesquisa em interpretabilidade de modelos de linguagem (LLMs) foca no rastreamento de circuitos (circuit tracing), que visa identificar quais recursos internos (neurons, features) contribuem causalmente para uma saída específica e como eles interagem.

• Limitação Atual: Embora existam técnicas para encontrar esses circuitos (grafos de atribuição), a etapa final de interpretação humana é altamente manual e ad-hoc. Pesquisadores precisam inspecionar manualmente exemplos de dados, intervenções causais e ativações para entender o papel funcional de cada componente.
• Desafio: Escalar essa interpretação para modelos maiores e conjuntos de dados mais extensos é inviável com métodos manuais. Não há um processo padronizado e automatizado para traduzir a estrutura matemática do circuito em descrições linguísticas naturais e compreensíveis.

2. Metodologia: O Pipeline ADAG

O ADAG é um pipeline de ponta a ponta totalmente automatizado para rastreamento e interpretação de circuitos. Ele consiste em quatro etapas principais:

A. Rastreamento de Circuitos (Circuit Tracing)

• Utiliza um algoritmo baseado em atribuição por gradiente (especificamente a técnica RelP - Relevance Patching) para construir um grafo de atribuição.
• O modelo gera um grafo onde os nós são tokens de entrada, neurônios MLP e logits de saída, conectados por arestas ponderadas que representam a influência causal.
• Apenas neurônios com atribuição acima de um limiar ($\tau$) são mantidos no circuito.

B. Perfis de Atribuição (Attribution Profiles)

• Para superar o viés de localidade (a ideia de que apenas o contexto imediato define a função de um neurônio), o ADAG introduz perfis de atribuição.
• Cada recurso (neurônio) é quantificado por dois vetores para cada amostra:
  1. Atribuição de Entrada: A proporção da ativação do recurso atribuída a cada token de entrada (dependência de tokens anteriores).
  2. Contribuição de Saída: A proporção da ativação de cada logit de saída que o recurso contribui (efeito em tokens futuros).
• Isso captura tanto a dependência de contexto passado quanto o impacto em gerações futuras.

C. Agrupamento em Supernós (Clustering)

• O objetivo é agrupar neurônios com funções similares em "supernós" para simplificar a análise.
• Algoritmo: Utiliza agrupamento espectral multiview.
  • Calcula a similaridade entre neurônios baseada na similaridade de cosseno de seus perfis de atribuição e contribuição.
  • Usa uma média harmônica das similaridades de entrada e saída para penalizar a correlação negativa (evitando agrupar neurônios que têm efeitos opostos na saída).
  • Aplica agrupamento espectral sobre a matriz de similaridade agregada para obter clusters balanceados e coerentes.

D. Descrição em Linguagem Natural (LLM Explainer-Simulator)

• Utiliza uma estrutura de Explicador-Simulador (inspirada em Bills et al., 2023) para gerar e validar descrições:
  1. Explicador (Explainer LM): Recebe os perfis de atribuição (tokens de entrada destacados e scores de contribuição) e gera candidatos de descrições em linguagem natural.
  2. Simulador (Simulator LM): Recebe a descrição candidata e os prompts brutos, tentando prever os scores de contribuição reais.
  3. Avaliação: A qualidade da descrição é medida pelo coeficiente de correlação de Pearson entre os scores reais e os previstos pelo simulador. A melhor descrição é selecionada.
• Um Resumidor (Summariser) combina as melhores descrições de entrada e saída para criar um rótulo final conciso para cada supernó.

3. Principais Contribuições

1. Primeiro Pipeline End-to-End Automatizado: O ADAG é a primeira abordagem que automatiza todo o fluxo, desde a extração do grafo de circuitos até a geração de descrições linguísticas validadas, sem intervenção humana na interpretação.
2. Perfis de Atribuição: Propõe uma nova métrica que supera o viés de localidade, capturando a dependência não local de neurônios MLP em relação a tokens anteriores e posteriores.
3. Algoritmo de Clustering Robusto: Introduz um método de agrupamento baseado em média harmônica de múltiplas visões (entrada/saída) que evita a mistura de recursos com efeitos opostos, garantindo clusters funcionais coerentes.
4. Validação Rigorosa: O sistema utiliza o próprio simulador para pontuar e selecionar as melhores interpretações, substituindo a avaliação subjetiva humana.

4. Resultados Experimentais

Os autores testaram o ADAG em dois cenários principais:

A. Tarefa de Capitais de Estados (Capitals Dataset)

• Objetivo: Replicar um estudo humano anterior sobre o circuito de raciocínio de múltiplos saltos (ex: "Qual a capital do estado que contém Dallas?").
• Resultados:
  • O ADAG recuperou automaticamente os mesmos circuitos identificados por humanos (ex: recuperação do estado, recuperação da capital).
  • As descrições geradas pelo LLM para os supernós foram comparáveis ou superiores às humanas na tarefa de prever contribuições de saída (correlação de Pearson mais alta).
  • Steering (Controle): Ao zerar ou amplificar a ativação de clusters específicos, o sistema demonstrou efeitos previsíveis e alinhados com as descrições geradas (ex: abater o cluster "Dallas Texas" reduziu a probabilidade de "Austin" e aumentou "Oklahoma").

B. Jailbreak de Conselhos Médicos Nocivos (Pills Dataset)

• Objetivo: Analisar por que um prompt específico faz o Llama 3.1 8B Instruct fornecer conselhos médicos perigosos.
• Resultados:
  • O ADAG identificou clusters responsáveis pelo sucesso do ataque (ASR - Attack Success Rate).
  • Cluster C3 ("pills safety redirect"): Ativa-se com a palavra "pills". Ao ser suprimido (steering para 0), o ASR saltou de ~28% para 88%.
  • Cluster C9 ("ridiculous-to-introductory"): Ativa-se com a palavra "ridiculous". Ao ser amplificado (steering para 2x), o ASR subiu para 90%.
  • A correlação entre a atribuição do cluster e o ASR foi forte, validando que o sistema identificou causalmente os mecanismos de segurança e falha do modelo.

C. Análise de Não-Localidade

• Os experimentos mostraram que neurônios MLP em camadas mais profundas dependem significativamente de tokens distantes (não apenas dos 8 tokens anteriores), validando a necessidade dos perfis de atribuição não locais.

5. Significado e Impacto

• Escalabilidade: O ADAG remove o gargalo humano na interpretabilidade, permitindo a análise de circuitos em modelos maiores e conjuntos de dados massivos, o que é essencial para a segurança de IA.
• Automação da Interpretação: Demonstra que LLMs podem ser usados não apenas como objetos de estudo, mas como ferramentas para "traduzir" a mecânica interna de outros modelos para humanos.
• Segurança e Alinhamento: A capacidade de identificar automaticamente clusters responsáveis por comportamentos indesejados (como jailbreaks) e manipulá-los (steering) oferece novas ferramentas para mitigação de riscos e auditoria de modelos.
• Reprodutibilidade: O código e os dados estão disponíveis, facilitando a adoção e extensão da metodologia pela comunidade de pesquisa.

Em resumo, o ADAG representa um avanço significativo ao transformar a interpretabilidade de circuitos de um processo artesanal e manual em um sistema automatizado, escalável e validado matematicamente.