Memory Mosaics at scale
Authors: Jianyu Zhang
Leon Bottou
Venue: NeurIPS 2025 oral ()
Keywords: In-context learning
memory mosaics
Primary area: deep_learning

Abstract: Memory Mosaics, networks of associative memories, have demonstrated appealing compositional and in-context learning capabilities on medium-scale networks (GPT-2 scale) and synthetic small datasets. This work shows that these favorable properties remain when we scale memory mosaics to large language model sizes (llama-8B scale) and real-world datasets. 

To this end, we scale memory mosaics to 10B size, we train them on one trillion tokens, we introduce a couple architectural modifications (*memory mosaics v2*), we assess their capabilities across three evaluation dimensions: training-knowledge storage, new-knowledge storage, and in-context learning. 

Throughout the evaluation, memory mosaics v2 match transformers on the learning of training knowledge (first dimension) and significantly outperforms transformers on carrying out new tasks at inference time (second and third dimensions). These improvements cannot be easily replicated by simply increasing the training data for transformers. A memory mosaics v2 trained on one trillion tokens still perform better on these tasks than a transformer trained on eight trillion tokens.

Memory Mosaies at scale
Jianyu Zhang
Léon Bottou
New York University, New York FAIR, Meta Inc., New York
FAIR, Meta Inc., New York New York University, New York
Abstract
Memory Mosaics [Zhang et al., 2025], networks of associative memories, have demonstrated appealing compositional and in-context learning capabilities on medium-scale networks (GPT-2 scale) and synthetic small datasets. This work shows that these favorable properties remain when we scale memory mosaics large language model sizes (Ilama-8B scale) and real-world datasets.
To this end, we scale memory mosaics to 10B size, we train them on one trillion tokens, we introduce a couple architectural modifications ("memory mosaics v2"), assess their capabilities across three evaluation dimensions: training- training-knowlede knowledge storage, new- new-knowledge storage, and in-context learning.
Throughout the evaluation, memory mosaics v2 match transformers on the learning of training knowledge (first dimension) and significantly outperforms transformers on carrying out new tasks at inference time (second and third dimensions). These improvements cannot be easily replicated by simply increasing the training data for transformers. memory mosaics v2 trained on one trillion tokens still perform better on these tasks than a transformer trained eight trillion tokens.
Introduction
In Machine Learning, compositional capabilities and in-contexttour-oo-dissriiiii bilities have been continuously pursued but remain challenging. Early attempts to achieve these goals include pursuing disentanglement via various statistical "independence" [Comon, 1994, Roth et al., 2022], pursuing mut-of-disttibutioonleamii
learning capa-
from the perspective of optimization on multiple environments [Finn al., 2017, Arjovsky et al., 2019, Bengio et al., 2019]. In contrast, transformer- based models demonstrate certain compositional capabilities and early in-context learning abilities. However, we still lack a clear understanding of how current transformers achieve these capabilities, and why earlier models were unable to.
Memory mosaics [Zhang et al., 2025], networks of simple key-value associative memories (without position encoding), offer a comparatively transparent way understand how composition or disen- tanglement occur. Trained and evaluated on medium scall networks and synthetic datasets, memory mosaies reveal promising superior in-context learning abilities. Therefore, we ask "To peruse strong and general new task learning capability, how can we scale memory mosaics to large networks and real-world datasets?" "
The contribution of this work: 1) We successfully scale up memory mosaics to llama-8B scale, using one trillion real-world training tokens. The resulting network is named as Memory Mosaics v2. Compared memory mosaies, memory mosaies v2 made three architectural modifications, including an adaptive bandwidth of associative memory, a gated time-variant key feature extractor, and a 3-level memory design. 2) We propose three evaluation dimensions comprehensively assess model ability (from i.i.d. to o.o.d. scenarios). 3) Our memory mosaics v2 demonstrate superior new-task learning capabilities (with fewer examples and less priori knowledge from human designers).
For clarity, "Memory Mosaies" refers to the version in Zhang et al. [2025], while "Memory Mosaics v2" refers to this scaled-up version.
/ htt.gh/twwooeeeeoeeeeeee///aaaaa
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

