Resumo Técnico: Atribuir e Somar: Um Estudo Mecanístico de Aritmética Composicional

Declaração do Problema

Grandes modelos de linguagem (LLMs) demonstram a capacidade de compor habilidades para realizar tarefas complexas não vistas durante o treinamento, mas os mecanismos específicos que permitem essa "generalização composicional" permanecem elusivos. Embora a generalização seja frequentemente atribuída à composição de mecanismos atômicos, a interação entre mecanismos distintos — especificamente a vinculação de variáveis (recuperação associativa via cabeças de indução) e a aritmética modular (representações estruturadas via características de Fourier) — não foi sistematicamente estudada em um ambiente controlado.

Este artigo aborda a questão: Como os transformers aprendem a compor a vinculação de variáveis e os mecanismos de adição modular para generalizar além da distribuição de treinamento, e como essas capacidades emergem durante o treinamento?

Metodologia

Os autores empregam uma configuração experimental controlada envolvendo um transformer pequeno e mecanisticamente interpretável, treinado em uma tarefa sintética que combina atribuição de variáveis e adição modular.

• Arquitetura do Modelo: Um transformer de 2 camadas e uma única cabeça com $d_{model} = d_{head} = 128$ e $d_{mlp} = 512$. A arquitetura deliberadamente omite vieses (biases), camadas de normalização e o MLP da Camada 1 para simplificar a descoberta de circuitos. Utiliza embeddings de posição absoluta no estilo GPT-2 e mascaramento causal.
• Definição da Tarefa: O modelo é treinado para realizar adição modular ($N=59$), onde os operandos podem ser constantes ou variáveis atribuídas dinamicamente no contexto (ex: c=1, b=17, a=42, b+19=?).
• Distribuição de Dados: As sequências são tokenizadas com preenchimento (padding) aleatório para evitar a dependência de posições absolutas. O conjunto de treinamento é particionado para avaliar a generalização:
  • 0-variáveis (0-var): Adição pura de constantes.
  • 1-variável (1-var) & 2-variáveis (2-var): Adição envolvendo variáveis atribuídas.
  • Conjuntos Excluídos (Held-out Sets): Aproximadamente 30% dos pares de adição são removidos do treinamento (ex: se 1+3 não é visto, todas as sequências contendo isso são excluídas). Além disso, variáveis específicas são restritas de certas posições no conjunto de treinamento para testar a generalização posicional.
• Técnicas de Análise: O estudo utiliza ferramentas de interpretabilidade mecanística (via biblioteca TransformerLens) para analisar padrões de atenção, vetores do fluxo residual e matrizes de pesos. Também emprega sondas lineares e análise de gradiente teórica para rastrear a dinâmica de treinamento.

Principais Contribuições

1. Modelo Tratável para Aritmética Composicional: Os autores apresentam um modelo simplificado que permite a interpretação mecanística de como a atribuição de variáveis se compõe com a aritmética modular.
2. Dinâmica de Treinamento Desacoplada: A análise revela que a aritmética modular é aprendida antes e independentemente da atribuição de variáveis. O modelo atinge precisão quase perfeita em tarefas de 0-var enquanto ainda está aprendendo a vinculação de variáveis.
3. Explicação Mecanística da Generalização: O artigo identifica que a generalização surge da composição de dois circuitos específicos:
   • Uma cabeça de indução modificada na Camada 1 que recupera variáveis atribuídas e roteia seus valores.
   • Um circuito de adição modular na Camada 2 (especificamente a matriz Output-Value e o MLP da Camada 2) que realiza a aritmética nos valores recuperados.
4. Estrutura Teórica: Os autores fornecem uma análise teórica demonstrando como a dinâmica de gradiente facilita a composição dos circuitos de vinculação de variáveis e de adição modular, mostrando que o gradiente negativo na matriz Query-Key da Camada 2 leva naturalmente à formação de um bloco de memória associativa.

Resultados

Dinâmica de Treinamento e Fases

O processo de treinamento exibe três fases distintas:

1. Aprendizado de Adição Modular: O modelo primeiro aprende a realizar adição em sequências sem variáveis (0-var), alcançando alta precisão antes que a vinculação de variáveis esteja totalmente estabelecida.
2. Aprendizado de Atribuição de Variáveis: O modelo subsequente aprende a estrutura necessária para a atribuição de variáveis. Um "pico" agudo na precisão (por volta do passo 5.300–5.400) coincide com a emergência do circuito de roteamento de atribuição de variáveis, onde a atenção da Camada 1 roteia corretamente as constantes para suas variáveis atribuídas.
3. Fase de Refinamento: Uma segunda fase (passos 14.500–16.000) envolve o refinamento do roteamento para sequências "difíceis" (ex: variáveis específicas em posições excluídas), corrigindo falhas iniciais de atenção a certas posições de variáveis.

Capacidades de Generalização

• Generalização Posicional: O modelo generaliza para novas ordenações de variáveis (ex: prevendo b+14 após ver 14+b).
• Generalização Fora da Distribuição (OOD): O modelo prevê com sucesso pares de adição excluídos (ex: 1+3) que nunca foram vistos em nenhuma forma durante o treinamento, alcançando 99,4% de precisão em conjuntos de teste curados.
• Requisitos de Dados: A generalização para pares de constantes não vistos requer treinamento em pelo menos ~25% da distribuição total ($f \approx 0.25$). A generalização para sequências de 2 variáveis requer uma frequência relativa de ~20% ($r \approx 0.2$) de sequências de 2-var nos dados de treinamento.

Achados Mecanísticos

• Representação Compartilhada: Os vetores do fluxo residual para adição com variáveis atribuídas e adição pura de constantes são quase idênticos (similaridade de cosseno $\approx 1$) antes do MLP da Camada 2, indicando um mecanismo computacional compartilhado.
• Características de Fourier: As pré-ativações dos neurônios do MLP da Camada 2 correspondem à forma teórica das características de Fourier ($\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$), confirmando o uso de mecanismos periódicos estabelecidos para aritmética modular.
• Modificação da Cabeça de Indução: A atenção da Camada 1 atua como uma cabeça de "token anterior" para constantes, atendendo à variável imediatamente precedente a elas. A matriz Query-Key da Camada 2 então forma um forte padrão de identidade nos embeddings das variáveis, permitindo que o token = atenda diretamente aos constantes corretos associados às variáveis.

Significância

O artigo afirma que a generalização composicional em transformers pode ser uma consequência natural da composicionalidade dos mecanismos internos. Ao demonstrar que mecanismos distintos (cabeças de indução e aritmética baseada em Fourier) podem ser aprendidos independentemente e depois compostos via roteamento de atenção, o trabalho fornece uma base mecanística para entender como os transformers generalizam para tarefas de raciocínio complexas.

Os autores mantêm a modéstia quanto ao seu escopo, observando que o estudo é limitado a um ambiente pequeno e controlado para facilitar a análise. Eles não afirmam que esses achados explicam diretamente todos os comportamentos de LLMs, mas sugerem que oferecem uma estrutura para entender como as capacidades de generalização surgem em transformers, potencialmente motivando investigações futuras em tarefas mais complexas e leis de escala.