Recordação Factual Geométrica em Transformers

Resumo

Como os modelos de linguagem baseados em transformadores memorizam associações factuais? Uma visão comum considera matrizes de peso internas como memórias associativas sobre pares de embeddings, exigindo contagens de parâmetros que escalam linearmente com o número de fatos. Desenvolvemos uma abordagem teórica e empírica de uma forma alternativa e geométrica de memorização na qual os embeddings aprendidos codificam diretamente a estrutura relacional, e o MLP desempenha um papel qualitativamente diferente. Em um ambiente controlado onde um transformador de camada única deve memorizar bijeções aleatórias de sujeitos para um conjunto compartilhado de atributos, provamos que uma dimensão de embedding logarítmica é suficiente: os embeddings de sujeitos codificam superposições lineares de seus vetores de atributo associados, e um pequeno MLP atua como um seletor condicionado pela relação que extrai o atributo relevante por meio de portas ReLU, e não como um mapeamento chave-valor associativo. Estendemos esses resultados para o contexto de múltiplos saltos — cadeias de consultas relacionais como "Quem é a mãe da esposa de x?" — fornecendo construções com e sem cadeia de pensamento que exibem um compromisso capacidade-profundidade demonstrável, complementado por um limite inferior teórico-informacional correspondente. Empiricamente, a descida de gradiente descobre soluções precisamente com a estrutura prevista. Uma vez treinado, o MLP transfere-se zero-shot para bijeções completamente novas quando os embeddings de sujeitos são adequadamente reinicilizados, revelando que ele aprendeu um mecanismo de seleção genérico, em vez de memorizar qualquer conjunto particular de fatos.

English

How do transformer language models memorize factual associations? A common view casts internal weight matrices as associative memories over pairs of embeddings, requiring parameter counts that scale linearly with the number of facts. We develop a theoretical and empirical account of an alternative, geometric form of memorization in which learned embeddings encode relational structure directly, and the MLP plays a qualitatively different role. In a controlled setting where a single-layer transformer must memorize random bijections from subjects to a shared attribute set, we prove that a logarithmic embedding dimension suffices: subject embeddings encode linear superpositions of their associated attribute vectors, and a small MLP acts as a relation-conditioned selector that extracts the relevant attribute via ReLU gating, and not as an associative key-value mapping. We extend these results to the multi-hop setting -- chains of relational queries such as ``Who is the mother of the wife of x?'' -- providing constructions with and without chain-of-thought that exhibit a provable capacity-depth tradeoff, complemented by a matching information-theoretic lower bound. Empirically, gradient descent discovers solutions with precisely the predicted structure. Once trained, the MLP transfers zero-shot to entirely new bijections when subject embeddings are appropriately re-initialized, revealing that it has learned a generic selection mechanism rather than memorized any particular set of facts.