This paper is organized as follows. Section 2 introduces background knowledge on associative memories. Section presents the architecture of memory mosaics v2. Then section 4 describes the training process, section 5 evaluates memory mosaics v2 and transformers across three dimensions training (persistent) knowledge storage, new knowledge storage, and in-context learning. Section discusses the failure of replicating memory mosaics v2 by simply increasing the training data (x8 more data) for transformers. Section studies the advantage of memory mosaics v2 in fine-tuning. Finally, section 8 provides discussion and further directions.
2 Background on Associative Memory
General speaking, memory mosaics architecture [Zhang al., 2025] replaces attention blocks in transformers [Vaswani et al., 2017] with associative memories. This section provides the background on associative memories, highlights the connection and differences between "associative memory" in memory mosaics and "attention" in transformers.
Associative Memory Associative memories have a long history in both psychology and computer science, referring relationships between unrelated items. In this work, follow the definition from Zhang et al. [2025], according to which an associative memory is a device that can store key-value pairs (Ra, 01) -.. (kn; vn)) and retrieve values given corresponding key:²
k + f(k:; [(k1, (kis vn)))
v1)
The key-value pairs are stored in set, and thus can be assumed to be permutation invariant. This exchangeability property suggests that we can view an associative memory as a device that estimates conditional probability distribution P(V K)) the basis of the sample (k1, 01) 01)... (kn, Un,) of key- value pairs. The retrieval function is then a conditional expectation over this estimated distribution:
[k; {(h1, 01) ... (kn; vn))) = B(V K = k))
This conditional expectation can be estimated by kernel regression, e.g. Gaussian kernel regression:"
11
f (K:; [(k1, ui) (kins 2n)))
- - U/
(3)
where 3 controls the bandwidth of Guassian kernel. Connection between associative memory and attention closely connected to attention [Bahdanau et al., 2015] when all key vectors k, share the same squared norm. That is, expression (3) becomes:"
Associative memory in Equation is
R
(k; {(k), 01) ... (km, bn)})
Σ
R
U4
1ド=1
Moreover, the size of associative memory (i.e., the number of key- examples) is analogous to the sequence length in attention.
Differences between associative memory and attention The associative memory viewpoint is conceptually simple and transparent. This simplicity contributes several key differences in associa- tive memory (compared with attention), including: 1) L2 normalized key vectors with an explicit bandwidth parameter 3, 2) a symmetric kernel with the same formula for keys as queries, and 3) the absence of explicit position encoding. These differences further contribute to the superior compositional capabilities and in-context learning capabilities in memory mosaics.
3 Memory Mosaics v2
Transformers show an interesting induction head mechanism [Olsson et al., 2022], that is, predict b after sequence [..., a, b, ...,...,. a). This mechanism contributes to their in-context learning ability. According to Bietti et al. [2023], position encoding and asymmetric query-key extractors (e.g.
both keys and values shall be vectors in Rdd
"This Gaussian kernel smoothing not only converges to the true conditional expectation B(K]V) when 72 - oo and 9 = /n [Nadaraya, 1964, Watson, 1964], but also makes it easy to compute gradients.
5 FFor the easy of reading, we reuse "attention score" notion
En xi associative memories.
#

---

4r = WGTT, kr = WEFT) are essential for transformers to achieve this induction head mechanism with at least two layers of attention.
Inspired by studies of induction head mechanism, Memory Mosaies [Zhang et al., 2025] construet associative memories using keys to represent the recent past and values to represent the near future (Figure left):
KT (o(TT, -I:::...) VT (oTTT1+, .T:...)
(5)
This simple designer allows memory mosaies to get ride of explicit position encoding, use the same key as query, perform induction head with only one layer. The resulting memory mosaics also reveal appealing in-context learning capabilities on small synthetic datasets.
Based memory mosaics, this section introduces Memory Mosaics v2, aiming at stronger and more general in-context learning ability on broader real-world tasks (without loss of performance on other common benchmarks). Compared to memory mosaics, memory mosaics v2 incorporates three architecture modifications, including an adaptive bandwidth in associative memory, a gated time-variant key feature extractor, and a 3-level memory design.
3.1 Adaptive bandwidth in Gaussian kernel smoothing
Memory mosaies use one fixed bandwidth parameter 3 for different sizes n of associative memory (Equation 1). It is well known that bandwidth controls the bias-variance trade-off [Hastie et al., 2009] of kernel regression (memory-based) methods. That is, for given distribution, the optimal bandwidth depends on the number of examples (key-value pairs in associative memory). Inspired by the asymptotic Mean Integrated Squared Error kernel bandwidth estimation approach where 1/vB x 1/(++4) [Garefa- Portuué,, 2024], memory mosaics v2 schedule / in Equation 4 as:
3in + +0,
where 3o 2 0, 31 > 0, 1 a > 0 are learnable parameters (Check Appendix Table for reparameterization and initialization details). Le., the more key-value pairs (examples), the smaller bandwidth 1/vB.
3.2 Gated time-varian key feature extractor
Memory mosaics employs a simple time-invariant leaky averaging to extract key features:
hy Norm (kr) with by = ky A by = W XT
(7)
The averaging weights in Equation 7 are fixed and independent of the semantic input x. As a result, semantically similar cases, such as "tom-and-jerry" and "tom-- -and- - -jerry", may receive different key features. Inspired by recurrent-style networks [Peng et al., 2023, Gu and Dao, 2023, Beck et al., 2025], memory mosaies v2 utilize the following gated vareanariant key feature extractor:"
{ kr = gr + A [e e WGrr E R, XT = = -larll E R
by = W IT
KT = Norm kr) with
(8)
where We g: Wx are learnable parameters, the averaging weights Xy E R and the exponential gate gr E R semantically depend on input øT. See Appendix Figure for graphical illustrations.
For key feature extractor, memory mosaics v2 reuses the same convolutional key extractor as in memory mosaics:
VT ag..Nrr (ūr) with Py 0r (1 -2) 0741
by W.IT, IT,
where Y, Oig. R and We learnable parameters.
3.3 3-level memory
Transformer architecture [Vaswani et al., 2017] consists of attention blocks and feedforward neural network blocks. The former handles local contextual information from an input sequence, while the latter stores global persistent information shared by different training sequences. Memory mosaics [Zhang et al., 2025] simplify the attention and the feedforward network in transformer as contextual associative memory and persistent memory, respectively. This simplification reduces the dependence between the "attention score" and the token position, as shown in Figure 1. Compared with transform-
'It worth noting that this work is neither a linearization of attention nor attention efficiency. The recurrent feature extractor in Eq. 7 is used to create keys, while associative memory in Eq. 1 still stores all key-value pairs.
3

---

RoPE
Memory Mosaic
10:1
avtrage over hest
head
10
awer iver head
20
10
10--
10:9
2:00
300
430
300
10-7 300
200
303
400
500
tolknn positions
taleen positions
Figure 1: Average attention scores of the last token attending previous tokens. Left: Transformer with RoPE position encoding, Right: Memory Mosaies [Zhang et al., 2025]. The (averaged) attention scores in transformer heavily depends on token positions (curly curves), while the attention scores in memory mosaics at far tokens (e.g. position 450) are almost invariant positions (flat curves).
ers (Figure left), the attention scores in memory mosaics (Figure right) exhibit structured pattern. That is, attention scores on near-tokens (positions) heavily depend on positions, while attention scores on far-tokens are almost invariant to token positions. Inspired by this experimental discovery, memory mosaies v2 replace each contextual associative memory in memory mosaics with two associative memories, short-term memory and long-term memory, using distinct parameters (as in Figure 2).
Short-term memory The short-term memory at position & only stores key-value pairs of near- tokens, ranging from h + to implementing Eq. (4) as:
e CBATA Ky i E
f (k: ((kg-14,,
(k=-1,0(-1)3)
E
B
Vt
(10)
Long-term memory In contrast, the long-term memory skips near tokens and only stores key-value pairs before position - m, implementing Eq. (4) as fk:; {(k1, 01) ... (kt-m vt })). By setting m < h, memory mosaics v2 create an overlap between long-term and short-term memory, resulting soft boundary between these two memories. Eventually, the outputs of many long-term memories and short-term memories are concatenated together, following by a linear projection W..
Persistent Memory Memory mosaics v2 implements persistent memory using dense two-layer neural networks with SwiGLU activation [Shazeer, 2020] due to computational efficiency concerns."
tradio without looking forward into (Pu) blut momeryy
Decocing layer
mik of copelsppaaeeeeoo.ocoo spanenetooos
Associativa Memory rstreveapaat valus paira n.]
memories
Inplencrtations
Pemsewnt momories
Assacistive memory
No prppoigg Procce
normalize
Key les
(k.)
(ty)
esents future the in Lime series. WWrro-ppooo
Reprnserchhe wetire
Feoture Extractor
lingaamm memorte
ahoot-llem memer es
kangtarm memoriee so
ahars-aam memories (estow
rummakew
(xx)
Inputtimc serie
Masi
Figure 2: Left: Memory Unit. Right: Memory Mosaics v2 architecture
Training
We train two memory mosaics v2 of difference sizes (small/large). Memory mosaics v2 small (Ilama-1.5 1.5B scale) contains 24 layers, 2048 hidden dimensions, and 16 heads, trained on 200 billion tokens of a diverse datamix. Memory mosaics v2 large (llama-8p scale) increases the number of layers to 32, hidden dimensions to 4096, and the number of heads to 32, trained on I trillion tokens of the same datamix. Both models are trained on 4,096 context length, followed by a fine-tuning process 768 context length. Other training details are provided in Appendix C.
Note that k v, and B in long-term and short- term memory are constructed with distinct parameters. AA two-layers feed-forwar forward network and key-value associative memory are interchangeable as shown in Sukhbaatar et al. [2019].

---

Stochastic long-term memory size During training, memory mosaics v2 samples the long-term memory delay step m from [64, 256], sets the short- memory window size h = 256. At inference, m is set to 64. This stochastic long- -ere-termm memory training setup encourages the allocation of position- invariant signals to long-term memory and rrsition-dependen signals to short-term memory (as shown in Figure 1). The experimental results in Appendix G Table 12 show that this training setup enhances context-length extrapolation ability by more than 15%.
Baseline We train two baseline transformers (small/arge) with the same configurations as their memory mosaics v2 counterparts. Unless otherwise specified, in this work, transformer models use Ilama architecture [Grattafiori et al., 2024] with multi-head attention.
Three evaluation dimensions
The evaluation design provides a means to assess the specific properties of a system. Memory mosaics v2 aims at the ability to learn new tasks with fewer examples and less task-specific priori knowledge [Zhang, 2025]. Thus, to fully assess this capability, this section adopts three evaluation dimensions.
Persistent-knowledge storage and retrieval, the ability of persistent-memont to store and retrieve knowledge of training dataset. This capability prepares knowledge that could be reused in other tasks during inference. We use common language benchmarks to access this aspect. New-knowledge storage and retrieval, the ability to store and retrieve new information of test dataset. It is a prerequisite for "learning" new tasks via memory-based methods. We employ imalti-unrelate-ddocume
storing and question-answerring tasks to evaluate this aspect. In-context Learning, directly evaluates the ability to learn new tasks with fewer examples and less spec-fe priori knowledge. We use multiclass classification to assess this aspect.
5.1 Persistent-knoowlddge storage and retrieval
Table 1 evaluates both memory mosaics v2 and baseline transformers on 19 commonly used language benchmarks, showing that they perform closely these benchmarks. This is expected since both models share the same persistent memory architecture.
Table 1: Memory mosaics v2 and transformers performance on 19 common language benchmarks.
model
comtext longts
doga
casy
are
wino- grande
are challenge
pige
booky
aswag
hell-
sq
sige
liga
gamsk
muiu alt
human eval-
squal
muth
udyy
midie
race
high
race
wanndomeer small
32k
35.2
61.0
60.1
31.4
73.6
6300
19.3
11.7
48.5
26.7
3.11
33.2
32.4
54.7
26.0
1,2
92
52.2
37.4
37.8
memory mosalcs Y2 small
32k
35.0
50.0
55.4
32.9
73.3
627
58.0
11.8
46.6
29.3
3.1
34.7
30.8
59.3
27.3
1.1
944
49.2
39.4
38.0
Inmafooemrr large
12k
77.3
21.
$2.6
80.8
726
79.2
31.
49.3
61.5
32.4
49.0
38.3
763
45.
62.6
45.6
52.21
memory mossics berg
326
78.0
71.
$1.8
80.4
73.1
78.6
50.9
486
62.0
27.4
45.0
78.
47.8
61.6
46.5
How do we know whether these benchmarks access knowledge ability? To answer this question, re-evaluate these benchmarks on memory mosaics v2 but with long-term memory being removed after training. The underlying reason is that if a task solely relies on the information stored in persistent memory and retrieved by short-term memory, removing long-term memory should not significantly affect performance.
ability rather than new-
Table 2 shows that removing long-term memory after training does not degrade the performance of 13 common benchmarks. This suggests that these 13 tasks are almost exclusively based on information stored in persistent memory and retrieved by short-term memory. In contrast, Appendix Table indicates that the other 6 benchmarks perform poorly when long-term memory is removed.
Based on these findings, we use the 13 tasks to evaluate persistent knowledge storage and retrieval capability. The results (Table 1) show that memory mosaics v2 and transformers perform similarly in this evaluation dimension, suggesting that both models are capable of effectively storing and retrieving persistent knowledge.
Table 2: Memory mosaics v2 performance on 13 common language benchmarks. Removing the "long-term memory" after training barely hurt the performance (56.6% vs 56.8%). Flops/token is estimated at context length 256 via tha approach of Casson [2023].
params Hopstoken
obqa
easy
arc
arc grande challenge
wizo-
piça booly aswag
mq
siqn
tggt
gsmsk
mmluu
alt
human
evalt
avg
Transforner lange
8.8B
16.78
15.8
77,3
72,3
52.6
80.8
72.6
79.2 31.9 49.3 61.5
32.4
49.0
38,3
57.1
memory mosnics v2 large
9.98
18.98
45.4
78.0
71.2
51.8
80.4
73.1
78.6
30.9
48.6
62.0
27.4
48.2
43.0
56.8
wilhout long-term memory
memory masaics v2 large
15.6
77.9
80.4
30.8
26.7
46.8
42.2
5

---

Computation and # parameters concerns Table 2 summarizes the size of parameters and compu- tation required for transformers and memory mosaics v2. Interestingly, removing long-terr memory from memory mosaics v2 after training achieves a comparable transformer performance on the 13 persistent-knowllee benchmarks, while using fewer parameters and computations.
5.2 New-knowledge storage and retrieval
The new-knowledge storage and retrieval ability is a prerequisite for learning new tasks via memory- based methods (e.g., Gaussian kernel regression), because the data of new tasks must be adequately "stored" before learning (Note that memory-based methods are lazy methods). To illustrate this point, consider poor goldfish with 7-second memory how can it possibly learn 90-minute movie? Similarly, a model with limited new- -nwwlegee storage ability will struggle to learn information that exceeds its storage (memory) capacity.
question- answering" tasks from the RULER benchmark [Hsieh et al., 2024]. These tasks involve multiple concatenated realistic articles followed by a question related to one of these articles, requiring the model to find the correct answer based on the correct article." 9 A prompt example is:
Task description To assess this ability, we employ two imulii-urrelate-ddccmeee
Answer the question based on the given documents. The following are given documents. Document 1: [...] Document2: [...] [...] Document 20: [...] Question: What religion were the Normans? Answer:
These tasks are notably more challenging than typical 'ncedle-in---haystace benchmarks [Kamradt, 2023], owing to their high information entropy. The typical
task is too easy, resulting in many models achieving near-perfect performance. See Table 13 Appendix for details.
Main results Table compares memory mosaics v2 and transformers, pretrained on context length, on these question-answer tasks. Memory mosaics v2 outperforms transformers on 4k task- length by 1.4%~~5.6%. Similarly, Table 4 presents the same comparison, but with both models fine-tuned at 32k context length. As task lengths increase to 32k, the "multi-unrelated-e setti-urrlaaedddccmmenn question-anas answering" tasks become more challenging. At this increased difficulty level, memory mosaics v2 significantly outperforms transformers by 12.3% to 14.8%.
Table 3: Comparison of memory mosaics v2 and transformer, trained on 4k context length, on RULER question- answer tasks. Memory mosaics v2 not only outperforms transformer on 4k task-length, but also successfully extrapolate the context length x4 ~ x8 times without any fine-tuning
model
context length
task-length
16k
32k
transformer small memory mosaics v2 small
4k
45.0
35.0
34.1
31.7
transformer large
4k
57.7
x
memory mosuics v2 large
4k
59.3
48.8
46.4
26.5
Table 4: Comparison of memory mosaies v2 and transformer, trained on 4k and fine-tuned on 32k context length, on RULER question-answer tasks. Memory mosaics v2 outperforms transformer by 12.3% 14.8%.
model
context length
16k
tusk-length 32k
64k
transformer small memory mosaies v2 small
32k
37.0
29.3
29.0
22.1
32k
44.3
39.3
39.4
36.9
25.3
transformer large memory mosaics v2 large
32k
$1.2
48.8
44.7
41.1
X
32k
53.4
46.4
Many memory compression algorithms, such RNNS, XLSTM [Beck et al., 2025], rwky [Peng et al., 2023], and state-space models [Gu and Dao, 2023], fail on this task by construction because they cannot store all articles before reading the question. Similarly, local-window memory approaches, such as Alibi position encoding Press et al. [2021] and winddoww attention Beltagy al. [2020], also struggle for the same reason." This incompetent
The failures of many potential baselines
"Smillarly to the process used in section 5.1 for verifying persistent-knowleee appendix Table 10 compares memory mosaics v2 with and without long- term memory these question- answering tasks, confirming the necessity of "long-term memory" for these tasks.
storage and retrieval tasks,
"One might argue to play around this shortage by reading the question before the multiple articles. However, this process involves task speciic priori knowledge from human designers. In the end, instead of proving the

---

of memory compression algorithms has also been experimentally demonstrated by Hsieh et al. [2024] and Li et al. [2024]. Also, see Appendix G for these experimental evidences.
Extrapolating context length (without fine-tuning) Context length extrapolation (without fine- tuning) not only is compulationally appealing, but also reveals the model's 's consistency in handling context. Unfortunately, transformers (with ROPE position encoding) struggle extrapolate context length, as shown in Table 3." In contrast, memory mosaics v2, trained on 4k context length, not only outperform transformers on 4k length, but also perform well after extrapolating context length x4 ~ x8 times without any fine-tuning or adaptation.
5.3 In-context learning
Having demonstrated the new-knowledge storage and retrieval ability of memory mosaics v2, this section takes a step further to evaluate its capacity to learn new tasks or distributions at inference time. This ability is also commonly referred in-context learning.
Tasks description To assess the in-context learning ability, employ classic multiclass classifica- tion problems, 12 adopted from Li et al. [2024]. The classification tasks include:
Banking77 [Casanueva ct al., 2020] is a banking-intenn classification task with 77 target categories. Each example has average length of 24 tokens.
Tacred [Zhang al., 2017] is relation classification task of two objects in a sentence, extracted from newswire or webtext, with 41 target categories. Each example has an average length of 77
tokens.
[Demszky et al., 2020] is an emotion classification task of Reddit comments with 28 target categories. Each example has an average length of 26 tokens.
Goemotion
To solely evaluate the ability learn new tasks (reduce the influence of training knowledge), we create an anonymous version with anonymous target labels (e.g. "class 1", "class 2") for each classification task. The original classification setup with semantic labels (e.g. "happy", "angry") is referred to as semantic version.
In this section, we adopt a few-shot learning setup where each "shot" consists of one (x, y) example from each possible target label category. By collecting multiple shots, create n-shot classifica- tion task. To encode these (z, y) examples for memory mosaics v2 and transformers, we serialize the (zc, y) pairs into a sequence followed by a test query Ztest- A prompt example is:
Given customer service query, please predict the intent of the query. [...] The examples are follows: query: Tshott,
instant: Yshori, [...], query: Tshot2, instant: Yshor2; [-..], query: Ttest: instant:
Main Results Figure compares the performance memory mosaics and transformers in three classification tasks with semantic target labels. The horizontal axis represents the number of shots, while the vertical axis represents the classification accuracy on Ttest- We can observe two phenomena: 1) memory mosaics v2 consistently improve classification performance as it sees more demonstration shots (blue curves). In contrast, transformers struggle to maintain their performance and exhibit counterintuitively degraded performance as more demonstrations are provided (red curves). 2) Memory mosaies v2 significantly outperform transformers by more than 10%. Appendix H provides
machine is intelligent, often proves that human designers are intelligent. Please recall that child does not prepare all questions before going to school.
"The comparison ignores many memory compression and local window approaches [Press al., 2021, Beltagy al., 2020], because they fail on this evaluation by construction.
We choose classic classification problems over other fancy benchmarks for two reasons. Firstly, the mecha- nisms underlying classification are well-studied, allowing to confidently attribute good poor performance to the system's properties. Secondly, classification tasks can be designed be arbitrarily different from the training set by changing the classification boundary, making it casier to measure the ability learn new distributions. In contrast, as of this writing, many fancy benchmarks may not offer the same level of control and fine- ine-grained analysis.
known be sensitive the prompt strategies [Gupta al., 2024, Mirzadeh al., 2024], such as the delimiter before æ and y. shufflingnnot-shufiin the (zz y) examples within each shot. To reduce the influence of prompt strategies, we evaluate each classification task with different delimiters ("[space]" and "\n"",, shuffted/non -hhuffdd/-o--uufeeed shuffled (z, y) examples. Then choose the best prompt strategy for each rh-hoot classification task. Check appendix I for prompt examples.
Transformmrs
7

---

similar comparison on smaller model size (~1.5B), with an even larger margin. Appendix E further summarizes the comparison under matched model size or computation (FLOPs).
Figure presents a similar comparison as Figure 3, but on anonymous target labels. Again, memory mosaics v2 significantly outperforms transformers on all classification tasks.
b7 with semantic label
tacred with semantic label
goemation with semantic label
1.75-
5.43
0.41
Mororyy Mosilca
3900
0:44
9.27
0.25
numbaa of Stots
FARTISS siiets
numba ofsteos
Figure 3: Semantic label in-context learning comparison between memory mosaics v2 and trans- former. Memory mosaics v2 significantly outperform transformers on in-context learning with a large margin (more than 10%). Meanwhile, memory mosaics v2 benefits from more demonstration shots (x-axis), unlike transformers.
banking77 with annrymous label
tacred with anonymous label
goemotion with ancnymous label
0:4h
0.1-
1.6
0
8.6
15.5
6.00
0:33
3:57
0.12
3.5
Vaniurrar lnge PereryPeatt
12
14
numberr Shope
number ohotss
rumber
Figure 4: Anonymous label in-context learning comparison between memory mosaics v2 and transformers, Memory mosaics v2 significantly outperform transformers on all classification tasks.
In summary, the experiments demonstrate that memory mosaics v2 not only outperform transformer by significant margin (more than 10%) on in-context learning, but also consistently improve performance as more demonstrations are provided. These results highlight the superior in-context learning ability of Memory Mosaies v2.
Augment transformer with long-short term attention Memory mosaics (v2) contains several unique components that are not applicable to transformers, such as the symmetric key and query. and the adaptive bandwidth. One seemingly applicable component for transformers the separation of long-term and short-term memories introduced in Section 3.3. However, Figure shows that augmenting a transformer with long-short-tent attention does not help it overcome the lim- itations of in-context learning. These phenomena imply that memory mosaics (v2) is not simply transformer variation but represents a different architecture.
tacred with anorymous label
0.3
Iraakima aral ceal
runber shes
Figure 5: Augmenting trans- former with long-short term atten- tion doesn't help in-context learning.
Computation and Parameter Concerns On the last two evaluation dimensions (new knowledge storage and retrieval,
and in-context learning), memory mosaics v2 outperform transformers by more than 10% with slightly more parameters. This 10% advantage holds even when comparing under the same number of parameters or the same computational budget. See Appendix Figure 12 for details.
6 Risk-reture trade-off of frontier-model--sizd memory mosaics v2
Having demonstrated the superior new tasks learning ability of memory mosaics v2 up to 9.9 billion parameters and I trillion training tokens, this section analyzes the "risk-return trade-off" further scale memory mosaics v2 to the size of the frontier model, unveiling potential benefits and challenges.
Two Approaches
To train a large frontier foundational model, one can either:

---

1) take a recipes (e.g. architecture), or
approach by investing more resources (GPUs and data) and reusing old
2) take a middle-risk middll-rris-hhih---euun approach by trying new smart techniques.
Taking the first approach, one can take advantage of existing software, hardware, experiences, and datasets to quickly "reproduce" huge foundational model. However, this approach is unlikely to result model that stands out from others, as it is based on shared recipes.
In contrast, taking the latter approach may require optimizing software and hardware, adapting techniques, a sharp sense of research direction, and possessing a keen sense of research direction along with strong problem solving abilities." 14 Despite the high requirements for personnel, this approach holds the potential for tremendous breakthroughs.
Ultimately, the decision between these two approaches depends on the available resources and per- sonnel. To aid in this decision decision-making process, this section provides simple and brutal comparison:
How much more data does the transformer recipe approach need to match the performance of memory mosaics v2?
6.1 Comparison of two approaches
To answer this question, we compare the new tasks learning ability of memory mosaics v2 and transformers trained on various amounts of data. Specifically, multiple transformer models are trained on 200B, IT, and 8T training tokens, while a memory mosaics v2 is trained on training tokens.
New-knowledge storage and retrieval Table shows the comparison on the new-knowledge storage and retrieval ability. Training on the same number of tokens (IT), transformers lag behind memory mosaics v2 by 12.3% (41.1% vs 53.4%). x8 times more training tokens (8T) improves the performance of transformers. However, the resulting transformer (trained on 8T tokens) still lags behind memory mosaies v2 (trained on tokens) by 6.5% (46.9% vs 53.4%).
Although further increasing training data may improve the performance of transformers in this evaluation dimension, it comes at the cost of significantly larger training cost (time and resource). Moreover, serious problem occurs: are running out of data!
Table 5: Comparison of memory mosaics v2 and transformers, trained on 4k and fine-tuned on 32k context length, on RULER question- -restion-answer answer tasks. ("transformer large*" uses group-query attention to reduce memory cost, increases training context length to 8k to boost long-context performance.)
model
context length
train tokens
4k
8k
16k
task-length 32k
64k
transformer large
32k
2008
48.6
42.9
40.7
33.8
transformer large
32k
IT
51.2
48.8
44.7
41.1
x
transformer large*
32k
8T
59.2
54,5
50.9
46.9
x
memory mosaies v2 large
32k
IT
58.9
55.5
54.9
53.4
46.4
banking77 with semantic label
tacred with samantiic label
gosmotion with semantic label
0.37
0.300
0.5
3330
827
0.00
4299
Lagt-merLggg--00 2008
1.1
Bundaneeringg at tuender Manory Monico Larga
num of 87ot
7um oTahoos
num ofalcta
Figure 6: Semantic label in-context learning comparison between memory mosaics v2 and trans- former. Memory mosaics v2 trained on IT tokens, while three transformers are trained on 200B, 1T, 8T tokens, respectively. Transformer with x8 times more training data (8T, dash red line) starts to match the performance of Memory Mosaics v2 (IT, solid blue line).
'These requirements, in turn, demand small group high-quality researchers and managers. In i..d. regime, such as persistent-knowledee storing and retrieval, of course, more data larger model better performance. This argument in i.i.d. scenario has been well studied three decades ago Vapnik [1991].

---

banking77 with anorymous
with anonymous label
goomation anonymous
2003 des
911
145
333
Monoyy Masalia Large tn
2.03
3.54
1:04
3.87
172
145
03
12 34
34
10
54
38
Figure 7: nonymous label in-context learning comparison between memory mosaics v2 and trans- formers. Sadly, transformers trained 8T (dash red line) still lag behind memory mosaics v2 trained on IT (solid blue line) by a large margin.
In-context learning Figures and 7 show the comparison on in-context learning ability. For semantic label tasks (Figure 6), x8 times more training data helps transformers (8T data) match the performance of memory mosaies v2 (IT data). However, for the more challenging anonymous label tasks, more training data cannot help transformers. Contour-intuitively, transformers trained more training data (8T) exhibit a degraded performance on anonymous label tasks (Figure 7).
In summary, x8 more training data helps transformers certain new task learning benchmarks. However, the resulting transformers (8T data) still lag behind memory mosaics v2 trained on IT data. More importantly, in anonymous label tasks that heavily rely on the new task learning ability, more training data cannot help transformers. These experiments answer the initial question: "How much data does the transformer recipe approach need to match the performance of memory mosaics v2?".
7 Fine-tuning speed: who can fine-tune with one minibatch?
Despite the strong in-context learning capability of memory mosaics v2 shown Section 5.3, may still be attractive fine-tune model for specific domain in order either reduce inference costs or improve in-domain performance. It is generally expected that such models can be efficiently fine-tuned for new domain using a comparatively small number of examples.
Figure 8 compares the fine-tuning speed (in terms of data size) of memory mosaics v2 and transformers. Both mod- els. pre-trained on 4k context windows, were fine-tuned to 32k context length using the recipe described in Section 4 and evaluated on the same RULER tasks (32k task-length) described in Section 5.2.
mu tntun decuments quedticn sneareding (13k task-lenghh lenght
+22
Surprisingly, single fine-tuning mini-batel (one opti- mization step) on memory mosaics v2 yields a 22% ac- curacy improvement. Two fine-tuning mini-batches memory mosaics v2 are sufficient to reach the optimal performance. In contrast, transformer fine-tuned with 800 mini-batches still lags behind memory mosaics v2 fine-tuned with single mini-batch.
lorg
Eomaning wiraacce (oporiaaicn slepe)
Figure 8: Fine-tuning speed comparison between memory mosaics v2 and trans- former.
Discussion and future direction
This work scales memory mosaies (named memory mosaics v2) llama-8B scale, demonstrating superior performance on new task learning, outperforming transformers by more than 10%. The three evaluation dimensions introduced in this work provide transparent and controlled assessment of model capabilities, particularly focusing the new task learning. The risk-return trade-off analysis reveals the weakness of the mainstream "more data more computation" belief, highlighting research opportunities on other smart techniques. One future direction is to reduce the computational cost for very long context lengths using fuzzy hashing Breitinger et al. [2014], Chen et al. [2024] and hierarchical memory Yuan et al. [2025], Lu et al. [2025] approaches.
Acknowledgments
Léon Bottou is CIFAR fellow. We thank Gabriel Synnaeve, Jade Copet, Badr Youbi Idrissi, and Ammar Rizvi for their considerable support with hardware, software, data, and baselines.
10

---

References
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprini arXiv:1907.0288 2019.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Güittr Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. Advances in Neural Information Processing Systemx, 37:107547-107603, 2025.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint vrXiv:2004.05/500 2020.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A transerr objective for learning to disentangle causal 2019.
mechanisms. arXiv preprint
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36: 3560-1588, 2023.
Frank Breitinger, Barbara Guttman, Michael McCarrin, Vassil Roussev, and Douglas White. Approx- imate matching: Definition and terminology. techreport nist special publication 800-168. national institute of standards and technology, 2014.
lñigo Casanueva, Tadas Temčinas, Daniela Gerz, Matthew Henderson, and Ivan Vulić. Efficient intent detection with dual sentence encoders. arXiv preprint arXiv:2003.04807, 2020.
Adam Casson. Transformer flops. 2023. URL httt:p///aadmaaasooccccoosssssee ffooss.
Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Chou,,Jaayyy Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, et al. Magicpig: Lsh sampling for efficient llm generation. arXiv preprint 2XXv:2410.16173 2024.
Pierre Comon. Independent component analysis, new concept? Signal processing, 36(3):287-314, 1994.
Dorottya Demszky, Dana Movshovitz-Attias Jeongwoo Ko, Alan Cowen, Gaurav Nemade, and Sujith Ravi. Goemotions: A dataset of fine-grained emotions. arXiv preprint arXiv:2005.00547, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leamnin for fast adaptation of deep networks. In International conference machine learning, pages 1126-1135. PMLR, 2017.
E. García- Garcfa-Portuguus. Notes for Nouparametric Statistics. 2024. URL https://bookdonn org/e Marpor/NP-UC3M, Version 6.9.1. ISBN 978 84--09--9537-11.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinay Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama herd of models. arXiv preprint arXhv:2407.2178. 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXi:22112.00752, 2023.
Vipul Gupta, David Pantoja, Candace Ross, Adina Williams, and Megan Ung. Changing Answer Order Can Decrease MMLU Accuracy, June 2024.
Trevor Hastie, Robert Tibshirani, Jerome Friedman, et al. The elements of statistical learning, 2009.

---

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.06654 2024.
Gregory Kamradt. Needle in a haystack pressure testing llms. Github, 2023. URL https: teaettacckk/ee//ma.....
//github. comm watarraadaeLeeeeeeeeeeeleelleae
Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. Long-context Ilms struggle with long in-context learning. arXiv preprint wrXiv:2404.020600 2024.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, et al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502. 13189, 2025.
Iman Mirzadeh, Keivan Alizadeh, Hooman Shahrokhi, Oncel Tuzel, Samy Bengio, and Mehrdad Farajtabar. GSM-Symbolic: Understanding the Limitations of Mathematical Reasoning in Large Language Models, October 2024.
E. Nadaraya. On estimating regression. Theory of Probability and Its Applications, 9:141-142, 1964.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 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 preprint arXiv:2305. 13048, 2023.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint 100::2208.12240, 2021.
Karsten Roth, Mark Ibrahim, Zeynep Akata, Pascal Vincent, and Diane Bouchacourt. Disentan- glement of correlated factors via hausdorff factorized support. arXiv preprint 17X01:2210.07347, 2022.
Noam Shazeer. GLU variants improve transformer. arXiv preprint arXiv:2002.05202. 2020.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory, 2019.
V. Vapnik. Principles of risk minimization for learning theory. In J. Moody, S. Hanson, and R.P. Lippmann, editors, Advances in Neural Information Processing Systems, volume 4. Morgan- Kaufmann, 1991.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Geoffrey S. Watson. Smooth regression analysis. Sankhyā: The Indian Journal of Statistics, Series A, pages 359-372, 1964.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502. / 2025.
Jianyu Zhang. Ai for the open-world: the learning principles. arXiv preprint arXiv:2504.14751, 2025.
Jianyu Zhang, Niklas Nolte, Ranajoy Sadhukhan, Beidi Chen, and Léon Bottou. Memory mosaics. In The Thirteenth International Conference on Learning Representations, 2025.
Yuhao Zhang, Victor Zhong, Danqi Chen, Gabor Angeli, and Christopher D. Manning, Position- aware attention and supervised data improve slot filling. Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing (EMNLP 2017), pages 35-45, 2017.
12

---

Memory Mosaics at scale
Supplementary Material
A Gated time-variant key feature extractor & convolutional value extractor
Figure illustrates how keys and values are constructed in memory mosaics v2.
K
ki
V
Ap
Att
A eap( (Waed)
gtkt
Üt Ut+1
9: 5()1)
ūt = W
h. Win:
L
Et I1+1
Figure 9: Left: Key he feature extractor: kg Norm(k)). 2 = Norm()
Right: Value feature extractor:
B Training data sequence length distributions
Figure 10 shows the distributions of the length of the training data sequence truncated to 4096 or 32,768 768 max length.
training data sequence le distribution Bruncated at 4k)
trsining data secuence length distribution (rurccated at 12k)
13
LOCK
2008
1530
35300
Wmigh
5mg
Wmyth
Figure 10: Training data sequence length distributions. For a given maximum sequence length during training (e.g. 4k), longer sequences are truncated to the maximum sequence length. This truncation results in the peaks at the end distributions.
13

