Uniform Memory Retrieval with Larger Capacity for Modern Hopfield Models

Introduction

We address the memory confusion problem in the modern Hopfield models by proposing a two-stage optimization formulation, termed $\mathtt{U\text{-}Hop}$, for the memory retrieval dynamics of modern Hopfield models. We construct the similarity measure of modern Hopfield models with a learnable kernel. The feature map of the kernel is trained by maximizing the separation among the entire stored memory set (Figure 2). This allows Hopfield models under $\mathtt{U\text{-}Hop}$ to distinguish different memory patterns with larger separation and hence achieve larger memory capacity.

For all $\mu\in[M]$, let $R\coloneqq{\frac{1}{2}}\mathop{\rm min}_{\mu,\nu\in[M];\mu\neq\nu}\norm{\bxi_\mu-\bxi_\nu}$ be the finite radius of each sphere $\mathcal{S}_{\mu}$ centered at memory pattern ${\bm{\xi}}_{\mu}$. We say ${\bm{\xi}}_{\mu}$ is stored if all $\mathbf{x}\in\mathcal{S}_{\mu}$ are generalized fixed points of $\mathcal{T}$, $\mathbf{x}^{\star}_{\mu}\in\mathcal{S}_{\mu}$, and $\mathcal{S}_{\mu}\cap\mathcal{S}_{\nu}=\emptyset$ for $\mu\neq\nu$. We say ${\bm{\xi}}_{\mu}$ is $\epsilon$-retrieved by $\mathcal{T}$ with $\mathbf{x}$ for an error $\epsilon$, if $\norm{\calT(\bx)-\bxi_\mu}\leq\epsilon$.

Let $\Delta_{\mu}\coloneqq\Braket{{\bm{\xi}}_{\mu},{\bm{\xi}}_{\mu}}-\mathop{\rm max}_{\nu\in[M],\nu\neq\mu}\Braket{{\bm{\xi}}_{\nu},{\bm{\xi}}_{\mu}}$ be the separation between a memory pattern ${\bm{\xi}}_{\mu}$ from all other memories in $\bm{\Xi}$, and $m$ be the largest norm among memory patterns. Ramsauer et al. 2020 gives the retrieval error bound:

for any $\mathbf{x}\in\mathcal{S}_{\mu}$. This bound is crucial not only for characterizing retrieval quality but also, in capacity analysis, as a necessary condition for pattern ${\bm{\xi}}_{\mu}$ to be stored in the model (Hu et al. 2023, Theorem 3.1). Yet, it depends on $\bm{\Xi}$.

$\Delta_{\mu}$ measures the distance from a given ${\bm{\xi}}_{\mu}$ to the nearest memory pattern in $\bm{\Xi}$. $R$ measures the minimal separation among all stored memories $\bm{\Xi}$. Hence, they are both $\bm{\Xi}$-dependent. This $\bm{\Xi}$-dependence in (1.3) results in potential fuzzy retrievals, namely metastable states caused by multiple nearby local minima in the energy landscape, especially when $\Delta_{\mu}-2mR$ is small. When this occurs, these fuzzy retrievals deviate the retrieval process from the ground truth, thereby hampering performance.

This fuzzy memory (memory confusion) issue is well-known in literature. The dense associative memory model (Krotov and Hopfield 2016) tries to solve this issue by using polynomial energy fucntion. The modern Hopfield models (Demircigil et al. 2017; Ramsauer et al. 2020; Hu et al. 2023; Wu et al. 2024; Hu et al. 2024b) try to solve this issue by using exponential energy functions. However, all these attempts still rely on the quality of $\bm{\Xi}$. In this work, we rethink the use of inner-product similarity measure (i.e. $\bm{\Xi}^{\mathsf{T}}\mathbf{x}$ in (1.2)), and consider it as primary source of the fuzzy memory problem. Specifically, due to its Euclidean nature, inner-product assigns equal importance to all dimensions of patterns and yields small $\Delta_{\mu}-2mR$ if they (${\bm{\xi}}_{\mu}$ and some ${\bm{\xi}}_{\nu}$) share similar direction. This motivate us to replace the overlap (inner product) construction of the energy function with a similarity measure utilizing this $\bm{\Xi}$-dependence.

We introduce a learnable feature map $\Phi$ that maps energy $E$ to a kernel space with kernel $\mathcal{K}(\cdot,\cdot)\coloneqq\Braket{\Phi(\cdot),\Phi(\cdot)}$. The resulting kernelized energy $E_{\mathcal{K}}$ , and its corresponding retrieval dynamics $\mathcal{T}_{\mathcal{K}}$ satisfy the defining properties of modern Hopfield models: convergence between local minima of $E$ and fixed points of retrieval dynamics $\mathcal{T}$. This allows us to construct a separation loss $\mathcal{L}_{\Phi}$ that distinguishes the local minima of $E_{\mathcal{K}}$ by separating stored memory patterns in kernel space.

Methodologically, we introduce Uniform Hopfield Memory Retrieval ($\mathtt{U\text{-}Hop}$). It is a two-stage optimization formulation. The first stage is separation loss $\mathcal{L}_{\Phi}$ minimization, distancing stored memory patterns in kernel space. The second stage performs energy minimization with the kernel-induced $\mathcal{T}_{\mathcal{K}}$. The first stage enhanced $E_{\mathcal{K}}$, making it able to relocate its local minima to a more separated coordinate. As a result, modern Hopfield models under $\mathtt{U\text{-}Hop}$ is able to obtain improved memory capacity.

Empirically, $\mathtt{U\text{-}Hop}$ improves memory retrieval outcomes by a large margin comparing to other baselines. When applied to deep learning scenarios, $\mathtt{U\text{-}Hop}$ significantly improves model’s memorization capacity, generalization and convergence speed. We show that $\mathtt{U\text{-}Hop}$ improves memory retrieval tasks by an average 30% margin even under a single iteration of separation minimization, and learning tasks by an average 3% margin.

Section 2 presents $\mathtt{U\text{-}Hop}$. Section 3 connects $\mathtt{U\text{-}Hop}$ to deep learning. Section 4 conducts extensive numerical experiments to support $\mathtt{U\text{-}Hop}$. Appendix includes proofs, experimental details, and additional experimental studies.

