Geometrische Feitelijke Recall in Transformers

Samenvatting

Hoe onthouden transformator-taalmodel feitelijke associaties? Een gangbare opvatting beschouwt interne gewichtsmatrices als associatieve geheugens over paren van embeddings, wat vereist dat het aantal parameters lineair schaalt met het aantal feiten. We ontwikkelen een theoretische en empirische beschrijving van een alternatieve, geometrische vorm van memorisatie waarin aangeleerde embeddings direct relationele structuur coderen, en de MLP een kwalitatief andere rol speelt. In een gecontroleerde setting waarin een transformator met één laag willekeurige bijecties van subjecten naar een gedeelde attribuutset moet onthouden, bewijzen we dat een logaritmische embeddimensie volstaat: subjectembeddings coderen lineaire superposities van hun bijbehorende attribuutvectoren, en een kleine MLP fungeert als een relatie-afhankelijke selector die het relevante attribuut extraheert via ReLU-poortwerking, en niet als een associatieve sleutel-waarde-afbeelding. We breiden deze resultaten uit naar de multi-hop-setting – ketens van relationele queries zoals "Wie is de moeder van de vrouw van x?" – waarbij we constructies bieden met en zonder chain-of-thought die een aantoonbare capaciteit-diepte-afweging vertonen, aangevuld met een corresponderende informatietheoretische ondergrens. Empirisch ontdekt gradiëntdaling oplossingen met precies de voorspelde structuur. Eenmaal getraind, draagt de MLP zero-shot over naar volledig nieuwe bijecties wanneer subjectembeddings op de juiste manier worden geherinitialiseerd, wat onthult dat het een generiek selectiemechanisme heeft geleerd in plaats van een specifieke set feiten te hebben gememoriseerd.

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.