---

C Training details
hyperparameters For all Memory Mosaics v2 and baseline Transformer models we use consistent set of hyperparameters. That is, batch size of 1024, a sequence length of 4096, an adamw optimizer with 81 = 0.9 and 32 = 0.95 accompanied by L2 weight decay of 0.1 and a gradient norm clip of I, learning rate warm-up of 2000 iterations followed by cosine learning rate scheduler that reduces the learning rate by a factor of 100 at the end. The initial learning rates (after warm-up) are set to 3e-4 for "small" models and le-3 for "large" models.
We also employ document-wise attention mask, where the attention scores are only computed within each sequence (document) in the training data, to reduce computation cost. Two special tokens, "<lbegin of textl>" and "<lend_of of textl>" appended at the begining and ending of sequence, respectively.
During training, memory mosaics v2 samples the long-term memory delay step m from [64, 256]. sets the short-term memory window size h 256. At inference, m is set to 64, as illustrated in Figure 11.
fixed (to 256/ throughout training inference
unifforlly sampled [64.256 throughout training fixed to 64 inference
Short-term memory Attention mask from to
Long form momory actention mask from
Figure 11: Randomly overlapped long-term & short-term memory.
is worth noting that these hyperparameters were originally searched and optimized for the baseline transformer models. We transfer these hyperparameters to memory mosaics v2 without further hyperparameter searching. Thus, it is possible that this hyperparameter setup is suboptimal for memory mosaics v2.
Parameter Initialization and reparameterization Table 6 summarizes the parameter initialization methods and reparameterization tricks. W1, W2, W3 refer the parameters in persistent memory that are implemented as two-layer dense neural networks, Wz (SiLU * W3(z)). SiLU x sigmoid(x) is an activation function. d E {2048, 4096} indicates the hidden dimension of Memory Mosaics v2 small and large. d' {6144, 14336) indicates the hidden dimension of the two-layer neural networks in persistent memory. indicates the depth of the Mosaics blocks, starting from 0.
Table 6: Parameter initialization methods and reparameterization tricks used in Memory Mosaics v2.
Parameter
Location
reparameterization
Initialization
udaptive bandwidth
6 1in(9.10)
0 1.5
8
adaptive bandwidth
mmnn(9,10) 10)
0 1.5
a
adaptive bandwidth
a min
= 1/3
feature extractor
0 0
feature extractor
U(0,
Ws. Wg. Wg, Ws, W.
long short memory
min (mxx(V(0..), -30)) 30). /2a01
W1.
persistent memory
min (mxx(V(0.r)
W5
persisteat memory
min (00)
[2P((=1)
W.,
embedding classifier
min max((((0
v28
16 hidden dimension of the two-layers neural network in persistent memory is set to 6,144 and 14,336 for small and large models, respectively.
14

---

D Failures of memory compression baselines
Many memory compression algorithms, such RNNs, XLSTM [Beck et al., 2025], rwkv [Peng al., 2023], and state-space models [Gu and Dao, 2023], fail new-task storage and retrieval and in-context learning evaluation dimensions by construction. The reason is that these memory compression algorithms lack the ability to store large amounts of information before getting command on how to process the information. One might argue to play around this shortage by reading the "command" before storing the large amounts of information. However, this process involves task task-specific priori knowledge from human designers. In the end, instead of proving the machine is intelligent, often proves that human designers are intelligent. Please recall that child does not prepare all questions before going to school.
This incompetent of memory compression algorithms has been experimentally demonstrated by Hsieh et al. [2024] and Li al. [2024] on both RULER benchmarks and in-context learning tasks.
Table 7 compares memory compression methods (rwkv-v5-7b and (aamaa-..bb-slimn. and non- compression method (llama--7b) on RULER long-context tasks. is clear that memory compression methods perform poorly as the required context length (i.e., required information storage space) increases.
Similarly, Table 8 compares memory compression methods (wwkv-5-world 7b and Mamba-2.88 (.8B) and non-compression method (qwen- 1enn-.5--b--bass and
on in-context learning tasks (Tacred few-shot classification [Zhang et al., 2017]). In this challenging in-context scenario, memory compression methods just don't 't work at all.
Please note that this section shouldn't be used to criticize or hinder the study of memory compression methods. Memory compression methods have their advantages. In gersistent-knowledee storage and retrieval evaluation dimension, they performs very well. For model efficiency, memory com- pression methods reveal charming computation complexity. The goal of this section to explain why this paper doesn't t choose memory compression methods as baselines.
Table 7: Comparison of memory compression methods (rwkv-v5-7b and mamba-2.88 2.8-ss--.. and non-compression method (llama2-7b) on RULER long-context tasks. memory compression methods perform poorly as the required context length increases. Numbers are copied from Hsieh et al. [2024] Figure 4.
model
task- Ik task-length 2k task-length 4k
Ilama2-7b
96.0
91.6
95.0
twky-v5-7b
87.5
73.7
51.4
mamba-2.88 2.b-sl
62.6
52.6
Table 8: Comparison of memory compression methods (rwky-5-world 7b and Mamba-2.8B. and non-compression method (qwen-1.5-7b-bl
and 7istral-7b-v0.2-bas. on in-context learning tasks (Tacred few-shot classification [Zhang al., 2017J). Memory compression methods fail on all cases.
Numbers are copied from Li et al. [2024] Table 4.
model
1-shot 2-shots 3-shots
4-shots
5-shots
gwen-1.5-7b-baas. 15-7b-base 7b
38.7
47.3
45.2
43.6
40.6
histral--bbv0.2-base
53.3
53.1
51.6
48.0
42.3
rwkv-5-world
2.3
1.0
1.2
M.mba-2.8B
0
15

---

E Model efficiency comparison
As we emphasized in main text, model efficiency (e.g. model service, throughput, VRAM, etc.) is the goal of this work. Many engineering works can be performed adapt memory mosaics v2 to a custom use case or hardware. To aid in these potential adaptations, Figure 12 provides model efficiency comparison in both computation (FLOPs) and model size (number of parameters) viewpoints. The results show that memory mosaics v2 outperforms transformer by more than 10% under either the same FLOPs or the parameters budgets.
RULLR QA accuracy task length)
Namory Nosalss v2 Trensformer
50
Mamory Mosaice v2 arga Trensfomerr
number of billion paremeters
number of billion FLPPs
TacBeed ICI accurscy
shots)
50
40-
40
Memory Mosalcs v2 large Transformer large
Memory Monaree v2 lang Vandfomer lnnge
billion parameters
Figure 12: Model efficiency (FLOPs or # parameters) comparison between Transformer large and Memory Mosaics v2 large. Top row shows the comparison on RULER "malti-unrelated--oommeel storing and question-answering tasks, while the bottom row shows comparison on the Tacred in- context learning task. Memory mosaics v2 large outperforms transformer large by more than 10%.
F Additional results on persistent-knowlede storage and retrieval
Table shows six language benchmarks which removing long-term memory from memory mosaics v2 after training degrades its performance.
Table 9:
Memory mosaics v2 performance 6 language benchmarks, where removing the "long- term memory" after training dramatically hurt the performance (42. 1% vs 34.9%).
params flops/token
squad bbh math mhpp
race middle
high
race
avg
transformer large
8.8B
16.7B
76.3
45.6
8.7
9.8
62.6
45.6
41.4
memory mosaics v2 large
9.98
18.9B
78.2
47.8
8.8
9.6
61.6
46.5
42.1
memory v2 large
8.3
15.6B
69.4
24.6
5.4
6.8
59.5
43.6
34.9
without long-terne term memory
16

---

G Additional results on new-knowledge storage and retrieval
Table 10 shows that removing long-term memory from memory mosaics v2 after training degrades the performance on the RULER question-anse annserr tasks by 20%~~30% This indicates that the ruler question-answer tasks rely on long-term memory to perform well.
Table compares memory mosaics v2 large and other public base models on RULER question-ansee answer tasks. Memory mosaies v2 large outperforms these models across all task lengths.
Table 12 illustrates the effect of the stochastic long-term memory size training setup introduced in Section C. This stochastic long-term memory size setup is used to encourage the allocation of position-invariant signals and position-dependen signals to long-term and short-term memories.
Table 13 compares memory mosaics v2 and transformers on typical 'needle-in -aselle---ahaystackk task from RULER [Hsich et al., 2024]. The typical 'needle-in-at in-a-haystack" is too easy such that many models can achieve a near-perfect performance.
Table 10: The effect of removing "long-term memory" of memory mosaics V2 large on RULER question-answer tasks.
model
context length
16k
32k
memory mosaics large memory mosaies large without long-term memory
32k
58.9
55.5
54.9
53.4
32k
38.5
22.2
20.0
20.2
Table 11: Comparison of Memory Mosaics v2 large (base model) and other public base models
(similar scale) RULER question-answer tasks. Memory Mosaics v2 large outperforms these models across all task lengths, despite that Memory Mosaics uses 1/4 generation lengths (32 tokens) of other public base models (128 tokens). The numbers in "** rows come from Hsich et al. [2024].
Model
claimed length
task-length 4k
8k
16k
32k
NemooryMosaicss-v22 (base)
32k
58.9
55.5
54.9
53.4
Llama2-7B (base)
4k
48.6
Mixtral-base (8x7B)*
32k
50.8
47.7
45.3
41.3
Mistral- (7B)*
32k
53.5
51.0
48.4
44.7
Together-base (7B)*
32k
47.5
44.6
33.6
0.0
LongLoRA- (7B)*
100k
34.5
32.1
33.6
29.4
Yarn-base (7B)*
128k
29.7
23.5
28.6
29.7
I.W--base (7B)*
1M
42.7
40.2
38.7
37.1
The effect of stochastic long-term memory size (during training) in memory mosaics v2 small model on RULER question-answer tasks. Both models trained on 4k context length, then evaluated on 32k context length without any fine-tuning. The stochastic long-term memory size setup boost context length extrapolation ability by more than 15%.
Table 12:
model
context length
stochastic long memory lask-length task-length
memory mosaics v2 small memory mosaics v2 small
4k
No
43.6
15.9 31.7 (+15.8)
4k
Yes
45.0
Table 13: RULER S-NIAH benchmark comparison between transformer and memory mosaics v2.
model
context length
4k
8k
16k
task- 32k
transformer small
32k
99.4
99.0
98.2
97.8
memory mosaics v2 small
32k
100.0
100.0
100.0
100.0
transformer large
32k
100.0
100.0
100.0
99.6
memory mosaics v2 large
32k
100.0
100.0
100.0
100.0
17

---

H Additional results for in-context learning
Figure 13 and 14 shows the in-context learning comparison between memory mosaics v2 small and transformer small (Ilama- 1.5B scale).
banking77 semantic
tacred with semantic label
semartic label
0.22
0.20
372
Mersoy srat
015
242
314
083
25 34 14
22
16
Figure 13: Semantic label in-context learning comparison between memory mosaics v2 and trans- former. Memory mosaies v2 significantly outperform transformers in-context learning with large margin (more than 10%). Meanwhile, memory mosaics v2 benefits from more demonstration shots (x-axis), unlike transformers.
banking?7 with and munrymous label
facred with anonymous label
goemetion with anonymous label Turabmeuuret
2.8
Menuay Mackis seat
37-
0.2%
3.99
3.5
2.5
3.07
010
3.4
rumber ohabens
nimber
number 27+hoo.
Figure 14: Anonymous label in-context learning comparison between memory mosaics v2 and transformers. Memory mosaics v2 significantly outperform transformers on all classification tasks.
18

---

I Prompt examples of multiclass classification tasks
I.1 Banking77 classification with semantic labels
We sweep the delimiter from "[retum]] and "[space]], leads to the following two prompts:
"Given a customer service query, please predict the intent of the query. The predict answer must come from the demonstration examples with the exact format. The examples as follows:
service query:
am still waiting my card?
intent category:
sty__rrivvll
service query: My card has been found. Is there any way for me to put it back into the app?
intent category:
city_linking
service query: Can I get card even if live outside the UK? intent category:
"Given customer service query, please predict the intent of the query. The predict answer must come from the demonstration examples with the exact format. The examples are as follows: service query: I am still waiting my card?
intent category: city_arrival
service query: My card has been found. Is there any way for me put it back into the app?
intent category: city _inking
service query: Can get a card even if live outside the UK? intent category:"
For each prompt with either "[return]" or "[space]" delimiter, we also try to shuffle the demonstration example (i.e., service query: [...], intent category:(...I) orders within each one shot. This shuffling process provides another two more prompts.
1.2 Banking77 classification with anonymous labels
Anonymous tasks the same of prompts except that anonymous tasks replace semantic labels (e.g. aity_arrival, city dit_kink with anonymous labels (e.g. class class ol).
19

---

1.3 Goemotion classification with semantic labels
We sweep the delimiter from "[return]" and "[space]], leads to the following two prompts:
"Given a comment, please predict the emotion category of this comment. The predict answer must come from the demonstration examples with the exact format. The examples are as follows:
comment: Her upper lip always looks terrible such an easy fix, can u believe she is so vain and never bothers
emotion category:
embarrassment
comment:
No problem. I'm happy to know it's not what you meant.
emotion category:
joy
comment: These refs have out for the colts. I didn't 't realize we traded our MVP 11 KC either. emotion category:
"Given a comment, please predict the emotion category of this comment. The predict answer must come from the demonstration examples with the exact format. The examples are as follows:
comment: Her upper lip always looks terrible such an easy fix, can u believe she so vain and never bothers wax
emotion category: embarrassment
comment: No problem. I'm happy to know it's not what you meant.
emotion category: joy comment: These refs have it out for the colts. didn't 't realize we traded our MVP 11 either. emotion category:"
For each prompt with either "[return]" or "[space]" delimiter, we also try to shuffle the demonstration example orders within each one shot. This shuffling process provides another two more prompts.
1.4 Goemotion classification with anonymous labels
Anonymous tasks use the same of prompts except that anonymous tasks replace semantic labels with anonymous labels (e.g. class 00, class class_0l), O1).
20

---

1,5 Tacred classification with semantic labels
We sweep the delimiter from "[return]" and "[space]], leads to the following two prompts:
"Given a sentence and a pair of subject and object entities within the sentence, please predict the relation between the given entities. The examples are as follows:
sentence:
But US and Indian experts say it has hesitated to take action against Lashkar-e- Barnkarr--Tabba Taiba, which means "The Army of the Pure. "believing that the Islamic militants could prove useful in pressuring its historic rival India.
the relation between Lashkar-e-Taila and Army of the Pure is:
org:alternate_ aeggaaaatteeaamee _ames
sentence:
The offer from ITW, the Glenview, IlI, diversified manufacturer of engineered products, repre- sents a premium of 85 percent to the Manitowoc bid.
the relation between ITW and Glenview is: orgicity of headquarters
sentence:
The statement from North Korea, carried by the country's official Korean Central News Agency, did not mention Kim by name, but South Korean Unification Ministry spokesman Kim Ho-nyeon said the North's 's state media has before used such wording to refer to him. the relation between Korean Central News Agency and North Korea is:
"Given a sentence and a pair of subject and object entities within the sentence, please predict the relation between the given entities. The examples are follows:
sentence: But US and Indian experts say it has hesitated to take action against Lashkar-e- Taiba, which means "The Army the Pure, "believing that the Islamic militants could prove useful in pressuring its historic rival India.
the relation between Lashkar-e- Taibaa and Army of the Pure is: org:alternate _ames sentence: The offer from ITW, the Glenview, III, diversified manufacturer of engineered products, represents premium of 85 percent the Manitowoc bid. the relation between ITW and Glenview is: org:city_or of headquarters
The statement from North Korea, carried by the country's 's official Korean Central News Agency, did not mention Kim by name, but South Korean Unification Ministry spokesman Kim Ho-nyeon said the North's state media has before used such wording to refer to him. the relation between Korean Central News Agency and North Korea is:"
sentence:
For each prompt with either "[return]" or "[space]" delimiter, we also to shuffle the demonstration example orders within each one shot. This shuffling process provides another two more prompts.
I.6 Tacred classification with anonymous labels
Anonymous tasks use the same set of prompts except that anonymous tasks replace semantic labels with anonymous labels (c.g. class_00 class class_07), _0)).
Computation resources
All experiments are conducted on H100 GPUs with 80GB VRAM.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction do accurately reflect the contribution and scope made by the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Check Section "Computation and # parameters concerns" for example. "Limitations" section does discuss the limitations of the work performed by the authors, including scope of the claims made, strong assumptions, and how robust the results are to violations of the assumptions.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
22

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
We do not include theorems.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: Check experimental setups Section 4, Section 5, and Appendix
Guidelines: The paper clearly explains the training setup for experiments run in the paper. These explanations are sufficient to reproduce the experiments. The code will be released upon acceptance of the paper.
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
23

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [No]
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurIPS code and data submission guidelines (https://nips.. cc/pu
wwt.:euwbbmssssoosioooo.oooo for more details. http//CoeeSuuiiieeee
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
ceappp.ccppuulcc/uutt gww.//////////eg///ee/// Seteee///ooddmdiiiiacaeaa wwtainssssonna.. ssionPoiccy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Check Experiment Section4 and Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The in-context learning result figures show the performance of different evaluation configurations together with a fitting curve. The many accuracy results of different evaluation configurations provides means estimate variance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative етrоr rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Check Appendix J.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / /twp//euuipp..eeee cc hpppulicc/thicsuiiiicoopooo
Answer: [Yes]
Justification: The research conducted in the paper does confirm, in every respsect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper discusses the potential impacts of such a works and the resulting effects introduced by these considerations.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
25

---

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML..
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: not applicable.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The creators and original owners of assets used in the paper properly credited and the license and terms of use explicitly mentioned.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assett are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paperswi thcode. com dataatsss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
26

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing or research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
27

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not non non- otanstaddard components. Please refer to our LLM policy (https:/ htop::://neurps.... (roonffereneee2225555 /LLM) for what should or should not be described.
involve LLMs as any important, original,
28


Reviews:
[{"review_id": "HACrrxVO0e", "reviewer": "Reviewer_vHo6", "summary": "This paper updates the \"Memory Mosaics\" architecture, an alternative to widespread autoregressive transformer language model architecture that replaces attention with an associative memory module. This paper scales up the architecture and benchmarks it on larger-scale language modeling tasks, introducing three architectural changes designed to improve language modeling performance. First, it makes the kernel smoothing in the memory adaptive to the number of examples stored in the memory. Then, it introduces a recurrent gate mechanism in front of the \"key\" computation, designed to allow the model to learn key representations that are flexible to extra tokens in the sequence (as you might want to do when moving from timeseries to language modeling). Finally, it uses attention masking to separate long-term and short-term memory modules.  It also sets aside the unified persistent associative memory in favor of going back to a traditional feedforward MLP module interleaved with the contextual memory instead of unified.\n\nThe paper then measures the resulting architecture in three ways. First, it does standard benchmark evaluations at two model sizes, comparing to a standard transformer at similar sizes, comparing on 19 standard benchmarks, including six that are appear to be very sensitive to long-term memory, showing comparable performance and interesting wins on the long-term tasks. Then it does a RULER \"needle in a haystack\" evaluation, showing very strong wins in both in-domain and extrapolated-length performance testing up to 8x or 16x the training context length. Finally it tests on a suite of ICL text classification tasks with both meaningful and anonymized target labels, and show significant wins over the similar transformers.", "questions": "[original questions below; after discussion my main concerns have been addressed.]\n\n* Would transformers with traditional attention achieve similar performance wins if attention heads were similarly split between short- and long-range attention through a similar masking scheme?\n* What is the empirical benefit of each of the innovations, and are all needed?  What is the performance benefit of adaptive smoothing?  What is the performance benefit of the gated recurrent key computation?  Why is long-term memory trained with randomly-changing masks while being tested with fixed masks? Can ablations or other measurements be shown to justify these added complexities?\n* Exactly where in the formulation is masking applied to distinguish between long- and short-range memories?  This should be clarified in the equations, as well as any other differences between the two types of memories.\n* After training, do the two types of memories learn different patterns of parameters or different kinds of behavior in a way that can be clearly visualized or quantified? For example, MMv1 noted that \"induction heads\" can be implemented in a single layer. In MMv2, do these induction heads show up in the short- or long-range memories, or both?\n* Can you clarify the goals of MM in the context of the other literature that also aims to replace attention. (Such as Mamba/RWKV etc which focus on inference performance, or other attention alternatives such as sparse attention.) Does MM have a different goal? What is the goal, and how does it relate to that other work?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: Memory mosaics is a very pretty simplification of the transformer attention architecture, and this paper shows that it holds promise in being able to perform just as well, or even better, on large-scale language modeling when compared to an autoregressive transformer. The introduced architectural modifications on top of MMv1 are all sensible and nicely motivated, the performance measurements are well-chosen, and the performance wins on ICL and needle-in-haystack tasks are noteworthy.\n\nWeaknesses: The argument is not fully convincing that the MM architecture is responsible for the observed ICL/long-context performance wins over transformers. In particular, the paper leaves the reader with the unanswered question: could the same benefits be achieved without changing the attention mechanism, but by simply enforcing long- and short-term masks on traditional transformer attention? No experiment to apply the simple separation between long/short attention to transformers is done. To make the case, a long/short attention baseline should be compared, to give transformers a fair shot.\n\nSome ablations are missing. Several innovations over MMv1 are introduced, but ablations are only shown for the long/short innovation, which leaves open the question \"are all three sets of extra parameters needed? what is the benefit of each?\"  Although the adaptive smoothing and the learnable gate for keys both seem like reasonable ideas, but it is not measured how much benefit each of them provides.  Even a small measurement would be helpful to justify and motivate the added complexities.\n\nAnother little trick is hidden, which is that the long-range memory is trained with a randomized mask but then tested at inference with a fixed mask. What motivates this trick, and how important is it?  Some ablation or evidence should be shared.\n\nThe presentation of the distinction between long vs short memories is not sufficiently precise, given its centrality to the method. It should be presented in formula form in addition to the Fig 2 picture. For example, the masking, which seems like the key to their operation, should be included in the equations, for example by showing precisely how the masks are used to modify eq 5, and then how the masks are different for long and short memories; if there are any differences besides masking, they should be explained.  Similarly, after training what are the effects of the change in masking? If it leads short- and long-term memory to learn something qualitatively different, insights or visuailizations on the different behavior or differently-learned parameters would be worthwhile. Or a toy-model motivation could help clarify.\n\nA paragraph is devoted to the weak RULER performance of linear alternatives to attention such as Mamba and RWKV making some (logical) claims about their weaknesses, but no measurements are shown; if this comment is made, at least the appendix should quantify the weak performance on this benchmark by recaping or reproducing. Or the paragraph should be rewritten to attribute the claims to a specific previous published result. Related to that: since a lot of research attention is spent on linear attention schemes that attempt to reduce inference scaling costs, the paper should clarify that that is not the goal of this method (if so). It seems like the goal is the opposite, which is to improve long-term behavior and ICL using a similar computational budget as trasnformers.", "was_revised": true, "final_justification": "The MM architecture provides some interesting insights over traditional transformers, and this work both shows that performance can be competitive while clarifying several aspects of the architecture.  Support accepting the paper.", "rebuttal": "Thank you for your thoughtful reviews and valuable suggestions! It was a pleasure to answer your questions.\n\n---\n\n## Q1: Could the same benefits be achieved without changing the attention mechanism, but by simply enforcing long/short-term masks on traditional transformer attention? \n\nThe long/short-term memory components in Memory Mosaics v2 are built based on the nice \"flat-then-peak\" attention distribution in Memory Mosaics (Figure 1 in our paper). Unfortunately, transformer attention lacks this nice attention distribution, due to its position-encoding and asymmetric query-key in transformers. Therefore, simply adding a long/short-term mask cannot help transformer attention to achieve the same in-context learning benefits.\n\nThis paper builds upon Memory Mosaics [1], which has already shown that the position-encoding (as well as corresponding asymmetric query-key) in transformer attention results in a position-dependent attention distribution (Figures 14 & 15 in [1]). This position-dependent behavior impairs transformer's ICL capabilities, as shown in Figures 8 & 9 of [1]. Simply adding long/short-term masks does not address this fundamental issue and thus fails to improve ICL performance.\n\nWe will visualize the attention distribution of MMv2 with long/short-term masks in the final version, addressing the concern of applying long/short-term masks to transformers. \n\n\n>[1] Zhang, Jianyu, et al. \"Memory mosaics.\" arXiv preprint arXiv:2405.06394 (2024).\n\n\n---\n\n## Q2: Motivation and Performance benefit of adaptive smoothing, gated recurrent key computation, long-/short-term memory? \n\nTraining Memory Mosaics v2 Large and Transformer Large on 1 trillion tokens requires over 4 GPU years with H100 GPU. Comprehensive ablations at this scale are not feasible within weeks. Fortunately, we can better address the concern using conceptual and practical evidence. \n\nThese three components are conceptually necessary and practically helpful:\n\n- **Adaptive Smoothing (bandwidth):** In memory-based methods like kernel smoothing, the optimal bandwidth depends on the number of examples, and must be adaptive. Rich literature [2-3] supports this need for an adaptive bandwidth (see lines 89-91 for details). Furthermore, the adaptive bandwidth in MMv2, $\\beta_0 + \\beta_1 * n^{\\alpha}$, degenerates to a constant bandwidth as in MMv1 when $\\alpha = 0$, providing a chance to learn a constant bandwidth. \\\nPractically, MMv2 learns positive $\\alpha$ (0.2~0.4) after training, achieves a significantly lower training loss and better performance on the common benchmarks (persistent-knowledge storage and retrieval).\n\n- **Gated recurrent key feature extractor**: the gated recurrent key feature extractor is time-variant, enabling the model to extract semantic meanings without being constrained by sequence structure. An example is shown in line 99,  where “tom-and-jerry” and “tom---and---jerry” share the same semantic meaning despite differing in sequence structure. Rich literature in recurrent-style networks (e.g., RWKV, mamba) has experimentally verified the benefit of the time-variant feature extractor (line 100). \\\nPractically, this gated recurrent key feature extractor helps MMv2 (small scale) on common benchmarks (+0.5%) and ruler benchmarks (+8%) when trained on 4k evaluated on 32k.\n\n- **Long/short-term associative memories**:\nThe motivation and benefits of long-/short-term memory have been discussed in previous answers. \n\nWe will include these clarifications and results in the final version. \n\n\n>[2] Trevor Hastie, Robert Tibshirani, Jerome Friedman, et al. The elements of statistical learning, 2009. \\\n>[3] E. García-Portugués. Notes for Nonparametric Statistics. 2024.\n\n\n---\n\n\n## Q3: Long-term memory mask is randomized during training while fixed at inference. Visualization of long/short-term memories. \n\nThe motivation for creating long/short-term memories comes from the \"flat-then-peak\" attention distribution as shown in Figure 1 (right plot). To achieve context length extrapolation, we need to allocate the \"flat\" part (position-invariant) to long-term memory, and the \"peak\" tail to short-term memory. As such, only the size of long-term memory increases during context length extrapolation. \n\n\nBecause we do not know the exact boundary between \"flat\" and \"peak\" parts, we use a randomized long-term memory mask length to further impair long-term memory from learning the tail of a sequence. Note that on average this randomization of mask length equals to a linearly-decayed importance. \n\nWe will compare and share the long/short-term attention distribution (visualization) using this randomized mask in the final version, explaining the motivation behind it. \n\n\n---\n\n## Q4: Formal equation for long vs short memories. \nThanks for the valuable suggestion! Following Eq. 4, we express long/short memories as: \n\nshort-term memory: \n$f(k; \\{(k_{max(t-h,1)}, v_{max(t-h,1)}) ... (k_{t-1}, v_{t-1})\\})$\n\nlong-term memory:\n$f(k; \\{(k_1, v_1) ... (k_{t-m}, v_{t-m})\\})$, where $t-m >= 1$\n\nThere are no other differences besides the mask of long/short-term memories. We will add these equations in the final version. \n\n---\n\n## Q5: Weakness of Mamba and RWKV on RULER benchmark \n\n**Weakness of memory compression algorithms on Ruler**:\n\nThe RULER paper [4] has already studied the weaknesses of model compression methods, such as Mamba and RWKV, via experimental comparison (Figure 4). Here is a summary:\n\n\n\n|model |task-length 1k| task-length 2k| task-length 4k|\n|----|----|----|----|\n|Llama2-7b|96.0|91.6|95.0|\n|rwkv-v5-7b|87.5|73.7|51.4|\n|Mamba-2.8b-slimpj|62.6|52.6|-|\n*Table 1: Weakness of Mamba and RWKV on RULER benchmark.*\n\n\nIn the original words of the RULER paper, \"*Finally, we show that non-Transformer architectures, such as RWKV and Mamba, still lag behind Transformer by large margins on RULER.*\" \n\nNote that the two question-answering RULER tasks used in our paper are notably more challenging than other \"needle in a haystack\" tasks in RULER, like multi-/single-key retrievals. The information entropy of these multi-/single-key retrieval tasks is extremely low, making them easy to compress, allowing many models to easily achieve >99% performance (see Appendix E and F in RULER [4]).\n\n**Weakness of memory compression algorithms on in-context learning**:\n\n[5] has already experimentally shown the weakness of model compression methods, such as Mamba and RWKV, in in-context learning (Note that our paper uses the same in-context learning datasets as in [5].). Here is a summary: \n\n|model |task-length 1 shots| task-length 2 shots | task-length 3 shots|task-length 4 shots| task-length 5 shots|\n|----|----|----|----|----|----|\n|Qwen-1.5-7B-base 7B|  38.7| 47.3| 45.2| 43.6| 40.6|\n|Mistral-7B-v0.2-base|  53.3| 53.1| 51.6| 48.0| 42.3|\n|RWKV-5-World 7B| 2.3 |2.6 |1.0| 0| 1.2|\n|Mamba-2.8B  |0 |0 |0 |0 |0|\n\n*Table 2:  Weakness of Mamba and RWKV on in-context learning benchmark (TacRed classification).*\n\n\nWe will refer these results to experimentally support the claim (weakness of memory compression methods) in the final version. \n\n>[4] Hsieh, Cheng-Ping, et al. \"RULER: What's the Real Context Size of Your Long-Context Language Models?.\" arXiv preprint arXiv:2404.06654 (2024). \\\n>[5] Li, Tianle, et al. \"Long-context llms struggle with long in-context learning.\" arXiv preprint arXiv:2404.02060 (2024).\n\n---\n\n## Q6: The goal of this work\n\nPage 3 Footnote 5 claims: *This work is neither a linearization of attention nor attention efficiency*. \n\nThe goal of this paper is to develop new architectures that can better learn new tasks. We neither focus on the performance of persistent-knowledge storage and retrieval (many common benchmarks), nor model efficiency on model serving, nor throughput, nor memory efficiency. Even though they can be future research directions. The three evaluation dimensions in our paper are, therefore, designed to assess the new-tasks-learning capability. \n\nThanks for the valuable suggestion! We will emphasize our goal more clearly in the final version. \n\n---\n\n## Summary\n\n**Q1–Q3 and Q6** concern clarification and explanation issues that could potentially impair readability and understanding. Thanks for pointing them out! We have provided more detailed explanations to address these concerns.\n\n**Q4** highlights the need for formal equations to make the paper scientifically sound; we have added the relevant equations accordingly.\n\n**Q5** points out the need for experimental evidence to support our claim regarding the \"weakness of memory compression methods.\" To address this, we refer to the rich experimental evidence provided by the Ruler benchmark [4] and the in-context learning benchmark [5]. Note that these two benchmarks are the ones used in this paper. \n\n**We hope our responses sufficiently address your concerns!**"}, {"review_id": "uBLKZBwg6i", "reviewer": "Reviewer_jbCa", "summary": "This paper builds on the key-value associative memory architecture known as Memory Mosaics, which addressed some of the inherent limitations of standard Transformers in storing and retrieving long-term knowledge. While the original Memory Mosaics showed promise in improving memory retention and compositional generalization, it was previously explored only in medium-scale settings and synthetic tasks. To overcome these constraints, this paper proposes Memory Mosaics v2, an enhanced version that scales the framework to large models and real-world tasks. Specifically, it introduces three key features: a multi-level memory hierarchy, an adaptive bandwidth, and a gated key extractor. Through extensive experiments, the paper shows that Memory Mosaics v2 consistently achieves performance comparable to standard Transformers while offering stronger memory retention and in-context learning abilities. Additionally, the ablation studies provide insightful analyses of each component’s contribution to the overall framework.", "questions": "In Section 5.2, the paper discusses the failures of many potential baselines. I would like to hear the authors’ thoughts from a slightly different perspective.\n\nFrom the Memory Mosaics architecture viewpoint, the model can be seen as consisting of contextual associative memory and persistent memory components. One could interpret many memory compression algorithms (which act like short-term models due to their difficulty recalling past information) as enhancing the long-term context understanding of the contextual associative memory, ultimately enabling more effective retrieval from the persistent memory.\n\nIf this interpretation aligns with the authors’ view, would it make sense to also compare the proposed model with memory compression algorithms in the experiments for Section 5.1 (persistent-knowledge storage and retrieval), in addition to the comparison in Section 5.2?", "limitations": "Some of the limitations are discussed in the Weaknesses and Questions sections.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### **Strengths**\n\n- **S1.** The paper is well-motivated, addressing the limitations of the existing Memory Mosaic v2 architecture.\n- **S2.** The proposed architecture successfully scales the Memory Mosaics framework to realistic large-model settings (e.g., LLaMA-8B), bridging the gap between small-scale memory-augmented models and production-scale large language models.\n- **S3.** The comprehensive experiments across persistent knowledge, new knowledge storage, and in-context tasks demonstrate that the model consistently matches or exceeds the performance of standard Transformers, with particularly strong results on tasks requiring long-context reasoning.\n- **S4.** The ablation studies and analyses thoughtfully examine the contributions of each architectural component, providing valuable insights into the practical impact and trade-offs of the proposed design.\n\n### **Weaknesses**\n\n- **W1.** Compared to a standard Transformer, the proposed model requires more parameters to achieve comparable performance, which may limit its efficiency in some scenarios.\n- **W2.** While the paper demonstrates the effectiveness of the model’s capacity for in-context learning, it would be helpful to include more detailed explanations, experiments, or analyses on how the capacity scales with the number of Mosaic v2 parameters. This would better clarify where the proposed memory structure provides its main advantages.\n- **W3.** In Section 5.2, the paper mentions comparisons with memory compression algorithms, but does not provide detailed quantitative results. Including concrete empirical comparisons would make the strengths and trade-offs of the Memory Mosaics v2 approach clearer.", "was_revised": true, "final_justification": "They provided additional experimental results showing that Mosaic v2 outperforms the Transformer while requiring only a modest increase in parameters (0.1B for the small model and 1.1B for the large model). They also clarified the design of Mosaic v2 and presented empirical evidence supporting their claims. Based on this, I have increased my score from 4 to 5. I encourage the authors to incorporate the key points from the rebuttal into the camera-ready version of the manuscript.", "rebuttal": "Thank you for your thoughtful reviews! After reading the weaknesses W1-W3, we find that all of them can be effectively addressed by improving clarity and providing more experimental evidence.\n\n---\n\n## W1: Compared to a standard Transformer, the proposed model requires more parameters to achieve comparable performance.\n\nActually, Memory Mosaics v2 achieves better performance under comparable #FLOPs/#parameters, or fewer #FLOPs/#parameters under comparable performance across all three evaluation dimensions.\n\n- **Persistent-knowledge storage and retrieval**: Table 2 in our paper shows that after removing long-term memory, memory moasics v2 contains fewer parameters (8.3B vs 8.8B) and Flops (15.6B vs 16.7B) than Transformer. Despite this reduction, Memory Mosaics v2 still achieves a comparable performance as Transformers. \n\n- **New-knowledge storage and retrieval**: This paper compares memory mosaics v2 and transformer on two scales in Table 4, allowing us to study the relationship between #params/#FLOPs and performance. Table A summarizes the #params, #FLOPs, and performance of these models. By plotting a curve between #params/#FLOPs and performance using the data in Table A, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under any #params or #FLOPs configuration. *(To the ease of reading, we will use a figure to showcase this result in the final version.)*\n| model  | # params | # Flops per token | # ruler question-answering (32k task-length)|\n|---|---|---|---|\n|transformer small |   1.5B  |4.4B| 22.1| \n|transformer large | 8.8B|25.2B |41.1|\n|Memory Mosaics v2 small |1.6B|4.7B|53.4|\n|Memory Mosaics v2 large |9.9B|27.4B|53.4| \\\n*Table A: Ruler question-answering performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.*\n\n\n- **In-context learning**: Similar to the comparison in New-knowledge storage and retrieval, Table B summarizes the #params, #FLOPs, and #performance comparison on in-context learning tasks (TacRed 10-shots). Again, by plotting a curve between #params/#FLOPs and performance using the data in Table B, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under the same #params or #FLOPs configurations.\n| model  | # params | # Flops per token | # in-context learning (tacred 10-shots)|\n|---|---|---|---|\n|transformer small |   1.5B |4.4B| 23.0| \n|transformer large | 8.8B|25.2B |50.0|\n|Memory Mosaics v2 small |1.6B|4.7B|45.5|\n|Memory Mosaics v2 large |9.9B|27.4B|63.0| \\\n*Table B: In-context learning (tacred 10-shots) performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.*\n\n\nEven though W1 is not a genuine weakness of this paper, it helps improve the writing of the paper. Thank you for pointing it out! We will add these results (using figures), as well as clarifications,  in the final version. \n\n---\n\n## W2.1: How does the capacity scale with the number of Mosaic v2 parameters?  \n\nThe previous answer in Tables A and B showcases how the capacity scales with the number of Mosaic v2 parameters and FLOPs, demonstrating that Memory Mosaics v2 outperforms Transformer by over 10% under the same #params or $FLOPs. We will add these results in the final version, addressing this concern. \n\n---\n\n## W2.2: Clarify where the proposed memory structure provides its main advantages.\n\nThe proposed long/short-term memories structure in Memory Mosaics v2 is built based on the \"flat-then-peak\" attention distribution observed in Memory Mosaics [1] (Figure 1 in this paper). To achieve context length extrapolation, we need to allocate the \"flat\" part (position-invariant) to long-term memory, and the \"peak\" tail to short-term memory. This ensures that only the size of long-term memory increases during context length extrapolation. \n\nFor handling in-context learning tasks (few-shot learning), a fair treatment is required for different shots/examples/demonstrations, regardless of their position in a sequence. The position-invariant nature of long-term memory provides this fairness. \n\nIn contrast, transformer attention lacks this \"flat\" attention distribution (Figure 1 left in our paper). Therefore, transformers struggle with context length extrapolation and perform poorly on in-context learning tasks. \n\nWe will visualize the attention distribution of Memory Mosaics v2 with long/short-term memories in the final version, addressing the concern of long/short-term memories. \n\n>[1] Zhang, Jianyu, et al. \"Memory mosaics.\" arXiv preprint arXiv:2405.06394 (2024).\n\n---\n\n## W3: Include experimental evidence to support the claim that \"Memory compression algorithms are weak in New-knowledge storage & retrieval and in-context learning tasks.\"\n\n**Weakness of memory compression algorithms in New-knowledge storage & retrieval**: \n\nThe RULER paper [4] has already studied the weaknesses of model compression methods, such as Mamba and RWKV, via experimental comparison (Figure 4). Here is a summary:\n\n|model |task-length 1k| task-length 2k| task-length 4k|\n|----|----|----|----|\n|Llama2-7b|96.0|91.6|95.0|\n|rwkv-v5-7b|87.5|73.7|51.4|\n|Mamaba-2.8b-slimpj|62.6|52.6|-| \\\n*Table 1: Weakness of Mamba and RWKV on RULER benchmark.* \n\nIn the original words of the RULER paper, \"*Finally, we show that non-Transformer architectures, such as RWKV and Mamba, still lag behind Transformer by large margins on RULER.*\" \n\nNote that the two question-answering RULER tasks used in our paper are notably more challenging than other \"needle in a haystack\" tasks in RULER, like multi-/single-key retrievals. The information entropy of these multi-/single-key retrieval tasks is extremely low, making these tasks easy to compress, allowing many models to easily achieve >99% performance (see Appendix E and F in RULER [4]).\n\n\n**Weakness of memory compression algorithms In-context learning**: \n\n[5] has already experimentally shown the weakness of model compression methods, such as Mamba and RWKV, in in-context learning (Note that this paper uses the same in-context learning data as in [5].).  Here is a summary: \n\n|model |task-length 1 shots| task-length 2 shots | task-length 3 shots|task-length 4 shots| task-length 5 shots|\n|----|----|----|----|----|----|\n|Qwen-1.5-7B-base 7B|  38.7| 47.3| 45.2| 43.6| 40.6|\n|Mistral-7B-v0.2-base|  53.3| 53.1| 51.6| 48.0| 42.3|\n|RWKV-5-World 7B| 2.3 |2.6 |1.0| 0| 1.2|\n|Mamba-2.8B  |0 |0 |0 |0 |0| \\\n*Table 2:  Weakness of Mamba and RWKV on in-context learning benchmark (TacRed classification).*\n\nMemory compression algorithms, such as Mamba and RWKV, completely fail on the in-context learning tasks. \n\nWe will refer these results to experimentally support the claim (weakness of memory compression methods) in the final version. \n\n> [4] Hsieh, Cheng-Ping, et al. \"RULER: What's the Real Context Size of Your Long-Context Language Models?.\" arXiv preprint arXiv:2404.06654 (2024). \\\n> [5] Li, Tianle, et al. \"Long-context llms struggle with long in-context learning.\" arXiv preprint arXiv:2404.02060 (2024).\n\n\n---\n\n## Q1: Perspective on memory compression algorithms\n\nFirst, we agree that the memory and computation efficiency are significant aspects of memory compression algorithms. Our gated recurrent key feature extractor is also inspired by memory compression algorithms. \n\nIn your perspective, you mentioned this: \"*One could interpret many memory compression algorithms [...] as enhancing the **long-term context** understanding of the contextual associative memory, ultimately enabling more effective retrieval from the persistent memory.*\" Note that memory compression algorithms don't contain a true \"**long-term context**\" analogous to either long-term memory in Memory Mosaics v2 or attention in transformer. (Otherwise, these algorithms would not be considered as memory \"compression\".)\n\nDue to the absence of a mechanism capable of storing contextual information over large amounts of new information, memory compression algorithms face difficulties with tasks like new-knowledge storage & retrieval and in-context learning. \n\nFor persistent-knowledge storage and retrieval tasks, we agree that memory compression algorithms perform well. Actually, rich literature (e.g., Mamba, RWKV, xLSTM, etc) has already experimentally verified this point. In our paper, Table 2 (memory mosaics v2 large without long-term memory) also supports this point. Note that \"memory mosaics v2 large without long-term memory\" is a kind of fixed-sized sliding window, and thus acts like a memory compression algorithm. \n\nThanks for the suggestion! We will clarify that *\"memory compression algorithms (Mamba, RWKV, xLSTM, etc) perform well on persistent-knowledge storage and retrieval\"* in the final version, using rich experimental evidence from literature. \n\n---\n\n## Summary\n\n**W1–W2** point out issues related to clarification. We have revised the text to improve clarity, added experimental details, and reorganized the experimental results accordingly.\n\n**W3** highlights the lack of experimental evidence supporting our claims about the weaknesses of memory compression algorithms. Thanks for pointing it out! We now refer to the extensive experimental evidence provided in [4] and [5] to support our argument.\n\n**Q1** raises a question regarding the perspective on memory compression algorithms. We have included a corresponding discussion in the revised version.\n\n**Does the response address your concerns?**"}, {"review_id": "EcXNhqryhD", "reviewer": "Reviewer_yHif", "summary": "This work investigates scaling up the Memory Mosaics model to LLM sizes, trained with real-world datasets. This work builds on the original Memory Mosaics model and introduces Memory Mosaics v2, with three modifications: (i) adaptive bandwidth for associative memory, (ii) gated time-variant feature extractor and (iii) 3-level memory design. Memory Mosaics v2 is trained on one trillion tokens and is thoroughly evaluated across three dimensions: (i) training-knowledge storage, new-knowledge storage, and (iii) in-context learning. The results indicate that Memory Mosaics v2 matches the performance of transformers in learning training knowledge, and outperforms transformers in acquiring new tasks during inference.", "questions": "- How would this model work on generative tasks and natural language modeling?\n- Does this architectural changes incur any computational overhead during inference? How does the performance compare under the same amount of FLOPs?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Demonstrates that the properties of Memory Mosaics model is retained at scale, when trained upto one-trillion tokens of real-world data \n- Demonstrates superiority in learning new-tasks compared to transformers, even when transformers are trained with substantially more data.\n- The new architectural changes are well motivated and offers a simpler way to understand the dynamics of the models\n- Comprehensive evaluation\n\nWeaknesses:\n- Lacks analysis on computational overhead compared to transformers", "was_revised": true, "final_justification": "I believed the work was already solid, with minor suggestions which had been addressed in a convincing manner during the rebuttal. The question I raised was based on my initial misunderstanding, and it was sufficiently clarified during the rebuttal.", "rebuttal": "Thank you for your thoughtful reviews and useful suggestions! It is a pleasure to discuss these questions.\n\n---\n\n## W1 & Q1: Lacks analysis on computational overhead compared to transformers. Do these architectural changes incur any computational overhead during inference? How does the performance compare under the same amount of FLOPs?\n\nThank you for pointing out the need for computation analysis of Memory Mosaics v2 and transformers. We have now included #FLOPs and #parameter analysis for Memory Mosaics v2 and transformers across all three evaluation dimensions. The results demonstrate that Memory Mosaics v2 achieves better performance under comparable FLOPs/parameters or fewer FLOPs/parameters under comparable performance.\n\n\n- **persistent-knowledge storage and retrieval**: Table 2 in our paper shows that after removing long-term memory, memory moasics v2 contains fewer parameters (8.3B vs 8.8B) and Flops (15.6B vs 16.7B) than Transformer. Despite this reduction, Memory Mosaics v2 still achieves a comparable performance as Transformer. *Flops/token is estimated at context length 256 via [1] approach.*\n\n- **New-knowledge storage and retrieval**: This paper compares memory mosaics v2 and transformer on two scales in Table 4, allowing us to study the relationship between #params/#FLOPs and performance. Table A summarizes the #params, #FLOPS, and performance of these models. By plotting a curve between #params/#FLOPs and performance using the data in Table A, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under any #params or #FLOPs configuration. *Flops/token is estimated at context length 32k via [1] approach.*\n| model  | # params | # Flops per token | # ruler question-answering (32k task-length)|\n|---|---|---|---|\n|transformer small |   1.5B  |4.4B| 22.1| \n|transformer large | 8.8B|25.2B |41.1|\n|Memory Mosaics v2 small |1.6B|4.7B|53.4|\n|Memory Mosaics v2 large |9.9B|27.4B|53.4| \\\n*Table A: Ruler question-answering performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.*\n\n- **In-context learning**: Similar to the comparison in New-knowledge storage and retrieval, Table B summarizes the #params, #FLOPs/token, and #performance comparison on in-context learning tasks (TacRed 10-shots). Again, by plotting a curve between #params/#FLOPs and performance using the data in Table B, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under the same #params or #FLOPs configurations. (A similar conclusion also holds for other in-context learning tasks and other numbers of shots). \n| model  | # params | # Flops per token | # in-context learning (tacred 10-shots)|\n|---|---|---|---|\n|transformer small |   1.5B |4.4B| 23.0| \n|transformer large | 8.8B|25.2B |50.0|\n|Memory Mosaics v2 small |1.6B|4.7B|45.5|\n|Memory Mosaics v2 large |9.9B|27.4B|63.0| \\\n*Table B: In context learning performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.*\n\n\nThanks for the useful suggestions! We will add figures for these comparisons and analyses in the final version.\n\n[1] Adam Casson. Transformer flops. 2023\n\n---\n\n## Q2: How would this model work on generative tasks and natural language modeling?\n\nThe **in-context learning** and **new-knowledge storage & retrieval** tasks are evaluated in a generative manner. Specifically, we allow Memory Mosaics v2 to generate texts freely. Then check whether the generated text correctly answers the question. \n\nMemory Mosaics v2 handles generative tasks and natural language modeling similarly to transformers. First Memory Mosaics v2 pre-fill prompts (e.g., a question), store KV-cache of the prompt, and then generate text using this stored KV-cache.  \n\nMeanwhile, other related techniques in transformers, such as KV hashing and quantization, can also be directly applied to Memory Mosaics v2. \n\n\n---\n\n## Summary\n\nThanks for your thoughtful reviews! \n\n**Q1** highlights the need for computational analysis (e.g., FLOPs). We address this point by adding FLOPs and parameter analysis on all three evaluation dimensions.\n\n**Q2** raises the concern about the text generation process of Memory Mosaics v2. We provide clarification, showing that the generation process of Memory Mosaics v2 is similar to that of transformers. Many evaluation tasks in this paper are assessed in the generative manner.  \n\n**We hope the response address your concerns!**"}, {"review_id": "67zZVdUBuI", "reviewer": "Reviewer_YUYJ", "summary": "This paper extends the work of Memory Mosaics, introducing the Memory Mosaics v2 architecture, which they train on 1 trillion tokens.  They demonstrate comparable performance to transformers (trained on the same data-mix and training recipe) on downstream language tasks, while showing superior results on the RULER (needle-in-a-haystack) benchmark, as well as in-context learning tasks.", "questions": "**Questions**\n* Could the authors provide throughput profiles for both the memory mosaics v2 architecture as well as the Transformer architecture (with flash-attention and all relevant optimizations).\n* Could the authors provide additional evaluations, perhaps on the full RULER suite including multi-key retrieval, as well as other long-context benchmarks (e.g. HELMET, BabiLong, PG-19, SWDE, SQA, SNQA)?\n* Did the authors run evaluations on other in-context learning benchmarks?\n* Re Figure 1: Do the different lines correspond to different layers?  Please edit the caption to make the figure self-contained.\n* Memory mosaics (informally) interpolate based on next value $(k_t, v_{t + 1})$, whereas transformers interpolate based on a value at the same position $(k_t, v_t)$.  Has an ablation ever been run where a transformer operates the same as memory mosaics (values shifted forward in time) and the diagonal of the attention mask is removed in order to preserve causality?  This seems sensible to try to see how well transformers perform under the most similar interpolation scheme.\n* If $m < h$, there is no computational saving from splitting into short and long-term memory, but rather the benefit is from separate parameters for long and short term memory?\n* In appendix: \"During training, memory mosaics v2 samples the long-term memory delay step m from [64, 256], sets the short-term memory window size h = 256. At inference, m is set to 64\", what is the intuition behind this?  Did you observe any loss of accuracy without the sampling and keeping m fixed at 64?\n* what is the number of shots for each benchmark in Table 1?  Since few-shot learning is analogous to in-context learning, it would be helpful to potentially vary the number of shots and compare with transformers.  This would be relatively low cost compared to the huge cost of the pre-training experiment, so would be worth adding to the paper.", "limitations": "The authors mention limitations in the \"Computation and # parameters concerns\" section.  They mention the difference in FLOPs/token relative to transformers in Table 2.  However, to be more convincing the authors could provide throughput profiles for Memory Mosaics v2 relative to Transformers.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n* The experiments are large scale (up to 8B model size and 1T tokens)\n* They achieve greater than 10% boost on RULER (NIAH) tasks compared to transformers at 32K context size\n* The memory mosaics architecture demonstrates superior performance on in-context learning tasks\n* The architectural modifications are relatively simple and easy to understand\n\n**Weaknesses**\n* The architectural benefits are demonstrated on a specific RULER task and 3 in-context learning tasks (Banking77, Tacred, Goemotion).  To be truly convincing I believe readers would need to see performance boosts on more benchmarks.\n* While they compare against transformers with a comparable FLOPs/token, there is no training + inference throughput profiling for the memory mosaics architecture relative to transformers.  Throughput is the most important performance metric for large scale training and model serving, thus comparable training and inference efficiency would need to be demonstrated in order for the Memory Mosaics architecture to be adopted at scale.", "was_revised": true, "final_justification": "The authors made efforts to address my concerns in the rebuttal.  While I still have lingering reservations regarding the method's efficiency and also would like to see more extensive evaluations on long-context benchmarks (e.g. HELMET) that are relevant for memory models, I believe this work still has merit and meets the acceptance bar all things considered, and thus I have raised my score from 3 to 4.  I hope that the authors will take reviewer feedback to improve the camera-ready version of the manuscript.", "rebuttal": "Thank you for your thoughtful reviews and insightful suggestions. We appreciate the opportunity to address your concerns. Most of the issues stem from insufficient clarification and misunderstanding regarding the primary goal of this paper. As suggested by reviewer [vHo6](https://openreview.net/forum?id=IfD2MKTmWv&noteId=J6zL2uXCjd), we would like to highlight footnote 5 in our paper, emphasizing that \"**this work is neither an attempt at linearizing attention mechanisms nor focused on model efficiency or model serving**\". \n\n**The goal of this paper** is to develop new architectures capable of effectively learning novel tasks. We do not focus on the performance metrics typically associated with persistent-knowledge storage & retrieval, nor model efficiency, model serving, or throughput. While model efficiency could be a future research direction once one wants to establish model serving at a product level, this paper is currently focused on the early stages of architecture exploration and new-task learning ability.\n\nTherefore, most concerns can be addressed by providing more clarity regarding our goals. We will emphasize the goal of this work in the final version!\n\n---\n\n## Q1: Additional ruler benchmarks \n\nWe chose two question-answering tasks in RULER, because they are notably more challenging than other \"needle in a haystack\" tasks like multi/single-key retrieval, where many models can easily achieve >99% performance (see Appendix E and F in RULER [4]). Note that the information entropy of multi/single-key retrieval tasks is extremely low, making these tasks easy to compress. \n\nOn the other hand, we are happy to share experimental results of Memory Mosaics v2 on these key retrieval tasks. Table A and B show that Memory Mosaics v2 outperforms Transformers.\n\n|model|task-length 4k| task-length 8k| task-length  16k|task-length  32k|\n|----|----|----|----|----|\n|transformer small    |  99.4  |  99.0  |  98.2  |  97.8 |\n|memory mosaics v2 small   |  100.0  |  100.0  |  100.0  |  100.0 |\n*Table A: RULER S-NIAH benchmark comparison between transformer small and memory mosaics v2 small.*\n\n|model|task-length 4k| task-length 8k| task-length  16k|task-length  32k|\n|----|----|----|----|----|\n|transformer large     |  100.0  |  100.0  |  100.0  |  99.6  |\n|memory mosaics v2 large  |  100.0  |  100.0  |  100.0  |  100.0 |\n*Table B: RULER S-NIAH benchmark comparison between transformer large and memory mosaics v2 large.*\n\nThank you for suggesting additional benchmarks such as HELMET, BabiLong, PG-19, SWDE, SQA, and SNQA. However, these benchmarks overlap with each other, RULER, and the three in-context learning tasks presented in this paper. \n\nFor instance: \n- HELMET includes HotpotQA, banking77, and some ruler tasks. \n- The HotpotQA is used in one RULER question-answering task, which is used in this paper. \n- Banking77 is one of 3 in-context learning tasks used in this work.  \n- BabiLong consists of \"needle-in-a-haystack\" tasks with low information entropy, similar to many RULER tasks. \n\nTherefore, this paper doesn't evaluate similar or redundant tasks repeatedly.\n\n> [4] Hsieh, Cheng-Ping, et al. \"RULER: What's the Real Context Size of Your Long-Context Language Models?.\" arXiv (2024).\n\n--- \n\n## Q2: The goal of this work\n\nAs mentioned previously, we would like to emphasize that Footnote 5 states: *This work is neither a linearization of attention nor attention efficiency*. \n\nThanks for the valuable suggestion! We will emphasize this goal more clearly in the final version.\n\n---\n\n## Q3: The number of in-context learning benchmarks\n\nThis paper includes 3 in-context learning datasets, each paired with target label strategies, 2 few-shot prompt strategies, 2 few-shot examples ordering strategies, and 10~16 possible number of shots. This results in over 300 different tasks. Despite that the more benchmarks the better understanding, we find that the 300 tasks in this paper are sufficient to support the in-context learning dimension within a limited space. Note that there are two other evaluation dimensions to explore within the limited pages.\n\nThanks for the suggestions. We will leave more in-context learning benchmarks as future work.\n\n--- \n## Q4: training + inference throughput profiling. \n\nThe primary concern regarding model efficiency and throughput appears to be based on the assumption that this paper targets model efficiency or serving, which is not the case. As mentioned earlier, *this work is neither a linearization of attention nor attention efficiency* (Page 3, Footnote 5). The goal of this paper is to develop new architectures capable of better learning new tasks.\n\nDespite these clarifications, we are also happy to share the training and inference costs for our current Memory Mosaics v2 implementation. \n\nCurrently, Memory Mosaics v2 is implemented based on FlexAttention [5] due to its flexibility. Compared to highly optimized transformer implementations over many years, like flash-attention or fused kernels, FlexAttention [5] achieves approximately 85%–90% of their speed in practice. Therefore, the training cost of Memory Mosaics v2 and baseline Transformer models remains similar. We do not observe any significant training or inference overhead.\n\nThanks for the suggestions on model efficiency. We leave it as a future direction when our model is ready for model serving. \n\n[5] Dong, Juechu, et al. \"Flex attention: A programming model for generating optimized attention kernels.\" arXiv (2024).\n\n---\n\n## Q5: What does the design of long/short-term memory target on, the separation of functionality or computational efficiency? \nAs emphasized by the goal of this paper,  the design of long/short-term memory targets the separation of functionality, rather than computational efficiency. \n\n---\n\n## Q6: Re Figure 1: Do the different lines correspond to different layers?\n\nThe figure is borrowed from [1], where different lines correspond to different attention/memory heads. Thanks for the suggestions! We will edit the caption to make the figure self-contained.\n\n> [1] Zhang, Jianyu, et al. \"Memory mosaics.\" arXiv preprint arXiv:2405.06394 (2024).\n\n---\n\n## Q7: Long-term memory mask is randomized during training while fixed at inference. Visualization of long/short-term memories. \n\nThe motivation for creating long/short-term memories comes from the \"flat-then-peak\" attention distribution in Figure 1 (right). To achieve context length extrapolation, we need to allocate the \"flat\" part (position-invariant) to long-term memory, and the \"peak\" tail to short-term memory. As such, only the size of long-term memory increases during context length extrapolation.\n\nBecause we do not know the exact boundary between \"flat\" and \"peak\" parts, we use a randomized long-term memory mask length to further impair long-term memory from learning the tail of sequence. Note that on average this randomization of mask length equals a linearly-decayed importance. \n\nWe will compare and share the long/short-term attention distribution (visualization) using this randomized mask in the final version, explaining the motivation behind it and showing its benefits. \n\n---\n\n## Q8: Add the number of shots in Table 1.\n\nThanks for pointing it out. Following the setups in Deepseek-V3 [8] and Llama3 [9], we set the number of shots for common benchmarks as shown in Table C. We will clarify the #shots setups in the final version.  \n\n||obqa|arc-easy|windogrande|arc-challange|piqa| boolq| hellaswag| nq| siqa|tqa| gsm8k|mmlu|human eval+|\n|----|----|----|----|----|----|----|----|----|----|----|----|----|----|\n|# shots|0|0|0|0|0|0|0|5| 0|5|8|5|0|\n\n||squad | bbh | math | mbpp | race middle| race high|\n|----|----|----|----|----|----|----|\n|# shots|1 |3 |4 |3 |0 |0 |\n\n*Table C: Number of shots for common benchmarks.*\n\n---\n\n## Q9: Suggestions on varying the number of shots for the common benchmarks in Table 1? \nFirst, we observe in [8-9] that people do not vary the number of shots for these common benchmarks. We conducted a study to explore different numbers of shots and found that performance on these common benchmarks (e.g., MMLU) usually does not increase with more shots. Note that few-shot learning is expected to see improved performance with additional shots.\n\nFurthermore, Deepseek-R1 [7] even observed a decreased performance as the number of shots increased. In their original words:  \"*When evaluating DeepSeek-R1, we observe that it is sensitive to prompts. **Few-shot prompting consistently degrades its performance**. Therefore, we **recommend** users directly describe the problem and specify the output format using a **zero-shot** setting for optimal results.*\" \n\nThese results reveal that these tasks are not classic \"few-shot learning\". The shots in these tasks provide the information of \"problem format\", rather than contextual knowledge. Therefore, this paper, as well as [8-9], doesn't vary the number of shots for these common benchmarks. \n\nThanks for the suggestions. We will add this clarification in the final version. \n\n\n> [7] Guo, Daya, et al. \"Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.\" arXiv  (2025). \\\n> [8] DeepSeek-AI. \"DeepSeek-V3 Technical Report\". (2024). \\\n> [9] Grattafiori, Aaron, et al. \"The llama 3 herd of models.\" arXiv  (2024).\n\n---\n\n# Minor suggested edits. \n\nThanks for the suggestions on editing! We will fix them in the final version. \n\n----\n\n## Summary \n\nThank you for your thoughtful reviews! \n\n**Q1-Q5** raise concerns mainly based on an incorrect assumption --- namely that \"this work focuses on model efficiency or model serving\". We provide details clarification to state our goal clearly, along with additional experimental results. These concerns should be well addressed after clarifying the goal of this work. \n\n**Q6-Q9** ask for details about the model, motivations, and experimental setups. We provide comprehensive explanations and experimental details. \n\n**Does the response address your concerns?**"}]

Author rebuttal:

Meta-review: The paper scales Memory Mosaics to 10B parameters on 1T tokens. It does this by introducing adaptive bandwidth, a gated key extractor, and a three-level memory. It matches transformers on standard benchmarks, and it has advantages for long-context retrieval and few-shot in-context learning. 

Strengths: The architecture is simple, principled, and results in improvements over the original version of Memory Mosaics. Evaluation seems to be done carefully to separate persistent knowledge, new memory, and ICL. The findings directly challenge the prevailing belief that scaling transformers and data is sufficient, and their architecture outperforms transformers by up to 10% (at least in their experiments).

Weaknesses: The paper lacks a transformer baseline with the same masking, making it hard to disentangle the benefits of associative memory. 

Discussion & rebuttal: Reviewers requested masking baselines, more ablations, and stronger efficiency profiling. The rebuttal provided FLOPs and parameter tables, clarified mask construction, and added extended results. While some issues remain, the main findings were strengthened, and most concerns seem to have been addressed.

Decision: Accept. The paper demonstrates an alternative to transformers with clear benefits in its experiments.
Decision: Accept (oral)