Dense associative memory for Gaussian distributions
Chandan Tankala, Krishna Balasubramanian
Abstract
Dense associative memories (DAMs) store and retrieve patterns via energy-function based fixed points, but existing models are limited to vector representations. We extend DAMs to Gaussian densities equipped with the 2-Wasserstein distance. Our framework defines a log-sum-exp energy over stored distributions and a retrieval dynamics aggregating optimal transport maps in a Gibbs-weighted manner. Stationary points correspond to self-consistent Wasserstein barycenters, generalizing classical DAM fixed points. We prove exponential storage capacity and provide quantitative retrieval guarantees under Wasserstein perturbations. We validate the method on synthetic and real-world image (CelebA and CIFAR-10 datasets) and text (text8 and NLI corpus) datasets. By generalizing from vectors to distributions, our work bridges classical DAMs with modern generative modeling and paves way for distributional storage and retrieval in memory-augmented learning.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 4617b5ed-5023-496a-9157-5decfd1a4c8eBuilds on11
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Hopfield Networks is All You NeedHubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl et al.ICLR 2021 · 620 citations
- On Sparse Modern Hopfield ModelJerry Yao-Chieh Hu, Donglin Yang, Dennis Wu, Chenwei Xu et al.NeurIPS 2023 · 52 citations
- Forward-Backward Gaussian Variational Inference via JKO in the Bures-Wasserstein SpaceMichael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, Adil SalimICML 2023 · 47 citations
- Uniform Memory Retrieval with Larger Capacity for Modern Hopfield ModelsDennis Wu, Jerry Yao-Chieh Hu, Teng-Yun Hsiao, Han LiuICML 2024 · 44 citations
Related papers
- Dense Associative Memory with Epanechnikov EnergyBenjamin Hoover, Zhaoyang Shi, Krishnakumar Balasubramanian, Dmitry Krotov et al.NeurIPS 2025 · 11 citations
- Multivariate Representation Learning for Information RetrievalHamed Zamani, Michael BenderskySIGIR 2023 · 7 citations
- Efficient Approximation Algorithm for Computing Wasserstein Barycenter under Euclidean MetricPankaj K. Agarwal, Sharath Raghvendra, Pouyan Shirzadian, Keegan YaoSODA 2025
- Semantically-correlated memories in a dense associative modelThomas F. BurnsICML 2024 · 8 citations
- Dynamical properties of dense associative memoryKazushi Mimura, Jun'ichi Takeuchi, Yuto Sumikawa, Yoshiyuki Kabashima et al.ICLR 2026 · 6 citations
