Memory Mosaics

Authors: Jianyu Zhang, Niklas Nolte, Ranajoy Sadhukhan, Beidi Chen, Leon Bottou
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: predictive disentanglement, Associative memory, language model

Abstract:
Memory Mosaics are networks of associative memories working in concert to achieve a prediction task of interest. Like transformers, memory mosaics possess compositional capabilities and in-context learning capabilities. Unlike transformers, memory mosaics achieve these capabilities in comparatively transparent way (“predictive disentanglement”). We illustrate these capabilities on a toy example and also show that memory mosaics perform as well or better than transformers on medium-scale language modeling tasks.

Decision: accept (poster)
Reviewer scores: 3, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1729208205392

thread_type:
official_review

<thread_object>
creator: reviewer_9krw

<content>
summary:
This paper proposes Memory Mosaics, a neural network architecture composed of associative memory units for language modeling tasks. The authors present the concept of "predictive disentanglement" to explain how the model decomposes prediction tasks. Experiments show comparable or better performance to transformers on language modeling benchmarks, with claimed advantages in interpretability and in-context learning.

strengths:
-  The paper addresses an interesting problem in developing more interpretable language models.
- The work is well-organized and clearly presented.

weaknesses:
- The core concept presented in Figure 3 bears similarities to Linear Transformers [1], particularly RetNet [2] (see RetNet Eq. 6). However, the paper fails to adequately discuss or cite these highly relevant works. Linear Transformers are only briefly mentioned in Figure 9, without proper citation or in-depth comparison. This oversight is particularly concerning as it appears the authors are aware of the similarities, given the inclusion of Linear Transformer, RetNet, and GLA [3] in Table 3 and Figure 9. The lack of a detailed explanation comparing their approach to the evolution of Linear Transformer methods (such as GLA/RetNet) raises questions about the paper's originality and positioning within the field.

- The authors introduce persistent memory units (P-mem) but do not clearly differentiate them from standard MLP layers. They state that P-mem can be seen as MLP layers (lines 328-330), but fail to elucidate the advantages or unique properties of P-mem that justify its introduction as a new concept. This raises concerns about the necessity and value of creating this new terminology.

- While the paper uses RegBench to assess in-context learning, this benchmark alone may not provide a comprehensive evaluation of the model's memory capabilities. The performance of state-space models like Mamba can be significantly impacted by the inclusion of short convolutions, as noted in DeltaNet [4] (See Figure 2). A more thorough evaluation, possibly including benchmarks like MQAR [5], would provide a more robust assessment of the model's capabilities relative to existing approaches. I understand that there may be not enough time for running such experiments in the rebuttal stage, so the author can just try to rerun run the experiment with/without the short convolution in RegBench.

- The paper focuses primarily on language modeling tasks but lacks evaluation on standard common reasoning benchmarks that are crucial for assessing language models' capabilities. Tasks such as ARC Challenge, ARC Easy, PIQA, and Hellaswag are widely used to evaluate smaller-scale language models and provide valuable insights into a model's reasoning abilities. The absence of these evaluations limits our understanding of the model's true capabilities and makes it difficult to compare with other approaches in a comprehensive manner. Including these benchmarks would provide a more holistic view of the model's performance across various aspects of language understanding and reasoning.

- The paper appears to present an incremental advance on linear transformers, repackaged with new terminology and concepts. While this approach creates an impression of novelty, its actual contribution to advancing the field seems limited. The work seems to complicate a relatively straightforward evolution of linear transformer models with complex conceptual packaging, which may hinder rather than help the development of the research community and the broader field.

If the authors can appropriately address my concerns or correct me if I am wrong, I would be willing to increase my score.

[1] Katharopoulos, Angelos, et al. "Transformers are rnns: Fast autoregressive transformers with linear attention." International conference on machine learning. ICML, 2020.

[2]Sun, Yutao, et al. "Retentive network: A successor to transformer for large language models." arXiv preprint arXiv:2307.08621 (2023).

[3]Yang, Songlin, et al. "Gated Linear Attention Transformers with Hardware-Efficient Training." ICML, 2024.

[4]Yang, Songlin, et al. "Parallelizing Linear Transformers with the Delta Rule over Sequence Length." NeurIPS, 2024.

[5]Arora, Simran, et al. "Simple linear attention language models balance the recall-throughput tradeoff." ICML, 2024.

questions:
See my weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729586529913

thread_type:
official_review

<thread_object>
creator: reviewer_k2dw

