Memory Mosaics

## Abstract

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 ways. We demonstrate these capabilities on toy examples and we also show that memory mosaics perform as well or better than transformers on medium-scale language modeling tasks.

## 1 Introduction

This paper presents a learning system architecture, *Memory Mosaics*, in which multiple associative memories work in concert to carry out a prediction task of interest. Such systems are closely related to memory networks (Weston et al., ; Sukhbaatar et al., ) and resemble transformers (Vaswani et al., ) despite significant differences. Like transformers, Memory Mosaics possesses some of the disentanglement and compositional capabilities that have long eluded machine learning systems (Lake and Baroni, ). Unlike transformers whose internal mechanism are hard to decipher (Olsson et al., ; Bietti et al., ), Memory Mosaics achieve these capabilities in comparatively transparent ways.

The three main contributions of this work are (a) recognizing and exploiting a similarity between smoothing associative memories and self-attention, (b) identifying and illustrating the predictive disentanglement principle which explains how training decomposes the overall task in interesting ways, and (c) showing that this comparatively transparent architecture matches the performance of decoding transformers on a language modeling task.

Section 2 describes the basic architecture and outlines its consequences. Section 3 illustrates the predictive disentanglement principle. Section 4 extends these ideas to fully formed memory mosaics. Section 5 reports on medium-scale language modeling experiments.

## 2 Memories

#### Associative memory

Generally speaking, an associative memory is a device that can store key-value pairs and retrieve values given a corresponding key. This definition omits important details about dealing with duplicate keys and approximate matches. For our purposes, both keys and values shall be vectors in $\mathbb{R}^{d}$. The retrieval process can then be represented as a function of the queried key $k$ and all the stored pairs $(k_{1},v_{1})\dots(k_{n},v_{n})$.

