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

Weakness

W1: the algorithm's main justification is a lower bound that is not shown to be tight, so optimizing the minimal overlap between memories is not well supported (here)....:

Response: Sorry for the confusion caused. We want to emphasize that the inequality in Def. 2.7 is a condition instead of a lower bound. We understand the submitted draft is a bit ambiguous about this. We have made the following modifications to clarify.

We decompose Def.2.7 into two separate definitions:

• Definition 1 (Well-Separation Condition). Given a set of memory patterns $\Phi(\mathcal{V}) = \{ \Phi(v_\mu)\}_{\mu=1}^M \subseteq \mathbb{S}^{D_\Phi-1}$, we say the memory pattern $\Phi(v_\mu)$ satisfies the well-separation condition is the following holds:

• Definition 2. (Memory Code). Given a finite set of points $\Phi(\mathcal{V}) = \{ \Phi(v_\mu)\}_{\mu=1}^M \subseteq \mathbb{S}^{D_\Phi-1}$, and $\beta > 0$. We say $\Phi(\mathcal{V})$ is a memory code if all points in $\mathcal{V}$ satisfies the well-separation condition, i.e.

Further, we denote $\Lambda$ as the set of all memory codes in $\mathbb{S}^{D_\Phi - 1}$, such that

We also added intuitive explanations:

• Note that this inequality is a desired property of the code/data which benefits memory storage. In the following section, we demonstrate methods to achieve such property for all different $\mathcal{V}$.

Moreover, our Algorithm 1 is actually built on an upper bound, given at the inequality in line 585.

We hope these clarifications have addressed your concern.

W2: If the main claim of the paper is improving the capacity to store memories, I would expect a numerical validatio...

Response: Thank you for pointing that out. We recognize that many reviewers have identified the same problem of lacking experiments. In response, we have added and multiple instance learning and memory retrieval experiments to support our main result. Please see Table 2 and Figure 3 of the attached pdf.

Importantly, these new experimental results indicate

• KHMs improves the retrieval outcome by a large margin.
• U-Hop+ improves the MHM based model in multiple instance learning tasks.

W3: Significance: the results are mostly of practical importance for practitioners who wish to implement MHM with an optimal kernel. For that crowd, it is not that important if the algorithm is based on a novel lower bound but that the algorithm works, which needs to be better established.

Response: The main purpose of this paper is to

1. Identify the quantity of optimal memory capacity
2. Justify a theoretical grounded algorithm to achieve optimal capacity.

We want to emphasize that, for practical usage, the method provides empirical improvement with or without our justification. We will better establish this in our experiment section. Please see attached pdf for more validations.

Questions

Q1: How does the lower bound on memory capacity (lemma 2.1) contribute to ...

Response:
Two main reasons for us to demonstrate Lemma 2.1:

• Showing the traditional approach in hopfield is all you need does not work under the flexibility provided by KHMs.
• Demonstrating the exponential scaling behavior in $D_\Phi$, which corresponds to the quantity $M^\star$. To emphasize this relationship, we added the following proposition and description.

Proposition. Following Theorem 2.1, let $\theta \in$ 0, \frac{\pi}{2} $$, we have

where $\theta = \arccos$\rho( C_{opt} )$$, and $\varphi(\theta) = -\log \sin{\theta}$.

And a remark

Remark. We have $M^\star = O(c^{D_\Phi})$, for some $c > 1$. This property echoes the exponential capacity lower bound in Lemma 2.1.

Q2: What is Lambda in Definition 2.8, Theorem 2.1?

Response: Lambda is defined in the last line of our Definition 2.7. It refers to the set containing all memory codes.

Q3: Is the lower bound from definition 2.7 tight? Can you establish it is not vacuous?

Response: The inequality in Definition 2.7 does not serve as a lower bound. Instead, it is a condition of data that we want to achieve with Algorithm 1 because the existence of this condition benefits memory storage. The goal of algorithm 1 is to make this condition hold for all memories in feature space.

We would also like to highlight that this is a commonly used storage condition in [1, 2, 3, 4], showing that it is a well established/recognized condition.

