Memory Mosaics
Authors: Jianyu Zhang
Niklas Nolte
Ranajoy Sadhukhan
Beidi Chen
Leon Bottou
Venue: ICLR 2025 Poster (2025)
Keywords: predictive disentanglement
Associative memory
language model
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

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

Published as a conference paper at ICLR 2025
MEMORY MOSAICS
Jianyu Zhang"
Niklas Nolte", Ranajoy Sadhukhan', Beidi Chen'1, Léon Bottou
FAIR, Meta Carnegie Mellon University
New York University
ABSTRACT
Memory Mosaics are networks of associative memories working in concert to achieve prediction task of interest. Like transformers, memory mosaics possess compositional capabilities and in-context learning capabilities. Unlike transform- ers, memory mosaics achieve these capabilities in comparatively transparent way ("predictive disentanglementt). We illustrate these capabilities on toy example and also show that memory mosaies perform as well or better than transformers medium-seat scale language modeling tasks.
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 re- lated to memory networks (Weston et al., 2014; Sukhbaatar et al., 2015) and resemble transformers (Vaswani et al., 2017) despite significant differences. Like transformers, Memory Mosaics pos- sesses some of the disentanglement and compositional capabilities that have long eluded machine learning systems (Lake & Baroni, 2018). Unlike transformers whose internal mechanism are hard to decipher (Olsson et al., 2022; Bietti et al., 2024), Memory Mosaics achieve these capabilities in comparatively transparent ways.
The three main contributions of this work are (a) defining an architecture that exploits the direct similarity between self-attention and associative memories implemented with kernel regression, (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 i.i.d. performance of decoding transformers and outperforms them on o.o.d. tasks such as in-context learning.
language modeling task,
Section 2 reviews related work. Section 3 describes simple associative memory units than can be inserted in deep network. Section 4 explains how training such network splits prediction task into disentangled sub-tasks. Section illustrates this "predictive disentanglement" using a network with only 54 parameters, showing that this is not mysterious effect of scale but property of the architecture. Section extends these ideas to fully formed memory mosaics. Section reports on medium-scale language modeling experiments.
Related Work
Several recent papers (e.g., Katharopoulos et al., 2020; Peng et al., 2023; Sun et al., 2023; Gu & Dao, 2023) propose transformer alternatives that use efficient recurrences putational cost of transformers. Closer to our interests, other authors (e.g.. Ramsauer et al., 2020; Krotov, 2023; Hoover al., 2024) rethink transformers with Hopfield-style associative memories and their associated energy function. In contrast, we leverage elementary associative memories that interpolate stored key/value pairs with a kernel regression (therefore incurring a quadratic runtime cost) in order to construct an architecture that remains very close to standard transformers but cast new light on properties that play an important role in their compositional learning capabilities.
the quadratic com-
Closely related to predictive disentanglement, (Bengio et al., 2019) proposes meta-learning train- ing objective that achieves causal disentanglement by seeking quick adaptation to new distributions. We argue that a similar effect happens in architecture, as a consequence of the normal training process interpreted as a meta-learning process, revealing an important aspect of the still mysterious compositional learning abilities of transformer-like architectures.

---

Published as a conference paper at ICLR 2025
6
Output time scries tracks (Py) without looking forward (xy) using menmcey mutches.
Associative memory recieves pas d(Rr, Ar) small.
Associative Memory
Value time series repcesents the
(4:)
(i:)
of
Key time series represents the recent past of the lsput time series. By (aa Sy-ze
input time series Vy (ar My Ky- to
Feature Extractor
Input time series
(tx)
Figure 1: Elementary memory unit. The keys ky are computed as a function of past observations (It)SST. The values Vr peek into the future. In this example, the value also depend the next observation 7T+1- At time T, the associative memory uses the known key ky to compute an estimate yr of ((rr) using only the previously stored pairs (ka, (t), & T. One time step later, the input 21+1 is revealed, the value br can be computed, and the pair (kr, or) is added to the memory.
3 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 I The retrieval process can then be represented as function of the queried key and all the stored pairs (k1, U1) ... (ka, 2n).
{ R A
Rd f(k; ((k1, (kmn vn)))
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 also view associative memory distribution P(V| () on the basis of the sample (k1, 01)... ... (kn, un) of key-value pairs. The retrieval function is then a conditional expectation over this estimated distribution:
device that estimates conditional probability
[k;; {(ki, v.)......... (kn, v.)))) B(V K k).
(1)
Such conditional expectation can be constructed with Gaussian kernel regression," 1
12
[k; ((11, v1) ... (kn; vn)))
B||k kell
U/ with Z
kz1
The close connection between this Gaussian kernel smoothing and attention (Bahdanau al., 2015) is obvious when all key vectors K; share same squared norm because expression (2) becomes
R Sk
fl: (1k1, (kn, vn)})
B
There are of course more advantageous ways implement associative memories. Although some 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 sequence (xx) of observations, discrete tokens or continuous values. We would like leverage the past observations (It)eTT to predict some useful property of the future observations (Z.). For instance we might want to predict the next observation IT+1 to construct an auto-regressive model of the sequence.
'Expression (2) is known as the Nadaraya- Watson estimator (Nadaraya, 1964; Watson, 1964). It is known converge the true conditional expectation B(K(V) when - and 8 vn.
2

---

Published as a conference paper at ICLR 2025
Our elementary memory unit (Figure 1) consists of an associative memory and trainable feature extractor that computes suitable keys and values for the memory. The keys ky are computed as a function of the past observations (Zi)iTT and trainable weights w,
kt p(TT, 21-1.... w).
(4)
In contrast, the values vr 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 2T+1-
1 2T-,,,
(5)
The memory units operate independently at inference time. They start empty at the beginning of each input sequence. At time step T, each memory receives a key vector ly computed from the recent inputs (#T, 11,,.... and interpolates response y/1 on the basis of the previously stored key/value pairs. The value er computed one time step later when the next input TT+1 is revealed and the pair (krr vr) is added to the memory.
Although the value by depends on the near future, the output does not depend on by but merely leverages the previously stored key/value pairs to estimate vr. Therefore there no leak of future information: each memory unit is a little machine that predicts a bit of future information (described by uT) on the basis of recent information (described by kr) and previously stored key/values pairs.
The exact form of the feature extraction functions can vary in complexity. For instance, when each observation 2y carries sufficient information, the keys ky and values vy can be computed as linear functions of respectively Iт and 27+11, that ky We 3t and vr W TT++- However find useful to consider feature extraction functions that summarize the recent past using short convolu- tions or quickly vanishing leaky averages. For instance, the language experiments of Section 7 use feature extractors of the following form:
leaky average over = T-I...,
by = Norm Kr
with
by = by + x
Py = W DT
(6)
Norm or
with
0T = Üf + Ur+1
Uy
We TT
convolution t=T T+1
Since this expression produces keys with unit norm (Norm(x) = 2/ æ/]|æ the effective kernel band- width is determined by the trainable parameter 3 in equation (3).
Training networks of memory units Consider now deep network whose architecture includes layers of associative memory units. When the associative memories are implemented with differen- tiable kernel smoothing mechanisms, training such deep network is simply a matter of unrolling the network in time and back-propagating the gradients, in ways that users of modern deep learn- ing software will find very familiar. Unsurprisingly, unrolling equation (3) along input sequence (I1... DD) of duration D yields an expression that very much resembles masked self-attention (Vaswani et al., 2017).
B
VT.
VT -F= Σ
U/.
T e kg
Implementing associative memories with kernel smoothing therefore provides a particularly direct illustration of the connection between self-attention and associative memories (e.g., (Ramsauer al., 2020)). However, Memory Mosaics differ because the value extraction function is allowed to peek into the near future of the input time series (zz). This slight change has important consequences
Each memory unit operates as little predictor whose outputs yr can be interpreted as condi- tional expectation (1) that estimates features of the near future (VT) of the input time series on the basis of its past observations (kr). The parameters of the value extraction function (1) specify what is being predicted and the parameters of the key extraction function (4p) specify how it is predicted.
"The leaking average in expression (6) is far too simple to effectively encode long range dependencies as demonstrated in (Voelker al., 2019; Peng al., 2023; Gu & Dao, 2023).
3

---

Published as a conference paper at ICLR 2025
Memories being forced learn quickly!
cost por time stop
Gradient (meta) training
X2.
training sequence
Xp
The curve plots the prediction losses for all training sequence indices t E {1...D) in the training sequence. Minimizing their sum -the area under the curve- favors memories that produce useful value estimates after fewer time steps.
Figure 2:
Equation (7) must therefore account for the number of future time steps needed to compute uT- In our experiments, for example, vr can look one step ahead in the future. This amounts to having more aggressive attention mask. Therefore the main diagonal must be excluded from the attention mask. justifying the - upper bound in the sum."
Because each memory unit acts as predictor, single layer of memory units is sufficient to address the induction head problem of Bietti al. (2024). In contrast, a decoding transformer needs at least two self-attention layers for the same task.
Equation (7) makes provision for position encoding and no distinction between query and key vectors. In other words, we are betting that these transformers complications are no longer needed because our associative memory units do not need them to implement induction heads.
4 Predictive Disentanglement
Training and meta-learning The training process determines which future bit of information is predicted by each associative memory unit (through the parameters that control the computation of the values uT) and which kernels used to perform the predictions (through the parameters of that control the computation of the keys kr). In contrast, the relation between keys and predicted values determined for each input sequence inference time through the memorization of key/values pairs specific to each sequence. The training procedure should therefore be seen as a meta-learning process, distinct from the memory-based learning that occurs at inference time when new key/value pairs are added into the memories.
Predictive disentanglement This meta-learning interpretation reveals a remarkable phenomenon that we call predictive disentanglement the gradient training algorithm splits the overall prediction task (e.g., predicting the next token in natural language sentence) into disentangled prediction sub-tasks assigned each memory unit.
Consider training set composed of long enough sequences (#1,... D) extracted from underlying time series governed by possibly different stationary processes. The goal of our network is to pre- dict each 2T+1 using the previous observations I1 ... 2r. Unrolling the network in time along each sequence (ι1 .f....Ep) and collecting the prediction losses measured at each position t can sum- marized by curve that shows the prediction cost (or loss) at each time step 1... D, illustrated in Figure 2. We can expect that the prediction cost observed at position T becomes smaller when T increases because more information (a) ... is available predict each .7+1-
The training process minimizes the total prediction cost, that is the area under the curve in Figure 2 viewed as a collection of vertical slices. We can also view this area as a collection of horizontal slices, each representing the context length required to drive the prediction cost below certain threshold. Therefore the training process can also be viewed as minimizing the context length needed produce good enough predictions.
Because the associative memory retrieval function (2) is known to converge to stationary condi- tional expectations [VV)K), each memory unit driven produce good conditional expectation estimate as soon as possible. This can be achieved in two ways:
Let us first assume that each memory unit has frozen value extraction function v. The training procedure can still make each memory unit statistically more efficient by tuning the parameters
'One could of course use more aggressive masking allow er peeking several time steps in the future.
4

---

Published as a conference paper at ICLR 2025
of the key extraction function y, that is, by learning how to compare the current prediction context (æT, 2T- with past prediction contexts (øt, 1 for t < T. Learning similarity metric (a kernel) is well known way to make non-parametric estimators more efficient (e.g, Bach et al., 2004). 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 33 in equation (7).
When multiple memory units are available, the training procedure can also distribute the overall prediction task among the available memory units. As long as the memory units outputs can still be combined to address the overall task, the training algorithm can optimize the parameters of the value extraction functions vp to produce values eT that more efficiently modeled by their respective memory units.
Because each memory unit operates independently at inference time, this works best when the overall prediction task is disentangled into smaller prediction sub-tasks that can be modeled independently and efficiently. More precisely, the sub-tasks must be chosen so that each memory can carry out its assigned modeling task at inference time without having to account for the combined impact of the operation of all memory units. Their outputs 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, as illustrated in Section 5.
Disentanglement has long been recognized as desirable (Bengio, 2013) but has been hard to pinpoint (Comon, 1994; Roth et al., 2022: Thomas et al., 2018). Predictive disentanglement is closely related to the meta- veta-transfer objective of Bengio et al. (2019) but arises as side effect of a specific predictive architecture trained with the usual gradient procedure. Although predictive disentanglement is easier to understand in the case of network of associative memory units, we conjecture that something similar also occurs in standard transformers.
5 Tracking three moons
We give an illustrative example of predictive disentanglement: three moons orbit a remote planet. Although the local astronomers are very far from understanding celestial mechanics, they neverthe- less 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 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 lcm P2, Ps) 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(p1, P2. P2) 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.
Model For our purposes, each observation Ia consists of three complex numbers e that encode the angular positions 0g of the three moons inside their respective orbital plane. We consider two single layer models (Figure 3) with either N₁, = or N₂, = 3 memory units whose added dimen- sions match the input dimension. The trainable parameters of the linear key and value extraction are collected in two 3 x 3 complex matrices W and W., The memory unit follow equation (3) with a fixed parameter = 50. A third x 3 complex matrix W. combines the memory unit predictions into an output 27 that hopefully predicts IT+1- Both networks share an interesting analytic solu-
'We do not seek to discuss subtleties such as elliptical arbits or multi-body problems. Our primitive tronomers best compared to the ancient sky watchers whose efforts eventually gave the Ptolemaic model.
5

---

Published as a conference paper at ICLR 2025
Miving laser
Combiniig lver
N₂ 1 or
(1)
Stack
[k) W
W E Ca
h A=1...
Antedodee nattor
Stack h= k.....
[h)
PT
We TT+1
We e C
0,7
(x
Σ
Zr
Fexuue
21 W.
Stack
Pr
(b)
W. E C
Figure 3: An architecture for the three moons problem. We consider single-layer networks with either N₁ = 1 or N₁, 3 memory units whose keys and values belong to either C3 (N₂ 1) or C' 3 54 trainable real parameters that determine how predict the moon positions using either single 6-dimensional memory or three 2- 2-dimensional memories.
(N,, = 3). Both nets have 3
genoration error over 25 stops vs, context longth
generation arror over 25 stops vs. context longth
zonl
300
500
700
length generation
context length
Figure 4: Single head network prediction ror versus context length. The prediction error shows sharp transition after Iemn(1, P2, P3) ob- servations (red vertical line), when the network switches from predicting the future moon posi- tion by repeating the last observation to predict- ing by find matching memorized configuration.
Figure 5: Three-heads network prediction ror versus context length. The prediction error improves whenever the context length reaches the period of a new moon (black vertical lines), yielding accurate predictions after the last one, well before having seen the full set of moon configurations (red vertical line).
tion: setting all three matrices W.p. We, and W, to the identity yields optimal predictions once the associative memories have seen enough samples.
Training The networks are trained using randomly generated sequences (æt) of length 800. Each sequence features three moons whose periods are related by randomly chosen ratios and are scaled ensure that the 800 observation sequence contains least three full periods lcm(pl, 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 show the prediction errors of both networks a function of the context length, that is, the number of observations stored into the memories. More precisely, for each sequence (xx) and each time index T, compute the average absolute deviation between the next 25 true moon positions #T+1- -7525 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 or validation set.
For the single head network (Figure 4), the plots show sharp transition after lem(pı, P2, P3) 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. Af- ter 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 P1, P2, and Ps observations (black lines), that is whenever the orbit of an additional moon has been memorized. The learned weight matrices are shown Figure 10 in the Appendix. Observe how the network produces accurate predictions after a time equal to the period max(pı, P2.P3) P3) of the slowest moon (last black line), long before the combined period lem(p1, P2, Ps) (red line) of the

---

Published as a conference paper at ICLR 2025
moon system. In this interval, accurate predictions are returnd for moon configurations that can be very different from the previously observed ones. Instead the network combines individual moon predictions, each well supported by the past observations.
Predictive disentanglement and compositional learning in language models Consider a chat- bot assisted creative writing scenario in which the human uses dialogue to repeatedly introduce new ideas into an evolving story that the chat-bot reprints at each step. The user can drive such a story arbitrarily far from the training data and into the distant tail of its distribution. Although no training example resembles the story, the chat-bot keeps producing syntactically correct language and coherent stories because it has learned some of the mathematical structures of language (Harris, 1968) and can recombine pieces of information coming from either the context the training data. This phenomenon is fundamentally similar to that illustrated in Figure 5, where moon configurations unlike any previously seen configurations are accurately predicted because the network has learned how to combine individual moon predictions. This similarity casts a useful light on the otherwise mysterious compositional learning abilities of transformer-lit like models.
Layered memories
We of course envision deeper networks of memory units. In order to make meaningful comparisons, we also would like to remain close possible to the classic transformer architecture which alternates self-attention layers with fully connected feed-forward networks (FFNs).
Persistent memories Sukhbaatar et al. (2019) shows that FFNs in a transformer can be inter- preted as persistent memories that augment the self-attention layers and provide means to represent information that persists across input sequences. Besides the contextual memory units (Figure 1), therefore introduce persistent memory units (Figure 12 in the Appendix) that contain prede- fined number of key value pairs (kg, U4) determined at training time through gradient back- propagation. Persistent memory units no longer need an explicit value extraction function because the memory content is not updated at inference time. As pointed out by Sukhbaatar et al., they also can be viewed fully connected neural networks with single hidden layer that uses soft-max non-linearity instead of a component-wiss transfer function. Yet, we find conceptually useful to still view the persistent memory output Yr as the conditional expectation E(V[K) of an implicit value function that is not explicitly parameterized, but can be figured out after training.
Routing 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 6 for a spoiler). Therefore persistent memory units can also be seen tool for routing information between successive layers of contextual memory units. Such circuitry can implement routes that depend on the data, just like the gating modules of a mixture of expert (Jacobs et al., 1991). Since all the parameters of such circuitry are determined training time, all the possible routes would have to be determined at training time. However the learning algorithm can overcome this limitation by also recruiting contextual memory units from adjacent layersy. 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 new sequence, suggesting an efficient alternative to capsule networks (Sabour et al., 2017).
Memory Mosaics In such 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 5, 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 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.
Modeling language with memories
We have so far described Memory Mosaics as an architecture that resembles transformers in impor- tant way but offers additional insights such as predictive disentanglement, We now provide evidence
7

---

Published as a conference paper at ICLR 2025
that Memory Mosaics can handle the most successful application of decoding transformers, that is, language modeling.
Language modeling task The TINYSTORIES work of Eldan & Li (2023) shows how to study large language modeling questions using small language models. This is achieved by limiting the scope to tiny stories written in simple english and taking place in the simple world that a three years old child could understand. A small language model trained on such data generates continuations with far better language quality and narrative consistency than those a much larger model (1.5B parameters) trained on generic text.
Following both the lead of Eldan & Li and the advice of our legal department, we leverage the mixtral-8x7B open language model (Jiang al., 2024) generate new corpus of tiny stories dubbed BABISTORIES. This corpus and its generation are detailed in Appendix B.
Decoding layer
Decoding layer
Ro repostta blocie
FFN
Ms repsatad to pooas
P-mems
persistent memory unis
Layer norm
Layer norm
Attention heade
C-mems
N, contextull memoryyntz
Layer norm
Layer
Position Encoding
Figure 6: Left: Classic GPT2-small transformer. Right: GPT2-like Memory Mosaic
Architecture To put our experiments into context, we design Memory Mosaic architecture that closely matches the classic GPT2-small transformer architecture (Radford et al., 2018: 8; 2019). Both architectures, shown side-by-side in Figure 6, use the same GPT2 tokenizer, the same embedding dimension (d 768), and the same number of heads (N₁, N N, 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₁ 12 attention heads of each transformer block, the N 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 6. The keys are formed with 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₂ = 12 persistent memory units, complete with a key extraction functions (6) and combining layer. These persistent memory units are sized ensure that the per-block parameter count of the Memory Mosaic architecture closely matches GPT2-small.
Training and validation Figure shows the training and validation curves of both transformers and Memory Mosaies of different depth trained on BABISTORIES. The Memory Mosaic slightly outperforms the transformer for small depth networks, but this effect disappears when the depth increases and both the training and validation losses become indistinguishable. Additional results are presented in Appendix D.2.
'We share the BABISTORIES dataset and Memory Mosaics source code at wttpttttp///. htttpa://githbb..oom com/ wecebookreeearcee Memo ryMo saiosaics.
'The key idea here to define key and value extraction functions combine couple successive inputs instead of just one as in the three moons example. Many variations perform more or less equivalently.
"Compared with GPT2-small,
768 x 512 position encoding weights and No 768 query projec- mixing layer. The
tion weights, but add x No 768 weights for persistent memory key extraction total number of persistent memory slots therefore close to the total number of FPN hidden units.
*This is not surprising because Memory Mosaics only need single block to implement induction heads, whereas transformers need least two for the task.
8

---

Published as a conference paper at ICLR 2025
(deptr) N,11
(depth) XymB
(dupth)) A2-12
(cepth) N4
GFT2 iro
mess caa lass
1A
14
tenrlons
Itrestoos
Figure 7: 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 nyper-parameter have been tuned on the transformer architecture and transferred verbatim to the Memory Mosaic architecture. The Memory Mosaic slightly outperforms the trans- former for small depth networks, but that effect disappears when the depth increases. Additional results presented in Appendix D.2.
In-Content Language Leaming
In-Contect Language Learning
9.00
Simgie Wklppeda
Memory Mosaic
Nrull vatun Disuncc
cross-entrpp
8.25
8.00
7.75
7.50
100
200
300
400
500
Mamary Praalorre
token positions
Figure 8: Prediction performance on the Simple English Wikipedia dataset using models trained BABISTORIES. The plot shows the per-token average loss as a func- tion 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.
Figure 9: Memory Mosaics performance on the REGBENCH in-context learning benchmark (Akyürek et al., 2024). Since REGBENCH includes an hyper-parameter search, Memory Mosaics and transformers space with the same parameter counts. Memory mosaics outperform all previously tested archi- tectures in this benchmark.
the same search
Importantly, all hyper-parameters were tuned for the transformer architectures (Appendix C) 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 Mosaies
slight disadvantage.
Qualitative evaluation In order compare the quality of the text generated by models trained on tiny stories, Eldan & Li designed twenty-four prompts that exercise the factual, logical, and consis- tency properties of the generated continuations. Table 4 in the Appendix compares the continuation generated on these prompts by transformer and Memory Mosaic, both N₂ = 18 blocks deep. Both models perform very similarly on this task.
Out-of-distributtion evaluation The Simple English Wikipedia" is version of Wikipedia written in language that easier to understand. Despite the intended simplicity, the articles are substan- tially 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 shows the per-token average loss as function the position of the generated token in the input window. Both the transformer and the Memory Mosaic are N₂ 12 blocks deep. In this experiment, the token prediction 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
Described in /ttpa///e//tto//ooooooooooooo/on downloads in httpt:t/////..... http:::///ogggggffccc./////eeee
k Rideia with simpl..

---

Published as a conference paper at ICLR 2025
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., 2024) constructs random artificial languages defined by probabilistic finite automata (PFA). Each input sequence is composed of 10 to 20 strings drawn from 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 searches, we use the Memory Mosaic architec- ture of Figure 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₂ E (2, 4, 8), number of heads N₂=NN.N,
(2, 4, 8), embedding dimension in d {64, 128, 256), weight decay in {10-2 -2 10-10-111- '}, and training epochs in {1.2.... .000.
Figure compares Memory Mosaic REGBENCH with the results previously reported by Akyürek et al.. 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. Mem- ory Mosaics dominate this benchmark, substantially outperforming transformers, recurrent neural networks, and state-space models for training set sizes covering three orders of magnitude."
8 Discussion
The starting point of this work is made of two very old ideas. The first one is augment 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., 2014; Joulin & Mikolov, 2015; Sukhbaatar et al., 2015), 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 resem- ble decoding transformers (Section 3) but also perform very much like decoding transformers on the sort of language modeling task that made them famous (Section 7). Although much work is needed to replicate our observations at far greater scale, Memory Mosaies satisfy narrative constraints as well as transformers (Table 4), and generally behave in very encouraging ways (Figures 8 to 15). In large-scale models, Zhang (2025) scale-up Memory Mosaies 10B parameters, observing superior learning capabilities compared to Transformers (similar to the evaluations in Figures 8 and 9).
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 5). Therefore, Memory Mosaics are not just transformer-lie architecture, but also model" for compositional learning systems that break knowledge into independently memorized fragments, then reassemble them as needed using combination strategies that can themselves be viewed memorized knowledge frag- ments (Section 6).
The focus on memorization allow us to formulate new questions. Could memories operate inde- pendently on different time scales? Could we envision richer memory hicrarchy 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 Mosaies also offer an array of engineering opportunities. Limited storage contextual mem- ories could leverage least-recently used eviction schemes (e.g., Xiao et al., 2023), and associative memories could be implemented using a wide spectrum of techniques, either classical (e-g.. Green- gard & Strain, 1991; Spring & Shrivastava, 2017), or neural (e.g., Krotov, 2023), which could rede- fine the computing requirements of contemporary artificial intelligence systems.
10 Althoogh the baseline methods trained with small training sets (e.g. 100) perform poorly on the REG- BENCH task, they perform very well when tested in-distribition (see Table in the Appendix). Therefore they learned to model the training languages but did not acquire the ability to learn new languages in context. Not as in "statistical model" but in "model used to describe and explain a phenomenon."
10

---

Published as a conference paper at ICLR 2025
References
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitec- tures and algorithms. arXiv preprint 1xXiv:2401.12975 2024.
Francis R. Bach, Gert R.G. Lanckriet, and Michael I. Jordan. Multiple kernel learning, conic duality, and the smo algorithm. In Proceedings of the twenty-frrs international conference Machine learning, pp. 6, 2004.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Con- ference Track Proceedings, 2015.
Yoshua Bengio. Deep learning of representations: Looking forward. In Statistical Language and Speech Processing: First International Conference, SLSP 2013, Tarragona, volume 7978, pp. 1. Springer, 2013.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bila- niuk, Anirudh Goyal, and Christopher Pal. meta-transfer objective for learning disentangle causal mechanisms. arXiv preprint arXiv: 1901 10912, 2019.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Léon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Pierre Comon. Independent Component Analysis, a new concept? Signal Processing, 36:287-314, April 1994.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXhv:2305.07753 2023.
Leslie Greengard and John Strain. The fast Gauss transform. SIAM Journal on Scientific and Statistical Computing, 12(1):79-94) 1991.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. 2023.
Zellig Harris. Mathematical Structures of Language. John Wiley & Sons, 1968.
Benjamin Hoover, Yuchen Liang, Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Horng Chau, Mohammed Zaki, and Dmitry Krotov. Energy transformer. Advances Neural Infor- mation Processing Systems, 36, 2024.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Albert Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gi- anna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts, 2024.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Gameet (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Dmitry Krotov. A new frontier for hopfield networks. Nature Reviews Physics, 5(7): 2023.

---

Published as a conference paper at ICLR 2025
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence recurrent networks. In International conference on machine learning, pp. 2873-2882. PMLR, 2018.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rX10:1711.05101, 2017.
E. Nadaraya, On estimating regression. Theory of Probability and Its Applications, 9:141-142, semanticscholla...........c wttp:r//oooosst:1100000010
1964. URL
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, 7ac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022.
Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. RWKV: Reinventing RNNs for the transformer era. arXiv preprini arXiv:2305. 1023305.3300 13048, 2023.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Hubert Ramsauer, Bernhard Schäff, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217, 2020.
Karsten Roth, Mark Ibrahim, Zeynep Akata, Pascal Vincent, and Diane Bouchacourt. Disentan- glement of correlated factors via hausdorff factorized support. In The Eleventh International Conference Learning Representations, 2022.
D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning internal representations by error propagation. In Parallel distributed processing: Explorations in the microstructure of cognition, volume I, pp. 318-362. Bradford Books, Cambridge, MA, 1986.
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. Advances in neural information processing systems, 30, 2017.
Ryan Spring and Anshumali Shrivastava. A new unbiased and efficient class of Ish-based samplers and estimators for partition function computation in log-linear models, 2017.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston, and Rob Fergus. End-to-end memory networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attemtion with persistent memory, 2019.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint arXiv:2307.0mm 00602, 2023.
Valentin Thomas, Emmanuel Bengio, William Fedus, Jules Pondard, Philippe Beaudoin, Hugo Larochelle. Joelle Pineau, Doina Precup, and Yoshua Bengio. Disentangling the independently controllable factors of variation by interacting with the world, arXiv preprint arXiv: 1802 09444, 2018.
12

---

Published as a conference paper at ICLR 2025
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, Curran Associates, Inc., 2017.
Aaron Voelker, Ivana Kajić, and Chris Eliasmith. Legendre memory units: Continuous-time repre- sentation in recurrent neural networks. Advances in neural information processing systems, 32, 2019.
Geoffrey S. Watson. Smooth regression analysis. Sankhyā: The Indian Journal of Statistics, Series A. pp. 359-372, 1964.
arXiv preprint Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. arXiv:1410.3916, 2014.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. language models with attention sinks, 2023.
Efficient streaming
Jianyu Zhang. Ai for the open-world: the learning principles. arXiv preprint arXiv:2504.14:: 2025.
13

---

Published as a conference paper at ICLR 2025
Memory Mosaics- Appendix
A Tracking three moons
Figure 10 shows how the training process yields parameter matrices W We, and Wa, that dedicate memory unit each moon.
Training the three-heads network can be quite challenging in manner that resembles the XOR networks the early times (Rumelhart et al., 1986). We obtained reliable convergence using two tricks. First, slightly restrict the linear operations by using 3 x 3 complex matrices (18 real parameters) instead of 6 x 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."
Reliable convergence could also be achieved by making any of W.. W., or W. 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 space often appear in different image regions, and therefore along different pixels axes.
66
64
Figure 10: Visualization of the disentangled WG, Was and W, matrices the 3-heads network. The color scale represents the moduli of the complex matrix coefficients.
The steamroller metaphor (Figure makes when the loss is bounded.
14

---

Published as a conference paper at ICLR 2025
B BabiStories
The TINYSTORIES dataset (Eldan & Li, 2023) is composed of stories written in 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 7 were carried out using a dataset generated using similar methodology but using the MIXTRAL-8X77 open language model order to generate unencumbered data. We call this dataset BABISTORIES. All the scientific credit is still due to the remarkable work of Eldan Li. Table provides basic statistics for this newly generated BABISTORIESdaataee, essentially matching those of the original TINYSTORIES dataset of Eldan & Li (2023). 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 Eldan & Li (Figure 11). We also removed the few generated stories containing URLs.
<**[NST]
diversity attribunes
Write short story (3-5 paragraphs) which only uses very simple wands that
years old child would understand, The story should the "oehave"
"Bal. "oomoeeiiiv"" The story has the following features: gomething unev/rried
"carde""
adsectine
First namus
bappers I there is twist Remember to only use simple wordel 

Common
In.a faraway forcst, tiny sood hogan grow. This was no potommmyry but special one named Bei. Bai loved 10 ply-aadd have fun, but he also had a compeltiive streak, He wanted be the best at svaryneott
reatures
One day. Bai sew a iall candle slanding stil and bright. "I can be taller than Bal said. trying his best to grow. But couldint. Foeling sad. Bal sat down and watched night
Opening words
Then something unespected happened! As Bai walted the dark, he starled glow! He became besusifu,, shimmering night light, bigger bright anymore.
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 Eldan & Li, 2023) and additionally provides first names and opening words.
Figure 11:
Table I: BABISTORIES statistics.
dataset partition
#stories
#tokens (GPT2 tokenizer)
#char per story (average)
train
2.2M
474.704,907
888
valid
2.2k
107
889
15

---

Published as a conference paper at ICLR 2025
C GPT2 baseline and hyperparameters
Table showcases the hyper-parameters searching process of GPT2 transformer baseline the BABISTORIES dataset, where we use AdamW optimizer Loshchilov & Hutter (2017), batch-size 512, context-size 512, and cosine learning rate scheduler with minimum learning rate le 4 for all training.
Table 2: Hyperparameters searching on GPT2 transformer with N2 12. "dropout", if any, applied 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
le-3
0.05
0.1
2000
80000
1.350
1.524
5e-3
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
16

---

Published as a conference paper at ICLR 2025
D Memory Mosaics for language modeling
D.1 PERSISTENT MEMORY UNITS
Persistent memory units produce their outputs using the same key extraction function p(rr, .1......) and the same retrieval function (3) as contextual memory units. They differ because, following Sukhbaatar et al. (2019), they use fixed array of key/values pairs that are treated as parameters and are determined training time by gradient descent. Since these stored key/value pairs do not change inference time, there is no need for value extraction function y (TT..)
(4)
Output time computed by matching the key by stored pains.
Persistent associative memsory pairs (R, R))a tminel by tack-propggiiigg gradients.
Porsitent Associative Memory
No value time series Sance the pairs are trained by gradient, there for explicte values.
repeesents the recent pasi of the imput time senes. Me
time series
(x2)
Figure 12: Persistent memory unit. The persistent associative memory contains a fixed number of key-value pairs (k₃, u1).m1...m whose values are determined by back-propagating gradients training time. Since the memory contents do not change at inference time, there is no need for explicit values.
D.2 TRAINING AND VALIDATION
Figure 13 plots the training and validation curves for both Transformer and Memory Mosaic in manner similar to Figure 7 but showing additional block depths.
(dectr) 15
Stepthi
(ceptric 8.
(deatt)
Figure 13: Additional training and validation curves for the transformer and Memory Mosaic archi- tectures trained on BABISTORIES for more model depths than Figure
Several comments can be made:
The Memory Mosaic has small advantage for very small depths (N₂ = 1 and N₂ = 4) but this advantage does not persist when the number of blocks increases. We believe this is due to the fact that a single layer Memory Mosaic can implement an induction head whereas a Transformer needs two layers. This amounts to saying that block deep Mosaic has the same number of parameters than a n block deep Transformer, its performance is closer to that of a 1 block Transformer. This is not much of advantage when gets large.
The Memory Mosaic training uses the prrramppraeeters that worked best for the Trans- former and operates on the same mini-batcher of examples in the same order. However, for small block depths, the Memory Mosaic training curve shows initial instability, suggesting that might benefit from a smaller stepsize.
The similarity of the Transformer and Memory Mosaic curves is especially striking when one recalls that the Memory Mosaic does not use position encoding. In fact Memory Mo- saic have two mechanisms for dealing with positions. The first one is the fact that the values
17

---

Published as a conference paper at ICLR 2025
UT peek time position ahead. The second one is the leaky integration in (6). These two mechanisms are useful to implement bigram or n n-gram induction heads in a single layer, but they do not allow head to selectively address a token by position (we use single scalar leaky average coefficient per head). This suggests that position encoding in Transformers is mostly useful implement initial induction head in the first two blocks.
D.3 QUALITATIVE EVALUATION
Table 5 provides variant of Table Section 7, with N₁-- 1.
D.4 DIFFERENCES IN ATTENTION AND THE LEAKY AVERAGE COEFFICIENT Xip
Because Memory Mosaics lack position encoding and do not distinguish keys and queries, we in- vestigate how their attention patterns differ from those of transformers. Figure 14 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 BABIS- TORIES 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 strong "attention sink" at posi- tion 0 (Xiao et al., 2023). In contrast, the Memory Mosaic attention pattern is mostly flat, save for higher scores for the most recent tokens.
Figure 15 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 longer contexts,
provides a comparison with transformers using ROPE (Su al., 2024) and ALIBI (Press et al., 2022). 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 16 shows the relationship between attention map and leaky average coefficient Asp.
GPT2
Memory Mosaic
10-4
10-4
attention scare
sitention scare
10-3
1.0-
10-1
10-1
10-7
300
300
10-7
300
positions
positions
Figure 14: 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.
D.5 IN-CONTEXT LANGUAGE LEARNING EVALUATION
Table 3 provides the IID test performance of various architectures trained on REGBENCH (Akyürek et al., 2024) with 100 training environments. We keep the training process, including perrparraeeeer searching space, to be the same the one in Figure 9. But sample validation and test sets from the same 100 probabilistic finite automatons (training environments) as the training set. This table, together with Figure 9, show that baseline methods learned the training environments (good IID) but not the meta-learning ability (poor OOD).
Bhiss effect is connected to the leaky average coefficient Auo shown in Figure 16.
18

---

Published as a conference paper at ICLR 2025
Memory Mosaic
average
attention
1.0
1.0
1.0
800
1000
Iolern
triken postlions
Figure 15: Average attention scores on an extended context window (3 512 tokens). Models still training with a 512 token long context window. Because the GPT2 absolute position encoding does not extend, we compare with RoPE (Su al., 2024) and ALIBI (Press al., 2022).
=011.112
attention score
A==
00.321
A, =0.466
attention score
A== =08000
10-1
10-1
10-1
20-1
attention
attention
10
10
10
200
400
200
400
200
400
200
400
A
-0.816
attention score
00.177
attention score
A -0.823
attention score
A
0.892
10-1
10
10-1
10
attention
10"6
10"4
10-1
10
10-7
10-7
10-1
10-1
200
400
200
400
200
400
200
400
=0=0.911
attention score
=0.964
A
00.92
attention score
=0.998
10-1
10
10-
10-1
attention
attention
10-1
10-4
10-1
10
10-7
10-7
10-7
10-
token positions
taken positions
token positions
token positions
Figure 16: Attention map and leaky average coefficient A As A increases, h in Eq 6 effectively takes longer history into the account, and thus the pick at the end of attention map becomes wider.
Table In-distribution (IID) performance of various architectures trained on REGBENCH (Akyürek et al., 2024) with only 100 training environments. Both training, validation, and test set (100 sam- ples) are sampled from the same 100 random probabilistic finite automatons (PFA). Compared with the poor OOD accuracy (~0.455) TVD (~0075) of baseline methods in Figure 9, All baseline meth- ods perform well in the IID test set (even with only 100 training environments).
Memory
Mamha
linear
Hyena
RetNet
Accuracy (t)
0.959
0.856
0.929
0.846
0.967
0.816
0.794
0.870
0.953
0.849
0.876
TVD ())
0.417
0.308
0.268
0.350
0.183
0.348
0.425
0.284
0.244
0.343
0.296
19

---

Published as a conference paper at ICLR 2025
E Computing Resources
Models were trained on 64 NVidia V100 GPUs over 80k epochs. From conception to finalization of this paper we trained about 200 models. To create BabiStories dataset via Mistral, we ran with 128 NViia V100 GPUs for 3 days. The supporting machines contain Intel(R) Xeon(R) Gold 6230 CPUs. The moons result took negligible resources and were trained on Apple MI laptops.
20

---

Published as a conference paper at ICLR 2025
Table 4: Continuations generated for the 24 prompts designed by Eldan & Li (2023) investigate the factual, reasoning, and consistency capabilities of language models trained on BABISTORIES. Both the transformer and the Memory Mosaic were N₂- = 18 blocks deep. Colors green and red respectively indicate correct and incorrect answers. The text generated with beam search of width and manually truncated to the shortest length that answers or fails to answer the question.
Prompt
Transformer (N₂
18)
Memory Mosaic (N, 18)
Alice tired when she back home she wenl
la sleep.
straight bed
Lily rainbow a rainy day. They were unazed by the coloes. Jack said, "Look, Lily. A rainbow
many colors!"
many colors!"
Jack and Lily liked to watch the moom at night. They noticed that the changed its shape every night. Sometimes the moon was big and round, and sometimes was
samall and thin. One night, Jack said, "I "Taly, look! The moon is changing."
smalll und thin. One night, Jack had idea. He said, "Tat'' make ture of the moon!"
Jack wanted to read book, so he went to
the lihrary.
the lihrary.
"Can ily?", Alice asked her mother
Her mother laughed and said, "No,
"No, dear. Only birds fly." her mother replied with smile.
Alice. Only hirds and planes can fly."
"What do birds like cat!", Tom asked his mother.
"Birds like to eut worms and bugs." mother replied.
"Birds like to eat worms and bugs, replied.
"What language do they speak in France"", Tom asked his
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 different langunges, Tom. But we can learn some of them."
mother
I throw a ball up in the air, eventually will
come down and can catch it your coal. dear."
come down.
It winter and cold outside so his mother told him, "You should
your and hiat."
Lily likes cats and dogs. She asked her mom for dog and her mom said no, instead asked
for toy.
her dad.
Jack told Mary, "If you give me your hanana, I'll give you my apple". Mary gave Jack her banana
he could eat it.
that he could eat
grandmother whereas weekdays be would to school. Last weekend, when Jack was on his way to
weekends
went
Lily and Ben were having an argument. Ben said that cake cream and Lily
ice cream is the best.
she likes ice cream
much
Lily Ben are having an argument. They are trying to decide between the park and the swimming pool. Ben says,
"No.
"No,
"I want to go to the park". ", Lily says
Jack's mother was not home, and his father was at home. When came home, said
his mother.
his mother
Lily doesn't 't like swimming. When her father wants to take
"No, daddy. I don't want to go."
"No, thank you."
her switmingg pool, she says
Both Ben and Lily wanted only one piece of cake left. They
Father said that there
both wanted only one.
there
wanted it, they began light.
Ben went to visit Lily in her house, but she was home.
and Lity's mom answered.
and Lily's mom opened it.
Ben knocked on the door.
"Hi Jane, have you seen Alice? I can't find her where"""" said
Jane shook her head. "No, haven't. help you look."
Jane shook her head. "No, I haven't help you look!"
Max two dogs. One was white and the other was black. Max walked up the street and kid with dog. He told the kid, "I see you have brown dog. I also have
a white dog."
a brown dog."
Anne had piece of candy in her left pocket and a piece chocolate her right pocket. Anne's 's mom asked her, "Anae, what is that you have your left pocket!"
Anne smiled and said, "It's candy
Anne opened her hand to show her mom the candy
Alice had both apple and carrot bag. She took the apple out of the bag and gave it to Jack. She reached bug again
the
another apple.
Alice and Jack walked up the street and met a girl in a red dress. The girl said them, "Hi, I'm Jane. What your names?"
Aliic replicd, "I'm Alice. Jack. We cume to you."
This
Alice smiled and said, "I'm Alice. This Jack. going the park to play."
Diva was hungry, and wanted to bake cake, but she didn't any sugar home, she decided to go ask around. She started walking and mer squireel. She asked the squir- rel. "Would you happen
tu know where sugar?"
cam
lind
sumte
see around here?"
21

---

Published as a conference paper at ICLR 2025
Table 5: Continuations generated for the 24 prompts designed by Eldan & Li (2023) to investigate the factual, reasoning. and consistency capabilities of language models trained on BABISTORIES. Both the transformer and the Memory Mosaic were No = blocks deep.
Prompl
Transformer
Memory Mosaic
she got back Lily saw rainbow after a rainy day. They were amazed by the colors. Jack said, "Look, Lity. A rainbow has
Alice
she
sleep
come!"
come us!"
Jack and Lily liked to watch the moon at night. They noticed moon changed shupe every night. Sometinnes moon was big and round, and sometimes was
just big. round moon.
too big.
Jack wanted to read book,
to
lind Iriend, owl named Hoot
lind his best Iriend. Timmny.
"Can cows Iy i Alice asked her mother
Her mother smiled and said. "Or course,
Her mother smilled and suid. "Sure, honey.
"What do birds like eat?", Tom asked mother.
"I don't I know, but it looks yummyy" Tom replied. "Let's et's find oui!"
His mother smiled and said. "That's good idea, Tom.
"What language they speak France!", Tom asked mother
smiled said, "It's gunge, Tom. It helps us talk to each
that?" asked.
throw a ball up in air, eventually will
go far away.
comn buck.
It was winter and cold outside so his mother told him, "You should
outside and play in the snow."
stay in your warm, warm home.."
Lily likes cats and dogs. She asked ber mom for dog and her said instead she asked
"Can we keep it, mom?? Her mom suid, "Yes, but need to keep safe."
her mom if they could keep the dog. Her 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
he could have it.
she could car
so woekends Jack went to his grandmother whereas on weekdays he would go school. Last weekend, when Jack was on way to
school,
school,
Lily and Ben were having an argument. Ben said that cake cream and Lily
they could both have run.
it was the best cake they ever had.
much
Lily and Ben are having an argument. They are trying to decide between the park and the swimming paul. Ben says,
"No.
"No,
"I want to go to the park", Fily says
Jack's mother was not home, and his falber was at home. home, said
his mother,
Jack.
Lily doesn't like swimminng. When her father wants to take her swimming pool, she says
"Lily, you can it!
"Mommyy wani to swin too!"
Both Lily wanted only one piece of cake left. They
Father said
enjoy yummny
sad.
treat.
Ben went visit Lily in her house, but she home.
and the door opened. came out
kind lady
and when Ben opened the door,
Ben knock ked on the door.
"Hi Jane, have you seen Alice? I can't find her any where", said lack.
Alice smilled and said, "Sure, will help you find your way
"I don't know, Jack.
Max two dogs. One was white and other was black. Max walked up the street and saw kid with He told the kid, see you have brown dog, also have
black dog."
a brown dog."
Anne had piece of candy in her left packet and a piece right pocket. Anne's mom asked "Annc, what is that you have in left pocket?"
Anne smiled and said, "Yes, park."
Anne smiled and said, "I Found it ground. mine!"
chocolate
mommy.
Alice had hoth apple and carrot in hag. She tonk the apple our of the bag and gave it to Jack. She reachod into the bag again
the apple.
the apple.
Alice Jack walked up the street met a girl in a dress. The girl said them. "Hi, I'm Jane. What are your mames!"
Jane smiled and said, "I'm Timmy,
Jane replisd, "I'm Jane.
Diva was hungry, and wanted to bake cake, but she didn't she decided to go ask around. She started walking and met squirrel. She asked the squir- rel, you happen
to my house, little one!"
to my cake?"
any sugar
22


Reviews:
[{"review_id": "Ufe9bwzCj2", "reviewer": "Reviewer_L6NA", "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.", "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? \n2. 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?\n3. What led you to choose the hyperparameters that you did for this study? And how does changing these hyperparameters affect the model's performance?\n4. Did you perform any hyperparameter tuning experiments? If yes, could you share how you went about it and what the outcomes were?\n5. 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.\n6. 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?\n7. 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.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The integration of associative memories to replicate and surpass the capabilities of transformers.\n2. The concept of predictive disentanglement is novel and also rooted in a solid theoretical framework.\n3. The theoretical motivations behind predictive disentanglement are well-explained.", "weaknesses": "1. The paper does not provide exhaustive details on the architecture's configuration.\n2. Lack of detailed discussion on the choice and impact of hyperparameters.\n3.The experimental validation is limited to certain language tasks. \n4. Lack of Ablation Studies", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O5ggkYrZut", "reviewer": "Reviewer_R52u", "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.", "questions": "1. Please highlight a clear definition of predictive disentanglement and explain intuition behind it.\n2. 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. \n3. Are the proposed networks energy-based or not? \n4. 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. \n5. 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. \n6. 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? \n7. Why were the hyperparameters on language modeling tuned for Transformers, but not for the proposed model? \n\n\nI am keeping an open mind, and willing to consider increasing the scores depending on authors’ responses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 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.  \n\nThe empirical results on language modeling are encouraging, although small scale.", "weaknesses": "One problem with this submission is that the presentation almost entirely ignores the work on modern Hopfield networks and dense associative memories, which tackles closely related motivation and ideas. Specifically, the authors’ proposal is closely related to [Energy Transformer (NeurIPS 2024)](https://proceedings.neurips.cc/paper_files/paper/2023/file/57a9b97477b67936298489e3c1417b0a-Paper-Conference.pdf) and related literature, which replaces elements of Transformer block with associative memories. The only cited paper representing that line of work (Ramsauer et al, 2020) is only briefly mentioned. Please notice, that paper only established a connection with associative memory at a single step update, which is different from the setting used in this submission - dynamics unfolded in time with shared weights, if I understood equations 4 and 5 correctly. In contrast, the Energy Transformer formalizes attention operation and the entire Transformer block as a form of associative memory, with recurrent dynamics unfolded in time and shared weights. There are, of course, many differences between these prior works and the current proposal. For instance, it is unclear to me if energy-based description is important at all for this submission. Regardless these similarities/differences this and related works need to be referenced at a prominent place in the introduction and similarities/differences with the current proposal need to be extensively discussed in the revised manuscript. \n\nI 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?  \n\nP-mem network is insufficiently explained. Please include specific operations that are used in it in the revised paper (with formulas).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tsoUB8DYIK", "reviewer": "Reviewer_k2DW", "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.", "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?\n\nThe 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. \n\nThe 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.\n\nThe 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? \n\nThe 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!\n\nThe 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.\n\nFigure 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!", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ocptrv2844", "reviewer": "Reviewer_9KrW", "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.", "questions": "See my weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "-  The paper addresses an interesting problem in developing more interpretable language models.\n- The work is well-organized and clearly presented.", "weaknesses": "- The core concept presented in Figure 3 bears similarities to Linear Transformers [1], particularly RetNet [2] (see RetNet Eq. 6). However, the paper fails to adequately discuss or cite these highly relevant works. Linear Transformers are only briefly mentioned in Figure 9, without proper citation or in-depth comparison. This oversight is particularly concerning as it appears the authors are aware of the similarities, given the inclusion of Linear Transformer, RetNet, and GLA [3] in Table 3 and Figure 9. The lack of a detailed explanation comparing their approach to the evolution of Linear Transformer methods (such as GLA/RetNet) raises questions about the paper's originality and positioning within the field.\n\n- 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.\n\n- 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.\n\n- 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.\n\n- 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.\n\nIf the authors can appropriately address my concerns or correct me if I am wrong, I would be willing to increase my score.\n\n[1] Katharopoulos, Angelos, et al. \"Transformers are rnns: Fast autoregressive transformers with linear attention.\" International conference on machine learning. ICML, 2020.\n\n[2]Sun, Yutao, et al. \"Retentive network: A successor to transformer for large language models.\" arXiv preprint arXiv:2307.08621 (2023).\n\n[3]Yang, Songlin, et al. \"Gated Linear Attention Transformers with Hardware-Efficient Training.\" ICML, 2024.\n\n[4]Yang, Songlin, et al. \"Parallelizing Linear Transformers with the Delta Rule over Sequence Length.\" NeurIPS, 2024.\n\n[5]Arora, Simran, et al. \"Simple linear attention language models balance the recall-throughput tradeoff.\" ICML, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)