Dual Path Attribution: Efficient Attribution for SwiGLU-Transformers through Layer-Wise Target Propagation

Título: Dual Path Attribution (DPA): Atribuição Eficiente para SwiGLU-Transformers através de Propagação de Alvo em Camadas

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) baseados em Transformers são dominantes na IA, mas seus mecanismos internos permanecem opacos. A atribuição (identificar quais entradas ou componentes internos são responsáveis por uma previsão) é crucial para a interpretabilidade, mas enfrenta dois desafios principais:

• Custo Computacional Proibitivo: Métodos existentes que oferecem alta fidelidade (como Activation Patching ou intervenções causais) exigem múltiplas passagens para frente (forward passes), uma para cada componente ou interação. Isso torna a atribuição densa (para todos os neurônios e cabeças de atenção) inviável em modelos grandes e sequências longas.
• Compromisso Fidelidade-Eficiência: Métodos mais rápidos (baseados em gradientes ou atenção) frequentemente sofrem com erros de aproximação, ruído ou saturação de gradientes, falhando em capturar interações não lineares complexas.

2. Metodologia: Dual Path Attribution (DPA)

O DPA propõe um novo framework de decomposição que rastreia o fluxo de informações em um único forward pass e um único backward pass, sem necessidade de exemplos contrafactuais. A abordagem é específica para arquiteturas SwiGLU (usadas em Llama, Mistral, Qwen, etc.), explorando sua estrutura bilinear.

Conceitos Chave:

• Propagação de Alvo Invertida (Target Propagation): Em vez de passar cada componente para frente para ver seu efeito, o DPA propaga o vetor de unembedding do token alvo (o objetivo) para trás, através da rede congelada, até a entrada.
• Decomposição em Dois Caminhos: O DPA decompõe analiticamente os módulos do Transformer (Atenção e GLU) em dois caminhos distintos que podem ser linearizados e invertidos:
  1. Caminho de Conteúdo (Content Pathway): Transporta as representações de dados (ex: projeções de valor e up-projection).
  2. Caminho de Controle (Control Pathway): Governa o roteamento e a transformação (ex: pontuações de atenção, gate scores, queries e keys).
• Linearização e Inversão:
  • Para o GLU (Gated Linear Unit), a transformação é dividida em caminhos de up-projection (linear) e gate-projection (não-linear, linearizado via razão de ativação).
  • Para o MHSA (Multi-Head Self-Attention), a propagação é decomposta nos caminhos de Valor (linear), Query e Key. Os caminhos de Query/Key exigem uma decomposição de Taylor da função softmax para lidar com a mistura de tokens de forma estável.
• Complexidade O(1): Ao propagar um único vetor de alvo efetivo recursivamente de cima para baixo, o DPA calcula a pontuação de atribuição para qualquer componente via um simples produto escalar com as ativações cacheadas. Isso elimina a dependência do número de componentes ($M$), resultando em complexidade temporal constante em relação a $M$.

3. Contribuições Principais

1. Novo Framework de Decomposição: Introdução do DPA, que atribui pontuações de logits de forma eficiente e fiel sem exemplos contrafactuais.
2. Decomposição Estrutural: Demonstra que o grafo computacional de Transformers SwiGLU pode ser decomposto em caminhos de controle e conteúdo distintos, permitindo a inversão analítica.
3. Eficiência e Fidelidade: Alcança o estado da arte em fidelidade (identificação precisa de mecanismos causais) com uma eficiência computacional sem precedentes, escalando para atribuição densa em sequências longas.

4. Resultados Experimentais

Os autores avaliaram o DPA em benchmarks padrão de interpretabilidade (Known 1000, SQuAD v2.0, IMDb, Indirect Object Identification) em diversos modelos (Llama-3.1-8B, Llama-2-7B, Qwen3, Qwen2.5, Mistral-7B).

• Fidelidade de Atribuição de Tokens:
  • O DPA superou consistentemente todas as baselines (incluindo Integrated Gradients, DePass, e métodos baseados em atenção) nas métricas de Compreensividade (disrupção) e Suficiência (recuperação).
  • Em tarefas complexas como análise de sentimento (IMDb), o DPA identificou tokens críticos com maior precisão do que métodos baseados apenas em atenção.
• Fidelidade de Atribuição de Componentes:
  • Na tarefa de Identificação de Objeto Indireto (IOI), o DPA isolou circuitos causais com precisão extrema, alcançando zero disrupção e alta recuperação, superando métodos como Attn-only e Gradient-based.
  • O DPA identificou mecanismos causais (cabeças de atenção e neurônios MLP) muito mais cedo na curva de recuperação do que as baselines.
• Eficiência Computacional:
  • Velocidade: O DPA foi 40x mais rápido que métodos baseados em Activation Patching (AP) e DePass em tarefas de contexto longo (SQuAD v2.0). Enquanto AP e DePass levaram ~1,5 horas, o DPA completou em ~2 minutos.
  • Escalabilidade: A complexidade temporal é $O(1)$ em relação ao número de componentes, enquanto AP é $O(M)$ e DePass é $O(M/B)$.
  • Memória: O uso de memória de pico é comparável às baselines, pois requer apenas o cache das ativações de uma passagem para frente.

5. Significado e Limitações

• Significado: O DPA preenche a lacuna crítica entre a fidelidade mecânica e a eficiência computacional. Ele permite a análise de "circuitos" internos em modelos grandes de forma prática, tornando a interpretabilidade densa viável para o desenvolvimento e depuração de LLMs.
• Limitações:
  • Especificidade Arquitetural: A decomposição matemática é otimizada para Transformers SwiGLU densos. Adaptá-lo para Mixture-of-Experts (MoE) ou modelos sem atenção (como State Space Models) requer a derivação manual de novas inversões.
  • Erro de Aproximação: A linearização local pode falhar em regiões de alta curvatura ou interações não lineares fortemente acopladas, embora funcione bem na maioria dos regimes.
  • Escopo de Logits: O método rastreia atribuição diretamente para o logit alvo, não capturando facilmente mecanismos indiretos onde a probabilidade aumenta devido à supressão ativa de logit concorrentes (pós-softmax).

Conclusão

O Dual Path Attribution representa um avanço significativo na interpretabilidade de LLMs. Ao transformar o problema de atribuição em uma propagação de alvo inversa e linearizada, o método oferece uma ferramenta robusta, rápida e precisa para desvendar os mecanismos de raciocínio interno dos modelos modernos, superando as limitações de custo e precisão das abordagens anteriores.