𝚄​-​𝙷𝚘𝚙\mathtt{U\text{-}Hop}: Retrieval as Two-Stage Optimization

In this section, Section 2.1 introduces a learnable feature map that maps patterns and the energy function into a kernel space, and demonstrate the fixed-point convergence property of kernelized modern Hopfield models. Section 2.2 presents $\mathtt{U\text{-}Hop}$ (Algorithm 1), a two-stage algorithm for the kernel learning with optimal theoretical guarantees. It maximizes pattern separation by minimizing a novel Separation Loss.

In this section, we first parameterize the similarity measure(s) in modern Hopfield model(s) with a learnable kernel (via feature map (2.1)), and then show the induced models (with energy (2.2)) satisfying the defining properties of modern Hopfield models (Theorem 2.1, Lemma 2.1).

with $\mathbf{z}^{\star}=\text{Sparsemax}\left(\beta,\mathbf{z}\right)$. Section D.1 includes the definitions of $\alpha\text{-EntMax}$ and Sparsemax.

With 2.1, the energy function $E(\mathbf{x})$ was monotonically decreased by the following retrieval dynamics:

where $\text{Sep}_{\alpha=1}\left(\cdot\right)=\text{Softmax}\left(\cdot\right)$, $\text{Sep}_{\alpha=2}\left(\cdot\right)=\text{Sparsemax}\left(\cdot\right)$ and $\text{Sep}_{\alpha\in}\left(\cdot\right)=\alpha\text{-EntMax}\left(\cdot\right)$.

By 2.1 and the convexity of $\mathcal{K}$, there exists an inverse map that transforms the CCCP results in kernel space back to the state space, where $\mathbf{x}$ and ${\bm{\xi}}_{\mu}$ are located. We then complete the proof using the Concave-Convex Procedure (CCCP) and the convex conjugate construction following (Hu et al. 2023; Wu et al. 2024). See Section E.2 for a detailed proof. ∎

The introduction of $\mathcal{K}$ releases similarity measure from Euclidean inner-product to a learnable form via the weight $\mathbf{W}$ of the features map $\Phi$. Moreover, the new Hopfield model ((2.2) and (2.4)) includes all deep learning compatible existing modern Hopfield models (Hu et al. 2023; Wu et al. 2024; Ramsauer et al. 2020). If we replace the kernel $\mathcal{K}(\cdot,\cdot)$ with inner-product $\Braket{\cdot,\cdot}$, then (2.2) reduces back to the general sparse model Hopfield model (Wu et al. 2024) Recall that the general sparse Hopfield model encompasses both dense (Ramsauer et al. 2020) and sparse (Hu et al. 2023) models as its special cases..

While Theorem 2.1 guarantees the monotonic minimization of energy using $\mathcal{T}$, the fixed point of $\mathcal{T}$ might not be the local minima of $E(\mathbf{x})$ according to Sriperumbudur and Lanckriet 2009. Therefore, we provide the next lemma to ensure their alignment, following (Hu et al. 2023; Wu et al. 2024; Ramsauer et al. 2020; Sriperumbudur and Lanckriet 2009).

Given the energy function $E(\mathbf{x})$ Equation 2.2 and retrieval dynamics $\mathcal{T}_{\mathcal{K}}(\mathbf{x})$ Equation 2.4, respectively. For any sequence $\{\mathbf{x}_{t}\}_{t=0}^{\infty}$ generated by the iteration $\mathbf{x}_{t^{\prime}+1}=\mathcal{T}_{\mathcal{K}}(\mathbf{x}_{t^{\prime}})$, all limit points of this sequence are stationary points of $E$.

By the monotonic energy minimization property of $\mathcal{T}$ (Theorem 2.1) along with (Hu et al. 2023, Lemma 2.2), we prove this through Zangwill’s global convergence theory (Zangwill 1969; Sriperumbudur and Lanckriet 2009). See Section E.1 for a detailed proof. ∎

In summary, with $\Phi$, the parameterized similarity measure $\mathcal{K}$ introduces an additional degree of freedom for us to relocate the minima of energy landscape $E_{\mathcal{K}}$. We show that the Uniform Memory Hopfield Energy (2.2) and its induced retrieval dynamics (2.4) satisfies the defining properties of modern Hopfield models (Theorem 2.1 and Lemma 2.1). Importantly, Lemma 2.1 states that minimizing the energy $E$ with $\mathcal{T}$ also leads to convergence to the fixed point of $\mathcal{T}$.

This is pivotal in motivating our next step: constructing a separation loss $\mathcal{L}_{\Phi}$. This loss distinguishes the local minima of $E_{\mathcal{K}}$ by separating stored memory patterns in the kernel space. With $\mathcal{L}_{\Phi}$, we then formulate the memory retrieval dynamics of the modern Hopfield associative memory model as a two-stage optimization, termed $\mathtt{U\text{-}Hop}$. This includes an additional stage of separation maximization (by learning the kernel), significantly enhancing memory capacity.

For all $\mu\in[M]$, let $R_{\Phi}\coloneqq{\frac{1}{2}}\mathop{\rm min}_{\nu\neq\mu;\nu,\mu\in[M]}\norm{\Phi(\bxi_\mu)-\Phi(\bxi_\nu)}$ be the finite radius of each (kernelized) sphere $\mathcal{S}_{\Phi,\mu}$ centered at (kernelized) memory pattern $\Phi({\bm{\xi}}_{\mu})$. We say ${\bm{\xi}}_{\mu}$ is stored if there exists a generalized fixed point of $\mathcal{T}_{\mathcal{K}}$, such that $\Phi(\mathbf{x}^{\star}_{\mu})\in\mathcal{S}_{\Phi,\mu}$, to which all limit points $\Phi(\mathbf{x})\in\mathcal{S}_{\Phi,\mu}$ converge to, and $\mathcal{S}_{\Phi,\mu}\cap\mathcal{S}_{\Phi,\nu}=\emptyset$ for $\nu\neq\mu$. We say ${\bm{\xi}}_{\mu}$ is $\epsilon$-retrieved by $\mathcal{T}_{\mathcal{K}}$ with $\mathbf{x}$ for an error $\epsilon$.

2 Separation Loss and 𝚄​-​𝙷𝚘𝚙\mathtt{U\text{-}Hop}

