# Learning Associative Memories with Gradient Descent

arXiv: https://arxiv.org/abs/2402.18724 · HTML: https://mlsanity.com/paper/2402.18724/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 39 references that are in the map.
- First posted on arXiv 28 Feb 2024.

## Builds on (most cited first, 28 of 28)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [What is the Effect of Importance Weighting in Deep Learning?](https://mlsanity.com/paper/1812.03372/) (2018)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Saddle-to-Saddle Dynamics in Deep Linear Networks: Small Initialization Training, Symmetry, and Sparsity](https://mlsanity.com/paper/2106.15933/) (2021)
- [The Dynamics of Sharpness-Aware Minimization: Bouncing Across Ravines and Drifting Towards Wide Minima](https://mlsanity.com/paper/2210.01513/) (2022)
- [Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stability](https://mlsanity.com/paper/2305.11788/) (2023)
- [Second-order regression models exhibit progressive sharpening to the edge of stability](https://mlsanity.com/paper/2210.04860/) (2022)
- [Beyond the Edge of Stability via Two-step Gradient Updates](https://mlsanity.com/paper/2206.04172/) (2022)
- [Scaling Laws for Associative Memories](https://mlsanity.com/paper/2310.02984/) (2023)
- [Large Stepsize Gradient Descent for Logistic Loss: Non-Monotonicity of the Loss Improves Optimization Efficiency](https://mlsanity.com/paper/2402.15926/) (2024)
- [Infinite-width limit of deep linear neural networks](https://mlsanity.com/paper/2211.16980/) (2022)
- [Outliers with Opposing Signals Have an Outsized Effect on Neural Network Optimization](https://mlsanity.com/paper/2311.04163/) (2023)
- [Learning Rate Annealing Can Provably Help Generalization, Even for Convex Problems](https://mlsanity.com/paper/2005.07360/) (2020)
- [Dual Training of Energy-Based Models with Overparametrized Shallow Neural Networks](https://mlsanity.com/paper/2107.05134/) (2021)
- [On the Benefits of Large Learning Rates for Kernel Methods](https://mlsanity.com/paper/2202.13733/) (2022)

## Built on by (most cited first, 16 of 16)

- [Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models](https://mlsanity.com/paper/2402.19449/) (2024)
- [Uniform Memory Retrieval with Larger Capacity for Modern Hopfield Models](https://mlsanity.com/paper/2404.03827/) (2024)
- [BiSHop: Bi-Directional Cellular Learning for Tabular Data with Generalized Sparse Modern Hopfield Model](https://mlsanity.com/paper/2404.03830/) (2024)
- [Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers](https://mlsanity.com/paper/2406.18400/) (2024)
- [Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers](https://mlsanity.com/paper/2406.03068/) (2024)
- [Emergence of Linear Truth Encodings in Language Models](https://mlsanity.com/paper/2510.15804/) (2025)
- [From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics](https://mlsanity.com/paper/2510.06954/) (2025)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)
- [Provable Benefits of In-Tool Learning for Large Language Models](https://mlsanity.com/paper/2508.20755/) (2025)
- [Sharp Capacity Thresholds in Linear Associative Memory: From Top-1 Retrieval to Tail-Average Learning](https://mlsanity.com/paper/2605.05189/) (2026)
- [The Features at Convergence Theorem: a first-principles alternative to the Neural Feature Ansatz for how networks learn representations](https://mlsanity.com/paper/2507.05644/) (2025)
- [Factual recall in linear associative memories: sharp asymptotics and mechanistic insights](https://mlsanity.com/paper/2605.10795/) (2026)
- [Fine-Tuning Dynamics of In-Context Factual Recall in Transformers](https://mlsanity.com/paper/2605.27774/) (2026)
- [Rethinking Associative Memory Mechanism in Induction Head](https://mlsanity.com/paper/2412.11459/) (2024)
- [The Stability of Singular Distribution: A Spectral Perspective on the Two-Phase Dynamics of Language Model Pre-training](https://mlsanity.com/paper/2605.26489/) (2026)
- [Muon Learns Facts Better: Understanding the Role of Spectral Orthogonalization](https://mlsanity.com/paper/2610.02798/) (2026)

## Abstract (arXiv)

This work focuses on the training dynamics of one associative memory module storing outer products of token embeddings. We reduce this problem to the study of a system of particles, which interact according to properties of the data distribution and correlations between embeddings. Through theory and experiments, we provide several insights. In overparameterized regimes, we obtain logarithmic growth of the ``classification margins.'' Yet, we show that imbalance in token frequencies and memory interferences due to correlated embeddings lead to oscillatory transitory regimes. The oscillations are more pronounced with large step sizes, which can create benign loss spikes, although these learning rates speed up the dynamics and accelerate the asymptotic convergence. In underparameterized regimes, we illustrate how the cross-entropy loss can lead to suboptimal memorization schemes. Finally, we assess the validity of our findings on small Transformer models.