Memory Mosaics

Please read our general responses, especially about the (non)-relation to Linear Transformers.

1. 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

As explained in the general response, our work is entirely orthogonal to work on Linear Transformer. In particular, equation (7) makes clear that the computational cost is quadratic, like that of an ordinary transformer. The only connection to Linear Transformer are the engineering opportunities mentioned as future work in the last paragraph of the conclusion.

2. 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.

We can affirm that this paper does not constitute an incremental advance on linear transformers, because it is not about linear transformers at all.

3. 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.

Please check our general response 3.

4. 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...

Apart from the Memory Mosaic results, all results in Figure 9 were obtained by Akyurek et al. in the RegBench paper. Their code uses the publicly available implementations in a very well organized way. Since our work is not related to Linear Transformer and variants, we do not see a reason to relitigate the Regbench results. Note however that there is a regime in which the Mamba and H3 performances exceeds that of transformers, but, even in this regime, both are much worse than Memory Mosaics (and that is our point).

5. 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.

Glad you asked. At the relatively small GPT-2 scale, we find that these benchmarks tell very little about the actual model performance at scale. See for instance the recent Apple paper (https://arxiv.org/pdf/2410.05229) revealing the brittleness of such "reasoning" benchmarks.

We like the TinyStories approach (https://arxiv.org/abs/2305.07759) because it provides the means to probe large language modeling concerns in small models. We find that a deep dive into the TinyStories/BabyStories performance is a much better predictor of the general language modeling performance of a similar model on a much larger scale. Our ongoing scaling work certainly agrees with this claim made in the TinyStories paper.

The reviewer appreciate authors' thorough and detailed responses to the concerns raised in the review. After carefully going through all your comments and clarifications, I find your rebuttals regarding the distinction from Linear Transformers and the other technical aspects to be well-reasoned and convincing.

As my primary expertise lies in model pre-training rather than the core focus of this paper, I acknowledge that I may have misinterpreted some of the technical aspects initially. However, I still have one remaining concern regarding the evaluation benchmarks. While the provided link discusses the limitations of mathematical reasoning benchmarks, my original comment referred to commonsense reasoning benchmarks (ARC, PIQA, HellaSWAG). These benchmarks have been shown to have a positive correlation with language modeling loss scaling laws [1], suggesting their validity even for smaller-scale models. I believe evaluating on these benchmarks would still provide valuable insights into your model's capabilities.

Given your thorough clarifications on the technical aspects and architectural decisions, I am revising my score upward.

[1]: Du, Zhengxiao, et al. "Understanding emergent abilities of language models from the loss perspective." arXiv preprint arXiv:2403.15796 (2024).

We thank you for accepting to re-examine your initial assessment and we commend you for your intellectual integrity doing so.

Regarding the benchmarks, please note that we do not claim that memory mosaics should replace transformers in LLMs at this time, we merely claim that they give enlightening insights in the how this kind models work, and that they offer opportunities that might prove useful. We believe that our limited benchmarks are adequate for this claim:

• First let's look at the minimum kind of evidence needed to claim that memory mosaics should actually replace decoding transformers for big natural language applications. Just to be competitive for the kind of standard benchmarks you are suggesting, we would need scale our models to a couple billion parameters trained on a generic language corpus. We would have to get serious about hyper parameter search, establish corresponding scaling laws, etc. We would also need to demonstrate substantial new capabilities, such as for instance show that our models have uncharted in-context learning ability (in the spirit of Fig 9). But that might not be suitable for all applications since superior combinatorial learning abilities could also make the models more willing to hallucinate, and we would have to sort that out. This an entirely new research effort, certainly worthwhile, but only possible as a second step. Setting such a bar for any research on transformer ideas would prevent a lot of research from happening.

• Since our paper focuses on the insights gained by our transformer-like setup, we only need to show when our model behaves like transformers (e.g. fig 7, tables 4 and 5) and when it differs in interesting ways (e.g. Fig 8, Fig 9, Fig 15). This is more a qualitative comparison than a quantitative one. These comparisons are not precise enough to be changed into scaling laws, but they are still informative enough for our claim.

As an aside, we carefully read the reference you suggest and in particular the way experiment distinguish "easy" benchmarks like HellaSwag from "hard" benchmark like GSM8k. We also notice that the "easy" benchmarks often involve selecting one of n preselected continuations which is far easier (and also far easier to measure) than letting the model produce free generations. In contrast, note how Tables 4 and 5 show free generations using quite small models (<160M). The quality of these free generations cannot be matched by much larger models (1.5GB params as discussed in the TinyStories paper) trained on a general language corpus, even though these larger models can perform well on HellaSwag-style benchmarks. All this to say that the tiny stories setup gives very interesting qualitative information, albeit not the kind you could use to obtain scaling laws for instance. Yet, for the objectives of our paper, this works.

Many thanks for your very detailed review.

1. From statistical physics it is well known that the normalizing factor (Z in Equation 2) is nearly always intractable for large enough systems.

The normalization factor in equation 2 is a sum of n terms where n is the number of key/values pair stored in the memory, which is equal to the number of past tokens visited in the input sequence, and always inferior to the maximal context length. This number of terms is perfectly tractable. The normalization of the key vectors does not change it.

2. The authors discuss "peeking" in the Memory Mosaic, [...] ahead would be considered equivalent to a bug.

This is indeed a key component of our architecture with nontrivial consequences. Please check our general response 2 for a rephrasing of this mechanism and its consequences. When presenting our work in front of an audience, we found it useful to spend more time on the explanation in lines 99-102 of the paper. Maybe these lines should be highlighted more prominently.

3. If my understanding of the paper is correct, training a neural network [...] 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,

This is exactly what we mean (see lines 155-159) and this interpretation is critical for our paper.

4. Does the Memory Mosaic always benefit from fast-learning memory units? [...] is it possible that slow-learning memory units result in more accurate predictions (for any number of reasons).

This is a good question. What we observe is that the (meta)-training process wants fast-learning memory units, and therefore constructs disentangled memory units. It might be that this is not the optimal way to model sequences that strictly follow the training set distribution. However, this property comes handy when the input sequence is out-of-distribution because it allows the network to quickly adapt to this new distribution. See Figures 8 and 9 for examples. This is key to in-context learning which itself plays an important role in the performance of transformer-like architectures.

5. ... 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

We actually discuss this point lines 361-367 in the beginning of the language Section. In addition, mining this sort of information from text might in fact be easier because, unlike moons, language was constructed to be easily learnable (an old argument in structural linguistics).

6. Is the correct interpretation here that the rows of these matrices could be permuted freely, and hence are equivalent to the identity? [In the three moons network]

This is correct. The network architecture is symmetric with respect to an arbitrary permutation of the three heads.

7. The main language-based task presented in this paper is based on BabiStories. [...] Due to the similarities to TinyStories, it seems that BabiStories must include some additional structure that makes it useful when working with Memory Mosaics.

We would have preferred to use the original TinyStories dataset. The unfortunate legal reason for constructing a new one is evoked line 376. To our knowledge, there is nothing in the BabiStories dataset that favors Memory Mosaics. Instead we went to some length to ensure that the BabiStories dataset behaves like the TInyStories dataset.

8. The performance gap is only substantial for a model depth of 1. [...] 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?

The gap at model depth 1 is explained in footnote 6. What these plots show is that Memory Mosaics (without hyperparameter search) essentially match Transformers for in-distribution performance. However, surprisingly strong differences appear with out-of-distribution inputs as in Figure 8 and 9.

9. The hyperparameter transfer from GPT-2 to Memory Mosaics is also terrifically interesting

We do not believe that the GPT2 optimal hyperparameters are optimal for Memory Mosaics. We simply find it too easy to manipulate the performance by working on the hyperparameter search. Therefore we chose a protocol that completely avoids this risk.

Thank you for your response! Your general responses above are well written and add to my understanding of your work. I acknowledge due to the length constraints placed by ICLR this information could not be placed into the paper, but it would be nice to have some of these explanations presented alongside the work you have done.

The concept of peeking is starting to come in to focus for me, especially with your general response #2. As you say, this section may need to be rephrased to come across clearly to a new reader.

I am still not sure if I agree with the terminology used in regards to "meta-learning" and "learning" in this paper. The convention in machine learning is of course to refer to the training of a model with data before inference time as "learning", although this work is somewhat different as this phase is "only" updating parameters and not memories. I appreciate there are two distinct phases of learning here, but I wonder if there is better terminology that could be applied to make it easier for a reader to understand what is happening in each phase.

It would be unreasonable to request that you investigate the usefulness of "slow-learning" memories in the Memory Mosaic considering the quantity of research already present, but certainly this is something to look at in the future!

I am not convinced that the same predictive disentanglement property discussed, and shown in the "three moons" toy problem, is present in language data. Lines 361 through 367 give some discussion on why the structure of some language tasks may be compatible with predictive disentanglement, but this does not constitute hard evidence to me. I could be lacking in my understanding of predictive disentanglement, but ideally it would be nice to see some kind of visualization or very concrete data on why language tasks have the same properties as the "three moons" problem.

I see Line 376 is the motivation behind BabiStories. How unfortunate that legal issues handicap the use of the dataset! I stand by my point that introducing a completely new dataset is perhaps excessive alongside the remainder of the author's work --- it is difficult to gauge how useful this dataset is without benchmarking on many different models, which of course could be a paper in itself.

I appreciate that the memory mosaic only requires one head to implement induction heads, whereas transformers require two. However, Figure 7 seems to imply (or at least suggest) that the Memory Mosaic outperforms transformer architectures for a depth up to 8. Is a depth of 8 when Memory Mosaics and transformers start to perform equally well, or does this occur at much smaller depths?

As for the hyperparameter choices, I appreciate you are showing that the Memory Mosaic performs well even without tuning (you are giving transformers the upperhand and so on), but it is eyebrow-raising that using the same hyperparameters gives nearly exactly the same performance in Figure 7. This would suggest to me either the models have some underlying equivalence (you are observing the same behavior in two different expressions) or the models are both reaching some upper limit of the dataset. Moreover, I would still be interested to see hyperparameter tunings for the Memory Mosaic, to observe how much further they can be pushed.

Thank you again for your reply.

Please check our revised paper (which is now accessible through the usual pdf link in open-review).

• We have added a section in the appendix that augments the material in Figure 7 and discusses it in detail. Note that we do not claim that memory mosaic outperform transformer in these in-distribution evaluations. We only claim that it closely matches them despite the absence of position encoding, for instance. Figures 8 and 9 then provide out-of-distribution situations in which memory mosaics perform better.

• In order to remind the reader that the compositional nature of language is essentially a mathematical property, we have added a reference to the 1968 book of Zellig Harris "Mathematical Structures of Language" which non only formalizes how sentences can be built through well-defined successive transformations, but also discusses how these transformations are discoverable from language data, consistent with the fundamental tenets of structural linguistics (which we believe very relevant to LLMs.)

There is no question that the scale of this happening in language is very different from the scale of the three moons problem. However the nature of the phenomenon is the same: when operating in the distant tail of the training distribution, far from any training example, the model can only rely on mathematical structures that were discovered in the body of the training distribution(s) and remain valid in the distant tail. This is the first time we have model able to discover such structures without cheating (that is without revealing bias introduced by the model designer.)

Please check our general response with respect to the work on the modern Hopfield networks. Although we find this work direction interesting, we also believe that it is only weakly related and does not provide the right context to understand our paper.

1. ​​Please highlight a clear definition of predictive disentanglement and explain intuition behind it.

Please see our general response 2.

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.

Please see our general response 1 and observe that our use of inference time memories is far more connected to the memory networks of Weston et al (which predate transformers) than to the modern Hopfield network (see our second sentence line 20). We'll make this clear.

3. Are the proposed networks energy-based or not?

They are not.

Remark: Substantial work on associative memories predates Hopfield's seminal work on memories defined through an energy functions, e.g. Anderson's associative memories (J. R. Anderson, "language memory and thought", 1976), Kohonen associative memories (T Kohonen, "Associative memory: a system-theoretical approach", 1977).

4. Associative memories with Gaussian kernel smoothing (equation 2) have been studied in End-to-end Differentiable Clustering with Associative Memories, ICML 2023. Please explain in the revised paper how your approach is different/similar to that prior work.

What we do is a direct application of the Nadaraya-Watson estimator of 1964, also known as kernel regression, applied to the list of stored key/value pairs (not using a distributed representation here.)

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.

Please see our general response 3. Persistent memories only play a supporting role in this paper. We call them “persistent” because they store key/value pairs that are determined during training and frozen during inference. In contrast the other memories are cleared at the beginning of each test sequence and are filled during inference.

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?

As rephrased in our general response 2, the output of the memory unit $y_T$ at time $T$ only depends on the past of the input sequence. It does not depend on $v_T$ which is only computed at time $T+1$ in order to store a new pair $(k_T,v_T)$ into the memory. There is no leak of information here. You can also see this in equation (7) where the summation stops at $T-1$ in order to exclude $v_T$.

7. Why were the hyperparameters on language modeling tuned for Transformers, but not for the proposed model?

We did this in order to be absolutely certain that our model matches or outperforms the transformer. It is far too easy to obtain good numbers with a subtly fancier hyper parameter search. Therefore we follow a protocol that totally avoids this risk.

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?

Please check our general responses as well as sections 2 and 6 in the paper. In particular there are no activation functions in our networks. The only non-linearities are the RMS normalizations and those of equation (7). The layer configurations are explicit and footnote 5 precisely compares the number of weights in our model with those of the GPT-2 model.

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?

Please check our general response 2. The memories operate at inference time. They’re cleared at the beginning of each sequence and acquire new key/value pairs at each time step. The gradient training process only determines the nature of the stored keys and values.

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?

In section 6, we chose to stay as close as possible to the GPT2 architecture, and accordingly chose to reuse the hyper-parameters that were optimal for the GPT2 baseline, without searching for better ones. Although this places Memory Mosaics at a slight disadvantage, it ensures that our positive results are not a consequence of a fancier hyper-parameter search.

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.

There is not much to ablate without making the system incoherent. We have experimented with various replacements for the key and value extraction functions (equation 6) and found that they made little difference as long as they each involved at least two successive inputs.

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?

Predictive disentanglement is not a feature that we can add or remove at will, but a phenomenon that can be understood when one realizes that the training process is in fact a meta-training process. Please refer to the general response 2.

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.

Since writing the paper, we have scaled the architecture to models with a couple billion parameters without particular problems. We also made further strides in leveraging the good properties suggested in figures 8 and 9. However this new research is not ready for publication at the current time.

Dear reviewer L6NA

Thank you for taking the time to review our paper. We have carefully considered your comments and have provided detailed responses to each of your points. We hope that our answers adequately address your concerns and provide a clearer understanding of our work.

If you have any further questions or concerns, please do not hesitate to leave a comment.

Persistent memories were introduced by Sukhbaatar in (2019) as a possible replacement for the fully connected layers features in classical transformer architectures. In our work, persistent memories compute their outputs exactly like the contextual memories (equation 7). The only difference is that the stored key/value pairs do not change at inference time. They have been determined by gradient during training, just like the weights of a fully connected layer. This is explained lines 324-328.

Although replacing the persistent memory layers by fully connected layers of equivalent size does not change the performance of the system, we find interesting that we can train a GPT2-equivalent system that only relies on memories (either contextual memories that store contents pertaining to the current test sequence, or persistent memories whose content have been determined at training time and persists across test sequences). The only non-linearities in the resulting system are those present in equation (7) and in the RMS Normalization steps.

Persistent memories only play a supporting role in our construction. We describe them because we found that they give a conceptually more interesting system without performance penalty.

The following is an attempt to concisely rephrase the main idea (because the details matter).

In a memory mosaic, each associative memories operates independently at inference time, starting empty at the beginning of each token sequence. At each time step $T$, each memory receives a key vector $k_T$ computed from the current and recent tokens. Each memory then interpolates a response $y_T$ on the basis of the previously stored key/value pairs. One time step later, it peeks into the future and computes the value $v_T$ (which contains a bit of information about the next token $x_{T+1}$) and stores the pair $(k_T,v_T)$ into the memory. This is explained lines 98-102, and this is formalized in equation (7), together with a bullet list of points that distinguish this process from the classical self-attention (lines 135-150).

Although the value $v_T$ depends on the future token $x_{T+1}$, the output $y_T$ does not depend on $v_T$ but merely leverages the previously stored key/value pairs to estimate $v_T$. Therefore there is no leak of future information. Each memory is simply used as a machine that predicts a bit of future information (described by $v_T$) on the basis of recent information (described by $k_T$) and previously stored key/values pairs.

Importantly, the gradient learning algorithm determines which future bit of information is predicted by each memory (through the parameters that control the computation of the values $v_T$) and which kernels are used to perform the predictions (through the parameters of that control the computation of the keys $k_T$). However the relation between keys and predicted values is entirely determined at inference time through the memorization of key/values pairs specific to each sequence. This is explained lines 135-139. This is why we write, lines 155-157, that the gradient training procedure is in fact a meta-learning process whose main purpose is to determine what the memories model at inference time and how they do it.

Through a process that we call predictive disentanglement, this meta-learning interpretation reveals an interesting consequence: the gradient training algorithm tends to split the overall prediction task (e.g. predicting the next token in a sentence) into disentangled prediction subtasks assigned to each associative memory unit. The steamroller figure (Figure 2) summarizes the intuition for this (see also lines 176-212) : the training cost is reduced when the memories are able to provide good predictions $y_T$ of their targets $v_T$ after seeing as few tokens as possible. This is best achieved when the prediction subtasks assigned to each of the memories represent disentangled components of the overall prediction problem, (an idea closely related to https://arxiv.org/abs/1901.10912).

Although similar phenomenon might happen in regular transformers, the memory mosaic perspective makes it considerably easier to understand (yet not fully obvious, unfortunately). Not only this observation is essential to understand why transformer-like architectures are able to learn models with compositional properties, but we also show that, contrary to common beliefs, this is not a consequence of scale since we are able to demonstrate it in a network with only 54 parameters.

If the reviewers find this rephrasing helpful, we'll modify the paper accordingly.