In this section, we first introduce a separation loss $\mathcal{L}_{\Phi}$ (Definition 2.2) over the stored memory set $\bm{\Xi}$. Minimizing $\mathcal{L}_{\Phi}$ results in the separation of stored patterns within any given $\bm{\Xi}$. Consequently, we incorporate this separation-maximization step into the standard memory retrieval process ((2.4)), leading to a novel two-stage formulation/algorithm, $\mathtt{U\text{-}Hop}$, for memory retrieval (Algorithm 1).

$\mathcal{L}_{\Phi}$ indicates the logarithm of average Gaussian separation of $\Phi$ vector pairs over $\bm{\Xi}$. Naturally, minimization of $\mathcal{L}_{\Phi}$ leads to an on-average dissimilarity among kernelized memory patterns, i.e., $\{\Phi({\bm{\xi}}_{\mu})\}_{\mu\in[M]}$. Notably, $\mathcal{L}_{\Phi}$ is convex by design and hence exists an optimizer $\mathbf{W}^{\star}$ the maximizes the average distance between all possible memory pattern pairs.

Now we introduce Algorithm 1, the Uniform Memory Retrieval $\mathtt{U\text{-}Hop}$, for learning a suitable kernel and then retrieving stored memory from the learned kernel space.

Algorithm 1 is a 2-stage optimization process. For the first stage, we run $N$ iterations of kernel learning to minimize the separation loss, thus resulting in larger $\Delta_{\mu}$ for $\mu\in[M]$. Next, we rescale each row of the affine matrix $\mathbf{W}$ to ensure the magnitude remains the same for memory patterns. For the second stage, we run $T$ update steps for the retrieval dynamics, thus resulting in Hopfield energy minimization. Note that the learned kernel results in a new energy landscape of $E$, and is expected to encode memory patterns into local minima that separates from all other memory patterns.

3 Exact Memory Retrieval

Let $\mathbf{x}^{\star}$ be fixed points of $\mathcal{T}$. By Definition 2.1 and Hu et al. 2023, the retrieval error exhibits a naive bound

The $\norm{\bx^\star-\bxi_\mu}$ term forbids the exact memory retrieval. Explicitly, exact memory retrieval requires the memory pattern to be the fixed point of $\mathcal{T}$, namely $\norm{\bx^\star-\bxi_\mu}=0$. With this observation, we deduce the condition of exact retrieval

where $\mathbf{e}_{\mu}$ is the one-hot vector with the $\mu$-th element as $1$. By plugging ${\bm{\xi}}_{\mu}$ into $\mathcal{T}(\cdot)$, we see it is a fixed point $\mathcal{T}({\bm{\xi}}_{\mu})={\bm{\xi}}_{\mu}$ and retrieves the target memory ${\bm{\xi}}_{\mu}$ only when (2.6) holds. In the standard modern Hopfield model (utilizing the $\mathop{\rm{Softmax}}$ Sep function), the inability of $\mathop{\rm{Softmax}}$ to satisfy (2.6) results in a lack of exact retrieval (Martins et al. 2023), thereby preventing the modern Hopfield network from converging to a single memory pattern.

To combat this, we show $\mathtt{U\text{-}Hop}$ achieves exact memory retrieval when $\alpha>1$, based on the sparse extensions of modern Hopfield model (Wu et al. 2024; Hu et al. 2023; Martins et al. 2023). Specifically, we study the application of $\mathtt{U\text{-}Hop}$ with $\alpha$-EntMax as separation when $\alpha>1$.

Let $\mathcal{T}_{\text{sparse}}$ be $\mathcal{T}_{\mathcal{K}}$ from Theorem 2.1 with $\alpha>1$. Let $\mathcal{T}_{\text{sparse}}$ a real-valued kernel $\mathcal{K}$ with feature map $\Phi$. Let $t>0,\beta>0$. Supposed the query $\mathbf{x}\in\mathcal{S}_{\Phi,\mu}$, $\Phi\left({\bm{\xi}}_{\mu}\right)$ is the fixed point of $\mathcal{T}_{\text{sparse}}$ if the following condition is satisfied:

From (2.7), minimizing the separation loss gives the benefit of having the memory pattern to be the fixed point of $\mathcal{T}_{\text{sparse}}$. As a result, Sparse and Generalized Sparse Hopfield models (Hu et al. 2023; Martins et al. 2023; Wu et al. 2024) under $\mathtt{U\text{-}Hop}$ further improves the retrieval accuracy. The next corollary is an extension of the above theorem where we observe the condition with respect to the lipschitzness of $\Phi$.

Let $L>0$ be the Lipschitz constant of $\Phi$. Following Theorem 2.2, $\mathcal{T}_{\mathcal{K}}$ achieves exact memory retrieval if

See Section E.3 for a detailed proof. Note that with $\Phi$ defined in (2.1), $\Phi$ is always $L$-Lipschitz. ∎

Connecting to Modern Deep Learning

To incorporate $\mathtt{U\text{-}Hop}$ into deep learning, we first introduce a kernelized Hopfield layer. Here we propose a deep learning compatible layer based on $\mathtt{U\text{-}Hop}$ as

Note that this is a kernelized version of $\mathtt{Hopfield}$ (Ramsauer et al. 2020), $\mathtt{SparseHopfield}$ (Hu et al. 2023) and $\mathtt{GSH}$ (Wu et al. 2024) layers, which serve as an alternative to attention mechanism variants.

Next, we introduce the average separation loss for deep learning compatible $\mathtt{U\text{-}Hop}$.

Specifically, $\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}$ is the loss minimizer of $\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}_{\Phi}\left(\bm{\Xi};t\right)$.

The empirical validation is in Appendix G. This theorem shows that with a suitable kernel, the expressiveness of $\mathtt{Hopfield}$ layers under $\mathtt{U\text{-}Hop}$ reaches its full potential. The main difference between this new loss function and the separation loss is the square on $\mathcal{K}(\mathbf{u},\mathbf{v})$. Note that this theorem requires $\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}(\mathbf{u},\mathbf{v})=0$ for any given $\mathbf{u},\mathbf{v}\in\bm{\Xi}$, $\mathbf{u}\neq\mathbf{v}$, which implies it is only possible when $d\geq M$. In the context of deep learning, the patch size must not be larger than the hidden dimension to realize this result. This theorem extends the representation theorem in (Bhojanapalli et al. 2020) to a practical setting, showing that $\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}$ overcomes the low-rank bottleneck of the attention mechanism and Hopfield layer as well.