• [1] Hopfield Netorks is all you need. ICLR 2021
• [2] On Sparse Modern Hopfield model. NeurIPS 2023
• [3] Uniform Memory Retrieval with Larger Capacity for Modern Hopfield Models ICML 2024.
• [4] Outlier-Efficient Hopfield Layers for Large Transformer-Based Models

===

Thank you for your suggestions and attention to details!

We hope that our responses adequately address your concerns, and we look forward to any further feedback.

Thank you again for your time and consideration!

W1: There is a lack of experimentation in this work. It would be great for the work to have comparisons, similar to those detailed in Uhop.

Response:

Thank you for pointing that out. We recognize that many reviewers have identified the same problem of lacking experiments. We have added additional memory retrieval and multiple instance learning experiments to support our main result.

The new multiple instance learning results and retrieval results are in Table 2 and Figure 3 of the attached pdf.

Q1: For figure 4, what is the gray dot in the center of the circle? Is it an image (of a class) in CIFAR or just a point indicating the center?

Response:

Yes, sorry for the confusion. And thanks for your attention in details! We meant to use the dot to indicate the center of the circle, it is not an image in CIFAR. We will remove it in our final version.

Q2: For the derivative of the RHS of the bound, …

Response: Yes, you are correct. Again, thanks for your attention in details! We have modified our proof accordingly.

Meanwhile, this correction does not affect our conclusion that the RHS is increasing in M.

Q3: Is it possible to see the loss curve for MNIST when $W$ optimizing? It would be great to see the stability of the new $\mathcal{L}$ just to get a sense of ...?

Response: Yes. Thanks for the suggestion. We provide the visualization of the loss curve for MNIST in Figure 4 of the attached pdf.

===

Thank you for your suggestions and attention to details!

We hope that our responses adequately address the reviewer's concerns, and we look forward to any further feedback.

Thank you for your time and consideration!

W1: The optimization procedure of U-Hop is similar U-Hop+ i.e., Gradient Descent is replaced by PGD. Due to this similarity U-Hop should have been added as baseline. Since in the standard case MHMs use learnable weight matrices (as can be seen in Eq. 10 of [Ramsauer et al., 2020] and other publications using MHMs) there should have been a comparison to these MHMs using gradient descent based on the average separation loss.

Response: Thank you for pointing that out. The short answer to this is — Equation 10 in [Ramsauer et al., 2020] does not ensure a well-defined retrieval update of an associative memory model.

For reasons:

• Equation 10 will only be a legal update rule for auto-associative memory if $W_V$ is excluded, with $W_K Y, W_Q R$ being memory and query, instead of $Y, R$. The consideration of hetero-associative memory is beyond the scope of this paper.
• Equation 10 without $W_V$ minimizes the energy function of $W_K Y$ and $W_Q R$, instead of real data $Y$, $R$. The retrieved pattern will be in $W_K Y$ space instead of $Y$ space. Thus, under MHM construction, we are not able to separate $W_K$ and $Y$.

W2: While memory capacity is an important property it is not immediately clear how much it helps for learning tasks. Therefore, experiments in a similar vein to [Wu et al, 2024a] should have been conducted on these tasks.

Response: We recognize that many reviewers have identified the same problem of lacking experiments. In response, we have added memory retrieval and multiple instance learning experiments to support our main results.

In our experimental result, U-Hop and U-Hop+ show similar results, but the theoretical justification of algorithm 1 makes this paper theoretically grounded.

For more experimental proofs, we have added the multiple instance learning results and retrieval results in Table 2 and Figure 3 of the attached pdf.

Q1: How do you make sure that in Algorithm 1 the weight matrix W keeps the full column rank property from Definition 2.7?

Response:

In our algorithm, we do not force W to be full-rank. Existing methods such as [1, 2] are able to satisfy this property in practice. We choose to only analyze the simple PGD case to keep the convergence analysis simple. Since $\mathbf{W}$ is in a high-dimensional space, it is almost surely full-rank (The probability of a randomly initialized matrix having full column-rank is 1). [Remark 2.1, Wu et al 2024] also explains this implication in detail. Thus we choose to not explicitly force this property in practice.