<content>
summary:
This paper introduces a new neural network architecture, the Memory Mosaics, which employ associative memories to complete prediction based tasks. The Memory Mosaic architecture consists of a number of "elementary memory units", which in turn consist of a feature extractor (outputting a time series of "key" vectors and a time series of "value" vectors), which feeds an associative memory. Prediction in the Memory Mosaic has the associative memories update based on new information from the task, rather than having weights frozen from some previous training, such as in traditional deep feed-forward neural networks. The authors link the Memory Mosaic architecture to the transformer architecture (in one specific example in Section 6, GPT-2) and the attention mechanism. This link is derived from Gaussian kernel smoothing and enforcing key vectors to have equal magnitude. To differentiate this work from other works linking associative memories and attention (e.g. Ramsauer et al. 2020) this paper focuses on the ability of the Memory Mosaic to "peek" one step ahead in prediction tasks. The paper asserts that the learning process of Memory Mosaics can be framed as a meta-learning process. The meta-learning process selects for associative memories that are faster at learning new tasks, and hence can produce useful value estimates in fewer time steps. The authors continue to discuss "predictive disentanglement"; a framework in which discrete elementary memory units may learn discrete aspects of the prediction task, allowing for interpretable networks with modules of definitive responsibilities (at least, for tasks that can be easily disentangled). The authors demonstrate predictive disentanglement with a toy predictive task based on modelling the orbit of three moons (represented as a time series of positions). In the presented example, the Memory Mosaic is capable of separating the independent orbits of each moon, allowing the network to make more accuracy predictions earlier, i.e. before all moons have finished an orbit. The next example tests the Memory Mosaic in language based tasks. The authors present an altered memory unit that has architectural similarities to the GPT-2 transformer. The authors create a new dataset based on TinyStories (Eldan & Li 2023) titled BabiStories. BabiStories is generated from the Mixtral-8x7B language model and is otherwise similar to the TinyStories dataset, except with an increased diversity of stories and requiring first names and specific opening words for each story. Memory Mosaics demonstrate improved losses compared to GPT-2 for small networks, but has strikingly similar performances for deeper networks. Out-of-distribution language-based experiments were conducted on the Simple English Wikipedia dataset, in which Memory Mosaics appear to perform much better than GPT-2. Further experiments are performed using RegBench, in which Memory Mosaics appear to outperform several state-of-the-art architectures in both accuracy and total variation distance.

strengths:
The paper presents an interesting, seemingly novel architecture for a prediction model based on associative memories. The architecture seems well justified from first principles and links nicely with existing works on the attention mechanism. This work is of good quality and presents its ideas reasonably well, including rigorous formalizations and numerous figures explaining the proposed architectures. The paper benchmarks the Memory Mosaic against a toy dataset to demonstrate properties of the network, as well as several language-based datasets which allow for comparisons with transform architecture models. These benchmarks add to the quality of the paper, and the newly introduced BabiStories benchmark is original.

weaknesses:
To my reading, the architecture is not differentiated strongly enough from existing, similar work on associative memories and the attention mechanism. The key property of the architecture, "peeking" in the value predictions, is not explained clearly enough for a reader to understand the significance. The meta-learning interpretation in Section 3 is somewhat confusing in how this relates to the broader field of meta-learning. The toy dataset, while very useful in understanding "predictive disentanglement", appears to be crafted specifically to show this property and it is not clarified if such a property is present in other, non-toy datasets. The paper introduces a new language-based dataset, BabiStories, which is similar to the existing TinyStories, but it is not immediately apparent (from the main text or the appendices) why a new dataset is required. Further, without benchmarking of the introduced BabiStories dataset with several models it is not clear if the dataset is useful in assessing the performance of the Memory Mosaic.

questions:
Equation 2 offers a nice, specific implementation of a conditional expectation set forth in Equation 1. The discussion here flows very smoothly and is easy to understand. However, the exact form of Equation 2 --- the Gaussian kernel smoothing --- bears striking similarity to the softmax distribution (and in turn the attention mechanism, as the authors discuss later) as well as the Boltzmann distribution. From statistical physics it is well known that the normalizing factor (Z in Equation 2) is nearly always intractable for large enough systems. In the context of Memory Mosaics, it is not immediately apparent what this would correspond to, but it would certainly be of interest to discuss where this implementation would stumble based on this factor. The authors do sidestep the intractability issue in Equation 3 by asserting that the magnitude of the key vectors (rather, the squared norm) are equal, which nicely reduces Equation 2. Again, it is not immediately clear if that is always going to be the case for Memory Mosaics, if it will always be possible to enforce this constraint in applications, or if this only occurs in the datasets presented in the paper. In any case, a brief discussion on the possibility of non-uniform squared norm key vectors may be of interest to readers and help ground the Memory Mosaic in the landscape of prediction models. The discussion on why this paper limits itself to only Gaussian kernel smoothing derived associative memories is excellent; clear and concise. Emulating this level of discussion for these other points may prove helpful. In short: is the calculation of the normalizing factor Z as intractable as it appears? What factor of the Memory Mosaic / dataset makes the normalization factor impossible to calculate? Does the uniform-squared-norm trick in Equation 3 work for all datasets, or is it a rare occurrence that makes the network otherwise untenable?