The next algorithm is the realization of searching for $\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}$ under supervised learning schema. Consider a supervised learning problem with input data $\mathcal{X}=\{\mathbf{X}_{1},...,\mathbf{X}_{n}\}$, label $\mathcal{Y}=\{\mathbf{Y}_{1},...,\mathbf{Y}_{n}\}$, model $F\coloneqq\mathcal{X}\rightarrow\mathcal{Y}$, where $F$ consists of one layer of “$\mathtt{U\text{-}Hop}$ + Hopfield layer”. The stage-I of $\mathtt{U\text{-}Hop}$ is parameterized by $\theta$, and $F$ is parameterized be $\theta_{F}$.

Experimental Studies

To validate the efficacy of $\mathtt{U\text{-}Hop}$, we test it on both associative memory retrieval task and deep learning task (image classification) with multiple real world datasets.

The memory retrieval task involves retrieving a memory pattern from a stored memory set. In particular, this experiment aim to reconstruct memories based on a query. The query is generated by randomly masking 50% of pixels in the target image. We compare our method against several modern Hopfield models (Hu et al. 2023; Ramsauer et al. 2020; Krotov and Hopfield 2016). We also vary the iteration number $N$ for the first stage in Algorithm 1. We use MNIST, CIFAR10 datasets for this task. Please see Appendix F for experimental details.

This experiment follows the same procedure as the memory capacity tasks, but with multiple levels of injected noise on the target image instead of masking out pixels to generate queries. We use Gaussian noise to contaminate the queries and vary the noise level by altering the mean of the Gaussian vectors. As the noise level increases, it becomes more difficult to retrieve the memory with low error. A higher noise level results in greater difficulty in achieving low-error retrieval. We use MNIST, CIFAR10 for this task. Please see Appendix G for experimental details.

We set $\beta=1,t=2$ across all the memory retrieval experiments. For the evaluation metric, we follow (Hu et al. 2023; Millidge et al. 2022) to use the Sum-of-Square pixel differences between the ground truth image and the retrieved image.

See Figure 3 for results of memory capacity and noise robustness, Figure 5 for results of “Stage I iteration improve retrieval error” and Section G.1 for the relationship between separation loss and retrieval error.

For memory capacity, $\mathtt{U\text{-}Hop}$ outperforms all other baselines by a large margin. This result shows the retrieval dynamics under $\mathtt{U\text{-}Hop}$ is near optimal across all memory set sizes. Next, we vary the iteration $N$ to observe how fast the retrieval error decreases as the $N$ goes up. In Figure 5, we show a strong correlation between $N$ and retrieval error.

For noise robustness, $\mathtt{U\text{-}Hop}$ shows strong performance against all baselines as well as showed in Figure 3.

2 Supervised Learning Tasks

For classification tasks, we compare our method against $\mathtt{Hopfield}$ (Ramsauer et al. 2020) and $\mathtt{SparseHopfield}$ (Hu et al. 2023). We test two settings:

$\mathtt{U\text{-}Hop}$ + Dense Modern Hopfield Model (Ramsauer et al. 2020), and

$\mathtt{U\text{-}Hop}$ + Sparse Modern Hopfield Model (Hu et al. 2023).

We vary the training sample size and observe model performance. We focus on (i) convergence speed (speed of loss decay), (ii) generalization power (test accuracy). We use CIFAR10, CIFAR100 and TinyImageNet for this task. Please see Appendix F for more experimental details.

We use the following $\mathtt{Hopfield}$ layer (Ramsauer et al. 2020) to replace the self-attention mechanism in Vision Transformer:

To verify Theorem 3.1, we evaluate how many samples a model can memorize in supervised learning task. We follow the image classification settings, and see how Hopfield models with and without $\mathtt{U\text{-}Hop}$ react to sample size growth.

We also use the $\mathtt{STanHop}\text{-}\mathtt{Net}$ (Wu et al. 2024) as our test-bed and observe the performance change with and without $\mathtt{U\text{-}Hop}$. For this task, we use ETTh1, ETTm1 and WTH datasets. We use the prediction horizon of $\{96,192,336,720\}$ for all datasets. Please see Appendix F for experimental and hyperparameter details.

We compare the performance of Modern Hopfield and Sparse Hopfield with and without $\mathtt{U\text{-}Hop}$. For image classification, we use Vision Transformer (Dosovitskiy et al. 2020) as test-bed and replace the attention mechanism with $\mathtt{Hopfield}$ (Ramsauer et al. 2020) and $\mathtt{SparseHopfield}$ layer (Hu et al. 2023). For time series prediction, we compare the performance of $\mathtt{STanHop\text{-}Net}$ (Hu et al. 2023) with and without $\mathtt{U\text{-}Hop}$.

See Table 1 for convergence results of image classification task, Figure 6 for expressiveness results (Theorem 3.1) and Table 8 for time series prediction.

For image classification, we observe that modern Hopfield models under $\mathtt{U\text{-}Hop}$ consistently outperform other baselines, and the performance gap increases with the sample size growth. Additionally, $\mathtt{U\text{-}Hop}$ models shows superior convergence speed comparing to other baselines on both training and test set. For model generalization, see Table 1, for convergence results, see Section G.5.

For model expressiveness, we observe that when the dataset size is small, $\mathtt{U\text{-}Hop}$ has similar memorization capability as $\mathtt{Hopfield}$. However, as the dataset size increases, $\mathtt{Hopfield}$ without $\mathtt{U\text{-}Hop}$ shows a sharp degeneration on training accuracy and struggles to converge well, as evidenced in Figure 6. For more detailed results, see Section G.4.

For time series prediction, our results (in Section G.6) demonstrate that even on SOTA Hopfield-based time series model, $\mathtt{U\text{-}Hop}$ delivers performance improvement across different datasets and prediction horizons.

3 More Discussions on Experimental Results

