Resumo Técnico: Como os Transformers Aprendem a Associar Tokens: Termos Líderes de Gradiente Trazem Interpretabilidade Mecanística

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) baseados em autoatenção demonstraram capacidades notáveis na captura de conhecimento factual e na geração de texto coerente, impulsionados em grande parte por sua capacidade de aprender associações semânticas (por exemplo, o vínculo entre "pássaro" e "voou" ou a intercambiabilidade entre "carro" e "caminhão"). Embora estudos recentes tenham identificado mecanismos específicos como cabeças de indução ou relações semânticas lineares, falta uma explicação principista e fundamentada teoricamente sobre como essas associações emergem durante o treinamento de transformers baseados em atenção em dados de linguagem natural do mundo real.

Trabalhos teóricos anteriores frequentemente dependem de suposições irreais, como linguagem estruturada sintética, arquiteturas simplificadas (sem codificações posicionais ou conexões residuais) ou regimes de treinamento não padrão (por exemplo, treinamento sequencial componente a componente). Esses desvios em relação à prática limitam a generalização de suas percepções para LLMs modernos. Este artigo aborda essa lacuna ao perguntar: Como as associações semânticas emergem durante o treinamento de modelos de linguagem baseados em atenção em dados de linguagem natural?

Metodologia

Os autores propõem um quadro teórico fundamentado na dinâmica de treinamento, analisando especificamente a aproximação de termo líder dos gradientes durante os estágios iniciais do treinamento.

1. Configuração do Modelo: O estudo foca em um transformer padrão de $L$ camadas baseado em atenção com codificações posicionais (relativas), mascaramento causal e fluxos residuais, treinado via descida de gradiente em lote completo padrão em corpora de linguagem natural.
2. Aproximação de Termo Líder: Reconhecendo que os transformers adquirem comportamentos centrais (incluindo relações semânticas) no início do treinamento, os autores derivam expressões de forma fechada para as matrizes de pesos (Saída $W_O$, Valor $V^{(l)}$, Consulta-Chave $W^{(l)}$ e Codificação Posicional $P^{(l)}$) aproximando as atualizações de gradiente. Eles mostram que, para um número suficiente de passos $s$ (onde $s \leq \eta^{-1} \min(\dots)$), os pesos permanecem próximos aos seus termos líderes de gradiente.
3. Decomposição em Funções de Base: A análise teórica revela que os pesos aprendidos podem ser expressos como composições simples de três funções de base fundamentais derivadas das estatísticas do corpus:
   • Mapeamento de Bigrama ($\bar{B}$): Captura dependências de próximo-token ($P(e_j|e_i)$).
   • Mapeamento de Intercambiabilidade ($\Sigma \bar{B} = \bar{B}^\top \bar{B}$): Captura similaridade funcional entre tokens baseada na similaridade de suas distribuições de tokens precedentes (por exemplo, sinônimos ou papéis gramaticais compartilhados).
   • Mapeamento de Contexto ($\bar{\Phi}$): Codifica co-ocorrência de prefixo-sufixo de longo alcance, representando o contexto médio no qual um token aparece.

Principais Contribuições

O artigo faz três contribuições primárias:

1. Caracterização Explícita de Pesos: Apresenta a primeira caracterização explícita e de forma fechada de pesos em transformers baseados em atenção treinados em corpora de texto do mundo real sob a perda de previsão de próximo-token. Os pesos são mostrados como composições das três funções de base mencionadas acima.
2. Interpretação Mecanística de Associações Semânticas: Os autores interpretam as características aprendidas nos pesos como composições específicas de mapeamentos de bigrama, intercambiabilidade e contexto. Eles demonstram como essas composições permitem que o modelo capture associações semânticas (por exemplo, como "peixe" se relaciona com "lagoa" via mapeamento de contexto, ou como "vermelho" se relaciona com "caminhão" via mapeamento de bigrama).
3. Validação Empírica: As previsões teóricas são validadas tanto em um modelo de atenção de brinquedo de 3 camadas treinado no conjunto de dados TinyStories quanto em um LLM prático e de grande escala (Pythia-1.4B) treinado no OpenWebText.

Resultados

• Alinhamento Teórico: Experimentos em um transformer de 3 camadas treinado no TinyStories mostram que a similaridade de cosseno entre os termos líderes previstos teoricamente e os pesos realmente aprendidos permanece extremamente alta (por exemplo, >0,99 para matrizes de Atenção e Valor) mesmo após 100 épocas de treinamento.
• Estrutura Semântica: A análise qualitativa confirma que as funções de base capturam estruturas linguísticas significativas. Por exemplo, o mapeamento de bigrama $\bar{B}$ correlaciona adjetivos com os substantivos que modificam, enquanto o mapeamento de contexto $\bar{\Phi}$ correlaciona entidades com seus ambientes típicos (por exemplo, "peixe" com "lagoa").
• Generalização para LLMs Reais: Quando aplicado ao Pythia-1.4B, as características líderes teóricas mostram forte concordância com os embeddings aprendidos e correlações de atenção, particularmente nos estágios iniciais do treinamento. A análise revela que, embora os pesos se desviem em direção a representações de conhecimento mais ricas à medida que o treinamento progride, eles retêm as características associativas fundamentais previstas pela teoria por um período significativo.
• Cooperação de Componentes: A análise do cálculo completo do modelo (Eq. 12) ilustra como o fluxo residual fornece uma previsão média de próximo-token, enquanto o bloco de autoatenção (impulsionado por $\bar{Q}$ e $\bar{\Phi}^\top \Sigma \bar{B}$) refina essa previsão focando seletivamente nos tokens mais indicativos do próximo token com base nas estatísticas do corpus.

Significância

O artigo afirma fornecer uma fundação mecanística para a aprendizagem de representação em modelos de linguagem de grande escala, conectando a dinâmica de treinamento de aprendizado profundo com a teoria linguística. Ao fundamentar a análise em distribuições de texto naturalistas e procedimentos de treinamento padrão, o trabalho reduz significativamente a lacuna entre a análise teórica formal e o uso prático de LLMs.

Os autores argumentam que suas descobertas oferecem uma maneira rigorosa de contextualizar a função de cada componente do transformer. Em vez de visualizar o modelo como uma caixa preta, a teoria sugere que o comportamento do modelo se desenvolve estabelecendo primeiro essas características associativas fundamentais (bigramas, intercambiabilidade, contexto), que servem como âncoras para aprender relações semânticas mais complexas. Este trabalho abre caminhos para pesquisa em interpretabilidade ao sugerir que as matrizes de pesos em diferentes componentes podem ser decompostas em funções simples de estatísticas compartilhadas do corpus, permitindo hipóteses mais amplas sobre como conceitos surgem em modelos além de mecanismos específicos.