$$ $\left\{\begin{array}[]{lcl}\mathbb{R}^{d}&\rightarrow&\mathbb{R}^{d}\\ k&\mapsto&f\big{(}k;\>\{(k_{1},v_{1})\dots(k_{n},v_{n})\}\big{)}\end{array}\right.$ $$

Except perhaps when duplicate keys are involved, an associative memory stores key-value pairs without consideration for their temporal ordering. Therefore the retrieval function can be assumed invariant with respect to any permutation of the stored pairs. This exchangeability property suggests that we can also view an associative memory as a device that estimates a conditional probability distribution $P(V|K)$ on the basis of the sample $(k_{1},v_{1})\dots(k_{n},v_{n})$ of key-value pairs. The retrieval function is then a conditional expectation over this estimated distribution:

$$ $f\big{(}k;\>\{(k_{1},v_{1})\dots(k_{n},v_{n})\}\big{)}\leavevmode\nobreak\ = \leavevmode\nobreak\ \mathbb{E}\/(V\>|\>K=k)\,.$ (1) $$

Such a conditional expectation can be constructed with Gaussian kernel smoothing,(^1^11This expression results from estimating the key-value distribution with Kernel Density Estimation using a Gaussian kernel. The Parzen windows convergence theorem then ensures that (2) converges to the true conditional expectation when $n\rightarrow\infty$ and $\beta=\sqrt{n}$.)

$$ $f\big{(}k;\>\{(k_{1},v_{1})\dots(k_{n},v_{n})\}\big{)}\leavevmode\nobreak\ = \leavevmode\nobreak\ \sum_{i=1}^{n}\frac{1}{Z}\leavevmode\nobreak\ e^{-\beta\| k-k_{i}\|^{2}}v_{i}\leavevmode\nobreak\ \quad\text{with}\quad Z=\sum_{i=1}^{n} e^{-\beta\|k-k_{i}\|^{2}}\,.$ (2) $$

The close connection between kernel smoothing and attention (Bahdanau et al., ) is particularly obvious when all key vectors $k_{i}$ share a same squared norm because expression (2) can be simplified as

$$ $f\big{(}k;\>\{(k_{1},v_{1})\dots(k_{n},v_{n})\}\big{)}\leavevmode\nobreak\ = \leavevmode\nobreak\ \sum_{i=1}^{n}\leavevmode\nobreak\ \frac{e^{\,\beta\,k^{ \!\top}k_{i}}}{\sum_{j=1}^{n}e^{\,\beta\,k^{\!\top}k_{j}}}\leavevmode\nobreak \ v_{i}\leavevmode\nobreak\ .$ (3) $$

There are of course more advantageous ways to implement such an associative memory, ranging from fast Gauss transforms (Greengard and Strain, ; Yang et al., ) to locality sensitive hashing (Spring and Shrivastava, ; Chen et al., ). Although these methods will certainly prove useful in the future, this paper only relies on associative memories implemented with Gaussian kernel smoothing, not least because that makes it easy to compute gradients.

#### Predicting with associative memories

Consider now a sequence $(x_{t})$ of observations, discrete tokens or continuous values. We would like to leverage the past observations $(x_{t})_{t\leq T}$ to predict some useful property of the future observations $(x_{t})_{t>T}$. For instance we might want to predict the next observation $x_{T+1}$ to construct an auto-regressive model of the sequence.

Figure: Figure 1: Elementary memory unit. The keys $k_{T}$ are computed as a function of past observations $(x_{t})_{t\leq T}$. The values $v_{T}$ peek into the future. In this example, the value also depend on the next observation $x_{T+1}$. At time $T$, the associative memory uses the known key $k_{T}$ to retrieve an approximation $y_{T}$ of the yet unknown value $v_{T}$ using the previously stored pairs $(k_{t},v_{t})$. At time ${T+1}$, the value $v_{T}$ can be computed, and the pair $(k_{T},v_{T})$ can be stored into the memory.
Refer to caption: x1.png

Our elementary memory unit (Figure 1) consists of an associative memory and a trainable feature extractor that computes suitable keys and values for the memory. The keys $k_{T}$ are computed as a function of the past observations $(x_{t})_{t\leq T}$
and trainable weights ${\mathbf{w}}$,

$$ $k_{T}=\varphi(x_{T},x_{T-1},\dots;{\mathbf{w}})\,.$ (4) $$

In contrast, the values $v_{T}$ are allowed to peek in the future because they represent what the memory module aims to predict. For instance, the systems described in this paper merely allow values to depend on the next observation $x_{T+1}$,

$$ $v_{T}=\psi({\color[rgb]{.75,0,.25}x_{T+1}},x_{T},x_{T-1},\dots;{\mathbf{w}})\,.$ (5) $$

At any given time $T$, the associative memory contains the previously observed pairs $(k_{t},v_{t})_{t\leq{T-1}}$ and retrieves an estimate $y_{T}$ of the yet unknown value $v_{T}$ by querying the freshly computed key $k_{T}$. One time step later, the observation $x_{T+1}$ becomes available, the value $v_{T}$ can be computed, and the pair $(k_{T},v_{T})$ is stored into the memory.

The exact form of the feature extraction functions can vary considerably in complexity. For instance, when each observation $x_{T}$ carries sufficient information, the keys $k_{T}$ and values $v_{T}$ can be computed as linear functions of respectively $x_{T}$ and $x_{T+1}$

$$ $\displaystyle k_{T}$ $\displaystyle=W_{\varphi}\,x_{T}$ (6) $\displaystyle v_{T}$ $\displaystyle=W_{\psi}\,x_{T+1}\leavevmode\nobreak\ .$ $$

However it is often attractive to consider more complicated feature extraction functions involving convolutions, leaky averages, layers, transfer functions, normalization, or even polynomial time-series summaries (Voelker et al., ; Gu et al., ). For instance, the language experiments (Section 5) use feature extractors of the following form, with $\mathrm{Norm}(x)=x/\|x\|$:

$$ $\displaystyle k_{T}$ $\displaystyle=\alpha_{\varphi}\mathrm{Norm}\big{(}{\color[rgb]{.75,0,.25}\bar{ k}_{T}}\big{)}$ $\displaystyle\text{with}\quad\overbrace{{\color[rgb]{.75,0,.25}\bar{k}_{T}}={ \color[rgb]{0,0,1}\tilde{k}_{T}}+\lambda_{\varphi}{\color[rgb]{.75,0,.25}\bar{ k}_{T-1}}\quad\quad{\color[rgb]{0,0,1}\tilde{k}_{T}}=W_{\varphi}\,x_{T}}^{ \text{leaky average over t = T, {T-1}\ldots, 1}}$ (7) $\displaystyle v_{T}$ $\displaystyle=\alpha_{\psi}\mathrm{Norm}\big{(}{\color[rgb]{.75,0,.25}\bar{v}_ {T}}\big{)}$ $\displaystyle\text{with}\quad\underbrace{{\color[rgb]{.75,0,.25}\bar{v}_{T}}={ \color[rgb]{0,0,1}\tilde{v}_{T}}+\lambda_{\psi}{\color[rgb]{0,0,1}\tilde{v}_{T +1}}\quad\quad{\color[rgb]{0,0,1}\tilde{v}_{T}}=W_{\psi}\,x_{T}}_{\text{ convolution over t=T and T+1}}$ $$

#### Training networks of memory units

Consider now a deep network that contains both glue layers and elementary memory units, with trainable weights in both the glue layers and in the feature extractors of the memory units. Together, these weights determine what each memory unit memorizes (through the key and value extraction functions) and how their outputs are combined to perform the overall prediction task.

When the associative memories are implemented with differentiable kernel smoothing mechanisms, training such a deep network is simply a matter of unrolling the network in time and back-propagating the gradients, in ways that users of modern deep learning software will find very familiar. Unsurprisingly, unrolling equation (3) along an input sequence $(x_{1}\dots x_{D})$ of duration $D$ yields an expression that very much resembles masked self-attention (Vaswani et al., ).

$$ $\displaystyle\forall\,T\in\{1\dots D\}\qquad y_{T}=\sum_{i=1}^{T-\delta} \leavevmode\nobreak\ \frac{e^{\beta\,k_{T}^{\!\top}k_{i}}}{\sum_{j=1}^{T- \delta}e^{\beta\,k_{T}^{\!\top}k_{j}}}\leavevmode\nobreak\ v_{i}\leavevmode \nobreak\ ,$ (8) $$

There are however notable differences between (8) and classical self-attention:

- •
A more aggressive masking involves the number $\delta\geq 1$ of future time steps needed to compute $v_{T}$.
- •
There is no need for position encoding of for distinction between key vectors and query vectors.
- •
The eventual purpose of each memory unit is observable because, at each time $T$, its outputs $y_{T}$ can be interpreted as a conditional expectation (1) that predicts the yet unknown but explicit quantity $v_{T}=\psi\big{(}(x_{t})_{t\leq{T+\delta}};\,{\mathbf{w}}\big{)}$ on the basis of past observations.

#### Learning and meta-learning

This training procedure should in fact be viewed as a meta-learning process, distinct from the learning that occurs at inference time when new key-value pairs are stored into the memories. Unrolling the network in time over a complete input sequence $(x_{1}\dots x_{D})$ exposes the operation of the memories. Since the overall prediction cost averages terms evaluated at each time step along the training sequences, the training process favors *statistically efficient* memories, able to output useful value estimates $y_{T}$ after storing as few key-value pairs as possible (Figure 2).

Figure: Figure 2: The curve plots the prediction losses for all training sequence indices $t\in\{1\dots D\}$ in the training sequence. Minimizing their sum —the area under the curve— favors memories that produce useful value estimates after fewer time steps.
Refer to caption: x2.png

- •
Let us first assume that each memory unit has a frozen value extraction function $\psi$. The training procedure can still make each memory unit statistically more efficient by tuning the parameters of the key extraction function $\varphi$, that is, by learning how to compare the current prediction context $(x_{T},x_{T-1},x_{T-2}\dots)$ with past prediction contexts $(x_{t},x_{t-1},x_{t-2}\dots)$ for ${t<T}$.
Learning a similarity metric —a kernel— is a well known way to make non-parametric estimators more efficient (e.g., Bach et al., ). For instance, the training procedure can construct keys that summarize the relevant contextual information, discarding noise factors that could increase the distance between keys associated with similar values. It can also adjust the effective kernel bandwidth, for instance, using parameter $\alpha_{\varphi}$ in equation (7).
- •
When multiple memory units are available, the training procedure can also optimize the value extraction functions $\psi$, *redistributing* the overall prediction work among the available memory units. For instance, the training procedures can construct values $v_{T}$ that are more efficiently approximated by their respective memory units, but whose approximations can still be combined to address the overall prediction task.
We argue in this contribution that this redistribution does not only make the memory units more efficient, but also disentangles the original prediction problem in elementary sub-problems that can later be recomposed in myriad ways.

#### Predictive disentanglement

We have argued that training the time-unrolled network of memory units is in fact a meta-training process that learns how to efficiently use memory units to address the task of interest.
Using multiple memory units rather than a single one is beneficial when the overall prediction task can be divided into smaller prediction tasks whose targets are expressed by the values $v_{T}$ and can be *more efficiently predicted in isolation than together*.
The disentangled memories can then be recombined to provide predictions for inputs that are globally very different from the training inputs, but whose disentangled components are individually predictable. Section 3 provides a toy example.

Disentanglement has long been identified as both desirable and hard to pinpoint (Bengio, ). Statistical definitions lack in robustness with respect to changing data distributions (Comon, ; Roth et al., ). Causal definitions cannot be tested without active experiments (Thomas et al., ; Bengio et al., ). Predictive disentanglement provides an appealing alternative that does not require a specific training algorithm but arises as a side effect of the meta-training process. Although predictive disentanglement is easier to understand in the case of a network of associative memory units, we can safely conjecture that similar phenomenon also occurs in transformers, albeit in a far less explicit manner.

## 3 Tracking three moons

Three moons orbit a remote planet. Although the local astronomers are very far from understanding celestial mechanics,(^2^22We do not seek to discuss subtleties such as elliptical orbits or multi-body problems. Our astronomers are best compared to the ancient sky watchers whose efforts eventually gave the Ptolemaic model.) they nevertheless observe periodic motions and debate how to predict future moon positions. A first astronomer proposes to compile a single table containing the daily positions of all three moons, arguing that if the current set of moon positions matches a previous observation, the future moon positions will match the following observations. A second astronomer suggests instead to make three tables, one for each moon, arguing that the future positions of each moon can be independently predicted by matching its current position with a previously observed one.

To make reliable predictions, the first astronomer needs a table that contains at least one record for each of the possible moon configurations. Our astronomer therefore needs to log the daily moon positions until all three moons return to their original configuration, after a number of days equal to the least common multiple $\mathrm{lcm}(p_{1},p_{2},p_{3})$ of the individual moon periods. In contrast, the second astronomer only needs to log daily moon positions until each of the moons returns to a previously observed position, for a number of days equal to the period $\max(p_{1},p_{2},p_{3})$ of the slowest moon.

One could argue that the proposal of the second astronomer is obviously superior because the three moons are distinct objects, well separated in space and time. One could instead argue that we view the moons as separate objects precisely because their respective futures can in general be independently predicted. Space and time separation merely suggests the possibility of independent predictions, as long as the moons do not collide.

Figure: Figure 3: An architecture for the three moons problem. We consider single-layer networks with either $N_{h}=1$ or $N_{h}=3$ memory units whose keys and values belong to either $\mathbb{C}^{3}$ ($N_{h}=1$) or $\mathbb{C}^{1}$ ($N_{h}=3$). Both nets have $3\times 3\times 2\times 3=54$ trainable real parameters that determine how to predict the moon positions using either a single 6-dimensional memory or three 2-dimensional memories.
Refer to caption: x3.png

#### Model

For our purposes, each observation $x_{t}$ consists of three complex numbers $e^{i\theta_{k}}$ that encode the angular positions $\theta_{k}$ of the three moons inside their respective orbital plane. We consider two single layer models (Figure 3) with either $N_{h}=1$ or $N_{h}=3$ memory units whose added dimensions match the input dimension. The key and value extraction functions follow (6) with trainable parameters collected in two ${3\times 3}$ complex matrices $W_{\varphi}$ and $W_{\psi}$. The memory unit follow (3) with a fixed parameter $\beta=50$. A third ${3\times 3}$ complex matrix $W_{z}$ combines the memory unit predictions into an output $z_{T}$ that hopefully predicts $x_{T+1}$. Both networks share an interesting analytic solution: setting all three matrices $W_{\varphi}$, $W_{\psi}$, and $W_{z}$ to the identity yields optimal predictions once the associative memories have seen enough samples.

Figure: Figure 4: Single head network prediction error versus context length. The prediction error shows a sharp transition after $\mathrm{lcm}(p_{1},p_{2},p_{3})$ observations (red vertical line), when the network switches from predicting the future moon position by repeating the last observation to predicting by find a matching memorized moon configuration.
Refer to caption: x4.png

Figure: Figure 5: Three-heads network prediction error versus context length. The prediction error improves whenever the context length becomes large enough to fully describe the orbit of each planet. The net yields accurate predictions after only $\max(p_{1},p_{2},p_{2})$ observations (last black vertical line), that is, before having seen the full set of moon configurations (red vertical line). This is possible because these unseen moon configurations are predicted by composing individual moon predictions.
Refer to caption: x5.png

Figure: Figure 6: Visualization of the disentangled $W_{\varphi}$, $W_{\psi}$, and $W_{z}$ matrices in the 3-heads network. The color scale represents the moduli of the complex matrix coefficients.
Refer to caption: x6.png

#### Training

The networks are trained using randomly generated sequences $(x_{t})$ of length 800. Each sequence features three moons whose periods are related by randomly chosen ratios and are scaled to ensure that the 800 observation sequence contains at least three full periods lcm($p1,p2,p3$) of the moon system. Validation sequences are constructed similarly using a set of moon periods that does not appear in the training set.

Figure 4 and 5 show the prediction errors of both networks as a function of the context length, that is, the number of observations stored into the memories. More precisely, for each sequence $(x_{t})$ and each time index $T$, we compute the average absolute deviation between the next 25 true moon positions $x_{T+1}\dots x_{T+25}$ and the next 25 auto-regressive predictions (in which the successive predictions are looped back into the network input.) The plots show curves averaged over 512 sequences sharing the same set of moon periods taken from either the training set or the validation set.

- •
For the single head network (Figure 4), the plots show a sharp transition after $\mathrm{lcm}(p_{1},p_{2},p_{3})$ observations, that is, when the memory contains a full set of moon configurations (red vertical line). Before this threshold, predictions are performed by repeating the last observation. After this threshold, predictions are performed by finding a matching moon configuration in the memory, just as suggested by the first astronomer.
- •
For the three-heads network (Figure 5), the prediction error curve drops after seeing exactly $p_{1}$, $p_{2}$, and $p_{3}$ observations (black lines), that is whenever the orbit of an additional moon has been memorized. The network produces accurate predictions after a time equal to the period $\max(p_{1},p_{2},p_{2})$ of the slowest moon (last black line), long before the combined period $\mathrm{lcm}(p_{1},p_{2},p_{3})$ (red line) of the moon system. During this interval, accurate predictions are returned for moon configurations that can be very different from the previously observed ones. This is possible because the network *combines individual moon predictions, each well supported by the past observations*.

Figure 6 depicts the learned weight matrices and shows how the network has successfully disentangled the three moons ($W_{\psi}$ and $W_{\varphi}$) and learned how to recombine individual predictions ($W_{z}$).

#### Training notes

Training the three-heads network can be quite challenging in a manner that resembles the XOR networks of the early times (Rumelhart et al., ). We obtained reliable convergence using two tricks. First, we slightly restrict the linear operations by using $3\times 3$ complex matrices (18 real parameters) instead of $6\times 6$ real matrices (36 real parameters) operating on the 3-dimensional complex vectors as 6-dimensional real vectors. Second, we clip the mean squared loss in order to prevent the training algorithm from trying to optimize the prediction error when the memories are nearly empty.(^3^33The steamroller metaphor (Figure 2) makes more sense when the loss is bounded.)

Reliable convergence could also be achieved by making any of $W_{\varphi}$, $W_{\psi}$, or $W_{z}$ equal to the identity. Doing so would of course bias the network toward the disentangled solution, something we wanted to avoid. Yet it is not unreasonable to believe that disentanglement can often be achieved in the canonical basis. For instance, objects well separated in space often appear in different image regions, and therefore along different pixels axes.

## 4 Layered memories

In order to motivate more complex networks of memory units, let us challenge our astronomers by adding a sun and a couple planets with their respective moons. Every so often, maybe only in their imagination, our astronomers are transported to a new planetary system, observe new skies, and must sort out a new set of celestial bodies.

#### Persistent memories

Various hints can suggest a structure for a new planetary system: the sun has a distinctive appearance; inner planets remain close to the sun; the moons of a distant planet never stray far away from their hub, etc. Our astronomers also need geometrical knowledge to determine what to tabulate and how to combine the tabulated quantities and predict future skies. We call this knowledge *persistent* because it applies to all planetary systems, in contrast with *contextual* knowledge such as the previously observed skies for a specific planetary system.

In the three-moons architecture of Figure 3, persistent knowledge is stored in the parameters $W_{\varphi}$, $W_{\psi}$ of the feature extractors and the parameters $W_{z}$ of the mixing layer. These parameters are determined at training time through gradient back-propagation and optimization. Such simple linear functions might no longer be sufficient to encode the required persistent knowledge. Both feature extraction functions and combination functions might have to perform nonlinear computations reflecting, for instance, the projection of actual planet positions onto the celestial sphere.

Following instead the approach of Sukhbaatar et al. (), persistent knowledge can be stored in *persistent associative memories* containing a predefined number of key value pairs $(k_{i},v_{i})_{i=1\dots N_{m}}$ whose values are determined at training time through gradient back-propagation, and remain fixed at inference time. A *persistent memory unit* (Figure 7) is therefore very similar to a *contextual memory unit* (Figure 1) but relies on a persistent associative memory. Persistent memory units no longer needs an explicit value extraction function because the memory content is not updated at inference time. As pointed out in Sukhbaatar et al. (), they also can be viewed as fully connected neural networks with a single hidden layer that uses a soft-max non-linearity instead of a component-wise transfer function. Interleaving layers of contextual and persistent memory units can then be understood as means to increase the effective complexity of either the feature extractors or the combining layers of contextual memories (see Figure 8 for a spoiler). However we also find conceptually useful to still view the persistent memory output $y_{t}$ as the conditional expectation $\mathbb{E}(V|K)$ of an implicit value function that is not explicitly parameterized, but can be figured out after training.

Figure: Figure 7: Persistent memory unit. The persistent associative memory contains a fixed number of key-value pairs $(k_{i},v_{i})_{i=1\dots m}$ whose values are determined by back-propagating gradients at training time. Since the memory contents do not change at inference time, there is no need for explicit values.
Refer to caption: x7.png

#### Routing

What kind of solutions can our astronomers design to predict the skies of a new planetary system? They may perhaps first tabulate the sun position, then the position of each planet relative to the sun, then the position of each moon relative to its planet. In order to implement such a solution, a network of memory units does not only need to compute all these relative positions, but also must route information between the memory units representing each of these tables. Furthermore, since the same network must handle sequences of observations associated with different planetary systems of different composition, these routes must change dynamically on the basis of the first observations of every new sequence of observations.

Although such dynamical routing problems seem overwhelming, we can first observe how networks of memory units can implement *static routing* circuits whose behavior is determined at training time. Static routing circuits are implemented using the feature extractors and the combining layer of each memory unit layer (Figure 3), both enhanced with nonlinear operations implemented by interleaved persistent memory unit layers. Just like the gating modules of a mixture of expert (Jacobs et al., ), such a routing circuitry can implement routes that depend on the data. However, we still call such routes static because all the possible routes are determined ahead of time during training.

In order to gain dynamic routing abilities, the training algorithm can recruit contextual memory units to replace or complement the persistent memory units involved in the routing circuitry. Because the contents of contextual memory units are updated at inference time, recruiting some of them into the routing circuitry provides the means to create new routes on the basis of the first observations of a new sequence, suggesting an alternative to capsule networks (Sabour et al., ).

#### Memory Mosaics

In such a complex network, the division of labor between contextual memory units is still determined by the predictive disentanglement principle. During training, the steamroller of Figure 2 pushes the contextual memory units towards functions that more easily memorized independently than in aggregation. This does not only hold for memory units that record primary pieces of information such as the moon positions of Section 3, but also for those that affect the routing circuitry and those that operate on the information produced by earlier memory units.

Therefore, *under the pressure of the predictive disentanglement principle*, a network of memory units does *not only memorize disentangled fragments of information, but also memorizes how they fit together and how their combinations can be again broken into new disentangled fragments and recombined in myriad ways*. This is why we call such networks *Memory Mosaics*.

## 5 From planetary systems to language modeling

We have so far described Memory Mosaics as an architecture that resembles transformers in important ways and whose essential properties are comparatively transparent thanks to the predictive disentanglement principle. In order to make this vision more compelling, we provide evidence that Memory Mosaics can handle the most successful application of transformers, that is, language modeling.

Let us first argue that predictive disentanglement is a meaningful concept for language modeling. In the case of a dialogue, predicting who is going to speak next can largely be predicted independently from what is being said. The same applies to predicting whether the next piece of dialog is expressed in English or German, whether the participants are in agreement or not, or whether their grammar is formal or informal.
Disentangling such predictions provide opportunities to leverage more training data. For instance, one can predict that a specific dialog participant will provide a tense answer because virtually any dialog participant placed the same a situation would provide a tense answer.

#### Language modeling task

The TinyStories dataset of Eldan and Li () provides the means to use relatively small language models to study language modeling phenomena that usually only appear in much larger models. This is achieved by constraining the text written in the simple language that a three years old child could understand and recounting stories that take place in the simple world that a three years old child could understand. Thanks to this limited scope, show how a small language model (33M parameters) trained on a corpus of such tiny stories generates continuations with far better language quality and narrative consistency than those of a much larger model (1.5B parameters) trained on a generic web corpus.

Following both the lead of and the admonitions of our legal department, we use the Mixtral-8x7B open language model (Jiang et al., ) to generate a new corpus of tiny stories dubbed BabiStories. This corpus and its generation are detailed in Appendix A.

#### Architecture

Figure: Figure 8: *Left*: Classic GPT2-small transformer. *Right*: GPT2-like Memory Mosaic
Refer to caption: x8.png

To put our experiments into context, we design a Memory Mosaic architecture that closely matches the classic GPT2-small transformer architecture (Radford et al., , ). Both architectures, shown side-by-side in Figure 8, use the same GPT2 tokenizer, the same embedding dimension ($d=768$), and the same number of heads ($N_{h}=N_{c}=N_{p}=12$). Both architectures are trained and tested using sequences of length 512, that is, one to three stories long.

There are three major differences between these two architectures. First, the Memory Mosaic does not use positional encoding. Second, unlike the $N_{h}=12$ attention heads of each transformer block, the $N_{c}=12$ contextual memory units in each block do not distinguish keys from queries (Figure 1) but instead use the key and value extraction functions described in Equation 7. The keys are formed with a leaky average of past inputs, and the values can peek one time step ahead. Accordingly, the attention mask excludes the main diagonal to avoid breaking causality. Finally, the feed forward networks (FFNs) of the classic transformers blocks are replaced by a layer of $N_{p}=12$ persistent memory units, complete with a key extraction functions (7) and combining layer, sized to ensure that the per-block parameter count of the Memory Mosaic architecture closely matches GPT2-small.(^4^44Compared with GPT2-small, we save ${768\times 512}$ position encoding weights and ${N_{b}\times 768^{2}}$ query projection weights, and add ${2\times{N_{b}}\times 768^{2}}$ weights for the persistent memory key extraction function and mixing layer. The total number of persistent memory unit slots is therefore close, but not exactly equal, to the number of FFN hidden units.)

Figure: Figure 9: Training and validation loss of the transformer and Memory Mosaic architectures trained on BabiStories for different model depths. The horizontal axis represents the number of training iterations. All hyper-parameters have been tuned on the transformer architecture and transferred verbatim to the Memory Mosaic architecture. The Memory Mosaic slightly outperforms the transformer for small depth networks, but that effect disappears when the depth increases.
Refer to caption: x10.png

#### Training and validation

Figure 9 shows the training and validation curves of both transformers and Memory Mosaics of different depth trained on BabiStories. The Memory Mosaic slightly outperforms the transformer for small depth networks,(^5^55This is not surprising because the Memory Mosaic architecture only needs one block to correctly handle the induction head problem of Bietti et al. (), whereas the transformer architecture needs two blocks for the same task.), but this effect disappears when the depth increases and both the training and validation losses become indistinguishable.

Importantly, all hyper-parameters were tuned for the transformer architectures (Appendix B) and transferred verbatim to the Memory Mosaics. This choice might explain why the training curves track each other so well. It also leaves the Memory Mosaics at a slight disadvantage in this comparison.

**Table 1: Continuations generated for the 24 prompts designed by Eldan and Li () to investigate the factual, reasoning, and consistency capabilities of language models trained on BabiStories. Both the transformer and the Memory Mosaic were $N_{b}=18$ blocks deep. Colors green and red respectively indicate correct and incorrect answers. The text was generated with a beam search of width 2 and manually truncated to the shortest length that answers or fails to answer the question.**
| Prompt | Transformer ($N_{b}=18$) | Memory Mosaic ($N_{b}=18$) |
| --- | --- | --- |
| Alice was so tired when she got back home so she went | to sleep. | straight to bed. |
| Jack and Lily saw a rainbow after a rainy day. They were amazed by the colors. Jack said, “Look, Lily. A rainbow has | many colors!” | many colors!” |
| Jack and Lily liked to watch the moon at night. They noticed that the moon changed its shape every night. Sometimes the moon was big and round, and sometimes it was | small and thin. One night, Jack said, “Lily, look! The moon is changing.” | small and thin. One night, Jack had an idea. He said, “Let’s make a picture of the moon!” |
| Jack wanted to read a book, so he went to | the library. | the library. |
| “Can cows fly?”, Alice asked her mother | Her mother laughed and said, “No, Alice. Only birds and planes can fly.” | “No, dear. Only birds can fly,” her mother replied with a smile. |
| “What do birds like to eat?”, Tom asked his mother. | “Birds like to eat worms and bugs,” his mother replied. | “Birds like to eat worms and bugs, Tom,” his mother replied. |
| “What language do they speak in France?”, Tom asked his mother | His mother smiled and said, “We all speak different languages, Tom. It’s what makes us unique and special.” | His mother smiled and said, “They speak in different languages, Tom. But we can learn some of them.” |
| If I throw a ball up in the air, eventually it will | come down and I can catch it. | come down. |
| It was winter and cold outside so his mother told him, “You should | wear your warm coat, dear.” | wear your warm coat and hat.” |
| Lily likes cats and dogs. She asked her mom for a dog and her mom said no, so instead she asked | for a toy. | her dad. |
| Jack told Mary, ”If you give me your banana, I’ll give you my apple”. Mary gave Jack her banana so | he could eat it. | that he could eat it. |
| On weekends Jack went to visit his grandmother whereas on weekdays he would go to school. Last weekend, when Jack was on his way to | school, | school, |
| Lily and Ben were having an argument. Ben said that cake is much better than ice cream and Lily said that | ice cream is the best. | she likes ice cream |
| Lily and Ben are having an argument. They are trying to decide between the park and the swimming pool. Ben says, “I want to go to the park”. Lily says | , “No, | “No, |
| Jack’s mother was not home, and his father was at home. When Jack came home, he said hello to | his mother. | his mother |
| Lily doesn’t like swimming. When her father wants to take her to the swimming pool, she says | , “No, daddy. I don’t want to go.” | , “No, thank you.” |
| Both Ben and Lily wanted cake. Father said that there was only one piece of cake left. They | both wanted to eat it, but there was only one. | all wanted to eat it, so they began to fight. |
| Ben went to visit Lily in her house, but she was not at home. Ben knocked on the door, | and Lily’s mom answered. | and Lily’s mom opened it. |
| “Hi Jane, have you seen Alice? I can’t find her anywhere”, said Jack. | Jane shook her head. “No, I haven’t. But I will help you look.” | Jane shook her head. “No, I haven’t seen Alice. But I can help you look!” |
| Max had two dogs. One was white and the other was black. Max walked up the street and saw a kid with a dog. He told the kid, “I see you have a brown dog. I also have | a white dog.” | a brown dog.” |
| Anne had a piece of candy in her left pocket and a piece of chocolate in her right pocket. Anne’s mom asked her, “Anne, what is that you have in your left pocket?” | Anne smiled and said, “It’s a candy | Anne opened her hand to show her mom the candy |
| Alice had both an apple and a carrot in her bag. She took the apple out of the bag and gave it to Jack. She reached into the bag again and took | out the carrot. | out another apple. |
| Alice and Jack walked up the street and met a girl in a red dress. The girl said to them, “Hi, I’m Jane. What are your names?” | Alice replied, “I’m Alice. This is Jack. We came to see you.” | Alice smiled and said, “I’m Alice. This is Jack. We are going to the park to play.” |
| Diva was hungry, and wanted to bake a cake, but she didn’t have any sugar at home, so she decided to go ask around. She started walking and met a squirrel. She asked the squirrel, “Would you happen | to know where I can find some sugar?” | to see any sugar around here?” |

#### Qualitative evaluation

In order to compare the quality of the text generated by models trained on tiny stories, designed twenty-four prompts that exercise the factual, logical, and consistency properties of the generated continuations. Table 1 compares the continuation generated on these prompts by a transformer and a Memory Mosaic, both $N_{b}=18$ blocks deep. Both models perform very similarly on this task.

Figure: Figure 10: Prediction performance on the Simple English Wikipedia dataset using models trained on BabiStories. The plot shows the per-token average loss as a function of the position of the generated token in the 512-token long input window. Memory Mosaics outperform transformers after about 50 tokens, suggesting superior in-context learning abilities.
Refer to caption: x11.png

Figure: Figure 11: Memory Mosaics performance on the RegBench in-context learning benchmark (Akyürek et al., ). Since RegBench includes a search over architectural hyper-parameters, Memory Mosaics and transformers use the same search space with the same parameter counts. Memory mosaics outperform all previously tested architectures in this benchmark.
Refer to caption: x12.png

#### Out-of-distribution evaluation

The Simple English Wikipedia(^6^66[https://simple.wikipedia.org/wiki/Simple_English_Wikipedia](https://simple.wikipedia.org/wiki/Simple_English_Wikipedia)) is a version of Wikipedia written in a language that is easier to understand. Despite the intended simplicity, the articles are substantially longer and more sophisticated than our BabiStories. Predicting Simple English Wikipedia articles using models trained on BabiStoriesis therefore a challenging out-of-distribution task.

Figure 10 shows the per-token average loss as a function of the position of the generated token in the input window. Both the transformer and the Memory Mosaic are $N_{b}=512$ blocks deep. In this experiment, the token prediction is expected to improve when the increasing context size reveals that the distribution is different. The transformer performance plateaus after 100 to 150 tokens, which is a bit shorter than a typical tiny story. Memory Mosaics substantially outperform transformers after about 50 tokens, suggesting superior in-context learning abilities.

#### In-context learning evaluation

In order to rigorously compare the in-context learning abilities of various architectures, the RegBench benchmark (Akyürek et al., ) constructs a large number of artificial languages defined by probabilistic finite automatons (PFA) constructed at random. Each input sequence is composed of 10 to 20 strings drawn from a same PFA and delimited separator tokens. The competing architectures are trained on a variable number of input sequences, then evaluated on their ability to predict the last token of testing sequences generated using held out PFAs.

Since RegBench performs a hyper-parameter search for each data point, we use the Memory Mosaic architecture of Figure 8 with the same search space as transformers, ensuring that both transformers and Memory Mosaics have the same parameter count for the same architectural hyper-parameters. We sweep over depth $N_{b}\in\{2,4,8\}$, number of heads $N_{h}{=}N_{c}{=}N_{p}\in\{2,4,8\}$, embedding dimension in $d\in\{64,128,256\}$, weight decay in $\{10^{-2},10^{-1}\}$, and training epochs in $\{1,2,\dots 200\}$.

Figure 11 compares Memory Mosaic on RegBench with the results previously reported by . The left plot shows the prediction accuracy for the test string last token. The right plot compares the predicted last token distribution with the exact distribution implied by PFA. Memory Mosaics dominate this benchmark, substantially outperforming transformers, recurrent neural networks, and state-space models for training set sizes covering three orders of magnitude.(^7^77In the low train environments case (e.g. 100), all baseline methods which perform poorly in the RegBench OOD test set can actually perform very well in a IID test set. Thus baseline methods learned the training environments (good IID) but not the meta-learning ability (poor OOD). Please check appendix table 4 for details.)

#### Attention differences

Because Memory Mosaics lack position encoding and do not distinguish keys and queries, we investigate how their attention patterns differ from those of transformers. Figure 12 shows attention scores for each head of either a one-block deep transformer using absolute position encoding (left plot) or a one-block deep Memory Mosaic (right plot). The scores are averaged on 5000 BabiStories sequences and show how the last position attends to earlier positions in the 512 token long context window. The transformer attention patterns are noisy, with a strong “attention sink” at position 0 (Xiao et al., ). In contrast, the Memory Mosaic attention pattern is mostly flat, save for higher scores for the most recent tokens.(^8^88This effect is connected to the leaky average coefficient $\lambda_{\varphi}$, as shown in Appendix (Figure 15).)

Figure 13 show the attention patterns for contexts extended to 1536 tokens, using models trained on 512 token long sequences. Because the absolute position encoding scheme cannot be extended to longer contexts, we provides a comparison with transformers using RoPE (Su et al., ) and AliBi (Press et al., ). The RoPE attention patterns do not extend nicely beyond the training context length. The AliBi attention patterns show the vanishing contribution of distant tokens. In contrast the Memory Mosaic attention patterns remain mostly flat.

Figure: Figure 12: Average attention scores of the last token attending previous tokens (evaluated on an in-distribution validation dataset). Each solid line indicates one head in either the transformer attention block or the Memory Mosaic contextual memory block. The dotted line averages the attention of all heads. All models are trained with context length 512.
Refer to caption: x13.png

Figure: Figure 13: Average attention scores on an extended context window ($3\times 512$ tokens). Models are still training with a 512 token long context window. Because the GPT2 absolute position encoding does not extend, we compare with RoPE (Su et al., ) and AliBi (Press et al., ).
Refer to caption: x15.png

## 6 Discussion

The starting point of this work is made of two very old ideas. The first one is augment a deep network with explicit memories. The second one is to let the learning process decide what gets memorized and how it gets retrieved. Although such ideas have been explored in memory networks (Weston et al., ; Joulin and Mikolov, ; Sukhbaatar et al., ), the importance of having lots of independent memories had not been fully appreciated.

This contribution focuses on networks of associative memories implemented with kernel smoothing, therefore amenable to gradient-based learning algorithms. Such learning machines not only resemble decoding transformers (Section 2) but also perform very much like decoding transformers on the sort of language modeling task that made them famous (Section 5). Although much work is needed to replicate our observations at far greater scale, Memory Mosaics satisfy narrative constraints as well as transformers (Table 1), and generally behave in very encouraging ways (Figures 10 to 13).

Most importantly, we understand what Memory Mosaics do far better than we understand what transformers do. First, the value extraction functions of the associative memory units precisely describe what each memory seeks to memorize. Second, the predictive disentanglement principle explains why training a Memory Mosaic breaks the overall prediction task into pieces that are more efficiently memorized when they are considered independently (Section 3). Therefore, Memory Mosaics are not just a transformer-like architecture, but also a model(^9^99Not as in “statistical model” but as in “model used to describe and explain a phenomenon.”) for compositional learning systems that break knowledge into independently memorized fragments, then reassemble them as needed using combination strategies that can themselves be viewed as memorized knowledge fragments (Section 4).

The focus on memorization allow us to formulate new questions. Could memories operate independently on different time scales? Could we envision a richer memory hierarchy than simply distinguishing persistent memories from contextual memories? Can intermediate memory tiers be trained like contextual memories, that is, without gradients? Can the persistent knowledge be then reduced to a compact high order bias?

Memory Mosaics also offer an array of engineering opportunities. Limited storage contextual memories could be implemented by evicting the least-recently used entry rather than by bounding the context size (see also (Xiao et al., )). Associative memories can be implemented using a wide spectrum of techniques, ranging from fast transforms (Greengard and Strain, ; Yang et al., ) to locality sensitive hashing mechanisms (Spring and Shrivastava, ; Chen et al., ) that could change the computing requirements of contemporary artificial intelligence systems.

### Acknowledgements

JZ and LB received partial funding from the CIFAR Learning in Machines & Brains program.

## Appendix A BabiStories

The TinyStories dataset [Eldan and Li, ] is composed of stories written in a simple language and taking place a narrow world. Such stories can be used to train relatively small language models that still must address some of the broader language modeling challenges such as obeying narrative necessity and maintaining logical consistency. This dataset is a wonderful way to study big problems with acceptable computation and quick turn around.

The experiments of Section 5 were carried out using a dataset generated using a similar methodology but using the Mixtral-8x7B open language model in order to generate unencumbered data. We call this dataset BabiStories.
All the scientific credit is still due to the remarkable work of .
Table 2 provides basic statistics for this newly generated BabiStoriesdataset, essentially matching those of the original TinyStories dataset of Eldan and Li []. We had to increase the diversity of the generated stories by expanding the prompt to specify first names and by providing opening words for the story, in addition to required words and story features used by (Figure 14). We also removed the few generated stories containing URLs.

Figure: Figure 14: Generation of BabiStories. In order to improve the diversity of the generations, each story is generated by a prompt that provides a list of required words and story features (as in Eldan and Li, ) and additionally provides first names and opening words.
Refer to caption: x18.png

**Table 2: BabiStories statistics.**
| dataset partition | #stories | #tokens (GPT2 tokenizer) | #char per story (average) |
| --- | --- | --- | --- |
| train | 2.2M | 474,704,907 | 888 |
| valid | 2.2k | 4,749,107 | 889 |

## Appendix B GPT2 baseline and hyperparameters

Table 3 showcases the hyper-parameters searching process of GPT2 transformer baseline on the BabiStories dataset, where we use AdamW optimizer Loshchilov and Hutter [], batch-size 512, context-size 512, and a cosine learning rate scheduler with minimum learning rate $1e-4$ for all training.

**Table 3: Hyperparameters searching on GPT2 transformer with $N_{b}=12$. “dropout”, if any, is applied on attention score, attention heads output (before combining layer), and FFN output.**
| learning rate | dropout | L2 weight decay | warm-up iters | training iters | train loss | valid loss |
| --- | --- | --- | --- | --- | --- | --- |
| 5e-3 | 0.05 | 0.1 | 2000 | 80000 | 1.336 | 1.494 |
| 1e-3 | 0.05 | 0.1 | 2000 | 80000 | 1.350 | 1.524 |
| 5e-3 | 0 | 0.1 | 2000 | 80000 | 1.281 | 1.556 |
| 5e-3 | 0.05 | 0.01 | 2000 | 80000 | 1.322 | 1.516 |
| 5e-3 | 0.05 | 0.1 | 200 | 80000 | fail | fail |
| 5e-3 | 0.05 | 0.1 | 2000 | 40000 | 1.325 | 1.532 |
| 5e-3 | 0.05 | 0.1 | 2000 | 160000 | 1.314 | 1.497 |

## Appendix C Memory Mosaics for language modeling

### C.1 Additional qualitative evaluation

Table 5 provides a similar continuation generation comparison as Table 1 in Section 5, but setting $N_{b}=1$.

### C.2 Attention map and leaky average coefficient λ φ subscript 𝜆 𝜑 \lambda_{\varphi} italic_λ start_POSTSUBSCRIPT italic_φ end_POSTSUBSCRIPT

Figure 15 shows the relationship between attention map and leaky average coefficient $\lambda_{\varphi}$.

Figure: Figure 15: Attention map and leaky average coefficient $\lambda_{\varphi}$. As $\lambda_{\varphi}$ increases, $k_{t}$ in Eq 7 effectively takes a longer history into the account, and thus the pick at the end of attention map becomes wider.
Refer to caption: x19.png

### C.3 In-context language learning

Table 4 provides the IID test performance of various architectures trained on RegBench [Akyürek et al., ] with 100 training environments. We keep the training process, including hyperparameter searching space, to be the same as the one in Figure 11. But sample validation and test sets from the same 100 probabilistic finite automatons (training environments) as the training set. This table, together with Figure 11, show that baseline methods learned the training environments (good IID) but not the meta-learning ability (poor OOD).

**Table 4: In-distribution (IID) performance of various architectures trained on RegBench [Akyürek et al., ] with only 100 training environments. Both training, validation, and test set (100 samples) are sampled from the same 100 random probabilistic finite automatons (PFA). Compared with the poor OOD accuracy ($\sim$0.45) / TVD ($\sim$0.75) of baseline methods in Figure 11, All baseline methods perform well in the IID test set (even with only 100 training environments).**
|  | Memory Mosaic | tf | Mamba | S4 | RWKV | linear tf | H3 | GLA | Hyena | LSTM | RetNet |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Accuracy ($\uparrow$) | 0.959 | 0.856 | 0.929 | 0.846 | 0.967 | 0.816 | 0.794 | 0.870 | 0.953 | 0.849 | 0.876 |
| TVD ($\downarrow$) | 0.417 | 0.308 | 0.268 | 0.350 | 0.183 | 0.348 | 0.425 | 0.284 | 0.244 | 0.343 | 0.296 |

**Table 5: Continuations generated for the 24 prompts designed by Eldan and Li [] to investigate the factual, reasoning, and consistency capabilities of language models trained on BabiStories. Both the transformer and the Memory Mosaic were $N_{b}=1$ blocks deep. Colors green and red respectively indicate correct and incorrect answers. The text was generated with a beam search of width 2 and manually truncated to the shortest length that answers or fails to answer the question.**
| Prompt | Transformer | Memory Mosaic |
| --- | --- | --- |
| Alice was so tired when she got back home so she went | to sleep | to bed |
| Jack and Lily saw a rainbow after a rainy day. They were amazed by the colors. Jack said, “Look, Lily. A rainbow has | come!” | come to us!” |
| Jack and Lily liked to watch the moon at night. They noticed that the moon changed its shape every night. Sometimes the moon was big and round, and sometimes it was | just a big, round moon. | too big. |
| Jack wanted to read a book, so he went to | find his friend, a wise old owl named Hoot. | find his best friend, Timmy. |
| “Can cows fly?”, Alice asked her mother | Her mother smiled and said, “Of course, dear. | Her mother smiled and said, ”Sure, honey. |
| “What do birds like to eat?”, Tom asked his mother. | “I don’t know, but it looks yummy!” Tom replied. “Let’s find out!” | His mother smiled and said, “That’s a good idea, Tom. |
| “What language do they speak in France?”, Tom asked his mother | France smiled and said, “It’s a language, Tom. It helps us talk to each other.” | “Mom, what is that?” France asked. |
| If I throw a ball up in the air, eventually it will | go far away. | come back. |
| It was winter and cold outside so his mother told him, “You should | go outside and play in the snow.” | stay in your warm, warm home.” |
| Lily likes cats and dogs. She asked her mom for a dog and her mom said no, so instead she asked | , “Can we keep it, mom?” Her mom said, “Yes, but we need to keep it safe.” | her mom if they could keep the dog. Her mom said yes, and they named the dog Sparky. |
| Jack told Mary, “If you give me your banana, I’ll give you my apple”. Mary gave Jack her banana so | he could have it. | she could eat it. |
| On weekends Jack went to visit his grandmother whereas on weekdays he would go to school. Last weekend, when Jack was on his way to | school, | school, |
| Lily and Ben were having an argument. Ben said that cake is much better than ice cream and Lily said that | they could both have fun. | it was the best cake they ever had. |
| Lily and Ben are having an argument. They are trying to decide between the park and the swimming pool. Ben says, “I want to go to the park”. Lily says | , “No, | “No, |
| Jack’s mother was not home, and his father was at home. When Jack came home, he said hello to | his mother, | Jack. |
| Lily doesn’t like swimming. When her father wants to take her to the swimming pool, she says | “Lily, you can do it! | , “Mommy, I want to swim too!” |
| Both Ben and Lily wanted cake. Father said that there was only one piece of cake left. They | all sat down to enjoy the yummy treat. | both felt sad. |
| Ben went to visit Lily in her house, but she was not at home. Ben knocked on the door, | and the door opened. A kind lady came out | and when Ben opened the door, |
| “Hi Jane, have you seen Alice? I can’t find her anywhere”, said Jack. | Alice smiled and said, “Sure, I will help you find your way home.” | “I don’t know, Jack. |
| Max had two dogs. One was white and the other was black. Max walked up the street and saw a kid with a dog. He told the kid, “I see you have a brown dog. I also have | a black dog.” | a brown dog.” |
| Anne had a piece of candy in her left pocket and a piece of chocolate in her right pocket. Anne’s mom asked her, ”Anne, what is that you have in your left pocket?” | Anne smiled and said, ”Yes, mommy. I found it in the park.” | Anne smiled and said, ”I found it on the ground. It’s mine!” |
| Alice had both an apple and a carrot in her bag. She took the apple out of the bag and gave it to Jack. She reached into the bag again and took | out the apple. | out the apple. |
| Alice and Jack walked up the street and met a girl in a red dress. The girl said to them, ”Hi, I’m Jane. What are your names?” | Jane smiled and said, ”I’m Timmy, | Jane replied, ”I’m Jane. |
| Diva was hungry, and wanted to bake a cake, but she didn’t have any sugar at home, so she decided to go ask around. She started walking and met a squirrel. She asked the squirrel, ”Would you happen | to my house, little one?” | to my cake?” |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
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.
soundness: 2
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.
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.
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

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 2
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.
soundness: 3
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.
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.
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).

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 3
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!
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: See my weakness part.
soundness: 2
strengths: -  The paper addresses an interesting problem in developing more interpretable language models.
- The work is well-organized and clearly presented.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)