For memory retrieval tasks, $\mathtt{U\text{-}Hop}$ delivers significant improvements on retrieval error by lowering separation loss over the memory set. For the epochs $N$ required for kernel learning, we demonstrate that low retrieval error has strong correlation with large size of $N$. This is expected as our separation loss is convex and guaranteed to obtain global optima with a rate of $\mathcal{O}(1/N)$. As showed in Figure 5, separation loss consistently decreased as $N$ goes up.

For classification tasks, $\mathtt{U\text{-}Hop}$ delivers significant improvements in predictive power of the underlying models. Comparing to contrastive self-supervised learning (Wang and Isola 2020; Chen et al. 2020), where they maximize pairwise distance between samples, $\mathtt{U\text{-}Hop}$ maximizes the pairwise distance between patches. As Saunshi et al. 2022 show maximizing the distance over samples improves class generalization and is beneficial to downstream tasks.

Our experiment results indicate 2 new insights that the separation on the patch/token level also leads to better generalization. Firstly, we hypothesize that the Stage I of $\mathtt{U\text{-}Hop}$ serves as a pre-training step for a better representation with more separated data geometry. Namely, tokens/patches projected to kernel space have higher quality of representation as $\mathtt{U\text{-}Hop}$’s first iteration leads to better patch separation. Secondly, though “ $\mathtt{Hopfield}$ layers with and without $\mathtt{U\text{-}Hop}$” and “expressiveness” experiments (Table 1 and Figure 6), we observe that solely increasing embedding dimension do not guarantee to escape from the low-rank bottleneck in attention- and Hopfield-based models (Bhojanapalli et al. 2020). We conclude that this is because these models do not utilize their full expressive power (as in Theorem 3.1), despite of high embedding dimension. This observation supplements the existing “high-dimensional embedding improves low-rank bottleneck” conjecture (Bhojanapalli et al. 2020) with an intuitive yet effective learning scheme.

Concluding Remarks

Algorithm 1 has a time complexity of $\mathcal{O}(N+T)$. Algorithm 2 has a time complexity of $\mathcal{O}(N_{o}N_{i})$. Although this increases the standard supervised learning training time by a factor of $N_{i}$, our experimental results demonstrate that models under $\mathtt{U\text{-}Hop}$ mitigate this issue with a faster convergence speed, requiring fewer epochs to converge. See Section G.5 for related empirical results.

One notable limitation is that the optimality of separation loss (Definition 2.2) does not guarantee maximal separation for $R\coloneqq\frac{1}{2}\min_{\mu,\nu\neq\mu\in[M]}\|\boldsymbol{\xi}_{\mu}-\boldsymbol{\xi}_{\nu}\|$ for any given $\boldsymbol{\Xi}$. This problem (maximizing $R$) is inherently a max-min (non-convex) problem and is less straightforward to analyze (Comparison between max and avg. loss is in Section G.2). To achieve provably optimal memory capacity, we plan to explore different loss functions or learning schemes in the future.

To address the above limitation, the follow-up work (Hu et al. 2024d) presents a provably optimal memory capacity bound for kernelized modern Hopfield models and introduces a sub-linear time algorithm, $\mathtt{U\text{-}Hop}$+, to achieve this optimal capacity.

Impact Statement

This research is theoretical and is not expected to have negative social impacts. As outlined in the introduction and related works, the primary goal of this study is to enhance our understanding of the underlying principles of large Hopfield-based and transformer-based foundation models from an associative memory perspective.

Acknowledgments

JH would like to thank Stephen Cheng, Shang Wu, Dino Feng and Andrew Chen for enlightening discussions, the Red Maple Family for support, and Jiayi Wang for facilitating experimental deployments. The authors would also like to thank the anonymous reviewers and program chairs for their constructive comments.

JH is partially supported by the Walter P. Murphy Fellowship. HL is partially supported by NIH R01LM1372201, NSF CAREER1841569, DOE DE-AC02-07CH11359, DOE LAB 20-2261 and a NSF TRIPODS1740735. This research was supported in part through the computational resources and staff contributions provided for the Quest high performance computing facility at Northwestern University which is jointly supported by the Office of the Provost, the Office for Research, and Northwestern University Information Technology. The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding agencies.

Supplementary Material

Section D. Supplementary Theoretical Backgrounds

Appendix B Related Work

Associative memory models (Willshaw et al. 1969; Kanerva 1988) have been widely discussed in both the neuroscience and machine learning fields. The main goal of these models are to store a set of memory patterns where those patterns can be retrieved with respect to a given query. Hopfield models represent a primary category within the class of computational associative memory models (Hopfield 1982). Starting from the classical Hopfield models (Hopfield 1982; Hopfield 1984; Krotov and Hopfield 2021), these models are able to store and retrieve binary patterns with guaranteed memorization capacity. Their biologically plausible designs provides significant insights to understand both human brains (Yampolskaya and Mehta 2023; Krotov and Hopfield 2021) and modern deep learning paradigms (Burns 2024; Cabannes et al. 2024b; Cabannes et al. 2024a; Kozachkov et al. 2023; Negri et al. 2023; Ramsauer et al. 2020). Recently, these Hopfield models regain interest in the deep learning field due to its connection to the attention mechanism in transformers. Notably, Ramsauer et al. 2020 propose the Modern Hopfield models (MHMs) whose single-step update is equivalent to the attention mechanism (Vaswani et al. 2017). As a result, this connection (starting from the dense associative memory model (Krotov and Hopfield 2016)) facilitates the integration of associative memory models into modern deep learning (Hofmann et al. 2024; Hu et al. 2024b; Xu et al. 2024; Wu et al. 2024; Burns and Fukai 2023; Auer et al. 2024; Widrich et al. 2020) and large foundation models (Hu et al. 2024a; Pan et al. 2024; Fürst et al. 2022).

Beside empirical success, Modern Hopfield Models (MHM) offer a low-assumption theoretical framework for analyzing transformer-based deep learning architectures. Toward their fundamental theory, Hu et al. 2023 and Wu et al. 2024 point out that the energy function of MHM and its sparse variants are actually tied to the convex conjugates of different entropic regularizers. This has led to the Sparse and Generalized Sparse HMHs, which are connected to attention mechanisms with various degrees of sparsity (Correia et al. 2019; Vaswani et al. 2017; Martins and Astudillo 2016). Extending this foundation, Hu et al. 2024b further complement this understanding with the principled construction of possible efficient variants from a nonparametric perspective. Furthermore, Hu et al. 2024c provide a detailed theoretical analysis of all possible efficient variants, through the lens of fine-grained complexity theory.

