Recuperación Factual Geométrica en Transformers

Geometric Factual Recall in Transformers

May 12, 2026

Autores: Shauli Ravfogel, Gilad Yehudai, Joan Bruna, Alberto Bietti

cs.AI

Resumen

¿Cómo memorizan asociaciones factuales los modelos de lenguaje basados en transformadores? Una visión común considera las matrices de pesos internas como memorias asociativas sobre pares de embeddings, lo que requiere recuentos de parámetros que escalan linealmente con el número de hechos. Desarrollamos un análisis teórico y empírico de una forma alternativa y geométrica de memorización en la que los embeddings aprendidos codifican directamente la estructura relacional, y el MLP desempeña un papel cualitativamente diferente. En un entorno controlado donde un transformador de una sola capa debe memorizar biyecciones aleatorias desde sujetos hacia un conjunto compartido de atributos, demostramos que una dimensión de embedding logarítmica es suficiente: los embeddings de sujetos codifican superposiciones lineales de sus vectores de atributos asociados, y un MLP pequeño actúa como un selector condicionado por la relación que extrae el atributo relevante mediante compuertas ReLU, y no como un mapeo asociativo clave-valor. Extendemos estos resultados al entorno de múltiples saltos —cadenas de consultas relacionales como “¿Quién es la madre de la esposa de x?”— proporcionando construcciones con y sin cadena de pensamiento que exhiben un trade-off demostrable entre capacidad y profundidad, complementado por una cota inferior informacional-teórica correspondiente. Empíricamente, el descenso de gradiente descubre soluciones con precisamente la estructura predicha. Una vez entrenado, el MLP transfiere en zero-shot a biyecciones completamente nuevas cuando los embeddings de sujetos se reinician adecuadamente, revelando que ha aprendido un mecanismo de selección genérico en lugar de memorizar cualquier conjunto particular de hechos.

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.