The authors discuss "peeking" in the Memory Mosaic, which allows the memory units to look one time step ahead in the "key" series to predict the next item in the "value" series. After several read-throughs, it is still difficult to understand what this means for the Memory Mosaic. In a more traditional prediction model looking a timestep ahead would be considered equivalent to a bug. In the Memory Mosaic this is seemingly not the case, but the reasoning is still hazy to me. Am I correct in my understanding that "peeking" is allowable in this model, and if so, why is this distinct from other associative memories? Is it possible to clear up the discussion on this point, particularly the subsection in Section 2, to explain the nuances of "peeking" more? Considering the importance of the "peeking" mechanism in this work, it is paramount it is easy to digest. 

The Memory Mosaic architecture has striking similarities to associative memories linked to attention mechanisms. In particular, Ramsauer et al.'s influential "Hopfield is All You Need" has almost the same derivations found in Section 2 of this paper. The authors discuss the Ramsauer paper at the end of Section 2, but the differentiation of the works does not appear substantial. In particular, the "peeking" of the Memory Mosaic is not well explained, which makes the reliance on it for the basis of novelty compared to Ramsauer et al. troublesome. With some further discussion of the "peeking" mechanism (as noted above), and potentially some more clarity on the section discussing Ramsauer et al., this paper would provide a much stronger foundation for Memory Mosaics in the context of existing literature.

The meta-learning process at the start of Section 3 is somewhat confusing when primed to think about both meta-learning as a field and traditional training of neural networks. If my understanding of the paper is correct, training a neural network (e.g. writing a training loop, feeding in training data, performing backpropagation, updating weights, and so on) is presented here as meta-learning, with the "regular" learning being the updates to the associative memories performed when prediction of a time series occurs. This is a fairly radical change in perspective, although an interesting interpretation, but again it seems that the paper could benefit from having this stated very directly as to guide the reader to this way of thinking about meta-learning in the Memory Mosaic. Is this interpretation of meta-learning in the Memory Mosaic correct? If so, the steam-roller metaphor (Figure 2) gives rise to another interesting question that is seemingly not addressed in the main body of the paper; does the Memory Mosaic always benefit from fast-learning memory units? Certainly I can see the benefit of having a predictor model that requires a smaller context window to start spitting out predictions, but is it possible that slow-learning memory units result in more accurate predictions (for any number of reasons). If that is the case, it would raise concerns about the gradient / meta-learning presented in the paper, and if that is not the case I have not been convinced by the text. Is it correct to say that faster-learning memory units in the context of Memory Mosaics are always better? 

The toy example presented --- the three moons problem --- is excellent at showing the application of "predicative disentanglement". However, it is less clear after reading through the main text and the appendices if this specific kind of "easy" disentanglement is present in other, non-toy datasets. Mining this sort of information from, say, a language-based dataset would be nigh impossible, but some discussion about whether the property is present in other datasets (or even if this is likely to occur) would make for a nice conclusion to the Section. Moreover, the authors mention several times that the weight matrices in the three moons problem could be replaced with the identity to improve results. In Appendix A these matrices are seemingly the anti-diagonal identity. Is the correct interpretation here that the rows of these matrices could be permuted freely, and hence are equivalent to the identity? Otherwise, these do not seem to be identity matrices, which makes the previous discussion odd!

The main language-based task presented in this paper is based on BabiStories. The new dataset is seemingly introduced without explanation as to why it is needed. Due to the similarities to TinyStories, it seems that BabiStories must include some additional structure that makes it useful when working with Memory Mosaics in particular, but this is not stated clearly in the main text. Appendix B may mention this briefly, but it is not clear if the dataset is created for this reason or not. Is it possible to use TinyStories to train Memory Mosaics and use these results in place of BabiStories? If not, perhaps a description of why not could be included in the main text to explain to the reader why a new dataset is required.