[1] McTorch, a manifold optimization library for deep learning [2] Trivializations for Gradient-Based Optimization on Manifolds (NIPS)

===

Thank you for your time and feedback! We look forward to further discussions!

W1: Prior work on KHNs already has analysis using linear kernels. It would be helpful to showcase the improvement brought by the current algorithm to improve clarity and contextualization relative to prior work.

Response:

The prior work in KHN mainly explains how separation maximization helps capacity but lacks theoretical justification for why their algorithm improves capacity. To fill this gap, this work provides a rigorous theoretical analysis of why and how separation maximization improves KHN memory capacity. Additionally, we introduce the notion of optimal capacity and prove the asymptotic conditions for achieving it.

We acknowledge that the current draft does not emphasize these points enough. We will update the final version accordingly.

W2: Moreover, the experiments only demonstrate the convergence property of the iterations. It would be nice to have more practical experiments on tasks such as multiple instance learning and retrieval. I could improve my evaluation if the authors would address the questions below.

Response: In response to more experiments, we have added new MIL results and retrieval results in Table 2 and Figure 3 of the attached pdf.

These additional experiments demonstrate clear improvements over standard MHN.

Q1: Since MHN is a special case of KHN, how do the current results reduce to the MHN results in [Ramsauer et al., 2020]? Are there any improvements?

Response:

The main difference of reducing to MHN is that MHN does not guarantee the well-separation condition will hold for all memories. Thus, not necessarily to store all the given points.

Second, MHN does not have an extra feature space to store memories. If the data dimension is too small, it is possible that the well-separation condition is impossible to achieve. For KHNs, since learning iteration of Algorithm 1 and $D_\Phi$ has no restrictions, one is able to select sufficiently large values of these two, to ensure the well-separation condition holds for all memories. We can also compare KHNs and MHNs with Lemma 2.1. With sufficiently large $D_\Phi$ and $N$, $R_\Phi$ will eventually surpass $R$, and thus obtain a higher lower bound.

Empirically, we conducted additional exps (in Table 2, Figure 3 of the attached pdf) showing improvements of UHop+ over MHN.

Q2: Definition 2.7 assumes that the patterns are normalized. This assumption is not explicitly stated elsewhere and is not well-justified.

Response:

We would like to highlight that this is a commonly used assumption in the analysis of [1, 2] and in experiments in [1, 2, 3].

A justification for this setup is its connection to transformer attention. The modern Hopfield update rule/retrieval dynamics are often compared to the attention mechanism, where the input query and memory set correspond to Q and K in attention. Since LayerNorm is commonly used in attention layers, this setup reflects real-world scenarios.

Analyzing non-normalized patterns is left for future work.

Lastly, we have added a highlight block for this assumption in line 92 of the final version.

Q3: Definition 2.8 is more like a statement rather than a definition. Why should the definition lead to the conclusion that "there is a feature map such that is a memory code"?

Response: Thanks for pointing this out. We recognize this potential confusion and have changed the definition to a Lemma, as it is a derived result from Definitions 2.2 and 2.7.

Q4: Is the convergence in Theorem 2.2 pointwise convergence or uniform convergence? For a given what's the error rate?

Response: The main purpose of Theorem 2.2 is to show one optimization problem converges to another. The convergence between two functions here is not our main focus. That being said, the convergence between $\mathcal{L}_0$ and HardMax loss is uniform convergence as $\tau$ goes to 0. We have made it clear in our manuscript by adding the following in line 589.

Given a $\tau > 0$, for the error rate $\epsilon$ between the convergence of $\mathcal{L}_0$ and HardMax Loss, we have $\epsilon \leq 2 \tau \log(M)$.

Q5: The proposed algorithm learns a linear map, what about non-linear maps?

Response: Yes, Algorithm 1 also works for non-linear maps.

