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

Summary

Suggest a method to optimize kernel versions of Modern Hopfield Models by mapping memories onto well-separated values on a sphere in feature space. Shows that the spherical arrangement of memories is optimal for retrieval (e.g., maximize capacity) and thus can improve upon the linear kernel [Wu et al. 2024a], which optimises memory separation.

Strengths

* Originality: the idea presented is quite simple and intuitive, but the perfected procedure is non-trivial and can potentially improve memory performance across a large family of relevant MNM-s. * Clarity: the submission is clearly written and well organized.

Weaknesses

* Quality: 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). * Quality: if the main claim of the paper is improving the capacity to store memories, I would expect a numerical validation that the algorithm improves the number of memories that may be stored or increases the amount of noise that may be tolerated. The included results (sections 3.1 and 3.2) are only qualitative and are achieved on toy problems. Such a comparison can be done relative to plain-vanilla dense MHM or to [Wu et al. 2024a]. * 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.

Questions

1. How does the lower bound on memory capacity (lemma 2.1) contribute to the derived algorithm? It’s dependance on R and D of the embedding, but those are not optimized for in the algorithm. Furthermore, it doesn’t depend on the minimal memory separation. 2. What is Lambda in Definition 2.8, Theorem 2.1? 3. Is the lower bound from definition 2.7 tight? Can you establish it is not vacuous?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have presented a theoretical approach based on a lower bound for capacity, which may not directly translate to improvement in task performance and has not shown such improvements through experiments.

Summary

This work provides a theoretical analysis of the optimal capacity on both Modern Hopfield Networks (MHNs) and Kernelized Hopfield Networks (KHNs). Specifically, it seeks to answer four fundamental problems: 1. How does memory capacity affect the gap between memories? 2. Given a fixed feature dimensionality, how to select an optimal separation function? 3. Is there a simpler and more efficient algorithm which optimize the kernel? 4. How to connect the gap between memories and the optimal capacity of KHNs theoretically? Through the viewing of the memory set as a spherical code, the work establishes the problem of storing memories as a hard-max problem of maximizing the number of patterns stored, while ensuring that the angle between them is at least a certain constant (which the work does show). Consequently, this portrayal of the problem enables a simpler algorithm to optimize the kernel in contrast to the previous work UHop. Furthermore, the work illustrates that the optimal memory capacity is strongly dependent on the projection function $\phi$, such that its optimality maximizes the gap (or angle) between each pair of memories.

Strengths

The formulation of the problem presented in the work is great while the proofs are easy to follow. Moreover, the work achieves its goals in answering its four fundamental questions while providing a nice associative memory framework, which could possibly also be applied to other MHNs.

Weaknesses

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](https://arxiv.org/pdf/2404.03827).

Questions

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? For the derivative of the RHS of the bound, $\nabla^\phi_\text{min} + \frac{1}{2\beta} \ln (\frac{1}{2} \nabla^\phi_\text{min}) \geq \frac{1}{\beta} \ln (2 (M - 1))$, is it not $\frac{1}{\beta (M - 1)}$ instead of $\frac{1}{2\beta M}$? Is it possible to see the loss curve for MNIST when optimizing $\mathbf{W}$? It would be great to see the stability of the new $\mathcal{L}$ just to get a sense of $O(\frac{1}{-4tR^2 - \mathcal{L}^*_\phi(\Theta)})$.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

There is a significant lack of experimentation (or comparisons) in this work. Other types of MHN might require modifications or a different framework to analyze their optimal capacity and compatibility with KHNs. The feature map is rather simple --- it is a linear affine function. Lastly, the spherical code framework only considers normalized points on the hypersphere.

Summary

The paper is a theoretical analysis of the memory capacity of MHMs and KHMs. It establishes a connection between memory capacity and spherical codes. A method for approximating optimal memory capacity is introduced. Experiments are conducted to validate the theoretical findings.

Strengths

The paper is well written and easy to follow. The memory capacity is an important problem that can have a significant impact on SOTA architectures like Transformers. The paper is an interesting theoretical discussion and shows a connection between memory and spherical codes in an original manner. The theoretical framework and parts of the proofs are based on [Ramsauer et al., 2020] with adaptions to KHMs. These previous concepts are extended to get results for optimal memory capacity and to introduce an objective for optimal memory capacity. This objective is relaxed by an "average separation loss" introduced in [Wu et al, 2024a].

Weaknesses

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. While memory capacity is an important property it is not immediately clear how much it helps for learning tasks. Therefore, experiments in similar vein to [Wu et al, 2024a] should have been conducted on these tasks.

Questions

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

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Summary

Modern Hopfield Networks (MHNs) are limited because they require sufficient minimal separation $\Delta_{\min}$ for theoretical guarantee. Kernelized Hopfield Models (KHNs) mitigate this limitation by storing the memories in the feature space. In particular, theoretical analysis in [Wu et al., 2024a] uses a linear kernel. The current paper proposes an algorithm to find the "optimal" feature map using spherical coding. By solving a surrogate optimization problem whose solution converges to the HardMax problem asymptotically with vanishing temperature, the algorithm improves the memory capacity.

Strengths

MHNs emerge as a powerful tool for theoretical understanding of neural networks such as Transformers. The current paper builds upon KHNs and gives analysis of necessary conditions for KHNs to achieve optimal memory capacity. Further, a sub-linear time algorithm is proposed to optimize the feature map.

Weaknesses

Prior work on KHNs already has analysis using linear kernel. It would be helpful to showcase the improvement brought by the current algorithm to improve clarity and contextualization relative to prior work. 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.

Questions

1. 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? 2. Definition 2.7 assumes that the patterns are normalized. This assumption is not explicitly stated elsewhere and are not well-justified. 3. 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 $\Phi$ such that $\Phi(V)$ is a memory code"? 4. Is the convergence in Theorem 2.2 pointwise convergence or uniform convergence? For a given $\tau>0,$ what's the error rate? 5. The proposed algorithm learns a linear map, what about non-linear maps? 6. 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. 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

Rating

5

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

The paper is of theoretical nature. A paragraph on limitation is given in the main content.

Response to rebuttal

I appreciate the author's improvements in the manuscript and believe my main concerns have been addressed. I'm especially impressed by the added experiments, which, to me, provide important empirical support to the mathematical justification. I will raise my score accordingly.