Figure 7 shows the performance gap between the Memory Mosaic and GPT-2 models when trained on the BabiStories dataset. These are promising curves and show excellent results. However, this performance gap is only substantial for a model depth of 1. Even at a model depth of 8, the second sub-figure, the gap between GPT-2 and Memory Mosaics has seemingly disappeared. It is still interesting that GPT-2 and Memory Mosaics perform equally well for larger depths, does this hint that perhaps the equivalent performance is due to the dataset rather than the architecture i.e. both architectures can saturate to the best possible accuracy on BabiStories? Also, is it possible to focus on specifically the region where there is a difference in performances? The jump from a depth of 1 to a depth of 8 in this Figure seems to leave out the most interesting region in the context of the paper. Does the Memory Mosaic continue to outperform GPT-2 at a depth of 2 or 3, or do the two architectures converge immediately after a depth of 1? The hyperparameter transfer from GPT-2 to Memory Mosaics is also terrifically interesting --- considering how fickle these networks can be, and the equal performance show in Figure 7, does this hint at some deeper connection between Memory Mosaics and the GPT-2 architecture, perhaps even a full equivalence of the models? This last question is less a review of the current manuscript and more of a potential lead into future research, so no need to address this point in your comments!
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669119702

thread_type:
official_review

<thread_object>
creator: reviewer_r52u

<content>
summary:
This submission proposes replacing parts of the Transformer block with associative memories. The notion of predictive disentanglement is discussed. The proposed architecture is evaluated on 3-moon problem and small scale language modeling tasks.

strengths:
In general, this is a though-provoking paper with some interesting ideas. Exploration of the relationship between associative memories, attention, and Transformer blocks is valuable, although the current presentation is heavily biased, and omits related ideas from prior work.  

The empirical results on language modeling are encouraging, although small scale.

weaknesses:
One problem with this submission is that the presentation almost entirely ignores the work on modern Hopfield networks and dense associative memories, which tackles closely related motivation and ideas. Specifically, the authors’ proposal is closely related to [Energy Transformer (NeurIPS 2024)](https://proceedings.neurips.cc/paper_files/paper/2023/file/57a9b97477b67936298489e3c1417b0a-Paper-Conference.pdf) and related literature, which replaces elements of Transformer block with associative memories. The only cited paper representing that line of work (Ramsauer et al, 2020) is only briefly mentioned. Please notice, that paper only established a connection with associative memory at a single step update, which is different from the setting used in this submission - dynamics unfolded in time with shared weights, if I understood equations 4 and 5 correctly. In contrast, the Energy Transformer formalizes attention operation and the entire Transformer block as a form of associative memory, with recurrent dynamics unfolded in time and shared weights. There are, of course, many differences between these prior works and the current proposal. For instance, it is unclear to me if energy-based description is important at all for this submission. Regardless these similarities/differences this and related works need to be referenced at a prominent place in the introduction and similarities/differences with the current proposal need to be extensively discussed in the revised manuscript. 

I have read this submission several times, but to be frank with the authors I am not sure I understand what they mean by predictive disentanglement, which seems to be the core concept here, but defined only vaguely. Unfortunately, I do not have a specific recommendation how one could improve the presentation in section 3. But, at present, it is insufficient for me. What architectural aspects of memory mosaics suggest that predictive disentanglement should hold?  

P-mem network is insufficiently explained. Please include specific operations that are used in it in the revised paper (with formulas).

questions:
1. Please highlight a clear definition of predictive disentanglement and explain intuition behind it.
2. Please make the discussion of associative memory based modifications of Transformers more scientifically balanced. At present, the submission leaves a strange impression that the work from FAIR is unreasonably promoted (e.g., Bietti et al, Weston et al., etc), but the work from elsewhere is ignored. 
3. Are the proposed networks energy-based or not? 
4. Associative memories with Gaussian kernel smoothing (equation 2) have been studied in [End-to-end Differentiable Clustering with Associative Memories, ICML 2023](https://proceedings.mlr.press/v202/saha23a/saha23a.pdf). Please explain in the revised paper how your approach is different/similar to that prior work. 
5. What aspect of P-mem makes it persistent? Please include a detailed description of this network with explicit formulas describing how it is implemented. The figure presented in (Fig 6 right) is insufficient for understanding how it is defined. 
6. Formula 5 confuses me. Wouldn’t such a definition of values trivially leak the information about the future to the output of the model? In other words, what prevents the model from simply copying that future token and output that copy as a next token prediction? 
7. Why were the hyperparameters on language modeling tuned for Transformers, but not for the proposed model? 


I am keeping an open mind, and willing to consider increasing the scores depending on authors’ responses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730813217202

thread_type:
official_review

<thread_object>
creator: reviewer_l6na

<content>
summary:
"Memory Mosaics" introduces an architecture where multiple associative memories are employed in unison to enhance predictive capabilities for tasks like language modeling. This architecture aims to combine the benefits of transformer models with greater transparency and efficiency through what the authors term "predictive disentanglement." The paper is technically rich, articulating a clear hypothesis, describing the underlying mechanisms in detail, and providing comparative analysis against existing models.

strengths:
1. The integration of associative memories to replicate and surpass the capabilities of transformers.
2. The concept of predictive disentanglement is novel and also rooted in a solid theoretical framework.
3. The theoretical motivations behind predictive disentanglement are well-explained.

weaknesses:
1. The paper does not provide exhaustive details on the architecture's configuration.
2. Lack of detailed discussion on the choice and impact of hyperparameters.
3.The experimental validation is limited to certain language tasks. 
4. Lack of Ablation Studies

questions:
1. Could you dive a bit deeper into the architecture of the Memory Mosaics for us? I'm particularly curious about the specifics of the associative memory units, what kind of activation functions you're using, and how the layers are configured. And also the comparison with other non-transformer models like rwkv and mamba? 
2. I'd love to understand more about how you set up the associative memories initially. How do you initialize and update the parameters during training?
3. What led you to choose the hyperparameters that you did for this study? And how does changing these hyperparameters affect the model's performance?
4. Did you perform any hyperparameter tuning experiments? If yes, could you share how you went about it and what the outcomes were?
5. Are there any ablation studies that you've conducted or plan to conduct? It would be really helpful to see how each component of the model contributes to its overall performance.
6. I'm interested in the predictive disentanglement feature of your model. Could you explain how altering or removing this feature affects the model's performance?
7. Have you conducted any scalability tests, particularly in comparison with transformers? I'm curious about how the model scales with larger datasets and how computationally efficient it is.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731688473614

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response 1 : About existing work on transformer alternatives.

comment:
We thank the reviewer for pointing out that it is useful to explain how our work differs from current work on transformer alternatives. 
We'll include a section that will make the following clear.

1)  The computational cost of self-attention grows quadratically with the sequence length. A number of recent papers [CITATIONS] propose transformer alternatives with linear complexity instead.  Our work is orthogonal to that line of research. As should be clear in equation (7), our models have quadratic complexity. We mention, in the last paragraph of the conclusion, that memory mosaics provide opportunities for future work in this direction (lines 520-524), but this neither is the main thrust of this paper nor is its main motivation. 

2) Another series of recent papers [citations] share our interest in associative memories but focus on the Hopfield style of associative memories which are governed by an energy function. Our associative memories simply store key/value pairs in the computer memory and use kernel regression to retrieve an interpolated value when provided with an approximate key. This is explained lines 40-47. The argument that justifies such memories is not an energy argument, but a well known statistical argument about kernel regression (see footnote 1. See also the Nadaraya-Watson estimator which is worth citing there.)  Other than relying on associative memories in a broad sense, we are not aware of a precise connection between our approach and the transformer alternatives based on modern Hopfield ideas.

