2000 character limit reached
Learning Associative Memories with Gradient Descent
Published 28 Feb 2024 in cs.LG, cs.AI, and stat.ML | (2402.18724v1)
Abstract: 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.
- Second-order regression models exhibit progressive sharpening to the edge of stability. In International Conference on Machine Learning (ICML), 2023.
- Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems (NeurIPS), 2019.
- Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
- The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research (JMLR), 24(316):1–36, 2023.
- On the benefits of large learning rates for kernel methods. In Conference on Learning Theory (COLT), 2022.
- Birth of a transformer: A memory viewpoint. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
- What is the effect of importance weighting in deep learning? In International conference on machine learning (ICML), 2019.
- Scaling laws for associative memories. In International Conference on Learning Representations, 2024.
- Beyond the edge of stability via two-step gradient updates. In International Conference on Machine Learning (ICML), 2023.
- On the global convergence of gradient descent for over-parameterized models using optimal transport. In Neural Information Processing Systems (NeurIPS), 2018.
- Infinite-width limit of deep linear neural networks. arXiv preprint arXiv:2211.16980, 2022.
- Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations (ICLR), 2020.
- Support-vector networks. Machine Learning, 1995.
- Dual training of energy-based models with overparametrized shallow neural networks. arXiv preprint arXiv:2107.05134, 2021.
- A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
- Transformer feed-forward layers are key-value memories. In EMNLP, 2021.
- Inequalities on the lambert w function and hyperpower function. J. Inequal. Pure and Appl. Math, 9(2):5–9, 2008.
- Hopfield, J. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the National Academy of Sciences of the United States of America, 1982.
- Neural computation of decisions in optimization problems. Biological Cybernetics, 1985.
- Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning (ICML), 2015.
- Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106.15933, 2021.
- The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory (COLT), 2019.
- Characterizing the implicit bias via a primal-dual analysis. In Algorithmic Learning Theory, 2021.
- Adam: A method for stochastic optimization, 2015.
- Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems (NIPS), 2016.
- Theories of associative recall. Quarterly Reviews of Biophysics, 1970.
- Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations (ICLR), 2019.
- A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
- Distributed representations of words and phrases and their compositionality. Advances in neural information processing systems (NIPS), 2013.
- Nakkiran, P. Learning rate annealing can provably help generalization, even for convex problems. arXiv preprint arXiv:2005.07360, 2020.
- In-context learning and induction heads. Transformer Circuits Thread, 2022.
- Hopfield networks is all you need. In International Conference on Learning Representations (ICLR), 2021.
- Outliers with opposing signals have an outsized effect on neural network optimization. arXiv preprint arXiv:2311.04163, 2023.
- Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:1805.00915, 2018.
- Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning (ICML), 2021.
- The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research (JMLR), 19(1):2822–2878, 2018.
- Non-holographic associative memory. Nature, 1969.
- Implicit bias of gradient descent for logistic regression at the edge of stability. In Neural Information Processing Systems (NeurIPS), 2023.
- Large stepsize gradient descent for logistic loss: Non-monotonicity of the loss improves optimization efficiency, 2024.
Paper Prompts
Sign up for free to create and run prompts on this paper using GPT-5.
Top Community Prompts
Collections
Sign up for free to add this paper to one or more collections.