We would like to comment further on the results of (Hu et al. 2024c). First, it observes that the magnitude of the patterns (i.e., the norms of queries and memories) not only affects retrieval accuracy (as seen in the linear $m$ scaling in (1.3)), but also determines the efficiency of a variant of the modern Hopfield model. This norm-based efficiency criterion, with precision guarantees, echoes the outlier effect in the attention heads of transformer models (Hu et al. 2024a). This outlier effect is well-known in pretraining large transformer-based models for its negative impact on model quantization performance (Sun et al. 2024; Bondarenko et al. 2023; Bondarenko et al. 2021). To address this, Hu et al. 2024a interpret the outlier effect as inefficient rare memory retrieval and propose the outlier-efficient Hopfield layer for transformer-based large models, demonstrating strong empirical performance and theoretical guarantees. The benefits of removing outliers in the attention heads of transformer-based large foundation models are also highlighted in (Gu et al. 2024a; Gu et al. 2024b; Alman and Song 2024a; Alman and Song 2024b; Alman and Song 2023; Gao et al. 2023) from various theoretical perspectives. In this work, the removal of outliers is achieved by the row-wise normalization in $\mathtt{U\text{-}Hop}$ (see line 4 of Algorithm 1).

Another line of research focuses on learning an associative memory model (Tyulmankov et al. 2021; Salvatori et al. 2021) that has the ability to “read” (retrieve) and “write” (store) memories. Particularly, this type of method contains a ”readout” network to retrieve/generate memories with a given query. Bartunov et al. 2019 propose a meta learning framework to learn a generative network that treats the retrieval error as their energy function. Yoo and Wood 2022 propose a hierarchical associative memory model that relaxes the requirement of meta learning. Salvatori et al. 2021 propose a hierarchical generative network trained with predictive coding. Instead of deriving the retrieval dynamics from the energy function, these methods normally use a generative model for memory retrieval. With the expressiveness of deep neural networks, such method showed great empirical performances. However, since the structure of the readout network does not connect or dependent on the energy function, they are not able to preserve appealing theoretical guarantees like Hopfield models.

Iatropoulos et al. 2022 propose a kernelized memory network In a similar vein, Schaeffer et al. 2024 bridge associative memory models and probabilistic modeling. . They formulate the modern Hopfield models with a recurrent SVM model. In particular, their kernel is a single layer feed forward network that is trained to memorize patterns. However, their framework consists of several high assumptions. In comparison, our proposed framework has mild assumption on parameters and pattern distributions. In addition, $\mathtt{U\text{-}Hop}$ has significant practical usage and was validated through extensive experiments in both memory retrieval and supervised learning tasks.

This work bridges two paradigms of associative memory models via a non-singular kernel, such that the kernelized energy function (2.2) still satisfies the defining properties of modern Hopfield models, i.e. attention-included retrieval dynamics (Theorem 2.1). A comparison between Uniform Memory Hopfield and similar models are shown in Table 3.

The usage of kernels and feature expansions in transformers has been extensively discussed in previous literature. One primary objective of these studies is to reduce the computational complexity associated with attention mechanisms. For instance, Chen et al. 2021b; Kitaev et al. 2020; Chen et al. 2021a demonstrate empirically and theoretically that these efficient algorithms can effectively approximate SoftMax attention. Song et al. 2021 provides a generalized framework for attention mechanism by decomposing it into two parts, RBF kernel as similarity measure and $L_{2}$ norm weighting on tokens. In our paper, we offer a distinct perspective aimed at enhancing memory capacity, drawing inspiration from the construction of the modern Hopfield model. Therefore, our approach differs from attempting to approximate the standard modern Hopfield association. Instead, we focus on relocating memory patterns to facilitate easier retrieval.

Appendix C Connection to Transformer Attentions

Suppose that $\mathbf{X}$ and $\bm{\Xi}$ are embedded from the raw query ${\mathbf{R}}$ and $\mathbf{Y}$ memory patterns, respectively, via $\mathbf{X}^{\mathsf{T}}={\mathbf{R}}\mathbf{W}_{Q}\coloneqq\mathbf{Q}$, and $\bm{\Xi}^{\mathsf{T}}=\mathbf{Y}\mathbf{W}_{K}\coloneqq\mathbf{K}$, with some projection matrices $\mathbf{W}_{Q}$ and $\mathbf{W}_{K}$. Then, taking the transport of $\mathcal{T}$ in (1.2) and multiplying with $\mathbf{W}_{V}$ such that $\mathbf{V}\coloneqq\mathbf{K}\mathbf{W}_{V}$, we obtain

This result enables that the modern Hopfield models are able to serve as powerful alternatives to the attention mechanism equipped with additional functionalities.

Appendix D Supplementary Theoretical Backgrounds

Here we quote some known results from (Hu et al. 2023; Martins and Astudillo 2016).

The variational form of $\alpha\text{-EntMax}$ is defined by the optimization problem

where $\Psi_{\alpha}(\cdot)$ is the Tsallis entropic regularizer given by (D.2).

D.2 Separation Loss

With the output vector of $\Phi$ is normalized, we have

D.3 Convergence Rate of Gradient Descent

where $f(x^{\star})$ is the optimal value of $f$. Intuitively, this means that gradient descent is guaranteed to converge and that it converges with rate $\mathcal{O}(1/k)$.

we must run $k=\mathcal{O}(1/\epsilon)$ iterations of gradient descent, which gives us a sub-linear

Appendix E Proofs of Main Text

Here we introduce a helper lemma from (Sriperumbudur and Lanckriet 2009, Lemma 5).

Following Theorem 2.1, $\mathbf{x}$ is called the fixed point of iteration $\mathcal{T}$ w.r.t. $E$ if $\mathbf{x}=\mathcal{T}(\mathbf{x})$ and is considered as a generalized fixed point of $\mathcal{T}$ if $\mathbf{x}\in\mathcal{T}(\mathbf{x})$. If $\mathbf{x}^{\star}$ is a generalized fixed point of $\mathcal{T}$, then, $\mathbf{x}^{\star}$ is a stationary point of the energy minimization problem in Equation 2.2.

