Rappel de faits géométriques dans les Transformers

Geometric Factual Recall in Transformers

May 12, 2026

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

cs.AI

Résumé

Comment les modèles de langage à base de transformeurs mémorisent-ils les associations factuelles ? Une vision courante considère les matrices de poids internes comme des mémoires associatives agissant sur des paires de plongements, nécessitant un nombre de paramètres qui croît linéairement avec le nombre de faits. Nous développons un compte rendu théorique et empirique d'une forme alternative et géométrique de mémorisation dans laquelle les plongements appris encodent directement la structure relationnelle, et le MLP joue un rôle qualitativement différent. Dans un cadre contrôlé où un transformeur monocouche doit mémoriser des bijections aléatoires allant de sujets à un ensemble partagé d'attributs, nous prouvons qu'une dimension de plongement logarithmique est suffisante : les plongements de sujets encodent des superpositions linéaires de leurs vecteurs d'attributs associés, et un petit MLP agit comme un sélecteur conditionné par la relation qui extrait l'attribut pertinent via un filtrage ReLU, et non comme un mappage clé-valeur associatif. Nous étendons ces résultats au cadre multi-sauts — des chaînes de requêtes relationnelles telles que « Qui est la mère de l'épouse de x ? » — en fournissant des constructions avec et sans chaîne de pensée qui présentent un compromis prouvable entre capacité et profondeur, complété par une borne inférieure information-théorique correspondante. Empiriquement, la descente de gradient découvre des solutions avec précisément la structure prédite. Une fois entraîné, le MLP se transfère en zéro-shot à des bijections entièrement nouvelles lorsque les plongements de sujets sont réinitialisés de manière appropriée, révélant qu'il a appris un mécanisme de sélection générique plutôt que d'avoir mémorisé un ensemble particulier de faits.

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.