Our contribution is much more closely related to the 2019 paper of Bengio et al (https://arxiv.org/abs/1901.10912) which discusses how a meta-learning objective that favors fast adaptation can disentangle causal factors. The is happening for free in memory mosaics (and this is what predictive disentanglement is about.)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732129079958

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision of the paper

comment:
We just posted a substantial revision of the paper with the intention to account for the comment of the reviewers.

Changes include:
* A new "related work" section listing the lines of work suggested by the reviewers (and the corresponding papers of course) and explaining how our work differs. We also comment on the paper of Bengio et al (2019) which presents ideas closely connected to predictive disentanglement and might help the readers.
* A rework of the paragraphs explaining the detailed operation of the associative memory units and presenting predictive disentanglement, along the lines discussed in the rebuttals.
* A new paragraph at the end of the three moon sections explaining how the moon disentanglement and recombinations is similar to phenomena that one can experience with modern LLMs (albeit on a different scale). We also cite (Harris 1968) to stress out that the compositional structure of language can be described as a strictly mathematical property.
* An additional section in the appendix augments the results presented in figure 7 and provides a substantially larger discussion, too large to fit in the main text.

All changes are shown in blue in the revised pdf (which is accessible using the usual pdf link on open review.)
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a neural network architecture that composes of associative memory units. Compared to transformers, it achieves a new capability, called predictive disentanglement, under a meta-learning interpretation of the training process. The paper shows better interpretability and in-context learning performance on language modelling benchmarks.

Pros:
- Novel use of associative memories with a strong theoretical motivation.
- A thought-provoking paper with interesting ideas.
- Encouraging empirical results, showcasing competitive performance.

Cons:
- Insufficient discussion of related works. (addressed in the rebuttal)
- Limited evaluation on relatively small networks.