Explaining Attention with Program Synthesis

Enunciado do Problema
Compreender os cálculos realizados por redes neurais profundas em termos algorítmicos continua sendo um desafio significativo no aprendizado de máquina. Abordagens anteriores de interpretabilidade basearam-se amplamente na atribuição de significado a neurônios ou características, seja de forma top-down (via sondagens para conceitos definidos por humanos) ou bottom-up (via resumos de entrada/saída). Embora poderosos, esses métodos frequentemente falham em fornecer uma descrição formal completa da computação neural. Além disso, rotular componentes intermediários com descrições de linguagem natural pode resultar em interpretações que são ambíguas ou difíceis de formalizar. Os autores postulam que há uma necessidade de um meio que preencha a lacuna entre a complexidade de modelos de bilhões de parâmetros e as explicações em linguagem natural — um meio que seja tanto legível por humanos quanto formalmente verificável.

Metodologia
O artigo propõe um framework baseado em síntese de programas para explicar os cálculos de cabeças de atenção de modelos de linguagem transformer. O objetivo central é encontrar um programa Python executável ($\pi$) que possa mapear uma sequência de texto de entrada ($X$) diretamente para uma matriz de atenção ($A$), aproximando efetivamente a lógica de uma cabeça de atenção específica.

O framework procede em quatro etapas distintas:

1. Extração de Mapas de Atenção: Os autores extraem matrizes de ativação de atenção de verdade fundamental de um modelo alvo através de um corpus de sequências de treinamento (especificamente usando o dataset TinyStories por sua simplicidade estruturada). Eles filtram esses mapas para isolar os 2,5% superiores de pesos de atenção para focar em interações salientes de pares de tokens.
2. Síntese de Programas: Um LLM auxiliar, atuando como um agente de síntese, é estimulado com os padrões de atenção extraídos e sequências de entrada. O agente tem a tarefa de gerar funções Python executáveis que reproduzam esses padrões. O agente de síntese tem acesso a bibliotecas como NumPy, spaCy e NLTK para processamento linguístico e numérico.
3. Refinamento e Classificação: Candidatos a programas são validados quanto à sintaxe e pontuados contra padrões de atenção reais usando a Distância de Jensen-Shannon (JSD). Os melhores candidatos passam por um loop de refinamento condicionado por feedback, onde o agente de síntese contrasta padrões reais e previstos para produzir programas revisados.
4. Seleção: O programa substituto (proxy) ideal para cada cabeça é selecionado maximizando a similaridade (medida via Interseção sobre União, ou IoU) em sequências de dados de teste (held-out).

Principais Contribuições

• Interpretabilidade Programática: Os autores introduzem um método para substituir cabeças de atenção neurais por código simbólico e executável, oferecendo uma alternativa formalmente verificável às descrições de linguagem natural.
• Validação Causal via Substituição: Ao contrário de trabalhos anteriores que apenas observam correlações, esta abordagem valida as explicações substituindo diretamente componentes neurais aprendidos por programas sintetizados durante o forward pass do modelo.
• Escalabilidade: O método é aplicado a quatro arquiteturas transformer distintas: BERT-Base, GPT-2-Small, TinyLlama-1.1B e Llama-3B, demonstrando a aplicabilidade da abordagem em diferentes tamanhos e arquiteturas de modelos (bidirecionais e causais).

Resultos

• Alto Alinhamento: Uma fração substancial de cabeças de atenção em todos os modelos testados pode ser aproximada por programas executáveis com alta precisão. Os programas de melhor desempenho alcançam pontuações médias de IoU de 69% para GPT-2, 74% para TinyLlama e 79% para Llama-3B.
• Escala e Arquitetura do Modelo: Os autores observam que as cabeças de atenção em modelos autorregressivos (decoders) são geralmente mais fáceis de ajustar do que aquelas em modelos bidirecionais (encoders) como o BERT. Além disso, a qualidade do ajuste aumenta com a escala do modelo, sugerindo que modelos maiores podem exibir maior especialização funcional entre suas cabeças.
• Fidelidade Causal: Ao substituir as cabeças de atenção por seus programas sintetizados, os modelos mantêm o desempenho funcional. Especificamente, até 25% das cabeças podem ser substituídas com um aumento de apenas 16% na perplexidade.
• Desempenho Downstream: Substituir de 30% a 40% das cabeças de atenção por seus melhores programas não degrada significativamente o desempenho em benchmarks de perguntas e respostas (incluindo HellaSwag, PIQA, SciQ, ARC-Easy, Social IQA e COPA).
• Correlação: Existe uma forte correlação negativa (Spearman $r > 0.9$) entre a similaridade de IoU de um programa e o aumento da perplexidade após a substituição da cabeça, indicando que o IoU é um proxy confiável para selecionar programas funcionalmente fiéis.

Significância e Alegações
O artigo alega que, mesmo em modelos de linguagem de última geração, uma fração substancial de padrões de atenção pode ser compreendida em termos simbólicos. A principal significância reside na demonstração de que componentes neurais treinados podem ser substituídos por substitutos simbólicos via código executável sem alterar substancialmente o comportamento do modelo. Isso sugere que os métodos modernos de síntese de programas impulsionados por LM oferecem um framework viável para abordar a questão mais ampla de como as redes profundas operam.

Os autores posicionam este trabalho como um passo em direção a uma caracterização simbólica completa dos transformers, onde o comportamento do modelo pode ser raciocinado, modificado e testado diretamente através de lógica, em vez de manipulação de pesos. No entanto, o artigo permanece modesto quanto às suas limitações, reconhecendo que nenhuma cabeça de atenção de modelo é totalmente caracterizada atualmente (muitas cabeças ainda alcançam baixos escores de IoU) e que a biblioteca de programas atual, embora eficaz, ainda não é exaustiva. O trabalho estabelece uma fundação para pesquisas futuras visando alcançar uma destilação simbólica completa das arquiteturas transformer.