🤖 AI Summary
Current large language models lack a factual knowledge storage mechanism that simultaneously achieves information-theoretic optimality, adaptability to arbitrary embedding geometries, and compatibility with the Transformer architecture. This work addresses this gap by analyzing the decoding boundaries of multilayer perceptrons (MLPs) and, for the first time, constructs a closed-form MLP formulation that satisfies three empirically grounded properties. The resulting construction reveals an intrinsic Hebbian-like learning behavior and enables modular factual editing. It achieves information-theoretic optimality in storing facts under both isotropic and arbitrary embedding geometries, reducing parameter count by 10–10⁴ times compared to existing approaches at equal factual capacity. When integrated into a Transformer, the method preserves optimal capacity scaling while lowering parameter overhead by 15–63 times.
📝 Abstract
Large language models (LLMs) store factual knowledge in their parameters. While recent work has shown that this knowledge resides in MLP layers, existing constructive and mechanistic interpretability models of fact-storage in LLMs fail to explain the surprising empirical phenomenon that they store facts at an information-theoretically optimal rate. In this work, we develop a theoretical account of this phenomenon. We develop the first Transformer-compatible fact-storing MLP closed-form construction that satisfies the following three properties empirically observed in LLMs: it (i) attains optimal fact storage scaling, (ii) handles arbitrary input/output geometries, and (iii) works inside Transformers. Key to our work is to analyze the decoding margin of MLPs, whereas prior work only studies MLP fact storage. Under isotropic embeddings, our construction achieves information-theoretically optimal storage capacity scaling and requires $10$-$104\times$ fewer parameters at matched fact count than prior constructions. For arbitrary key and value embeddings, we show that our construction attains the same storage capacity scaling, up to penalization factors depending on the embedding geometries. Moreover, we demonstrate that our constructed MLPs can be used within Transformer blocks for factual recall tasks at optimal capacity scaling, requiring $15$-$63\times$ fewer parameters at matched fact count than prior constructions. Finally, as a proof-of-concept, we show that fact-storing MLPs enable modular fact editing by swapping a Transformer's MLP with a new one.