Memory Layers at Scale
Authors: Vincent-Pierre Berges
Barlas Oguz
Daniel HAZIZA
Wen-tau Yih
Luke Zettlemoyer
Gargi Ghosh
Venue: ICML 2025 poster (2025)
Keywords: memory
scaling laws
llm
moe
factuality
Primary area: deep_learning->large_language_models
TL;DR: Provides scaling laws for LLMs with memory layers, showing imroved performance and inference speed at similar training FLOP and memory cost compared to MOE layers.

Abstract: Memory layers use a trainable key-value lookup mechanism to add extra parameters to a model without increasing FLOPs.  Conceptually, sparsely activated memory layers complement compute-heavy dense feed-forward layers, providing dedicated capacity to store and retrieve information cheaply.  This work takes memory layers beyond proof-of-concept, proving their utility at contemporary scale.  On downstream tasks, language models augmented with our improved memory layer outperform dense models with more than twice the computation budget, as well as mixture-of-expert models when matched for both compute and parameters.  We find gains are especially pronounced for factual tasks.  We provide a fully parallelizable memory layer implementation, demonstrating scaling laws with up to 128B memory parameters, pretrained to 1 trillion tokens, comparing to base models with up to 8B parameters.

Memory Layers at Scale
Vincent-Pierre Berges
Barlas Oğuz"
Daniel Haziza
Wen-tau Yih
Luke Zettlemoyer'
Gargi Ghosh
Abstract
Factual QA Accuracy vs. Memory Size
Memory layers use trainable key-value lookup mechanism to add extra parameters to model without increasing FLOPs. Conceptually, sparsely activated memory layers complement compute-heavy dense feed-forward layers, pro- viding dedicated capacity store and retrieve information cheaply. This work takes memory layers beyond proof-of-concept, proving their utility at contemporary scale. On downstream tasks, language models augmented with our im- proved memory layer outperform dense models with more than twice the computation budget, as well as offexeeeof-expprt models when matched for both compute and parameters. We find gains are especially pronounced for factual tasks. We provide fully parallelizable memory layer im- plementation, demonstrating scaling laws with to 128B memory parameters, pretrained to trillion tokens, comparing to base models with up to 8B parameters.
NQ
60
TQA
Accuracy
20
1024
2048
40962
8192²
Memory Parameters (x (x2048)
NLL vs. Memory Size
Negative Log Likelihood
12
NQ
TQA
10
1. Introduction
Pretrained language models encode vast amounts of infor- mation their parameters (Roberts et al. 2020), and they can recall and use this information more accurately with in- creasing scale (Brown et al. 2020). For dense deep neural networks, which encode information primarily weights of linear matrix transforms, this scaling of parameter size is directly coupled to an increase in computational and energy requirements. It is unclear if this is the most efficient so- lution to all information storage needs language models. important subset information that language models need learn are simple associations. For example, LLMs learn birthdays of celebrities, capital cities of countries, or how one concept might relate to another. While feed- forward networks can in principle (given sufficient scale)
0
1024
2048
40962
81922
Memory Parameters (x2048)
Figure 1. Scaling the size the memory for a 1.3 billion param- eter model (zero memory parameters corresponds to dense model), trained to I trillion tokens. On the left, factual QA ac- curacy (exact match on NaturalQuestions and score Trivi- aQA), on the right task NLL (lower is better). Dashed lines show performance of a 7B model trained on trillion tokens with 10x more FLOPs.
learn any function (Hornik et al. 1989,, including lookup tables of associations, using an associative memory for this purpose would be both more efficient and more natural.
Equal contribution 'Meta FAIR. Correspondence Vincent- Pierre Berges <vincentpierre @ meeta.omm>. Barlas Oguz <bar- ais@mmet..com..
Such memory layers can be implemented with a simple and cheap key-value lookup mechanism where both keys and values are encoded embeddings (Weston al. 2015).
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Memory Layers at Seale
Earlier works introduced end-to-end trainable memory lay- (Sukhbaatar et al., 2015) and incorporated them as part of neural computational systems (Graves et al. 2014). De- spite early enthusiasm however, memory layers have not been studied and scaled sufficiently to he useful in modern AI architectures. There are distinctive challenges one en- counters when attempting to scale memory layers, which we touch upon in Section 3 In contrast to dense lay- ers which are predominantly FLOP-bound memory lay- ers with their sparse activation pattern are almost entirely memory bandwidth bound. Such components are rarely used in modern architectures and have not been optimised for hardware accelerators. In addition to, and partly as a result of this, little research was done to improve their per- formance. Instead, the field focused on alternatives such as mixture-of-experts (Shazeer et al., 2017), which more closely resemble dense networks and are thus casier to scale, but are challenging to optimize for inference.
haviour and phase transitions (Wei et al. et al. 2022). Nevertheless, given well behaved metric (such as task likelihood loss), most tasks exhibit smooth improvements with scaling (Schaeffer et al. 2023).
2022 Ganguli
Kaplann et al. 2020) showed that performance scales log- linearly with compute and parameter size across a wide range of architecture hyper- paper-parmeters, such as model depth and width. It has been difficult to find architectures which substantially deviate from these laws. Mixture-of- experts (MOE) (Shazeer et al. 2017 Lepikhin et al. 2020) is a notable exception. MOE adds extra parameters to the model without increasing the computation budget. While scaling laws for MOE also mostly focus on training per- plexity, gains transfer well to downstream applications, as evidenced by the popularity of MOE architectures in recent state-of- ootte-of-thheart model families (Jiang et al.] [024]] OpenAl et al. 2024 Team et al.) [024]. Nevertheless, scaling laws for specific task families and capabilities like factuality re- main understudied.
In this work, we show that memory layers, when im- proved and scaled sufficiently, can be used to augment dense neural networks to great benefit. We do so by replac- ing the feed-forward network (FFN) of one or more trans- former layers with memory layers (we leave other layers unchanged). These benefits are consistent across a range of base model sizes (ranging from 134 million to 8 bil- lion parameters), and memory capacities (up to 128 bil- lion parameters). This represents a two orders of magni- tude leap in memory capacity compared to previous mem- ory layers reported in the literature. Our results (Section ( indicate that memory layers improve the factual accuracy of language models by over 100% as measured by fac- tual QA benchmarks, while also improving significantly on coding (HulanEval, MBPP) and general knowledge (Hellaswag, MMLU). In many cases, memory augmented models can match the performance of dense models that have been trained on 4x more compute. They also outper- form mixture-of-expert architectures with matching com- pute and parameter size, especially on factual tasks. Given these findings, we strongly advocate that memory layers be integrated into all next generation AI architectures.
Like MOE, memory augmented models also aim to aug- ment the parameter space of the model without adding sig- nificant computational cost. Memory networks were pro- posed initially in (Weston al.
2015), and with end-to-end
training in (Sukhbaatar et al. 2015) Neural Turing Ma-
chines (Graves et al. 2014 2016) combine external train- able memory with other components to build a neural train- able computer. Product-key networks (Lample et al. 2019) were introduced to make the memory lookup more effi cient and scalable. The recent PEER (He 2024) builds on this work, replacing vector values with rank-one matrices, forming bridge between memory architectures and MOE, Memory networks make use of form of sparse attention using a top-k lookup mechanism similar to Top-k Atten- tion (Gupta et al., 202T). Moreover, the idea of replacing MLPs with attention mechanisms is similar to what was done in "Augmenting Self-attention with Persistent Mem- ory" (Sukhbaatar et al.; (201)).
Factual text generation has long been considered funda- mental capability for generative models, typically bench- marked through factual open domain question answer- ing (Chen et al. 2017 Chen & Yihl (020) and other Enowledge-intensive tasks (Petroni et al. 2021). Being able to memorize the facts in the training corpus enables the model to answer fact-seeking, knowledge intensive tasks more factually and accurately. Indeed larger models have been shown to be more factual (Roberts et al.) 2020 Brown et al. 2020), but even modern LLMs are known to struggle with hallucination (Ji et al. 2023). A tested way of ensur- ing more factuality through retrieval augmented genera- tion (Lewis et al. 2021
2. Related work
Language model scaling laws (Kpllan et al. 2020) study the empirical performance of language models as they are scaled in compute, data, and parameter size. Scaling laws are typically formulated in terms of training/test log like- lihood, which is generally believed to correlate well with downstream performance. Scaling plots on downstream tasks are also not without precedent (Brown et al. 2022). but have sometimes been shown to exhibit non- linear be-
2020 Lee al.) Khandelwal et al. 2020). We use
Karpukhin et al.
2019 Guu et al. 2020
OOur implementation is available at https: / http:/gittttt com/ http//acebockeesearcchc////
short-form QA tasks in this work to demonstrate the effec-
2

---

Memory Layers at Scale
tiveness of memory layers and leave the long-form gener- ation tasks for future work. Recently, wide literature has emerged in mitigating LLM hallucinations through data re- lated methods, architecture variants, pre gre-training and infer- ence time improvements. We refer to (Ji et al. 2023) sec- tion 5 for a comprehensive survey.
Let I1, and $1, 82 be the top-k indices and scores ob- tained from the respective key sets K1, K2. Since there are only VN keys in each set, this operation is efficient. The overall indices and scores can be found by taking argmax i. 81 [&1] + s2[i2].
3.1.2. PARALLEL MEMORY
3. Memory Augmented Architectures
Memory layers are naturally memory-intensive mostly due to the large number of trainable parameters and associated optimizer states. To implement them at the scale of several millions of keys, we parallelize the embedding lookup and aggregation across multiple GPUs. The memory values are sharded across the embedding dimension. At each step, the indices are gathered from the process group, each worker does a lookup and then aggregates the portion of embed- dings in its own shard. After this, each worker gathers the partial embeddings corresponding to its own portion of the indices. We take care to keep activation memory manage- able at this stage, by making sure each GPU only gets its own portion, and does not need to instantiate the entire em- bedding output. The process is illustrated in Figure [ The implementation is independent of other model parallelism schemes such as tensor, context or pipeline parallelism, and operates on its own process group.
Trainable memory layers work similarly to the ubiqui- tous attention mechanism (Bahdanau et al. 2016). Given query q E IR" set of keys K E RN and val- ues V E <R-x", the output is soft combination of val- ues, weighted according to the similarity between q and the corresponding keys. Two major differences separate memory layers from attention layers as they are typically used (Vaswani et al. 2023). First, the keys and values in memory layers are tranable parameters, as opposed to acti- vations. Second, memory layers typically have larger scale in terms of the number of keys and values, making sparse lookup and updates a necessity. For example, in this work, we scale the number of key-value pairs to several millions. In this case, only the top-k most similar keys and corre- sponding values take part the output. A simple memory layer can be described by the following equations:
I = TopkIndices(Kg).
Softmax y= sVs (1)
Here / is a set of indices, E RR Kr, E Bkxn, and the output E R.. Each token embedding (for us, the output of the previous attention layer) goes through this memory lookup independently, similar to the FFN operation that we replace.
3.1.3. SHARED MEMORY
Deep networks encode information at different levels of ab- straction across different layers. Adding memory to mul- tiple layers may help the model use its memory in more versatile ways. In contrast to previous work (Lample et al.] 2019), we use a shared pool of memory parameters across all memory layers, thus keeping parameter count the same and maximizing parameter sharing. We find that multi- ple memory layers increase performance significantly over having a single layer with the same total parameter count, up to a certain number of layers (in our case, 3). Beyond this point, replacing further FFN layers degrades perfor- mance, showing sparse and dense layers are both needed and likely complementary (see Section 5.4) 5.44.
3.1. Scaling memory layers
Being light on compute, and heavy on memory, memory layers have distinet scaling challenges. We detail some of these challenges and how we address them in this section.
3. PRODUCT-KEY LOOKUP
One bottleneck which arises when scaling memory layers is the query-key retrieval mechanism. A naive nearest- neighbour search requires comparing each query-key pair, which quickly becomes prohibitive for large memories. While fast approximate vector similarity techniques (John- son et al. 2019) could be used here, it's a challenge to in- corporate them when the keys are being continually trained and need to be re-indexed. Instead, we adopt trainable produet-quantized keys from (Lample et al.) 2019). Prod- uct keys work by having two sets of keys instead of one, where K1, K2 E R R The full set of keys of size N x n, which is never instantiated, consists of the prod- uct of these two sets. The top-k lookup on the full set of keys can be efficiently done by searching the much smaller set of keys first, saving compute and memory. To per- form the lookup, we first split the query as qı, 42 E R
1.1.4. PERFORMANCE AND STABILITY IMPROVEMENTS
The main operation in the memory layer is to compute the weighted sum of the top-k embeddings: it is implemented in Py PoTohh's EmbeddingBag operation. As the number of floating-poin operations is negligible, we expect this operation to be solely limited by the GPU memory band- width but find multiple inefficiencies in PyTorch's imple- mentation in practice. We implemented new and more ef- ficient CUDA kernels for this operation. Our forward pass optimizes memory accesses and achieves 3TB/s of mem- ory bandwidth, which is close to our H1 specification of 3.35TB/s (compared to less than 400GB/s with Py Torch''s
3

---

Memory Layers at Scale
GPU-0
Share Indexes
idi
Embeddings
18x
IBk
18x
18k
GPU-1
ICW
X
ISK
Idx
IBX E
Embeddings
Figure 2. Illustration of the parallel EmbeddingBag implementation for "Memory Group" of two GPUs. Each GPU performs the EmbeddingBag operation on all of the indices of the group, but on half-dimension embeddings it has access
implementation). The backward pass is more complicated as multiple output gradients have to be propagated to the same weight gradient. We benchmarked multiple strate- gies: using in-memory atomic additions, using locks, or computing the inverse mapping from embedding.id id to to- ken_ids to make the backward pass embarrassingly parallel. We exclusively used the atomic strategy in our trainings because the best performing, but the other ones were competitive, and could be used when the hardware does not support atomics in the given data type.
the activated parameters. Even with small batch size, MoE models can activate large portion of their param- eters, making the decoding time many times slower than that of a dense model. Models with memory layers, on the contrary, only activate a small subset the entire embed- ding table for every token in the batch, hence the number of parameters activated remains roughly constant in the small batch decoding regime. This makes decoding of memory transformers up 5x faster than MoE models in the small batch size regime (see Figure 4
Overall, our custom kernels make the embedding bag operation end-to-end 6x faster compared to PyTorch's EmbeddingBag for our use cases.
4. Experimental setup
For our base model architecture, we follow closely the Llama series of dense transformers (Touvron et al. 2023 Dubey et al. 2024), which also serve our dense base- lines. We augment the base models by replacing one or more of the feed-forward layers with a shared memory layer. For scaling law experiments, we pick base model sizes of 134m, 373m, 720m, and 1.3b parameters. For these models, we use the Ilama2 tokenizer with 32k tokens, and train IT tokens with a pretraining data mix that is similar to that of Llama2 (Touvron et al.) 2023). For experiments the 8B base model scale, we use the Llama3 (Dubey et al.) 2024) configuration and tokenizer (128k tokens), and a bet- ter optimized data mix similar to Llama3.
3.2. Architectural improvements to memory layers
We improve training performance of the memory layer by introducing input-dependent gating with silu non- linearity (Hendrycks & Gimpel 2023). The output
in Equation 1 then becomes
output = (y @ silu(a W)))
(2)
where silu(x) sigmoid(x) and @ is the element-wise multiplication(see also Figure 3 We find that for large memory layers, training can become unstable, especially for small base models. We use nk-normalization (Team
2024) when needed to alleviate this issue.
4.1. Baselines
3.3. Decoding inference efficiency
In addition to the dense baselines, we also compare to other parameter augmentations including mixture-of- experts (MOE) (Saazeer et al.) 2017) and the more recent PEER (He) 2024) model. MOE, each FFN layer is com- posed of multiple "experts", only a subset of which partici- pate in the computation for each input. The PEER model is conceptually similar memory layer, but instead of trieving single value embedding, it retrieves pair of em- beddings, which combine into a rank-1 matrix. Several of these are assembled together into a dynamic feed-forward layer. PEER works similarly to memory layers in prac-
Memory layers have important advantages over alternatives like Mixture-of-Experts (MoE) when it comes to inference efficiency. Inference involves two phases: prefilling and decoding. Prefilling typically processes many tokens, and in the same compute- -ouud regime training, where memory layers only incur a small overhead. However, dur- ing decoding, GPU typically processes only few tokens at a time, due to either latency requirements, or memory limitations. This process is memory bandwidth bound: the decoding time for single token depends on the size of
4