We would like to highlight that only linear feature maps have been discovered for KHNs. The explicit form of non-linear maps is not yet known. The desired (retrieval dynamics with) non-linear map should satisfy fixed-point convergence and monotonic energy minimization to be well-defined under the KHN framework. Discovering its explicit form is left for future work.

Q6: In line 597, since you are referring to a known upper bound in a book, it would be better to provide more details for the reader to find the result.

Response: Thank you for pointing this out. We have modified the citation to [Thm 1, Moore 1974].

Q7: There are some typos: Line 48: does -> is Line 105: lower bounded by what? Definition 2.5: It seems that there should be an arg in front of minmax Definition 2.7: satisfies -> satisfy

Response: Thank you very much for pointing these out! For Line 105, we change it to: “lower bounded by the following lemma:” For other typos, we have conducted another round of proofreading and fixed all typos.

===

• [1] Uniform Memory Retrieval with Larger Capacity for Modern Hopfield Models. ICLR 2024
• [2] Sparse and Structured Hopfield Networks. ICLR 2024
• [3] Universal Hopfield Networks: A General Framework for Single-Shot Associative Memory Models
• [4] Trivializations for gradient-based optimization on manifolds. NIPS 2019
• [5] Vector Packing in Finite Dimensional Vector Spaces. 1974
• [6] On kissing numbers and spherical codes in high dimensions. 2018
• [7] New lower bounds on kissing numbers and spherical codes in high dimensions. 2023

We hope these responses have addressed your concerns. Thank you!

Thank you very much for the response.

Q1: In Fig. 3 of the attached...

Response: It is almost overlapped with the Generalized Sparse Hopfield + U-hop+ line.

Q2: Could you elaborate on how it is derived from Definitions 2.2 and 2.7?

Response: We apologize for the confusion. The logic behind this is that Def 2.2 leads to Def 2.7, and then Def 2.7 leads to the maximal capacity as we refer to the solution of the $\max_{ \Phi(V) } | \Phi(V) |$, and $\tilde{\Phi}$ is the solution of $\argmax_{ \Phi } \max_{ \Phi(V) } | \Phi(V) |$.

Q3: I don't think my question Q4 is addressed.

Response: With the same $\Xi$ and $\Phi$, $\mathcal{L}_0$ converges uniformly to the HardMax loss with same error rate of $\epsilon = O(\tau \log M)$. However, since Thm 2.2 considers two optimization problems, it is indeed $\Gamma$-convergence between these two minimization problems (convergence for functionals). With the property of $\Gamma$-convergence, the minimizers of one minimization problem will also converges to another. We will address that in the proof of our manuscript.

Thanks again for the response and your insightful review, we hope our responses address your concerns.

Glady!

Q1: Could you provide any intuition on why the Sparse Hopfield + U-hop+ line is almost overlapped with the Generalized Sparse Hopfield + U-hop+ line?

Response: Sparse Hopfield is a special case of Generalized Sparse Hopfield (GSH) with the sparsity parameter $\alpha=2$. Notably, $\alpha=2$ represents the most sparse GSH. Therefore, for datasets with high sparsity, the two may overlap.

Q2: In my previous question, I was wondering why Definition/Lemma 2.8 leads to the conclusion that "there is a feature map such that is a memory code". It seems to me that there are some important premises for Memory Code (Definition 2.7). One is the assumption that the patterns are normalized. Another is that the feature map is linear.

Response: Thanks for pointing this out. We notice our expression causes confusion. We have changed the paragraph into the following:

"From this definition, with a fixed $D_\Phi$, for some set of patterns $\mathcal{V}$ that has size $M^*$, there might exist a function $\tilde{\Phi} : \mathbb{R}^d \rightarrow \mathbb{R}^{D_\Phi}$, such that it pushes memories further from each other to the level where $\tilde{\Phi} ( \mathcal{V} )$ is a memory code. In other words, if one hopes to store all points in $\mathcal{V}$ with KHNs, the objective is to find an appropriate $\tilde{\Phi}$."

This paragraph is meant to be an intuitive explanation rather than a theoretical result. Thank you again for pointing this out, we have revised our manuscript accordingly.