Research map / Optimization / Implicit bias and generalization
Learning Associative Memories with Gradient Descent
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes gradient dynamics of a single linear associative memory trained with cross-entropy by reducing it to a system of interacting particles, deriving logarithmic margin growth, oscillations and loss spikes from token imbalance and correlated embeddings, and checks these on small Transformers.
Problem
Understanding training dynamics of an associative memory (linear layer with fixed input/output embeddings) trained with cross-entropy loss, including loss spikes, oscillations, and memorization under over- and under-parameterization.
Approach
Views the gradient dynamics of the associative memory as a system of interacting particles (scores w_ij), with interactions governed by embedding correlations and class structure; derives closed-form margin growth for orthogonal embeddings and a two-particle analysis for correlated embeddings.
Key results
- In the orthogonal binary case, margins grow as log(c_x t) with loss bounded by 1/(t eta) times sum p(x)/c_x.
- For two correlated particles with alpha<=0 a single gradient step reaches perfect accuracy; with strong imbalance and correlation, a single gradient update gives loss at least eta(alpha p1 - p2) p2, which can be arbitrarily large (loss spikes).
- In the under-parameterized d=M=2 case with N=3, minimizing cross-entropy can forget the most frequent association, showing the cross-entropy loss is not calibrated to 0-1 loss for this model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a larger gradient-descent learning rate | Associative memory model · Overparameterized regimes | steps to perfect accuracy | Mixed |
| Use a larger gradient-descent learning rate | Associative memory model · Overparameterized regimes with correlated embeddings | training loss: \cL(W_1) \geq \eta(\alpha p_1 - p_2) p_2 | Mixed |
| Increase embedding dimension | Associative memory model with random embeddings · N = 5 tokens | margin convergence speed | Helped |
| Train with limited capacity (d < N) | Associative memory model · Underparameterized regimes | 0-1 accuracy | Hurt |
| Train all three matrices rather than only W_O^2 | Simplified two-layer Transformer · Bigram in-context learning task | training loss | Mixed |
| Increase embedding dimension | Simplified two-layer Transformer · Bigram in-context learning task | convergence speed | Helped |
| Use orthogonal rather than correlated input embeddings | Associative memory model · Overparameterized regimes | training dynamics | Helped |
| Increase token-frequency imbalance | Associative memory model with correlated embeddings · Overparameterized regimes | steps to perfect accuracy | Hurt |
Surprising findings
- Cross-entropy minimizers in the limited-capacity regime can forget the most frequent association, so the loss is not calibrated to 0-1 error.
- Large learning rates speed convergence but cause oscillations and loss spikes, even though gradients always increase the max-margin direction.
Limitations
- Analysis restricted to a single linear associative memory with fixed embeddings; ignores factorized parameterizations, normalization, adaptive optimizers and noisy data.
- Multi-class and gradient descent results are partly heuristic or stated without full proof (indications only).
- Transformer experiments are simplified and only empirical.
Experimental setup
Datasets: synthetic bigram in-context task (16,384 sequences of length 256, N=64 tokens) synthetic random embeddings on sphere
Metrics: training loss classification margins 0-1 loss sharpness (Hessian operator norm) number of steps to perfect accuracy
Builds on: induction heads (Elhage et al. / Olsson et al.) Hopfield network storage scaling
Closest papers by meaning
- Scaling Laws for Associative Memories (ICLR 2024 · similarity 0.90)
- Factual recall in linear associative memories: sharp asymptotics and mechanistic insights (NeurIPS 2026 · similarity 0.88)
- Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory (NeurIPS 2026 · similarity 0.88)
- Semantically-correlated memories in a dense associative model (ICML 2024 · similarity 0.87)
- Dense Associative Memory Through the Lens of Random Features (NeurIPS 2024 · similarity 0.86)
- Why Diffusion Models Don’t Memorize: The Role of Implicit Dynamical Regularization in Training (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.