Provably Optimal Memory Capacity for Modern Hopfield Models: Transformer-Compatible Dense Associative Memories as Spherical Codes

TL;DR

This paper establishes the first tight, optimal asymptotic memory capacity for modern Hopfield models by linking their configuration to spherical codes, and introduces an efficient algorithm to achieve this capacity.

Contribution

It provides a rigorous analysis connecting Kernelized Hopfield Models to spherical codes, deriving the optimal capacity and proposing a sub-linear time algorithm to reach it.

Findings

Optimal memory capacity matches exponential lower bound.

Proposed $ ext{U} ext{-} ext{Hop}$+ algorithm achieves capacity efficiently.

Theoretical analysis of feature dimension scaling.

Abstract

We study the optimal memorization capacity of modern Hopfield models and Kernelized Hopfield Models (KHMs), a transformer-compatible class of Dense Associative Memories. We present a tight analysis by establishing a connection between the memory configuration of KHMs and spherical codes from information theory. Specifically, we treat the stored memory set as a specialized spherical code. This enables us to cast the memorization problem in KHMs into a point arrangement problem on a hypersphere. We show that the optimal capacity of KHMs occurs when the feature space allows memories to form an optimal spherical code. This unique perspective leads to: (i) An analysis of how KHMs achieve optimal memory capacity, and identify corresponding necessary conditions. Importantly, we establish an upper capacity bound that matches the well-known exponential lower bound in the literature. This provides the first tight and optimal asymptotic memory capacity for modern Hopfield models. (ii) A sub-linear time algorithm $\mathtt{U}\text{-}\mathtt{Hop}$+ to reach KHMs' optimal capacity. (iii) An analysis of the scaling behavior of the required feature dimension relative to the number of stored memories. These efforts improve both the retrieval capability of KHMs and the representation learning of corresponding transformers. Experimentally, we provide thorough numerical results to back up theoretical findings.