Based on Zangwill’s global convergence theory (Zangwill 1969), a set of limit points of $\{x_{t}\}^{\infty}_{t=0}$ are all generalized fixed points if the retrieval dynamics and energy function satisfies the following conditions:

For any sequence $\{x_{t}\}^{\infty}_{t=0}$ with $x_{0}\in S_{\mu}$ as starting point, all points in $\{x_{t}\}^{\infty}_{t=0}$ are in the compact set $S_{\mu}$.

$E(\mathbf{x})$ is monotonically decreased by $\mathcal{T}(\mathbf{x})$, where $E(\mathbf{x}_{t+1})\leq E(\mathbf{x}_{t}),\forall\mathbf{x}_{t+1}=\mathcal{T}(\mathbf{x}_{t})$.

For all $t$, if $E(\mathbf{x}_{t+1})\leq E(\mathbf{x}_{t})$, $\mathcal{T}$ is closed at $\mathbf{x}_{t}$.

From Definition 2.1, since radius $R$ is bounded and closed, $S_{\mu}$ is a compact set. Thus satisfies the first condition. CCCP (Yuille and Rangarajan 2001) studied the monotonic decreasing property. With our definition of $E_{\text{convex}},E_{\text{concave}}$, we have $E_{U}(\mathbf{x},\mathbf{y})=E_{\text{convex}}(\mathbf{x})+E_{\text{concave}}(\mathbf{y})+\braket{(\mathbf{x}-\mathbf{y}),\nabla_{\mathbf{x}}E_{\text{concave}}(\mathbf{y})}$ is continuous in $\mathbf{x},\mathbf{y}$. As a result, by (Hu et al. 2023, Lemma E.1), condition (iii) holds due to the non-empty assumption on the point-to-set map $\mathcal{T}$. Thus, by Zangwill’s global convergence theory, all limit points are also the stationary points of the energy minimization problem in (2.2). By the results in Lemma E.1, these fixed points are also the stationary points of the minimization problem. Thus, (2.2) is guaranteed to converge to local minimum. ∎

E.2 Proof of Theorem 2.1

Since the $\mathop{\rm{lse}}$ function is non-decreasing and convex, and $\mathcal{K}$ is convex, the composited function $\mathop{\rm{lse}}\left(\mathcal{K}(\bm{\Xi},\mathbf{x})\right)$ is convex. Thus, the energy function is the sum of a convex function: $\Braket{\mathbf{W}\mathbf{x},\mathbf{W}\mathbf{x}}/2$ and a concave function: $-\mathop{\rm{lse}}\left(\mathcal{K}(\bm{\Xi},\mathbf{x})\right)$.

Therefore, we have $E(\mathbf{x})=E_{\text{convex}}(\mathbf{x})+E_{\text{concave}}(\mathbf{x})$. With the Concave-Convex Procedure (CCCP) (Yuille and Rangarajan 2001), the energy function $E(\mathbf{x})$ is guaranteed to monotonically decrease the energy $E$ as a function of time with the following update rule:

E.3 Proofs of Theorem 2.2 and Corollary 2.2.1

Here we use the same proof strategy in (Martins et al. 2023, Proposition 2).

If we have exact memory retrieval of pattern ${\bm{\xi}}_{\mu}$, the following equation holds:

This is also equivalent to ${\bm{\xi}}_{\mu}$ itself being the fixed point.

With the assumption of normalized patterns, we are able to reduce the above condition to

E.4 Proof of Theorem 3.1

By construction, any two columns in $\mathbf{X}^{\prime}$ satisfies:

With $\mathbf{D}_{\widetilde{\mathbf{W}}_{KQ}}=\mathbf{D}_{0}$, and using (E.4), we obtain

As a result, to construct $\mathbf{W}_{K}$ and $\mathbf{W}_{Q}$ such that they satisfy Theorem 3.1, any two matrices must satisfy

Appendix F Implementation Details

MNIST. It is a hand written digits image recognition dataset (LeCun et al. 1998) consists of 60000 training samples and 10000 test samples. Each image has the size of $28\times 28$. The label contains digits from 0 to 9.

CIFAR10. It is an image recognition dataset (Krizhevsky et al. 2009) consists of 50000 training samples and 10000 test samples. Each image has the size of $32\times 32$. The dataset contains 10 categories with 6000 samples for each.

CIFAR100. It is an image recognition dataset (Krizhevsky et al. 2009) consists of 50000 training samples and 10000 test samples. Rach image has the size of $32\times 32$. The dataset contains 100 categories with 600 samples for each.

TinyImageNet. It is an image recognition dataset (Le and Yang 2015) contains 100000 images of 200 classes. Each image is downsized to 64×64 colored images. Each class has 500 training images, 50 validation and test images.

ETT (Electricity Transformer Temperature). ETT (Zhou et al. 2021) records 2 years of data from two counties in China. We use two sub-datasets: ETTh1 (hourly) and ETTm1 (every 15 minutes). Each entry includes the “oil temperature” target and six power load features.

WTH (Weather). WTH records climatological data from approximately 1,600 U.S. sites between 2010 and 2013, measured hourly. Entries include the “wet bulb” target and 11 climate features.

F.2 Memory Capacity

For memory capacity experiment, we follow the settings in (Hu et al. 2023; Wu et al. 2024).

We randomly mask $50\%$ of the pixels in the image, using the masked image as a query for a single-step update with various Hopfield models. In the case of $\mathtt{U\text{-}Hop}$, we trained the kernel with different numbers of epochs on the memory set and then used it for memory retrieval. We reported the sum-of-square pixel difference between the retrieved image and the ground truth. In each run, we repeated this process for every image in the memory set, conducting the experiment 20 times for each baseline. The range of kernel learning epochs, memory set size can be found in Table 4. For Figure 26, we use $N=100$ for MNIST and $N=500$ for CIFAR10.

F.3 Noise Robustness

For noise robustness experiment, we follow the settings in (Hu et al. 2023; Wu et al. 2024).