---

Memory Layers at Scale
Regular Memory Block
Values
softmax
top-k
FFN input
FFN output
Sub Key2
Key1
Query MLP
Memory+ Block
Regular Memory Block
W2
FFN input
FFN output
W1
silu
Figure 3. On the top, the regular memory layer. On the bottom, the Memory+ block, with the added projection, gating and silu non- linearity
Forward time (dim=2048, n layers=222
Forward time (dim=8192, n n_layers=48)
0.005
dense
dense
mem-layers (1M entries, top 128) (4 experts)
0.14
mem-layers (16M entries, top 128)
0.004
experts)
0.12
runtime forward
runtime per forward
0.10
0.003
0.08
0.002
0.06
0.04
0.001
0.02
0.000
0.00
10
10
10a
10
10
102
103
number of tokens in batch
number of tokens in batch
Figure 4. Roofline analysis of the decoding time for a dense model, an model, and a Memory transformer. models have the number of activated parameters per token. The MoE and Memory models have the same number of total parameters.
tice, but requires twice the number of parameters for the same number of keys. Like memory layers, these methods increase the number of parameters in the model without significantly increasing FLOPs. We pick the number of ex- perts in MOE and the number of keys in PEER to match the number of parameters our memory-augmented mod- els as closely as possible. MOE models are trained with expert choice (Zhou et al. 2022). and evaluated with top- routing. PEER layers share the same configuration and hyper-parameters as memory layer implementation.
4.2. Evaluation benchmarks
Our evaluations cover factual question answering (Natu- ralQuestions (Kwiatkowski al., 2019). TriviaQA (Joshi)
et al. 2017)). multi-hop question answering (Hot-
potQA Yang et al., 2018)), scientific and common sense world knowledge (MMLU (Hendrycks et al. 2021). Hel- laSwag (Zellers et al. 2019), OBQA (Mihaylov et al.] 2018), PIQA (Bisk et al., 2019)) and coding (Hu- manEval (Chen et al. 2021). MBPP (Austin et al. 2021 We try report the most commonly used accuracy ries (exact match Fl score for QA benchmarks, pass- at-l for coding). For some benemarks, the performance of small models can be very low, and accuracy numbers
5

---

Memory Layers at Seale
noisy. Therefore use negative log-likelihood (nll) of the correct answer for model ablations.
NaturalQuestions
15
Memory+
5. Scaling results
MOE
Dense
We compare Memory models to baselines in a compute- controlled setting.
Accuracy
5.1. With fixed memory size
5
First, we fix the size of the memory, and therefore the number of extra parameters, and compare with the dense baseline, as well as roughly parameter matched MOE and PEER models. Models with the same base model configu- ration have negligible differences in FLOPs. For Memory models, we fix the number of half keys to 2 and thus the number of memory values to 2 (roughly 1 million). For the PEER baseline, we pick the number of half-keys to 768, resulting in slightly more total parameters than Memory. For MOE models, we pick the lowest number of experts such that the parameter count exceeds that of Memory. This corresponds to 16, 8, 6, and 4 experts for the 134m, 373m, 720m and 1.3b sizes respectively.
134m
373m
720m
1.3b
Base Parameters
TriviaQA
40
fl-score
20
The vanilla Memory model has a single memory layer, which we pick to replace the middle FFN layer of the trans- former. Our improved Memory+ model has 3 memory lay- ers, placed centered with a stride of for the 134m mod- els and for the others. Additionally it includes a custom swilu non-linearity, and optimized key dimension (set to equal half of the value dim). As noted earlier, memory lay- ers share parameters, thus have identical memory footprint to single memory layer.
Memory+ MOE
Dense
0
124m
373m
720m
1.3b
Base Parameters
Figure Accuracy vs. Base Parameters for NaturalQuestions and TriviaQA (Memory+ models use I million memory embeddings.)
We can see from Table [ that Memory models improve drastically over the dense baselines, and generally match the performance of models with twice the number of dense parameters on QA tasks. Memory+ improves further over Memory, with performance falling generally between dense models with 2x-4x higher compute. The PEER archi- tecture performs similarly to Memory for the same num- ber of parameters, while lagging behind Memory+. MOE models underperform the memory variants by large mar- gins. Figure shows the scaling performance of Memory, MOE and dense models QA tasks across various base model sizes.
10x more FLOPs. (see also Table 2
5.3. Results at 8B scale
Finally, we scale our Memory+ model with an 8B base model and 4096² memory values (64B memory parame- ters). We use the Llama3 8B (Dubey et al. 2024) archi- tecture and tokenizer, and train on a data mix similar to Llama3 (Dubey et al. 2024). We report results at 200 bil- lion and trillion tokens of training Table 2 On an ex- panded set of benchmarks, including general scientific and world knowledge and coding, we see that memory aug- mented models significantly outperform dense baselines.
5.2. Scaling memory size with fixed base model
Next, investigate scaling behaviour with respect to the memory size for a fixed base model. In Figure D we see that factual QA performance for a Memory+ model keeps increasing predictably with increasing memory size. At 64 million keys (128 billion memory parameters), 1.3b Memory model approaches the performance of the Llama2 7B model, that has been trained on 2x more tokens using
The gains are more pronounced earlier in training (200B tokens), suggesting that memory helps models learn facts faster. At only trillion tokens of training, our Memory- model approaches the performance of Llama3.1 8B, which was trained on 15 trillion tokens.
6

---

Memory Layers at Scale
Table 1. Comparing memory augmented architectures with baseline models on QA tasks. Memory models have million value embed- dings unless otherwise specified in the model configuration column. Metrics are accuracy for NQ. PIQA, OBQA and FI score for TQA, HotpotQA.
Base
Model
Total
Params
Tokens
FLOPS
Configuration
Params
NQ
TQA
PIQA
OBQA
HotPot
Dense
134m
0.91
7.7
62.13
16.40
5.18
MOE
984m
2.49
13.08
65.78
18.80
7.80
134m
IT
7.9e20
PEER
1.037b
2.46
16.34
67.25
17.40
8.82
Memory
937m
2.1
16.31
66.65
17.80
9.28
Memo ry+
937m
3.16
18.77
65.94
17.60
9.35
Dense 373m
2.58
17.68
67.47
18.80
10.06
MOE
1.827b
3.99
19.94
68.88
22.20
12.50
373m
IT
2.6c21
PEER
1.575b
5.1
26.39
70.19
21.60
12.96
Memory
1.441b
4.95
24.24
69.37
20.40
12.53
Memory+
1.434b
5.76
28.10
71.22
22.00
13.34
Dense 720m
3.77
24.85
71.33
22.60
12.90
MOE
2.768b
7.04
28.08
70.08
20.80
14.10
720m
IT
4.9e21
PEER
2.517b
7.92
33.26
71.98
25.00
14.03
Memory
2.316b
7.2
34.8
71.82
24.40
14.94
Memory+
2.316b
9.39
36.67
72.42
24.00
14.92
Dense 1.3b
7.76
32.64
72.74
23.40
13,92
MOE
3.545b
8.14
31.46
73.72
25.20
15.15
PEER
3.646b
12.33
42.46
73.34
26.60
15.39
Memory
3.377b
9.83
39.47
72.29
25.80
15,46
1.3b
IT
8.5e21
Memoryt
3.377b
13.68
42.89
75.35
26.80
16.72
Memory+ 4m
9.823b
14.43
51.18
75.03
27.80
18.59
Memory+ 16m
35.618b
20.14
58.67
76.39
26.80
20.65
Memory+
64m
138.748b
20.78
62.14
77.31
30.00
20.47
llama2 7B
2T
9.1e22
Dense 7b
25.10
64.00
78.40
33.20
25.00
Table 2. Results with an 8B base model. Memory+ models have 16 million memory values (64 billion extra parameters). Metrics are accuracy NQ. PIQA, OBQA, HellaSwag, MMLU; F1 score TQA, HotPoiQA; pass for HumanEval, MBPP. The number training tokens for each model denoted in parenthesis.
Model (8B)
Tokens
FLOPS
HellaS.
Hotpot
HumanE.
MBPP
MMLU
NQ
OBQA
PIQA
TQA
I
1ST
6.8e23
60.05
27.85
37.81
48.20
66.00
29.45
34.60
79.16
70.36
dense
200B
9.1c21
53.99
20.41
21.34
30.80
41.35
18.61
31.40
78.02
51.74
Memory+
200B
9.1e21
54.33
21.75
23.17
29.40
50.14
19.36
30.80
79.11
57.64
dense
IT
4.6c22
58.90
25.26
29.88
44.20
59.68
25.24
34.20
80.52
63.62
Memory+
1T
4.6e22
60.29
26.06
31.71
42.20
63.04
27.06
34.40
79.82
68.15
5.4. Model ablations
formance initially increases. However, as we're effectively removing dense parameters from the model for each added memory layer, eventually the model performance degrades, revealing sweet spot at around memory layers (Table top). Moreover, we experiment with the placement of these layers, modifying the centring and spacing. We find that centred placements with larger strides are better, and we adopt this for our Memory architecture.
In this section, we present results which motivate our mod- elling choices for the Memory+ architecture.
Memory layer placement Since the memory pool is shared, we can replace more FFN layers with memory lay- ers without increasing either the memory or the compute budget. We see that as we add more memory layers, per-

---

Memory Layers at Seale
Memory layer variants We experiment with minor modifications to the memory mechanism (Table bottom). We try 1. gating the memory with the input using a linear projection, 2. adding a custom swilu non-linearity (Fig- 3). 3. adding random key-value pairs in addition to the top-k during pre-training to unbias key selection, 4. adding a single fixed key-value pair (softmax sink) to the top-k se- lected values during pre-training to serve as "anchor". We find that the swilu non-linearity consistently improves re- sults, and we adopt this improvement into our model. Sim- ple gating improves performance only in some cases, and swilu already covers this behaviour to some extent, so we decide not to do additional gating. For key sampling im- provements, including the random keys and the fixed (sink) key, we see minor improvements, however these have some negative impact on training speed in our implementation, and the gains were not consistent for larger model sizes, therefore we excluded them from our experiments, leaving this direction open for future exploration.
ters the model, and thus we cannot increase it indefinitely without breaking fair comparisons. We pick key dimen- sion of half the base model dim for our experiments.
nll
NQ nll
TQA nll
v_dim
#values
64
16m
2.15
12.86
8.75
256
4m
2.14
12.63
8.49
1024
1m
2.11
12.13
8.34
2048
512k
2.14
12.49
8.53
nll
NQ nll
TQA nll
key dim
256
2.11
12.13
8.34
512
2.12
12.32
8.15
1024
2.11
12.37
8.25
2048
2.09
11.98
7.83
Table 4. Ablation studies: on the top, varying the value embed- ding dim while keeping total parameter count the same, on the bottom varying key dim. Metrics are all log likelihood, training set, NQ answers and TQA answers. These were ran on the 373m model size, which uses latent dimension of 1024. key dim is the sum of the dimension of the sub-keys.
nll
NQ nll
TQA nll
layer #
the
12
2.11
12.13
8.34
12,16,20
2.08
11.60
7.54
1,12.16
2.07
11.79
7.64
4,12,20
2.06
11.32
7.20
5,8,11,14,17,21
2.11
11.79
7.73
6. Implications and shortcomings of the work
nll
NQ all
TQA nll
Scaling of dense transformer models has dominated progress in the AI field in the last 6 years. As this scal- ing is nearing its physical and resource limits, it's useful to consider alternatives which might be equally scalable with- out being as compute and energy intensive. Memory layers with their sparse activations nicely complement dense net- works, providing increased capacity for knowledge acqui- sition while being light on compute. They can be efficiently scaled, and provide practitioners with an attractive new di- rection to trade-off memory with compute.
Model
PK base
2.11
12.13
8.34
+gated
2.11
12.24
8.17
+swilu
2.11
12.05
8.09
+random values
2.11
12.36
8.09
+softmax sink
2.11
12.19
8.04
Table 3. Ablation studies: on the top, number of memory layers with shared memory, on the bottom different memory architecture variations. Metrics are all log likelihood, the training set, NQ answers and TQA answers.
While the memory layer implementation presented here is orders of magnitude more scalable than previous works, there still remains substantial engineering task to make them efficient enough for large scale production uses. Dense architectures have been optimized for and co- evolved with modern GPU architectures for decades. While we believe it's in principle possible to make mem- ory layers as fast, or even faster than regular FFN layers, acknowledge that this needs non-trivial effort.
Key and value dimension By default, the memory value dimension is chosen to be the same as the base model di- mension. However, we can potentially trade-off the value dimension with the number of values in the memory with- out changing the total parameter size of the memory using an extra projection after Memory. We present this ablation in Table 4 top, and find that the default configuration is optimal. We can also independently increase the key em- bedding dimension, which we do in Table g bottom. We find unsurprisingly that increasing the key dim is beneficial. However, increasing the key dim adds more dense parame-
We have so far presented only high level empirical evi- dence that memory layers improve factuality of models. However, we believe the sparse updates made possible by memory layers might have deep implications to how mod- els learn and store information. In particular, we hope that
8

---

Memory Layers at Scale
new learning methods can be developed push the effec- tiveness of these layers even further, enabling less forget- ting, fewer hallucinations, and continual learning.
et. al., A. S. The Ilama 3 herd of models, 2024, URL
https: httpss///xxvvvoorgg.o wtt.ooggg/bb///0/0 22177.22188
Ganguli, D., Hernandez, D., Lovitt, I.,, Askell, A., Bai, Y. Chen, A., Conerly, T., Dassarma, N., Drain, D., El- hage, N., and EI Showk, S. a. Predictability and sur- prise in large generative models. In 2022 ACM Con- ference on Fairness, Accountability, and Transparency, FAccT '22. ACM, June 2022. doi: 10.11453333 3533229. URL http::/dx..dee http:///..oii..o/.....1..01 3531146.533227
Impact Statement
This goal of this work is to advance the field of ML. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Graves, A., Wayne, G., and Danihelka, I. Neural turing
Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., and Sutton, C. Program synthesis with large language mod- http::/:///xxx.x..../.o//
machines, 2014. URL
org/abs/ abs/
1401
cls. 2021.
Graves, A., Wayne, G., Reynolds, M., Harley, T., Dani- A., Colmenarejo, S. G., Grefenstette, E., Ramalho, T., Agapiou, J.. Badia, A. P., Hermann, K. M., Zwols, Y., Ostrovski, G., Cain, A., King, H., Summerfield, C., Blunsom, P., Kavukcuoglu, K., and Hassabis, D. Hybrid computing using neural network with dynamic external memory. Nature, 538 7666))471-476 October 2016. ISSN 00280836. URL hetp://dx.doi.. www.doo....gg..... org/ www.//0/attre20000
07732
helka, I.,
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate, 2016. URL
0473
Bisk, Y., Zellers, R., Bras, R. I.., Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural language, 2019. URL
Gupta, A., Dar, G., Goodman, S., Ciprut, D., and Berant, J.
11641
Memory-efficient transffrres via top-k attention, 2021. URL httttttss///aattea org/ abs////////// 220./1110..0 06899.6.66699
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.. Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G. and et. al., A. A. Language models are few-shot learners, 2020. URL https::/arxiv.org//eet 2005. 14165]
Guu, K., Lee, K., Tung, Z., Pasupat, P., and Chang, M.- W. Realm: Retrieval-augmented language model pre- training, 2020. URL https::/arxiv..rrt /bbs/
Chen, D. and Yih, W.-t. Open-domain question answer- ing. In Savary, A. and Zhang, Y. (eds.), Proceed- ings of the 58th Annual Meeting of the Association for Computational Linguistics: Tutorial Abstracts, pp. 34-37, Online, July 2020. Association for Computa- tional Linguistics. doi: 10. 1011863//11000..ac//oonaa 8. URL httpps:/accaattol..ee ort.a//tthoooooo..oo//0//00.... acl -tuutoiaals..
He, X. O. Mixture of a million experts, 2024. URL /ttp///abbb2240004153
Hendrycks, D. and Gimpel, K. Gaussian error linear units
(gelus), 2023. URL https:://rxxvv.org/
1606 1606.08415
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. 2021. htLpst
Chen, D., Fisch, A., Weston, J., and Bordes, A. Reading wikipedia to answer open-domain questions, 2017. URL htttps://arxiv
wttp:///rgg/b//////a/o/o 2tbb/220903300 03300
awwp////b///0/./..//// 00051
Homik,, K., Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. Neu- ral Networks, 2(5): 1989 ISSN 0893-6080. doi: URL httpss//www..cciennceeeeee.. siiene//rticee/////3338888999
Chen, M., Tworek, J., Jun, H., Yuan, Q., de Oliveira Pinto, H. P., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., and et. al., H. K. Evaluating large language models trained on code, 2021. URL htt.tttss///xxi..orrr abs/2107. 03374
Ji, Z., Lee, N., Frieske, R., Yu, T., Su, D., Xu, Y., Ishii, E., Bang, J.J. Madotto, A., and Fung, P. Survey of halluci- nation in natural language generation. ACM Computing Surveys, 15(12):1-38, March 2023. ISSN 1557-7341
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., and

---

Memory Layers at Seale
doi: 10.1145/3571730. URL http: /ttttp:/dxx.di..... org/
Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Küttler, H., Lewis, M., tau Yih, W., Rocktäschel, T., Riedel, S., and Kiela, D. Retrieval- augmented generation for knowledge-imlensive nlp tasks, 2021. URL /tttps://arxvv.orgg
10. 15.14533357730
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G. Bour, G., and et. al., G. L. Mixtral of experts, 2024. URL https: / /arxiiivv....... www.i.owg/aaa...aa...t......aaee
2005. 11401
2401 04088
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering, 2018. URL https:// abs http://b//11890 /w0./8809.000 02789
Johnson, J., Douze, M., and Jégou, H. Billion-scale similar- ity search with GPUs. IEEE Transactions on Big Data, (13):535-547, 2019.
OpenAI, Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L.. Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., Avila, R., Babuschkin, I., and et. al., S. B. Gpt-4 technical report, 2024. URL https: :
Joshi, M., Choi, E.,, Weld, D. S., and Zettlemoyer, L. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. 2017. URL
08774
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language mod- els, 2020. URLitttss: / //trxvaaavv////////a//eaaeeawavl org/ abs/ /2011.
Petroni, E. Piktus, A., Fan, A., Lewis, P., Yazdani, M., Cao, N. D., Thorne, J., Jernite, Y., Karpukhin, V., Mail- lard, J., Plachouras, V., Rocktäschel, T., and Riedel, S. Kilt: a benchmark for knowledge intensive lan- guage tasks, 2021.
08361
URL
Karpukhin, V., Oğuz, B., Min, S., Lewis, P., Wu, L., Edunov, S., Chen, D., and tau Yih, W. Dense passage re- trieval for open-domain question answering, 2020, URL https: wwtpi/iwwigoaaaaaaaaaoo////oooaa abo///as./2002///// 04906
Roberts, A., Raffel, C., and Shazeer, N. How much knowl- edge can you pack into the parameters of a language model?, 2020. URL https://arxiv.or abs/
Khandelwal, U., Levy, O., Jurafsky. D., Zettlemoyer, L., and Lewis, M. Generalization through memorization: Nearest neighbor language models, 2020. (tthttps:
2002. 08910
Schaeffer, R., Miranda, B., and Koyejo, Are emergent abilities of large language models mirage?, 2023. URL 15004
//arxiv. wrr.iivvvrrr//oo abs/ 1911. 00172
Kwiatkowski, T., Palomaki, J., Redfield, O., Collins, M., Parikh, A., Alberti, C., Epstein, D., Polosukhin, I., Devlin, J., Lee, K., Toutanova, K., Jones, L., Kel- cey, M., Chang, M.-W., Dai, A. M., Uszkoreit, J., Le, Q. and Petrov, S. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466, 2019. doi: 10.1162/tacl a 00276.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-expertt layer, 2017. URL https:
06538
Sukhbaatar, S., szlam, a. Weston, J., and Fergus, R. End-to-end memory networks. In Cortes, C., Lawrence, N., Lee, D., Sugiyama, M., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 28. Curran Associates, Inc.. 2015. URL httpp::/pprcceeenggg.o
URL https://
acl aclantholon wowh. org. 0111//09911100
Lample. G., Sablayrolles, A., Ranzato, M., Denoyer, L., and Jégou, H. Large memory layers with prod- uct keys, 2019. URL
1907. 05242]
cc hc//appet..files/pppee////////11
Lee, K., Chang, M.-W., and Toutanova, K. Latent re- trieval for weakly supervised open domain question an- swering, 2019. URL
pdf
Sukhbaatar, S., Grave, E., Lample, G., Jegon, H., and Joulin, A. Augmenting self-attention with persistent
1906. 00300]
memory, 2019. URL
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang. Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scal- ing giant models with conditional computation and auto- matic sharding, 2020. URL
1907.01477 01470
Team, C. Chameleon: Mixed-modal early-fusion foun-
http:://xxxvvrrrr///
dation models, 2024.
URL httts://arxiv..ozz
abs/ 2006. 6668
abs/2405. 240.2.955888
10

---

Memory Layers at Scale
Team, G., Georgiev, P., Lei, V. I., Burnell, R., Bai, L.,, Gu- lati, A., Tanzer, G., Vincent, D., Pan, Z., Wang, S., Mar- iooryad, S., Ding, Y., Geng. X., Alcober, F,, and et. al., R. F. Gemini 1.5: Unlocking multimodal understand- ing across millions of tokens of context, 2024. URL. httpss://axxyy: 24tp://.b///22....555
Touvron, H., Martin, L., Stone, K., Albert, P., Alma- hairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., and et. al., C. Llama Open founda- tion and fine-tuned chat models, 2023. URL https: 7/arxiv. orq/ abs/ 2307 09288
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, Atten- tion is all you need, 2023. URL
org/ abs/ /1766 03762
Wei, J. Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., Chi, E. H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W. Emergent abilities of large lan- guage models, 2022. URL https://arxiv..rr org/
abs / 221//2266.07 07682
Weston, J., Chopra, S., and Bordes, A. Memory net- works, 2015. URL
1410. 3316
Yang, Z., Qi, P., Zhang, S., Bengio, Y., Cohen, W. W., Salakhutdinov, R., and Manning, C. D. Hotpotqa: A dataset for diverse, explainable multi-hop question an- swering, 2018. URL https::/aaxxv.org///
1809.09600
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sen- tence?, 2019. URL https:
1905.07830
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A., Chen, Z., Le, Q. and Laudon, Mixture-of- experts with expert choice routing, 2022. URL https: 7/arxiv verxiiior www.vvorrrggaaoo htt.ab////22292222 09368

---

Memory Layers at Scale
A. Training hyperparameters
Here are the hyperparameters used to train our transformer models.
Table Model Configurations
Model Size
Embedding Dim.
Number of Layers
Attention Heads
Learning Rate
134m
768
12
10 4
373m
1024
24
x 10
720m
1536
22
12
3 10
1.3b
2048
22
16
3 10
8b
4096
32
32
1 x 10-4
For all model sizes, Memory and Memory+ experiments use 4 heads and 32 top-k values for the memory embedding lookups. We did ablations on these values, but largely came to the same conclusions as the original product-keys pa- per (Lample et al. 2019). Overall, varying the number of heads or top-k while keeping their product (128) the same affects results minimally. Increasing the total keys improves little beyond 128, but has substantial memory lookup and gpu memory costs. On the other hand, decreasing to has non-negligible accuracy degradation.
B. Distribution of keys during training
We monitored the distribution of selected keys during training and noticed that the distribution tends to become more uniform as training progresses. The distributions starts heavily skewed and progressively improves as training continues.
Train step:
10-3
Train step: 1000
Train step: 2000
10
Train step: 4000
Train step: 8000
10
Train step: 13000
10
10-7
10
10
10
0.0
Figure 6. Evolution of the distribution of Memory+ selected keys during training. Keys on the x-axis sorted by decreasing occurrence.
12


Reviews:
[{"review_id": "VfxhkruQmY", "reviewer": "Reviewer_3YhT", "summary": "The authors conduct LLM scaling experiments in which the dense FFNs in a transformer are replaced with \"Memory Layers\".  A Memory Layer uses the attention operation to attend over a block of trainable parameters.  The advantage of attention over a traditional MLP, is that it is possible to implement sparse variations of attention; in this case the authors use product-quantized keys (Lample et al., 2019) to implement a sparse Top-K lookup.  \n\nTop-K attention with memory layers is similar in concept to sparse mixture-of-experts (MoE) architectures.  In both cases, the number of trainable parameters in the FFN can be dramatically increased, because those parameters are now sparsely activated, thus allowing more parameters to be utilized without a corresponding increase in FLOPs.  \n\nThe authors compare Top-K Memory Layers against standard (dense) transformers, sparse MoE models, and the newer PEER architecture, which similarly uses product-quantized keys.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and thoughtful comments!  We will try to address the feedback in order of priority, starting with the discussion of performance as this was deemed to be the “biggest weakness” of the paper.\n\n## Discussion of performance\n> The biggest weakness of this paper is that the discussion of performance with respect to implementation is lacking, especially regarding parallelism and batching\n\nThe reviewer’s intuition and observations in this section about the performance of memory layers vs. MOE or dense models is correct. Memory layer models have a large advantage compared to MOE at inference time, where decoding batch sizes are small.  At large batch sizes, the three models are in principle equivalent.  This is the typical setting for training.  In practice, memory layers do have some additional communication overhead, however we have been able to realize training throughput that is within 10% of the dense model training speed using our optimized implementation.\n\n> I would have liked to see a chart that compares the lookup speed of Dense vs. MoE vs. Top-K for different sized batches of query vectors\n\nWe provide such a chart in figure 4, where we plot the decoding time (dominated by memory lookup for small batch sizes) of the 3 architectures against the batch size. \n\n## MOE results\n > I am surprised by the poor performance of the MoE models\n\nWhile we cannot fully explain the poor performance of MOE models (especially for the 1.3b setting), we can provide the following additional information:\n- Our MOE implementation reuses code from a recent SoTA MOE architecture paper (reference not given to protect anonymity), which we believe to be well optimized for MOE.\n- We have independent evidence from other internal teams that match our MOE results\n\nThat being said, the 1.3b MOE results do not fit the scaling trends.  We don’t know the cause of this, but it might be because the MOE configuration which is parameter matched to memory layer at this scale only has 4 experts, which is quite small.\n\n> If there are 2^20 values, then I'm a bit confused about why there are only 984m total parameters\n\nThe value dimension in the memory layer is equal to the model dimension, which is different for each dense model scale.  For the 134m model, this is 768.  This is indeed confusing, and we will provide a table of model parameters for each scale in the appendix.\n\n## Theoretical claims\nThe reviewer is mostly correct, however the way product keys are implemented, we have 2 sets of half-keys, and the full set of keys is defined as the product of these sets.  As a result, if A is in the first set and B is in the second set, then AB is by definition part of the key set.  For the example given by the reviewer, if A’B’, AX and BY are in the keys, then so are A’B, AB’, which both have higher scores than A’B’.\n\n## Additional references\nThank you for providing additional references, we will include them in the final version.", "claims_and_evidence": "I am surprised by the poor performance of the MoE models in Figure 5.  MoE models are quite similar in many ways to product keys, where the first \"key\" of the pair determines the choice of expert.  Thus, for equal numbers of parameters, I would expect the MoE model to be very close to Memory, perhaps lagging slightly behind. Why is there such a big difference between them?  And why does the performance of MoE drop so sharply at the 1.3b parameter mark?  That seems odd.\n\nIn fact, looking at Table 2, MoE and Memory may be relatively close, it's just that Figure 2 shows Memory+ rather than Memory.  Perhaps you could put curves for both Memory and Memory+ into Figure 2? \n\nTable 1 should also specify the dimension of the values stored in memory.  If there are $2^{20}$ values, then I'm a bit confused about why there are only 984m total parameters.  In general, full architectural details (number of layers, embedding size, etc.) should at least be in the appendix.", "methods_and_evaluation_criteria": "In general, the experiments seem to have been well designed, and make fair comparisons between architectures.  \n\nThe biggest weakness of this paper is that the discussion of performance with respect to implementation is lacking, especially regarding parallelism and batching.  When training a model with a sequence length of 8k, which is common, there are batch_size x 8k query vectors that must be processed in parallel.  Both dense MLPs and MoEs can easily group the queries into large batches.  The large batch size means that attention can be done with a matrix-matrix multiply, and the keys and values can be read from memory in blocks, which is very efficient.  \n\nIn contrast, a sparse Top-K algorithm must select a *different* set of values for each query vector, which requires a sparse read from memory.  \n\nAs a result, I would expect Top-K lookup to be very fast for inference, when the batch size is small, because the values are sparsely activated.  (The authors do in fact make this claim.)  However, I would expect it to be significantly slower during training, when large numbers of queries have to be processed in parallel.  I would have liked to see a chart that compares the lookup speed of Dense vs. MoE vs. Top-K for different sized batches of query vectors.  The authors do mention some scaling challenges in Section 6, but do not give details.", "theoretical_claims": "It should be noted that, contrary to the claims in (Lample et al., 2019), product keys are an approximate Top-k lookup algorithm; it is not guaranteed to return the actual Top-k keys.  To see why, assume that A,B,... are half-keys of dimension d/2, which have been normalized to unit length.  Assume we are searching for the key AB, and the memory contains the following:\n\n* 500 keys of the form AX, where X$\\cdot$B $\\leq$ 0.\n* 500 keys of the form YB, where Y$\\cdot$A $\\leq$ 0.\n* A'B' where A$\\cdot$A' = 0.9 and B$\\cdot$B' = 0.9. \n\nThe closest matching key to the query for AB is A'B' (AB $\\cdot$ A'B' = 1.8).\nHowever, with K=128, the product key algorithm will fail to find it.  \n\nPlease update the text to mention that product keys are an approximate Top-k algorithm.", "experimental_designs_or_analyses": "The experiments are well designed, and properly attempt to compare architectures by balancing number of parameters.", "relation_to_broader_scientific_literature": "The authors cite the appropriate literature.", "essential_references_not_discussed": "Gupta et. al., \"Memory-efficient transformers via top-k attention,\" also uses a top-k lookup mechanism. \nSukhbaatar et. al.  \"Augmenting Self-attention with Persistent Memory\" replaces MLPs with dense attention over trainable keys and values.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "The term \"Memory\" is rather overused at this point, and can mean many different things.  Instead of using the word \"Memory Layer\", I would use the term \"Top-K Memory Layer\", especially in the title, abstract, and introduction.", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kEqPGKJzKp", "reviewer": "Reviewer_orhQ", "summary": "This paper proposes to replace the feed-forward layer in LLM with a memory layer. A memory layer consists of a key and a value matrix. Similar to the attention mechanism, each token representation will attend to the top-k selected values. Since it's sparsely activated, the computation cost will be much lower than the original feed-forward layer, and is more merory-bound.\n\nThrough dedicated engineering design, the authors can successfully scale up the memory layer up to 128B parameters for LLMs in the size of 134M-8B. Compared to baselines (dense LLM and MoE), LLM with memory layer performs significantly better across various tasks, while requiring less compute than the dense LLM and comparable compute to MoE.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for the positive feedback! We will do our best to improve based on the reviews. Here is how we plan on addressing the comments:\n\n## Experimental details\nWe will add an appendix with the details about the experimental setup including model dimensions and training hyper-parameters. Here are some of these details.\nFor 134m base models, we use dim of 768, 12 layers and attention with 12 heads.\nFor 373m base models, we use dim of 1024, 24 layers and attention with 16 heads.\nFor 720m base models, we use dim of 1536, 22 layers and attention with 12 heads.\nFor 1.3b base models, we use dim of 2048, 22 layers and attention with 16 heads.\nFor 8b base models, we use dim of 4096, 32 layers and attention with 32 heads.\nIn all Memory and Memory+ experiments, we use 4 heads and 32 top-k values.\nWe used a learning rate of 3e-4 for the 134m to 1.3b models and 1e-4 for the 8B models.\n\n## Code release for EmbeddingBag\nOur code is already open source (withholding reference due to anonymity), however here is the relevant file with the kernel implementation for your review: https://justpaste.it/cb6xz .\n\n## Unbalanced key usage\nRegarding unbalanced usage of keys, we notice that there can be a few steps at the very beginning of training during learning rate warmup where the keys become quite unbalanced. This issue goes away on its own and the distribution naturally smooths out during training.  In general, we did not need to do any regularization to enforce balance.  We will add a plot showing how this distribution evolves during the first few steps of training in the appendix.", "claims_and_evidence": "Yes, all claims are well supported.", "methods_and_evaluation_criteria": "Yes, all evaluation criteria make sense for me, with enough experimental results to support the claims.", "theoretical_claims": "This is an engineering paper without theoretical proof.", "experimental_designs_or_analyses": "Yes, the experimental designs are very sound and valid.", "relation_to_broader_scientific_literature": "This paper is related to efficient training and inference of LLM, and offers very promising results. It is a scale-up work of previous memory layer [1], also making it more efficient.\n\n\n[1] Sukhbaatar, S., szlam, a., Weston, J., and Fergus,R. End-to-end memory networks.", "essential_references_not_discussed": "No, the paper thotoughly discuss all closely related works.", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is well-written.\n2. The experimental results are thorough and promising, clearly demonstrating the effectiveness of memory layer.\n\n\nWeaknesses:\n1. Missing experimental details. The experimental setup is not very detailed, most training hyper-parameters are not shown.\n2. Since this is an engineering paper, and the authors do a lot contributions to the optimization side (like kernel), it would be better to include such materials for review.", "other_comments_or_suggestions": "None", "questions_for_authors": "1. I would like to review the implementation of the custom EmbeddingBag, since this is the key contribution for the speed up. But there is little implementation detail in the paper. I'm willing to raise my score, if the authors can show the implementation and explain it in details for my validation, because it's the key for memory layer at scale.\n2. Do you observe any unbalance unsage of the key and value? Since the key and value matrices are in large N, I wonder whether such a memory layer will have unbalance problem, i.e. only a limited amount of keys and values are used. It would be good to show the usage rate of the key and values.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9qZ7M2Y4Mn", "reviewer": "Reviewer_5Zym", "summary": "This paper describes a scaling analysis of memory layers and a comparison with alternative sparsely activated layers like MoEs and PEER. The main claims that the paper makes are:\n1. Performance improves by increasing the size of the memory layers.\n2. Memory layers significantly outperform dense layers.\n3. Memory layers outperform mixture-of-experts architectures with matching compute and parameter size.\n4. Memory layers are faster than mixture of expert layers at low batch sizes during decoding.\n\nThere are several additional claims made via ablation studies:\n5. Replacing more than three FFN layers degrades performance.\n6. The Memory+ block with added projection, gating and silu improves performance", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful review.  Please find our answers and clarifications below:\n\n## Evidence for claims\n> it would be helpful to include a similar plot for non-factual tasks (e.g. HumanEval or MMLU)\n\nWe opted to not provide numbers for these benchmarks for the small model sizes, since they are often not meaningful (e.g. MMLU is ~25% which is chance level).  However, here are the nll values for the two mentioned benchmarks at 1.3b scale for reference (lower is better):\n| Model    | HumanEval nll    | MMLU nll |\n| ---------- |   -------------------- | -------- |\n| dense   |          53.52          |  1.23  |\n| MOE     |          52.29         |  1.17  |\n| Mem+   |          52.02          |  1.14  |\n\n> I can’t find a compute-matched comparison anywhere in the paper\n\nAll of our experiments are compute (FLOPs) matched within the same scale category.  E.g. for the 1.3b scale, the dense, MOE and memory layer models all have (almost) identical flop cost.  For the MOE-memory layer comparisons, we also attempted to match parameter count.  We will make it more clear in the paper that the comparisons are FLOP-controlled, as this is one of the main claims of the paper.\n\n> How would varying the number of experts and size of experts affect the performance relative to memory layers\n\nFor the main experiments, we set the number of memory values at 2^20 for the memory models, and tried to set the MOE experts to match the parameter count.  This choice results in varying number of MOE experts for different scales (e.g. 16 experts for the 134m model, 4 experts for the 1.3b model).  While we agree that having more experiments at various parameter / expert counts would be ideal, we did not have the needed resources to run them.  We would guess that having more but smaller experts would behave more similarly to a memory layer.  In fact, the PEER work which we compare against, investigates the limit of this, where they have up to a million rank-1 experts.  Our work performs comparably or better, while being simpler and more compute efficient.\n\n## Other questions and comments\n> How does the efficiency of a dense model compare to the memory model at training and inference?\n\nWe provide such an analysis in figure 4, plotting latency against batch size for dense, MOE and memory models.  Memory models have similar efficiency with dense models in both the small batch-size (typical inference) and large batch-size (typical training) regimes, while MOE inference latency is much higher.  This is a roofline analysis.  In practice, we observe ~10% lower training throughput compared to a dense model during training at 8B scale, due to communication overheads.\n\n> What is the value of k chosen and how does it affect performance?\nThank you for catching this oversight!  We used 4 heads and k=32 (total keys 128) for all experiments, we will add this to the paper.  We also did ablations on these values, but largely came to the same conclusions as the original product-keys (Lample et al) paper, so decided not to repeat those results here.  Overall, varying heads or k while keeping their product the same effects results minimally.  Increasing the total keys improves little beyond 128, while having substantial memory lookup and gpu memory costs, while decreasing to 64 has non-negligible accuracy degradation.  In the end we stick to 128, which was also the default for the original paper.\n\n## Typos and nits\nThank you for pointing these out, we will edit this in the final version of the paper.", "claims_and_evidence": "1. This claim is nicely supported by the results in Figure 1. However, it would be helpful to include a similar plot for non-factual tasks (e.g. HumanEval or MMLU).\n2. This claim is well supported.\n3. I some have concerns with the quality of evidence for this claim. The results in Table 1 appear to be a memory-matched comparison, but I can’t find a compute-matched comparison anywhere in the paper. Also, it looks like a comparison was made only for a single hyperparameter configuration.  How would varying the number of experts and size of experts affect the performance relative to memory layers. \n4. The evidence for this claim is in Figure 4. A limitation is that the MoE baselines are implemented with very low expert counts. Would there still be a significant speedup, if more experts were used in the MoE layer?\n5. This result is well-supported. \n6. This result is well-supported.", "methods_and_evaluation_criteria": "See claims and evidence.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "See claims and evidence.", "relation_to_broader_scientific_literature": "This work provides a scaling analysis of a well-known technique from the literature. Section", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "There are several important open questions, which are not addressed in the paper.\n- How does the efficiency of a dense model compare to the memory model at training and inference?\n - What is the value of k chosen and how does it affect performance?", "other_comments_or_suggestions": "I am quite interested in the relationship between MoE and memory layers. For certain choices of hyperparameters, they appear to be basically equivalent. So, my main concern with the presentation and experiments in this paper center around the categorical nature of the MoE claims (*i.e.* that MoEs are worse than Memory layers). \n\nI think the the paper would be much strengthened by carefully analyzing these hyperparameters and understanding how these techniques relate. \n\nTypos and nits:\n\n- In Figure 1, the caption should read top/bottom not left/right.\n\nTypos and nits:\n- In Figure 1, the caption should read top/bottom not left/right.", "questions_for_authors": "No questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8XeStmJF1x", "reviewer": "Reviewer_9R67", "summary": "The paper proposes a trainable key-value look-up to embed within the transformer architecture as an inductive bias for memory. To make this computationally efficient, they propose ways to parallelize the search across GPUs and show that this architectural change helps factuality.\n\n### Update after rebuttal\nWhile there remains questions on generalization to unseen or new domains at inference-time and how the memory layer could actually be detrimental, I believe the paper proposes a novel architecture that acts as an inductive bias for trainable memory and demonstrate that it reliably improves performance on especially factuality-based tasks. I believe this paper should be accepted and maintain my score of 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the supportive feedback!  For benchmarks beyond factuality, we tried to include a variety of standard benchmarks on table 2 for the 8B models.  In addition to this, here are results for the GSM8K math benchmark that we recently ran (8B, 1 trillion tokens):\n| GSM8K exact match |  |\n| --- | --- |\n|  dense    |       35.8  |\n| Memory+  |    43.4  |\n\n> I am curious how this may affect out-of-distribution tasks or reasoning in new domains unseen during pretraining\n\nWhile this is a valid and interesting question, at this time we do not have any data to help us speculate on this issue.  However, we are continuing research with memory models, including in new domains of knowledge, and hope to contribute in this direction in future work.", "claims_and_evidence": "Yes, the claims seem to be sound and well supported.", "methods_and_evaluation_criteria": "Yes, the methods are sound and the evaluation across different QA benchmarks is comprehensive.", "theoretical_claims": "There are no substantial theoretical claims in this paper.", "experimental_designs_or_analyses": "The experimental designs seem to be sound.", "relation_to_broader_scientific_literature": "The contributions of the paper are related to general machine learning and specifically in the field of investigating newer and different architectures that separate memory from reasoning, such as Neural Turing Machines.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths\n- The paper is very clearly written and the contributions of the paper are very clear. The authors show that they can in general improve the performance for factuality-related tasks upon other model architectures or training paradigms with similar or larger parameter count. Also, they show that they can maintain the runtime efficiency of dense models (over MoEs).\n- Evaluations are comprehensive across different QA benchmarks.\n\nWeaknesses\n- Although the results on the knowledge-based benchmarks are comprehensive and convincing, the generality of the finding across different types of tasks like math or reasoning benchmarks hasn't been tested and seems important given that they are proposing an architecture-level change for the future generation of models to explore.", "other_comments_or_suggestions": "None", "questions_for_authors": "One of the obvious weakness of embedding these trainable memory layers, in lieu of retrieval-based systems like RAG, is that it is difficult to swap knowledge at inference-time (although RAG can also be attached on top of this model). But, because of this inductive bias of the memory layers serving as \"memory retrieval\" within the model, I am curious how this may affect out-of-distribution tasks or reasoning in new domains unseen during pretraining.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sLItXrKCs2", "reviewer": "Reviewer_Sxwo", "summary": "This paper proposes an improved Memory layer design, which adds extra parameters to the model without increasing FLOPs. Following Memory layer design of previous works, this paper optimizes embedding bag kernels and improves architectures by gating networks for performance and QKNorm for stability. \n\nIn the experiments, the paper shows that Memory+ shows an decoding advantage than MoE where the activated parameters are relatively small under small batch size. Besides, the performance of Memory+ also outperforms Dense and MoE baselines. The model is scaled to 8B to show the effectiveness of Memory layer.\n\n## Update after rebuttal\nMy score keeps as Accept. The author solved most of my concerns.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the kind review! We will try to address the comments in the order they were raised.\n\n## MoE Results\nWhile we cannot fully explain the poor performance of MOE models (especially for the 1.3b setting), we can provide the following additional information:\n- Our MOE implementation reuses code from a recent SoTA MOE architecture paper (reference not given to protect anonymity), which we believe to be well optimized for MOE.\n- We have independent evidence from other internal teams that match our MOE results\nIt is quite surprising that MoE performs slightly worse than the dense model for NQ and TQA on Figure 5. MoE does perform better than dense at the 1.3b scale on all other tasks of Table 2, PIQA, ObQA and HotPot.\n\n## Memory+ layout\n> Are there any ablation studies, showing that whether layout plays a role on the performance?\n\nTable 3 (top) provides some insights into how the memory layers layout (number of layers, centering and spacing) affects performance. In the table, adding more Memory+ layers beyond 3 decreases performance. We believe this is because for every memory layer we add, we remove the FFN layer, resulting in effectively less computation.  Also, we share the value parameters across all the Memory+ layers of the model, thus the total parameter count does not increase.  As a result, there is a sweet spot after which adding more memory layers doesn’t help.  It is possible to add memory layers in addition to or parallel to an FFN layer, however we did not consider this as it would break our compute-controlled comparisons.\n\n## Decoding Efficiency of Memory+\n> I'm not fully convinced that the decoding efficiency of Memory+ layer is consistent with dense baseline across all the batch size. When the batch size is increasing, the activated memory ratio will increase, where the memory access demand also increases.\n\nYou are indeed correct that Memory+ layers memory access grows with batch size.  However, for the top-k value that we pick (4 heads, 32 top-k, 128 total), the memory demand remains comparable to the dense model for small batch sizes (note that the analysis for one forward pass through the entire model, not just a single memory layer).  For large batch sizes, all models are compute bound, thus the curve follows the dense models curve closely for the full range of batch sizes.  However, MOE models quickly need a lot more memory bandwidth even at small batch sizes, and are thus severely memory bound in this regime.\n\n## Model Design Philosophy\nOur design is built on top of the work of Large Memory Layers with Product Keys (https://arxiv.org/abs/1907.05242) in which they use two independent sets of sub keys that they then combine into a complete key.  This is done for efficiency reasons, since otherwise it would be infeasible to compute top-k over millions of keys.  This likely causes some degradation in key lookup (similar to product quantization).   On the other hand, the value parameters are kept in a flat table and are independent of each other. \n\n## Relation between performance and sparsity ratio\n> Is there any experiment about the relation between performance and sparsity ratio? In other words, when the sparsity ratio is further increasing, can the performance also improve consistently?\n\nIn Table 2 at the 1.3b scale, we study the relation between increasing the Memory+ parameter count (which increases sparsity) and performance. At the scales studied, it seems increasing the parameter count increases performance. However, we were unable to scale further in this research work and could not verify if further increase always results in better performance.\n\nAnother way to increase sparsity would be to decrease k.  We find that doing this (e.g. go from top-128 to 64) hurts performance.  We also did ablations on these values, but largely came to the same conclusions as the original product-keys (Lample et al) paper, so decided not to repeat those results here.  Overall, varying heads or k while keeping their product the same effects results minimally.  Increasing the total keys improves little beyond 128, while having substantial memory lookup and gpu memory costs, while decreasing to 64 has non-negligible accuracy degradation.  In the end we stick to 128, which was also the default for the original paper.", "claims_and_evidence": "This paper makes a systematic experiment to support their claims.", "methods_and_evaluation_criteria": "The proposed method aims at adding extra parameters to enhance model's capability without increasing FLOPs. The experiments are standard under LLM training scenarios.", "theoretical_claims": "No Theoretical claims.", "experimental_designs_or_analyses": "The experimental designs are sound in general. However, I have some concerns in the comparisons with MoE models:\n\n1. There are many improved MoE baselines recently, for example, fine-grained and shared experts in DeepseekMoE. Usually, MoE models are at least better than Dense baselines. However, in Figure 5, the MoE models are worse than Dense in 1.3b dense. Therefore, I'm curious at the implementation details of MoE.\n\n2. Besides the module difference, there is also an important difference, where Memory+ activate Memory layers every several layers while MoE model replace FFN with MoE in every layers. Are there any ablation studies, showing that whether layout plays a role on the performance?", "relation_to_broader_scientific_literature": "This paper is an improved work of Memory network, which is a kind of sparse activation modules. Besides, it also has connections with Mixture-of-Experts.", "essential_references_not_discussed": "All the essential references are discussed in my point of view.", "other_strengths_and_weaknesses": "Strength:\n1. I think this network design will make an influence on hardware design. The computation patterns are different from current dense matrix multiplications.\n2. LLM needs to go beyond current MoE designs for further sparse ratio.\n\nWeaknesses:\n1. Following Experimental Designs Or Analyses, I think there should be more detailed ablations and comparisons with MoE and layout design.", "other_comments_or_suggestions": "Overall, this is a valuable paper and shows a promising direction for LLM sparsity. However, I have some doubts about the details in the experiments. I list them in different sections of this review. I hope a useful discussion can help us better understand and position this paper.", "questions_for_authors": "1. I'm not fully convinced that the decoding efficiency of Memory+ layer is consistent with dense baseline across all the batch size. When the batch size is increasing, the activated memory ratio will increase, where the memory access demand also increases. Compared with a linear projection, how can the curve almost the same?\n\n2. For model design philosophy, the same sub key2 is used among different sub key1, which means there are some shared knowledge in the same row or column. Is there any explanation for that?\n\n3. Is there any experiment about the relation between performance and sparsity ratio? In other words, when the sparsity ratio is further increasing, can the performance also improve consistently?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper investigates Memory Layers—a sparse, trainable key-value mechanism that replaces standard feedforward layers in transformers—and introduces a refined variant called Memory+, which incorporates projection, gating, and normalization. Through a comprehensive scaling study (from 134M to 8B parameters trained on 1T tokens), the paper demonstrates that Memory Layers scale effectively and outperform both dense and Mixture-of-Experts (MoE) baselines on a range of factual QA and reasoning benchmarks. The authors also present decoding latency and efficiency comparisons using roofline analysis.

The reviewers generally praised the paper for its architectural clarity and thorough experimental validation, highlighting several strengths:

- Memory Layers provide sparse computation with lower decoding latency than MoEs, particularly in low-batch inference scenarios.

- Compared to dense models, Memory Layers achieve better accuracy under equal-FLOP budgets.

- The empirical study includes comparisons across dense and MoE baselines, layer placement ablations, factual QA performance, and some reasoning tasks.

Meanwhile, the paper contains a number of limitations that were noted by some reviewers but not treated as major concerns, and which, in my view, deserve greater emphasis:

- The highlighted benefit of inference efficiency over MoEs is demonstrated only through a roofline analysis, without a practical implementation or end-to-end latency measurement.

- There is no discussion of training efficiency for Memory Layers compared to MoEs or dense models. As noted by Reviewer 3YhT, both dense MLPs and MoEs can group queries into large batches, enabling efficient memory access and compute via block-based operations. In contrast, Memory Layers introduce fine-grained, dynamic sparsity that is not well-supported on standard accelerators. This type of sparsity is usually leveraged for gaining efficiency during inference instead of training (e.g., [a]).

- While Memory Layers outperform dense models under equal-FLOP comparisons, the paper does not examine how FLOP reductions translate into wall-clock benefits. This is a critical omission when claiming practical efficiency improvements.

- Much of the claimed benefit of Memory Layers derives from activation sparsity, but the paper omits discussion of highly relevant prior work showing that intermediate activations in transformers are either naturally sparse (e.g., [b]) or can be made sparse through architectural or training modifications (e.g., [c]). These works imply that FLOP-based comparisons may be unfairly biased against dense models, as the actual FLOP count for dense models with sparse activations is often significantly lower than nominal estimates.

[a] PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU, SOSP 2024
[b] The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers, ICLR 2023
[c] ProSparse: Introducing and Enhancing Intrinsic Activation Sparsity within Large Language Models, 2024

**Recommendation**: Acceptance. This paper introduces Memory Layers, demonstrating through a comprehensive scaling study and strong empirical results that this novel sparse architecture outperforms dense models at iso-FLOPs and shows advantages over MoEs on key benchmarks. While limitations regarding practical latency data, training efficiency details, and full context within sparsity literature exist, these are areas for refinement. The core contribution—validating a viable, scalable, and high-performing alternative architecture—warrants acceptance, expecting authors will address limitations in the final version.
Decision: Accept (poster)