Year: 2024
Type: preprint
Abstract: Large language models have demonstrated an impressive ability to perform factual recall. Prior work has found that transformers trained on factual recall tasks can store information at a rate proportional to their parameter count. In our work, we show that shallow transformers can use a combination of associative memories to obtain such near optimal storage capacity. We begin by proving that the storage capacities of both linear and MLP associative memories scale linearly with parameter count. We next introduce a synthetic factual recall task, and prove that a transformer with a single layer of self-attention followed by an MLP can obtain 100% accuracy on the task whenever either the total number of self-attention parameters or MLP parameters scales (up to log factors) linearly with the number of facts. In particular, the transformer can trade off between using the value matrices or the MLP as an associative memory to store the dataset of facts. We complement these expressivity results with an analysis of the gradient flow trajectory of a simplified linear attention model trained on our factual recall task, where we show that the model exhibits sequential learning behavior.
Source: arXiv (Cornell University)
Institutions: Princeton University, Flatiron Institute
Language: English
Topic: Topic Modeling
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: Collaborative Research: CIF: Medium: MoDL:Toward a Mathematical Foundation of Deep Reinforcement Learning, CAREER: Towards a Theory of Deep Learning, CIF: Medium: Collaborative Research: Theory of Optimization Geometry and Algorithms for Neural Networks, G6671297155, RI: Medium: Provable Reinforcement Learning with Function Approximation and Neural Networks
OpenAlex ID: W4405918571