For the noise robustness experiment, we randomly sampled a Gaussian noise vector for each image, varying the norm of the sampled noise to adjust the noise level. We then added the noise to the query image and performed a single-step update with different Hopfield models. For $\mathtt{U\text{-}Hop}$, we trained the kernel with N iterations on the memory set and then used it for memory retrieval. We set $N=100$ for MNIST and $N=200$ for CIFAR10. We reported the sum-of-square pixel difference between the retrieved image and the ground truth. In each run, we repeated this process for every image in the memory set, conducting the experiment 20 times for each baseline.

F.4 Classification

CIFAR10 and CIFAR100. For these two datasets, we consider four different Hopfield layers as encoder:

$\mathtt{Hopfield}$ (Ramsauer et al. 2020)

$\mathtt{SparseHopfield}$ (Hu et al. 2023)

$\mathtt{Hopfield}$ +$\mathtt{U\text{-}Hop}$

$\mathtt{SparseHopfield}$ + $\mathtt{U\text{-}Hop}$.

We use a fully connected layer right after the encoder for classification. For each image, we follow the same process as introduced in (Dosovitskiy et al. 2020). We split an image into patches and add an additional $\mathtt{CLS}$ patch for classification. We send patches into the $\mathtt{Hopfield}$ layer with the $\mathtt{CLS}$ patch as query and other patches as memory. We then send the output to a fully connected layer for prediction. We use the CrossEntropy loss and Adam optimizer for training. Hyperparameters are in Table 5.

Tiny ImageNet. For this dataset, we use a 3 layered Vision Transformer as backbone (Dosovitskiy et al. 2020), and use $\mathtt{Hopfield}$ variations to replace attention mechanism in . Other processes are the same as introduced in the above paragraph. For kernel learning, we learn all kernels in different layers by passing a full forward pass. We then send the output to a fully connected layer for prediction. Hyperparameters are in Table 6.

F.5 Hopfield-Based Time Series Prediction with STanHop-Net (Wu et al. 2024) (Table 8)

For this task, we use STanHop-Net (Wu et al. 2024) as backbone, and equip it with $\mathtt{U\text{-}Hop}$. In addition, we use Algorithm 2 to minimize both separation loss and the MAE loss. For prediction horizon of $96,192$, we use one layered STanHop-Net, for $336,720$, we use a two layered STanHop-Net. We use the same settings for with and without $\mathtt{U\text{-}Hop}$ We thank the authors of (Reneau et al. 2023) for their helpful comments on this part..

Appendix G Additional Numerical Experiments

This section is a visualization of relationship between separation loss and retrieval error on MNIST and CIFAR10. The result shows that the retrieval error is highly correlated with respect to the value separation loss.

G.2 Max. Loss v.s. Avg. Loss

In the main paper, we discuss the differences between minimizing the maximum separation loss and the average separation loss.

Ideally, minimizing the maximum separation loss directly contributes to $R_{\Phi}$. However, as stated in the main text, such a loss is a max-min problem, which is challenging to optimize. Moreover, it entails an additional quadratic time complexity due to the max operation. On the other hand, the average loss is more time-efficient. It is also convex, thereby guaranteeing convergence to the global optimum at a rate of $\mathcal{O}(1/N)$ under gradient descent, where $N$ is the number of iterations. However, the average loss does not guarantee maximizing $R_{\Phi}$, nor does it ensure an optimal $\Delta_{\Phi,\mu}$ for any $\mu\in[M]$. Therefore, its theoretical impact on memory capacity and retrieval error bound is difficult to quantify.

As a result, we conduct an analysis comparing the performance of each loss function. We vary the memory size and the kernel learning iteration $N$ and perform memory retrieval on MNIST and CIFAR10 datasets.

The results demonstrate that minimizing the average loss yields a lower retrieval error, and this advantage grows with the size of the memory set. Additionally, the retrieval error decreases more rapidly with respect to $N$ under average loss than under maximum loss. This outcome is anticipated, as minimizing the maximum loss is a non-convex problem and does not guarantee reaching global optima. This empirical finding indicates that in practice, minimizing the average loss leads to better efficiency and retrieval outcomes.

We observe that with the average loss, $\mathtt{U\text{-}Hop}$ achieves almost perfect retrieval outcomes with $N=200$. However, $\mathtt{U\text{-}Hop}$ using the maximum loss struggles to reach its global minimum during optimization, thus hindering its ability to achieve optimal memory capacity. This is observed in Figure 8.

With CIFAR10 being more difficult comparing to MNIST, $\mathtt{U\text{-}Hop}$ under average loss still outperforms $\mathtt{U\text{-}Hop}$ under Max. loss.

G.3 Memory Retrieval

Here we again show the memory retrieval results with higher resolution. The result demonstrates with $\mathtt{U\text{-}Hop}$, modern Hopfield models obtain significant improvement on both datasets.

G.4 Model Expressiveness

Here we present the model expressiveness on training data. This is a empirical validation for Theorem 3.1. We observe that without $\mathtt{U\text{-}Hop}$, baseline models suffer from sharp performance drop, which also lead to generalization degradation as show in Section G.5.1. In contrast, with $\mathtt{U\text{-}Hop}$, models show better robustness against sample size increase.

G.5 Classification

Here we conduct empirical analysis on the correlation between dataset size and model convergence. In general, it is more difficult to memorize all samples for larger dataset. However, learning from more samples might lead to a better generalization performance, results in higher test accuracy.

The result demonstrates $\mathtt{U\text{-}Hop}$ significantly improves model’s performance on 3 aspects:

The improvement also became more obvious when the dataset size increases. This is reasonable as the standard $\mathtt{Hopfield}$ layer is powerful enough to memorize small sample size with and without $\mathtt{U\text{-}Hop}$.

The model behavior was similar comparing to what we observe from CIFAR10. The performance improvement under $\mathtt{U\text{-}Hop}$ became stronger with the increase of dataset size. With $\mathtt{U\text{-}Hop}$, the model improves on both convergence speed, memorization capacity (training accuracy) and generalization power (test accuracy).

Models under $\mathtt{U\text{-}Hop}$ continue to show strong performance against baselines on Tiny ImageNet dataset. Notably, we use a 3 layer encoder for this dataset, which provides additional insights ensuring that $\mathtt{U\text{-}Hop}$ works well under deep neural network architecture.

G.6 Time Series Prediction

Here we report the results of our time series prediction experiment. From Table 8, we observe that $\mathtt{U\text{-}Hop}$ obtains improvement in most datasets and prediction horizons.

References