Dense associative memory on the Bures-Wasserstein space
This work bridges classical associative memory with generative modeling, enabling distributional storage and retrieval, which is significant for memory-augmented learning but is incremental as it extends an existing framework to a new domain.
The paper tackled the problem of extending dense associative memories (DAMs) from vector representations to probability distributions, specifically using the Bures-Wasserstein space for Gaussian densities, and demonstrated exponential storage capacity with quantitative retrieval guarantees under perturbations.
Dense associative memories (DAMs) store and retrieve patterns via energy-functional fixed points, but existing models are limited to vector representations. We extend DAMs to probability distributions equipped with the 2-Wasserstein distance, focusing mainly on the Bures-Wasserstein class of Gaussian densities. 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, provide quantitative retrieval guarantees under Wasserstein perturbations, and validate the model on synthetic and real-world distributional tasks. This work elevates associative memory from vectors to full distributions, bridging classical DAMs with modern generative modeling and enabling distributional storage and retrieval in memory-augmented learning.