Resumo Técnico: Recordação Factual Geométrica em Transformers

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) possuem a capacidade de armazenar e recuperar vastas coleções de associações factuais. Uma questão central para compreender essa capacidade é como as arquiteturas de transformer memorizam essas associações. Estruturas teóricas existentes frequentemente modelam isso como um problema de memória associativa, onde matrizes de pesos internos (especificamente MLPs ou mecanismos de atenção) armazenam pares de embeddings de entrada e saída como vetores aleatórios, quase ortogonais. Nessa visão, a contagem de parâmetros necessária para memorizar $N$ fatos escala linearmente com $N$.

No entanto, observações empíricas sugerem um mecanismo diferente: os transformers podem utilizar memorização geométrica, onde embeddings aprendidos codificam diretamente a estrutura relacional dos dados. Este artigo investiga se os transformers podem memorizar associações factuais de forma mais eficiente ao estruturar o espaço de embeddings geometricamente, em vez de depender de armazenamento associativo por força bruta. Os autores focam em um cenário controlado envolvendo $N$ sujeitos, $R$ relações e um conjunto compartilhado de atributos, onde o modelo deve prever o atributo associado a um par sujeito-relação, ou compor essas relações em consultas de múltiplos saltos (por exemplo, "Quem é a mãe da esposa de $x$?").

Metodologia

O artigo combina construção teórica com verificação empírica em dados sintéticos e do mundo real.

Estrutura Teórica

Os autores analisam a capacidade dos transformers de resolver tarefas de recordação factual quando os embeddings são tratados como parâmetros aprendíveis. Eles distinguem entre dois regimes:

1. Recordação de Único Salto: Prever um atributo dado um sujeito e uma relação.
2. Recordação de Múltiplos Saltos: Compor $k$ relações sequencialmente sem saídas intermediárias (previsão direta) versus com Cadeia de Pensamento (CoT).

A análise teórica baseia-se na construção de soluções explícitas para transformers com embeddings aprendíveis. Fundamental nessas construções é o uso de superposições lineares dentro dos embeddings de sujeitos. Em vez de armazenar fatos como vetores distintos e ortogonais, um embedding de sujeito é construído como um vetor em bloco contendo as representações de todos os seus $R$ atributos possíveis. O MLP atua então não como uma tabela de consulta, mas como um seletor condicionado à relação que extrai o bloco de atributo relevante por meio de um portão (gating) ReLU.

Os autores também derivam limites inferiores teóricos da informação sobre a contagem de parâmetros e a dimensão do embedding necessária para a recordação de $k$ saltos, utilizando argumentos de contagem sobre grafos coloridos por arestas regulares-$R$.

Configuração Empírica

Para verificar se o descenso de gradiente descobre essas estruturas teóricas, os autores realizam experimentos extensivos em dados sintéticos:

• Conjunto de Dados: $N=4096$ sujeitos com $R \in \{2, \dots, 16\}$ relações, mapeando para um pool compartilhado de atributos.
• Arquitetura: Transformers de camada única com atenção uniforme ou aprendida e MLPs de duas camadas.
• Controles: Os experimentos comparam modelos com embeddings treináveis contra um controle onde os embeddings de entidades de entrada são congelados na inicialização aleatória (imitando o regime de embedding fixo de trabalhos anteriores).
• Análise: Análises pós-treinamento incluem leituras lineares para verificar estruturas de superposição, intervenções causais para testar a seletividade do MLP e experimentos de transferência "congelar-e-trocar" para testar a generalidade do mecanismo aprendido.
• Sondagem do Mundo Real: Os autores sondam LLMs pré-treinados (por exemplo, Qwen, Llama, Phi) para ver se os estados ocultos de sujeitos exibem codificação relacional linear semelhante às suas previsões teóricas.

Contribuições Principais

1. Memorização Logarítmica com Atributos Compartilhados

O artigo prova que um transformer de camada única pode memorizar $N \times R$ fatos usando uma dimensão de embedding $d = O(R \log N)$, desde que o MLP seja pequeno e condicionado à relação.

