Dense associative memory for Gaussian distributions

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.

The Hopfield network is one of the oldest ideas in machine learning. It stores a collection of patterns and recalls the complete one when you give it a corrupted or partial version. Show it half a face, and it fills in the rest. Then researchers showed that a modern version of the Hopfield network does almost the same computation as the attention mechanism in Transformers, the architecture behind modern AI systems like ChatGPT. A memory model from the 1980s turned out to be closely related to the engine driving today's AI. That link is what motivated our work: Hopfield networks, and the attention layers based on them, store plain data points. But a lot of modern AI represents a word or an image not as a point but as a probability distribution which combines a best guess with a measure of how uncertain that guess is. We built a version of these memory networks that stores and retrieves these uncertain representations instead of points. We proved it can hold a large number of them and recover the right one from a noisy query, and we tested it on real images, words, and sentences.