• Mecanismo: Os embeddings de sujeitos codificam superposições lineares de seus $R$ vetores de atributos associados.
• Papel do MLP: O MLP atua como um seletor que extrai o vetor de atributo específico correspondente à relação consultada por meio de um portão ReLU, em vez de armazenar o mapeamento em seus pesos.
• Significado: Isso reduz a dimensão de embedding necessária de linear em $N$ (como em modelos de memória associativa) para logarítmica, sem inflar o tamanho do MLP.

2. Compensação Capacidade–Profundidade para Recordação de Múltiplos Saltos

Os autores estabelecem uma compensação fundamental para a recordação factual de $k$ saltos sem Cadeia de Pensamento (CoT):

• Sem CoT: Resolver consultas de $k$ saltos requer ou uma dimensão de embedding exponencial em $k$ ($d = \tilde{O}(R^k)$) ou uma largura de MLP linear no número de fatos ($O(N \cdot R)$). Isso é provado por um limite inferior teórica da informação (Teorema 4.2) mostrando que, sem saídas intermediárias, o modelo deve codificar toda a árvore de avaliação dentro de um único estado oculto.
• Com CoT: Permitir que o modelo emita etapas relacionais intermediárias (Cadeia de Pensamento) quebra esse gargalo. Um transformer de camada única com $d = \tilde{O}(R + k)$ e largura de MLP $\tilde{O}(R)$ é suficiente para qualquer $k$. O CoT efetivamente reduz o problema de múltiplos saltos a uma sequência de problemas de único salto, contornando a necessidade de armazenar a árvore completa de $k$ saltos no espaço de embeddings.

3. Verificação Empírica da Estrutura Geométrica

Os experimentos confirmam que o descenso de gradiente descobre naturalmente as estruturas geométricas previstas:

• Limiar de Capacidade: O limiar $d = \Theta(R \log N)$ é observado empiricamente para embeddings treináveis, enquanto embeddings congelados requerem $d = \Theta(N^R)$ (ou escala polinomial similar), confirmando a vantagem teórica de embeddings estruturados.
• Superposição: Leituras lineares mostram que os embeddings de sujeitos são bem aproximados como superposições lineares de vetores de atributos.
• Seleção Genérica: Intervenções causais confirmam que o MLP atua como um seletor específico da relação. Crucialmente, quando os embeddings de sujeitos são re-inicializados para codificar novas bijeções aleatórias, o MLP congelado alcança alta precisão zero-shot. Isso demonstra que o MLP aprendeu um mecanismo de seleção genérico, não uma memorização de fatos específicos.

Significado e Alegações

O artigo alega fornecer uma conta teórica e empírica da memorização factual geométrica em transformers, oferecendo uma alternativa à visão predominante de memória associativa.

• Eficiência: Ao estruturar embeddings para codificar diretamente a topologia relacional, os transformers podem alcançar a recordação factual com significativamente menos parâmetros (logarítmicos em $N$) em comparação com modelos de memória associativa de embedding aleatório.
• Mecanismo de Recordação: As descobertas sugerem que a recordação factual é implementada por meio de representações geométricas estruturadas onde o MLP serve como um seletor genérico e condicionado à relação. Isso explica por que os modelos podem generalizar para novos fatos se o espaço de embeddings for re-inicializado adequadamente, pois o mecanismo de seleção é desacoplado dos valores específicos dos fatos.
• Papel da Cadeia de Pensamento: O trabalho fornece uma explicação teórica para a eficácia da Cadeia de Pensamento no raciocínio de múltiplos saltos. Demonstra que o CoT contorna os gargalos de capacidade inerentes à previsão direta de $k$ saltos ao distribuir o cálculo através de passos autoregressivos, permitindo que modelos rasos resolvam tarefas de raciocínio profundo.
• Relevância do Mundo Real: A observação de codificação relacional linear em LLMs pré-treinados sugere que esses mecanismos geométricos não são meramente artefatos teóricos, mas estão presentes em modelos treinados em linguagem natural, potencialmente explicando fenômenos como a linearidade da decodificação de relações observada em trabalhos empíricos anteriores.

Os autores mantêm uma postura modesta, observando que seu cenário trata relações como tokens opacos e atributos como entidades arbitrárias, enquanto o conhecimento do mundo real envolve estruturas semânticas ricas. Eles sugerem que estender essa conta geométrica para dados com tais regularidades é uma direção para trabalhos futuros.