Ultra-Sparse Memory Network
Authors: Zihao Huang
Qiyang Min
Hongzhi Huang
Yutao Zeng
Defa Zhu
Ran Guo
zhou Xun
Venue: ICLR 2025 Poster (2025)
Keywords: Large language model
sparse model
scaling law
Primary area: foundation or frontier models, including LLMs
TL;DR: Ultra-sparse memory network significantly enhances the efficiency and scalability of large language models while maintaining performance. Compared to Mixture of Experts, it has a significant advantage in inference speed.

Abstract: It is widely acknowledged that the performance of Transformer models is logarithmically related to their number of parameters and computational complexity. While approaches like Mixture of Experts (MoE) decouple parameter count from computational complexity, they still face challenges in inference due to high memory access costs. This work introduces UltraMem, incorporating large-scale, ultra-sparse memory layer to address these limitations. Our approach significantly reduces inference latency while maintaining model performance. We also investigate the scaling laws of this new architecture, demonstrating that it not only exhibits favorable scaling properties but outperforms MoE. In experiments, the largest UltraMem we train has \textbf{20 million} memory slots. The results show that our method achieves state-of-the-art inference speed and model performance within a given computational budget, paving the way for billions of slots or experts.

Published as a conference paper at ICLR 2025
ULTRA-SPARSE MEMORY NETWORK
Zihao Huang, Qiyang Min; Hongzhi Huang, Defa Zhu, Yutao Zeng, Ran Guo, Xun Zhou Foundation Team, ByteDance
(huangzihad not aota,bo, mingiy huanghongzhi, 51, zhudefa, yutao. zeng, guoran 94, Jebutxaneeyyeeaacee. com
ABSTRACT
It is widely acknowledged that the performance of Transformer models is logarith- mically related to their number of parameters and computational complexity. While approaches like Mixture of Experts (MoE) decouple parameter count from com- putational complexity, they still face challenges in inference due to high memory access costs. This work introduces UltraMem, incorporating large-scale, ultra- sparse memory layer to address these limitations. Our approach significantly reduces inference latency while maintaining model performance. We also inves- tigate the scaling laws of this new architecture, demonstrating that it not only exhibits favorable scaling properties but outperforms MoE. experiments, the largest UltraMem we train has 20 million memory slots. The results show that our method achieves sfte-of-thheatt inference speed and model performance within a given computational budget, paving the way for billions of slots or experts.
INTRODUCTION
2.7
MoE: 1.6.1.66.20 2in34
80
NoE-1.68-210
68-2in34
UltraMem-1.66e 1.68-1122
70
UtraMmm-1.68-12
UltraMem-1.6e
1,68-112
2.6
Dense-1.68
60
Momany accees (GR)
30
Dense-1.68
Dense-6.58
25
Validation
2.5
20
xx17
2.4
6,0x
5,1x
2.3
0.06
10
12.2x
4,5x
256
1024
10:
10
10:
107
Consumed (Billion)
Batch Size
(a) Validation loss
(b) Inference
(c) Memory access
Figure 1: We ensured that three models have the same computation, and MoE and UltraMem have the same parameters. The x-axis is plotted on logarithmic scale. In (b) and (c), the sequence length because during decoding time, we can only predict one token a time, and the key/value cache length is 2048. The experiments in (b) and (c) are conducted on the A 100-SXM-80GBB
Recent advancements in natural language processing (NLP), driven by Large Language Models (LLMs) (Radford et al. 20199 Brown, 2020), require exponentially more computational resources they scale, posing challenges in resource-limiee environments like real- sal-time applications. To address computational issues, the Mixture of Experts (MeE) et al. [022 Jiang et al.) (024) and Product Key Memory (PKM){Lampl) parameters, boosting training efficiency but impairing inference time due to increased memory access. PKM maintains consistent memory access with fewer value embeddings but its performance is significantly worse than MoE.
al. 2019) have been introduced. Moł selectively activates
As shown Figure Tb), MoE model, despite having the same computational cost and twelve times more parameters than dense model, runs 2 to 6 times slower in inference, varying by batch size.
Equul contribution.

---

Published as a conference paper at ICLR 2025
This slowdown, as depicted in Figure stems from high memory access demands, highlighting its inefficiency in inference scenarios. The primary challenge is how to match or even surpass the effectiveness of the MoE model while maintaining memory access levels comparable to those of dense models.
In this paper, we introduce UltraMem, an architecture that builds upon and extends the concepts from PKM. UltraMem incorporates large-scale, ultra-sparse memory layers that significantly enhance computational efficiency and reduce inference latency while maintaining or even improving model performance across various benchmarks. This architecture not only supports the deployment of highly effective language models in resource-constralee constraiedd environments but also opens up new avenues for constructing even larger models without the previously associated prohibitive costs.
In summary, we make the following contributions:
1. UltraMem is greatly enhanced compared PKM, and outperforms MoE same scale. Compared to PKM, UltraMem truly possesses the prerequisites for training large-scale mod- els on extensive computational resources and has undergone comprehensive experimental validation.
2. UltraMem has significantly lower memory access cost during inference compared MoE. Under common inference batch sizes, it can be up to 6 times faster than MoE with the same parameters and calculations. The inference speed of UltraMem is almost identical to that of dense model with equivalent computational resources.
3. We have verified the scaling ability of UltraMem. Similar to MoE, UltraMem has strong scaling ability, and we have observed stronger scaling ability than MoE.
RELATED WORK
Mixture of Expert. Shazeer et al. (017)) proposed MoE and Fedus et al. (2022) introduced the MoE in large language models, where each token selects one expert for inference each time, thereby increasing model parameters without increasing computation. |Rajbhandari et al. (2022) introduced the concept of shared experts, where each token utilizes some fixed experts along with some unique experts. Subsequent research has focused on improving the gating functions of MoE, including token choice (Chi al., 2022). non-train non-trainable token choice (Roller et al.) a 2021) and expert choice (Zhou) et al. 2022), primarily to address the issue of expert imbalance. [iu et al. (2024); Dai et al. (0024) opted to slice the experts into smaller segments while activating more experts per token, achieving significant performance improvements. Concurrent study (Krajewski et al. 2024) meticulously explored the benefits of granularity and increasing the number of experts, alongside investigating the scaling laws associated with MoE. In this paper, use fine-grained MoE our baseline, wherein
the granularity of the MoE is set to 2. This means that each expert is half the size of the original MultiLayer Perceptron (MLP). with two experts activated per token.
Large Memory Layer Lample et al. 2019) first introduced the concept of large memory layer, called PKM, which can be seen as slicing the MoE experts the smallest possible configuration. Kim & Jung 2020) introduced a concept similar to shared experts in MoE, allowing PKM and MLP to operate parallel. Csordás et al. (2023) made a slight modification to PKM by removing the Softmax operation. PEER (He 2024) improved the activation of values in PKM to activate a small expert with an inner dimension of 1, achieving significant performance gains. However, current research on PKM is limited to smaller models, and even the latest improved versions of PKM only outperform MoE in certain scenarios. Additionally, current PKM do not possess characteristics suitable for large-scale training. We address these issues in this paper.
Tensor decomposition breaks down a tensor into series of small matrices or tensors. In deep learning research, such methods are commonly used to approximate a large tensor during training, aiming to save on computation and parameters. Product quantization (Jegou et al. 2010) breaks a vector into multiple sub-vectors, allowing us to reconstruct the original vector using smaller number of sub-vectors, thereby reducing the model parameters. Bershatsky et al. (2024) initializes several matrices and
tensor, trains these parameters during the fine-tur tinettunig phase, and reconstructs the original large tensor in a manner of Tucker Decomposition at the end of training to reduce training costs. We borrow this insight to improve PKM's key retrieval.
2

---

Published as a conference paper at ICLR 2025
row keys
Σ
Weiigttdd sum poaling
top m
values
fetch values
query
fetch values
output
keys
col
keys
GeLU
values
output
top m
fetch values
col query
query
(a) Multilayer perceptron
(b) Large memory layer
Figure 2: An overview of multilayer perceptron (MLP) and large memory layer (LML). For the sake of brevity, we omit the third top-m operation from memory layer. An MLP typically consists of two linear layers and GeLU activation. We consider the weights of the first linear layer as keys, and those of the second linear layer as values. LML uses row and column keys to determine the 2-D logical address index memory values, whereas MLP uses I-D logical address. "fetch value" refers to retrieving values based on the indices with higher scores.
ULTRAMEM
3.1 PRELIMINARY
Here we firstly introduce the origin large memory layer (LML.) based on product keys, which serves as the foundation for our proposed approach. The concept of product key-based memory layer (PKM) was first explored in prior work (Lample et al. 2019). In their approach, the authors incorporated an external memory module into language models, with the goal of expanding the model's parameters while maintaining similar level of computational complexity. The overall structural diagram is depicted in Figure bb.
memory layer generally consists of parts: keys K RNxxs and values V NND. To retrieve information from memory values, a query vector q E RD finds most relevant values by multiplying keys obtain scores. The higher the scores are, the better impact should the values have. Consequently, this process can be formulated as:
= (Kg) o= V
where s is the scores, a is non-linear activation, is the output. Attention layers, who memorize context contents, and MLP layers, who memorize world knowledge, also follow the above formulation with being SoftMax in attention layers and GeLU in MLP layers (Geva et al. 2020) (see Figure[(a)).
Product-key memory layers scale up the memory size with N > 10 while activating only a few values with top-m scores. Here, is a hyper-parameter controlling sparsity. Though values are sparsely accessed, the keys, which are as large as values, must fully computed to obtain scores before top-m activation following equation [ To alleviate the computation complexity for keys, product keys are proposed. Borrowing the idea Product Quantization, utilizes a 2-D logical address (see Figure (b)) typically a x n grid where n /N, for memory value retrieval. Specifically, 2-D logical address (z. j) is used to index memory value at physical address i j. With such strategy, logical scores are then represented as matrix, which is further decomposed as an addition of row and column scores:
Srow TTopM Krowwoow (x)),
Scot Tropm (TTopn(Koottooo (eottens (x)),
S Sgrid (Topp(SSoo TT));
VT x SoftMax(ve:((srii))
where Krow, Keot E R Di. grow, Acol : R D. Du convert input hidden x E R Di to row and column query, (-) preserves top-m largest elements in the input and set the rest to negative infinity, and the matrix addition with unmatched matrix shape is implemented by element broadcasting. It should be noted that removing Ozopn from equation does make any difference. The only
3

---

Published as a conference paper at ICLR 2025
reason for applying top-m to the row and column scores is to reduce the computation for the last top-m operation on gridd As Syow: Scol have only m activated scores, Syrid has only m² candidates for top-m operation rather than N, i... top-m complexity reduces from O log m) to O((vN + log log..
Note that the Sgrid undergoes SoftMax operation akin to the one employed in the self-attention mechanism. Moreover, PKM adopts the multi-head mechanism from the self-attention module, wherein it utilizes multiple key sets to retrieve the shared values, we denote H as the number of PKM heads.
3.2 STRUCTURE IMPROVEMENTS
Improve PKM with bag of tricks. We first studied the structure of PKM and found that series of minor adjustments can steadily improve the model's performance:
1) We remove the operation Softmax in equation
which is well-established in the stud-
ies (Shen et 2023 Csordás et al., 2023).
2) We conduct Layer Normalization (LN) (Ba et al. 2016) on query and keys for stability of training.
3) PKM suggests using constant learning rate of 0.001 learn the values, which is much higher than the learning rate for other parameters. We found that gradually decaying the value learning
provides further benefits.
4) PKM uses a linear layer to generate query, we add a causal depthwise convolutional layer (Howard. 2017) before this linear layer enhance query.
5) Similar to Group Query Attention (Ainslie et al.) 2023), we share query in two key sets. This can reduce the computational cost of generating the query by half, with little performance impact.
6) By halving D.. we double the number of values. Under the condition of keeping the activation value parameter unchanged, we increased the diversity of activated values, and the model effect is further improved. In order to make the output consistent with hidden dimension, we add linear layer on the aggregated output.
UltraMem Overall structure. We then take a deeper investigation into the model structure and propose UltraMem. Figure shows Shows the PKM and our improved UltraMem structure, based Pre-LayerNorm Transformer architecture. PKM replaces MLP or operates in parallel (Kim & Jung 2020) with MLP the one of deeper layers with memory layer. We notice three drawbacks to PKM:
1. As value size N significantly increases, queries can harder find correct values.
Transfomer byer
Transtomerr layer
Product key decomposition introduces bias retrieval topology. For example, let (i.j) be the logical address for the top- I score, then top-2 score must be located on row i or column j, which significantly limits the diversity of top-m selection.
UltraMem block
Transormer layer
UltraMeem block
Tronsformer byer
PK3
MIP
Transformer
3. There are issues with unbalanced multi- GPU computation and communication during large-scale parameter training, the full model parameters cannot be placed single GPU.
UlraMee
MIL
Artention
LN
Aitertion
To alleviate problems and we decompose this large memory layer into multiple smaller mem- ory layers distributed at fixed intervals across the transformer layers. Additionally, this skip-layer structure allows us to overlap the execution of the memory layer and the transformer layers, the memory layer is predominantly memory-bound during training.
Transfomer layer
LN
Transformer
Transforme- byyer
Transformet
Figure 3: Overall of PKM and UltraMem.
4

---

Published as a conference paper at ICLR 2025
Step
Step 3
Step
tucker
row score
merged row score
merged
rDW
col score
merged col score
filtered col
rOW score
filtered row score
merged index
Stop
key
merged row score
top-m
top m index
top
index
merged col score
col index
col keys
top
Step7
final
Step
filtered
merged score
col query
Top
row score
Step2
top-m row
Terch
top m hdee
tucker core (2x2]
SVD
Score
merged index
tog-t col index
Figure 4: Flow of Tucker Decomposed Query-Key Retrieval (TDQKR), here r 2. The term "fetch" refers to the action of retrieving scores based on a given index (corresponding to koorch.gater" TDQKR replacing Product Quantization, serves as more precise retrieval module for recalling value indices UltraMem. Each step of the TDQKR process is meticulously referenced within the main text for understanding.
Tucker Decomposed Query-Key Retrieval (TDQKR), We explore a more complex multiplicative approach alleviate problem and 2, where tucker decomposition (Malik Becker, 2018) is adopted in place of product quantization. The whole process of TDQKR is illustrated in Figure Specifically, tucker decomposition estimates grid scores with rank-r matrix multiplication:
Srow K Krow Grow (x),
S Scol Keot (x),
(4)
S Syrid Trepx (S
x x Scot),
(5)
where Srow Scot E x and C E R is the tucker core, which is a learnable parameter with random initialization. To produce x r shaped row and column score, the dimensions of the query and key are reshaped, resulting in Krow, Koot (Dxxxx(D//r) and grow, Gcot E (Da//Dr corresponding to Figure step 1.
However, equation 5 inefficient be directly applied in practice, the top-m operation cannot be simplified with an equivalent two-phase top-m technique like product quantization can. As consequence, we propose an approximated top-m algorithm tackle this problem. The key is to do rank-1 approximation for the tucker core, so that the overall top-m can be approximated by:
C R ut
(Sppea(TToN C Scot) drapss(( Srow)] (S Scol))
(6)
where u, t Brxx1 Note that Srow),, T Scot) E R2*= are row vectors, then the two-phase top- technique pertains to the approximated objective Trope(( " Srow)) (t Scot)), corresponding to Figure step 3. Overall, we conduct approximated top-m on row and column scores, filtering out non-top elements, then index scores precise:
the concrete objective in the final top-m operated on Sgrid, keeping
C CuutTT
(7)
Sent =T Swnt) + Smat
(9)
Srow =Izope(u" Srow) (roww) Srow
(8)
S Syrid = Tom C Scot),
(10)
where ITope() is binary value function, which converts top-m elements to and otherwise to 0. Equation correspondd to Figure [ggure4step step 4&5,and Equation I0 corresponding to Figure 4 step 6&7. As for the rank-1 approximation, 2007) to factorize the tucker core with u. t be the left and right singular vectors corresponding to the leading singular value, corresponding to Figure step 2.
leverage Singular Value Decomposition (SVD) (Abdi

---

Published as a conference paper at ICLR 2025
Virtual block1
2
Weighttd sum pooling
Virtual block2
Fetch value
Linearl
Virtual block3
Linear2
Virtual block4
Physical Memory table
L Linear3
Shuffie
Given
Lineard
and
Virtual memory address table
Figure 5: Flow of Implicit Value Expansion (IVE), here = 4, m = 16. IVE reduces memory access and scales up memory size by expanding the memory table virtually. Fach virtual block is reparameterization of the physical memory table. Every virtual memory address corresponds to a physical memory address and projector index. The weighted sum pooling is grouped by the virtual blocks, followed by linear layer to produce the final output.
Last but not the least, the approximation enror should be concerned when non-maximum singular values are as large as the maximum one. To mitigate this, an auxiliary loss that manages approximation error is introduced during training by constraining non-maximum eigenvalues:
C == UAT (by SVD)
(11)
Laux
mmax (0, - -))
(12)
where, A denotes the singular values for C in descending order, with T serving as margin to prevent C from degenerating into a rank-1 1 matrix, and a is the coefficient for the loss.
Implicit Value Expansion (IVE). Though sparsely used, maintaining a large memory table is still costly during training due to the large amount of memory access. To reduce memory access as well as scale up the memory size, we propose virtual memory as implicit value expansion. Given virtual expansion rate > I, virtual memory expands the memory table into E times size. design virtual memories multiple reparameterizations of the original memory value V, which we denote physical memory. Then, E' linear projectors {Wplp E], Wp P.x.) } are utilized, and virtual memory block V, corresponding the p-th reparameterizztion be defined
VW
(13)
Then the overall virtual memory is concatenation of the virtual blocks V VV VT.... T E
Note the dimension of virtual values D' is not necessarily consistent with the dimension of physical values D..
To apply the virtual memory is intuitive, where memory table can be replaced from V to V. And to fit virtual memory size, the key size is expanded by EE times. Moreover, we suggest random shuffle for virtual memory to eliminate some unnecessary index topology prior introduced by row and column scoring. Concretely, if the virtual memory tables are unioned by concatenation, each memory value and its expansions would be located the same column logical address, and thus can be potentially more frequently chosen simultaneously.
naive reparameterization for virtual memory still introduces lots of computations, which is N D. D'., and E times GPU memory access. A better idea is to compute reparameterization demand. That is, we expand the logical address to triplets (i.j.P) p)) where (i, .) is the original logical address and p is index for the virtual memory block, and then simultaneously conduct sum pooling and compute virtual memory value. Consequently, equation 3 rewritten as:
$ ==Shuff 3=ShuffLe(vec(S grid)),
(14)
o = T
x
$ VV. xS = = (VT x Sp)
(15)
where š, represents the scores corresponding to p-th virtual memory block. With equation 15] we can firstly lookup and pool values according to the virtual block index and then transform the reduced

---

Published as a conference paper at ICLR 2025
physical values directly into reduced virtual values. This trick reduces extra computation from E-N. De - D to E B- Dv D',, where B is the number of tokens in batch, and has nearly no extra GPU memory access except for the linear projectors.
[ the flow of IVE.
Multi-Core Scoring (MCS). PKM shares a single score across dimension Do for each value. Empirically, assigning multiple scores single value has shown enhance performance. Thus, we rewrite the tucker core C series of component cores C Σx C This allows employing
{CO)) ===
to generate individual score maps bucker = -STow Scots Obviously,
(i)
h
h
Stucker S rown Σ C( = EST
row C(i) = - Shucker
(16)
We keep top-m conducted on aggregated score Stucker, while applying individual scores ((1) vertically split value table V [V(). V(h), ]] with V(i E RNNx(D./h), i.e.
Pucker on
(
s(h) TV.
(17)
When this technique incorporates with IVE, we split physical memory values instead of virtual memory values to keep the equivalence in equation 15
Improved initialization. PKM initializes values with a Gaussian distribution W(0, D.). Since PKM applies Softmax to the scores, the variance of the pooled outputs is 1/D,- We argue that LMI should be considered as a component similar an MLP and, therefore, should use an initialization method akin to that of MLPs. Before training. the output of an MLP typically follows a Gaussian distribution (0, 21. Broww 2020p, where the total number of layers. We initialize value E with (0, 2m.H ). where m is the activated value number, is the head number, E is the value expansion times. To ensure that the output distribution of UltraMem is (O, 2 We need to confirm that the mean of top-m score details see Appendix
QUANTITATIVE ANALYSIS WHY ULTRAMEM INSTEAD OF MoE
The most effective method for enhancing model capacity without significantly raising computational costs is MoE. This strategy employs a set of specialized sub-models, known as "experts", which work together to tackle complex problems. However, the MoE model poses challenges for inference processes.
Consider the Transformer hidden dimension D. the inner dimension of MLP is 4D, given the inference batch size as B. Using the MoE with 2in (choose 2 in Nmme experts per token) as an example, where the inner dimension of expert is 2D. Assuming the expert chosen is fully balanced, we can get the memory access of single MoE layer as min(2B, Nmoe) x 2D2 For the UltraMem, assuming value dimension D/2, and each token activates the top-m values, then its memory access is min( Bm, N) x D/2. As the batch size increases, the memory access of Moï grows rapidly until reaches an upper limit where all expert parameters need to be accessed. access of UltraMem increases very slowly, only reaching parity with MoE when the batch size is in the tens of thousands. However, in inference scenarios, the batch size typically not very large.
contrast, the memory
Figure shows the inference time and memory access of a 1.6 billion parameter Transformer with 2in34 MoE and x UlraMemm. For larger batch sizes, see Figure 7 in Appendix. Compared to MoE, UltraMem achieves the maximum acceleration of x6 at a batch size of 64, and also shows significant acceleration at other batch sizes.
EXPERIMENTS
In this section, we demonstrate the scaling capabilities of UltraMem, showing that it outperforms MoE. We additionally show how the performance of UltraMem varies with different top-mn values and the number of parameters, and perform an ablation study measure the impact of each part of UltraMem.
'The number of parameters in UltraMem is 12 times the number of parameters in the dense layer. In this case, the total parameters and total computation of UlrrMem the same the 2in34 MoF.
7

---

Published as a conference paper at ICLR 2025
5.1 SETUP
Datasets. Training data comes from RedPajama (Computer, 2023), containing trillion tokens. RedPajama represents clean-room, fully open-source version of the LLaMa (Touvron et al. 2023) dataset. Validation data includes the validation set (Raffel al. 2020 derived from the Common Crawl web corpus. The C4 training set is also incorporated within the RedPajama training data.
Tokenizer based on the GPT-NeoX (Black et al. 2022) tokenizer, which uses the Byte-Pair Encoding (BPE) (Sennrich et al. 2015) algorithm and has vocabulary size of 50,432.
Evaluation. We conducted a comprehensive evaluation all models across ten benchmark datasets. These datasets included MMLU, Trivia-QA- GPQA, and ARC for assessing the models' knowledge capabilities; BBH, BoolQ, HellaSwag, and WinoGrande for evaluating reasoning skills; DROP for testing reading comprehension abilities; and AGleval for measuring overall model performance.
The decoding hyperparameters are aligned with those of LLaMA3 (Dubey et al. 2024). Detrails see Appendix B
Training details. We used a standard pre-norm transformer (Xiong et al. 2020) with rotary embed- dings (Su al. 2024) for our dense models, which have 151M, 680M, 1.6B, and 6.5B parameters For sparse models, including UltraMem, PKM and MoE, we expand the sparse parameters twelvefold from the 15 MM. 680M, and 1.6B dense models. In MoE models, two experts are activated per to- ken QJiag et al. 2024), using balance loss (Fedus et al. 2022) weight of 0.01 to ensure even expert selection. We slightly increased the width of MoE's experts match UltraMem's 's computational and parameter costs. In UltraMem models, the auxiliary loss weight is
0.001 and margin т 0.15. The learning rate for values ten times other parameters and decays linearly. For model structure and hyperparameters details, see Appendix E and for large-scale training optimizations, Appendix C
5.2 EVALUATION ON LANGUAGE MODELING DATASETS
We evaluate models of various sizes, the results are shown in Table where FLOPs is the computa- tion cost of single token, the curves showing changes over the course of training are provided in the Figure i Appendix. We observe that as the model capacity increases, UltraMem can outperform PKM and MoE with the same parameter and computation. On 1.6B dense model, an UltraMem model with 12x the parameters match the performance of a 6.5B dense model.
Table 1: Performance metrics of various models.
Model
Param (B)
FLOPs
Val. losx]
GPQA
TriviaQA?
BBH col
Hella SwagT
Wino Grandet
DROP1
Avg
(G)
Dense-151M
0.15
35.07
PKM-15IM-x12
42.25
MoE-151M-2in32
2.04
0.35
2.63
17.30
33.27
23.24
48.44
55.96
18.57
33,20
2.03
0.35
2.67
19.42
28.97
22.65
43.96
50.83
14.08
29.99
48.83
PKM-680M-x12
8.95
1.50
2.46
26.97
57.32
61.72
25.20
39,70
MoE-680M--2n333
8.95
1.50
2.39
20.54
34.19
26.63
62.71
$9.98
26.54
UltraMem 680M-x12
8.93
1.49
2.37
21.99
$5.17
26.62
65.15
60.54
25.14
42.27
Dense-L66
1.61
58.6
22.63
PKM-1.6B-x12
21.13
MoE- 1.68-2in34
21.36
3.52
2.30
21.32
59,56
29.46
67.34
63.93
28.81
UltraMem 1.6B
21.41
3.50
2.24
24.66
66.38
30.63
71.52
66.38
29.99
48.26
Dense-6.5B
6.44
12.88
$7.28
31.14
69.73
65.9
33.12
46.19
FExcludes tokenizer vocabulary embedding and prediction head parameters. 2This table only includes evaluation results where the metrics have steadily increased with training. For all results, the Table 7 Appendix.
8

---

Published as a conference paper at ICLR 2025
160
3.6
Mcl-1519-2m32
MoE 630M 2m33
McE-1 24--148
2.95
320K
640K
top-m=16
top-m-8
3.4-
PCM-1.6
PCM 15LM x12
PKM B1M-880N-x12
2.90-
80K
160K
top
top-m=32
Validation Loss
3-2
3.0
Udrattem-1
VirrMtom
Dense/A5LM
Validation Loss
2.885
40K
Time (ms)
80
tap-m-128
top-m-512
top-m-256
2.8
Donse 1:48
Deire-AARM
2.80
MoE
Denserb
40
2.75
2.4
2.70
2.2
256
1024
16
64
29
16 32 64
Consumed toke ns (Billian)
Ratio of Sparse v.s. Dense Param.
Ratio of Sparse v.s. Dense Param.
(a) Scaling validation loss
(b) Loss across varying sparsity
(c) Speed across varying sparsity
Figure 6: (a). C4 validation loss of different models at different scale. (b). Scaling curves at different sparsity with 151M activated parameters. Each line represents the model sparsity; e.g. 20K indicates that approximately one out of every 20,000 values will be activated. The loss decreases linearly as the sparse parameters increase exponentially. (c). Inference time for UltraMem and MoE with 1. 1.6B activated parameters. The batch size is 512, sequence length is 1, and key/value cache length is 2048. With fixed activation parameters, UltraMem's inference time remains nearly constant as sparse parameters increase, while MoE's inference time increases significantly.
5.3 VALUE NUMBER AND TOP-m
In most sparse LLMs. such as MoE and UltraMem, there is a clear positive correlation between sparsity and model performance. Therefore, in this section, we conduct a series of scaling experiments by varying selected top-m and value number, i.c., the parameters of the sparse modules, to verify the changes in model performance with respect to sparsity. The result is shown in Figure (b).
is evident that, at the same level sparsity, the validation loss decreases the number of parameters increases and can maintain certain degree of decline. Additionally, the smaller the sparsity, i.c., the larger the proportion of activated parameters, the better the model performance. However, this also results in higher memory access overhead. Thus, there is trade-off between memory access volume and scaling efficiency. our final experiments, we selected a sparsity ratio of 80K as the default model configuration.
As sparse parameters increase, Figure ((c) shows that UltraMem maintains stable inference time despite exponential growth in parameters, long activated parameters (top-m) stay constant. In contrast, MoE's s inference time rises significantly under similar conditions. Additionally, Figure [[b) demonstrates that with smaller batch sizes, MoE's inference speed deteriorates to UltraMem.
compared
5.4 ABLATION
We conduct comprehensive ablation studies based on the 151M dense model. In the baseline, the PKM is version that operates in parallel with the MIP,, making it stronger baseline. For this group of experiments, the learning rate (LR) set to 1. 1.2e-4, with training on 500B tokens and evaluating the cross entropy loss the training and C4 validation sets. We ensure that the parameter count and computational cost of the final version of the model were essentially at the same level.
Table shows the ablation results. We identify 6 changes that significantly improved performance:
1. Doubling the number of values while halving their dimension, and simultaneously double the top-m selections keep the active parameters consistent.
2. Splitting single UltraMem into multiple smaller units evenly across the transformer layers, with outputs skipping several blocks. This arrangement keeps the total parameter count, computational cost, and sparse parameter activation below pre-split levels.

---

Published as a conference paper at ICLR 2025
Table Ablation study of model improvements
Train Loss
Valid. Loss
Dense Param.(M)
Sparse Param.(G)
FLOPs (M)
PKM x10
2,604
2,828
173.01
1.534
346.06
softmax
2.570 -0.034
2.822 -0.006
173.01
1,534
346.06
+half' vdim+proj
2.556 -0.014
2,800 -0.022
178.47
1,529
356.98
+share query
2,560 +0.004
2,803 +0.003
173.46
1.529
346.96
*split big mem&sskip
2.584 -0.006
2,788 -0.015
161 64
1,536
323.32
+query/key LN
2,789
1.536
323.54
+IVE
1,536
344.98
+TDQKR
2,538 -0.006
2.764 -0.008
172.37
1,536
344.98
+MCS
2,521 .0.017
2.761 0.003
172.37
1,536
344.98
+improved init
2,518 00.03
2,758 0.003
172.37
1,536
344.98
+value It decay
2.494 -0.024
2.736 -0.022
172.37
1.536
344.98
+query conv
2.493 -0.001
2.736 -0,000
172.38
1.536
345.02
Total DifT'
-0.111
-0,092
-0.64
+0.002
-1.04
3. Tucker Decomposition Query-Key Retrieval introduces negligible additional parameters
while reducing computation, here r 2.
4. Multi-Core Scoring significantly reduces training loss, and slightly reduces validation loss, here 2.
5. Implicit Value Expansion slightly increases both the parameter count and computational cost, but the improvement is significant, here E = 4.
6. The LR for the value parameters starts at ten times that of the other parameters and linearly decays to match them by the end of training.
Among other changes, sharing the query helps cut computational costs with minor trade-off in performance. Normalizing the query/key greatly reduces spikes in training perplexity and enhances training stability, shown in Figure (a). Improved initialization prevents score and output variance explosions in the early to middle training stages, detailed in Figure [0] (b) and (c). Additionally, employing convolution further limits the variance divergence UltraMem outputs(Figure) (0) The above results are based on incremental ablation. Results from independent ablation can be found in the Table in in Appendix, and they align with our expectations.
Beside, We conduct another ablation studies on IVE, TDQKR, and MCS with different configurations. which are documented Table 3 For IVE, as increases, there is consistent improvement in model performance alongside a notable increase in computational cost. However, the marginal gains decrease as rises, leading us to recommend E. 4. For TDQKR and MCS, increasing r and does not significantly change the computational load, but the effectiveness longer shows marked improvement, hence we suggest using r = 2 and h = 2.
Table 3: Ablation different config on IVE, TDQKR, and MCS
IVE
TDQKR
MCS
Baseline
E=4
E=9
E=16
Baseline
r=2
1=3
r=4
Baseline
h=2
h=4
h=
Training loss.
2.553
-0.019
2.544
-0.0065
-0.0063
-0.017
-0.017
Validation loss.
2.789
-0.027
2.772
-0.0082
-0.003
FLOPs(G)
323.54
+26.4%
344.8
+0.003%
+0.001%
+0.003%
+0.007%
CONCLUSION
In this paper, introduce UltraMem, which, compared to MoE, has minimal memory access and therefore achieves up to a sixfold speed advantage. Concurrently, in terms of performance, UltraMem surpasses MoE with the same parameters and computation model capacity increases, indicating its superior scaling capability. This work presents promising direction for developing more efficient and scalable language models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We extend our deepest gratitude to Pingshuo Ma and Wenda Liu for their invaluable assistance in optimizing the early stage training of large-scale UltraMem. We also appreciate the inference optimization work for the UltraMem carried out by Siyan Chen, as well as Fan Xia's efforts in assessing the inference speed of MoE.
REFERENCES
Hervé Abdi. Singular value decomposition (svd) and generalized singular value decomposition. Encyclopedia of measurement and statistics, 907(912):44, 2007.
Joshua Ainslie, James Lee-Thore Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245 2023.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv: 1007.06000004500 2016.
Daniel Bershatsky, Daria Cherniuk, Talgat Daulbaev, Aleksandr Mikhalev, and Ivan Oseledets. Lotr: Low tensor rank weight adaptation. arXivv preprint arXiv:2402.013376 2024..
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, et al. Gpt-neox-2066 An open-source autoregressive language model. arXiv preprint 3xXx:22004006445. 2022.
Tom B Brown, Language models are few-shot learners. arXiv preprint ArXiv:2005. /4165, 2020.
Zewen Chi. Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song. Xian-Ling Mao, et al. On the representation collapse of sparse mixture of experts. Advances Neural Information Processing Systems, 35:34600-34611 2022.
Christopher Clark, Kenton Lee, Ming-
ding-WWii Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 1905. /0044, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1003::v8800005555 2018.
Together Computer. Redpajama: open source recipe to reproduce llama training dataset, 2023.
URL httph:ttp//:tthtot http://wwwwo..ooooeeeeeeeeneeeee ReeeeereeeeeaPaaaaeaeeea. Data
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. Approximating two-layer feedforward networks for efficient transformers. arXiv preprint arXiv:2310.10837, 2023.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-expsrts language models. arXiv preprint arXiv:2401.06066 2024.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. arXiv preprint arXiv: 1rXiv:1903.00161, 2019.
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.2 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23 (3(120)1139, 2022.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forwart layers are key-value memories. arXiv preprint 2xXiv:20/2.14913. 2020.

---

Published as a conference paper at ICLR 2025
Xu Owen He. Mixture of million experts. arXiv preprint arXlv:2407.04153, 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.000 00t::2009003300, 2020.
AG Howard. Mobilenets: Efficient convolu-tional neural networks for mobile vision applications. arXiv preprini arXiv: /tt:::/704.0000 048661 2017.
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2010.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint S.XXi:24401000050 2024.
Mandar Joshi, Funsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: (XKv::700.0355// 2017.
Gyuwan Kim and Hwan Jung. Large product key memory for pretrained language models. arXiv preprint arXiv:2010.0388)) 2020.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóżdž, Piotr Sankowski, et al. Scaling laws for fine- one-grained mixture of experts. arXiv preprint arXiv:2402.07871, 2024.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. Deepseek-v2: A strong, economical, and efficient mixture-of- experts language model. arXiv preprint arXiv:2405.04434 2024.
Osman Asif Malik and Stephen Becker. Low-rank tucker decomposition of large tensors using tensorsketch. Advances in neural information processing systems, 31, 2018.
Deepak Narayaman, Mohammad Shocybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient acage-scale language model training gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15, 2021.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1 1667, 2020.
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jefí Rasley, and Yuxiong He. Deepspeed-moe: Advancing mixture-of-expert inference and training to power mext-generation ai scale. In International conference on machine learning, pp.
PMLR, 2022.
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. Gpqa: A graduate-level google-proof q&a benchmark. arXiv preprint arXiv:2311. 12022, 2023.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566, 2021.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106. 2021.
12

---

Published as a conference paper at ICLR 2025
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint wrxv:11508.0790,, 2015.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- maatee mixture-of-experts layer. arXiv preprint arXiv:/701.066 066388 2017.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer. arXiv preprint arXiv:2302.0646// 2023.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Img Training multi-billion parameter language models using model parallelism. arXiv preprint N00::/99090005555 2019.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Nenrocomputing. 568: 127063, 2024.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe V Le, Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533 PMLR, 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint
2019.
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin Wang, Amin Saied, Weizhu Chen, and Nan Duan. Agieval: A human-centric benchmark for evaluating foundation models. arXiv preprint MXXv:2304.06664 2023.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mfxture-of-expers with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114 2022.
13

---

Published as a conference paper at ICLR 2025
A ULTRAMEM INITIALIZATION
We initialize value with (0, E
2KHL where k is the activated value number, is the head number, E is the value expansion times. To ensure that the output distribution of UltraMem is N(0, 2 We need to confirm that the mean of top-m score
Assuming the candidate score follows (0, 1), and k < K. We can simplify the problem as follows: Given standard Gaussian distributed random variables X1,... X,, and the random variable Y (rrran((opm((X))) ..., Xn)), find the expected value B(Y). It is difficult to obtain an analytical solution for E(Y), so we approximate E(Y) by sampling M times points from a Gaussian distribution and calculating the mean of the top-ın values.
Then initialize the query layer norm weight as 1/ EEY), the keys layer norm weight as 1/vDa to ensure the expected of candidate score is
B INFERENCE TIME AND MEMORY ACCESS
Figure shows that UltraMem has a much slower growth in memory access compared to MoE, only aligning with MoE in terms of memory access when the batch size reaches 131,072, and it continues to have an advantage inference speed.
MoE-1.68-2nn34
McE-1,688 18--1.68-2h334
UltraMem-1166.
Ultra
10a
Dense-1.68
Memory acçess (GB)
Dense-1.68
Times (ms)
10
10a
10a
10a
10:
10a
10*
10a
10a
10a
10'
10f
10a
Batch Size
Batch Size
(a) Inference time
(b) Memory access
Figure 7: Inference time and memory access of Transformer, MoE and UltraMem. We ensured that three models have the same computation, and MoE and UltraMem have the same parameters. The x-axis and y-axis are both plotted on a logarithmic scale, The sequence length is because during inference, we can only predict one token a time, and the key/value cache length 2048. The modes run on the A100-SXM.
C MEGATRON SUPPORT FOR TRAINING EFFICIENCY
As memory table scales towards billions even trillions of parameters, model parallelism becomes essential to distribute model parameters and optimizer states across multiple devices to ensure they fit into device memory and trainable within reasonable time frame. We leverages Megatron's (Shoeybi et al. 2019 Narayanan et al. 202T) 3D parallelism (pipeline parallelism, data parallelism, and tensor parallelism) for training. However, several parallelism modifications are required to support the memory table effectively. Because pipeline parallelism cannot address scenarios where single layer's 's parameters exceed the memory capacity of single device, and tensor parallelism typically limited to relatively small group of GPUs, making it insufficient to meet the memory table's memory requirements. Consequently, we propose sharding the memory table across a combination of data parallel and tensor parallel groups its subgroups, ensure efficient distribution and scalability.
The memory table can be partitioned either number-win wise or dimension-win wise. The entire process of number-win and dimension -iins partitioning, along with their communication volume analysis and guidance on how to choose the appropriate partitioning method, is detailed in Appendix Di In our structural improvements, halving dim simultaneously reduce the communication overhead for
14

---

Published as a conference paper at ICLR 2025
memory table partition in dedkers
memory parther deskedd
deweno
dewenO
algther
dencect
dewee:
Viluo
wolues
(a) Number-wist partitioning
(b) Dimension-win partitioning
Figure 8: Process of Number-wise partitioning and Dimensionnwvis--arrtttinnng The weighted sum pooling step is omitted in the diagram.
both number-wise and dimension-wise partitioning. However, increasing top-m will proportionally increase the communication overhead. Additionally, Implicit Value Expansion, due to the increase in the size of values after weighted sum pooling, will further impact the communication volume for dimension wise partitioning.
To further augment performance, several key modifications have been implemented:
Fused Lookup-Reduce Reduce Operator: This newly introduced operator accelerates computations and reduces memory usage by combining the lookup and weighted sum pooling operations into single, more efficient step.
Asynchronous Execution Strategy: Recognizing the benefits of cross-layee utilization of the memory layer, we have adopted an asynchronous execution strategy. This strategic choice allows for the concurrent processing of memory calculations alongside dense network operations, substantially enhancing the overall system performance.
These enhancements demonstrate the efficacy of our parallelism strategy within the Megatron framework, paving the way for more efficient training of large-scale models.
D NUMBER- WISE AND DIMENSION WISE PARTITION DETAILS
Figure 8 shows the process of number-wise and dimension-wise partition. For number-wise partition- ing, first perform all-to-all indices distribute them the corresponding devices. After the lookup operation, the results are sent back to the original devices, then do weighted sum pooling. For dimension-wild partitioning, we need to perform an all all-gather operation on indices obtain all indices across devices. The lookup operation is then performed, dimension wiseessiioiiiseeeeee partitioning allows the results be sent back each device after completing the weighted sum pooling.
15

---

Published as a conference paper at ICLR 2025
tspmin2 Espming 1opm-226
3300
1000
1000
500
230
500
900
1020
1200
1400
Figure 9: Relationship between P and dim for communication volume of number-wise / dimension- wise equals 1, the shaded area is number-wise dimension wise greater than
Assuming the memory table is distributed across P processors, the communication volume can be described as follows:
Number-wise Partitioning Communication Volume (not considering indices deduplication):
All-to-all transmission of indices: sizeof(int) x bs x topm x (P-1)/2
All-to-all transmission of embeddings after lookup: sizeof bs topm x v did x (P-1)/P P)
Dimension-wise Partitioning Communication Volume:
AllGather indices: sizeof(int) x bs x topm x (P --1) 1)
AllGather scores: sizeof(bfloati) x bs x topm (P 1)
All-to-all transmission of embeddings post-lookup reduction: fizeof(blloat16) x x v_dim (P- 1)/P
Here v_dim is the value dimension, bs is the batch size times sequence length. Figure showw the relationship between P and v dim for communication volume of these two partitioning methods, helping choose the appropriate partitioning method under fixed configuration.
16

---

Published as a conference paper at ICLR 2025
E EXPERIMENT SETTING
Table displays common hyper-parameter settings for all experiments. "LR" stands for Learning Rate, corresponding to the val- ues 6e-4, 2.5e-4, 2e-4, and 1.2e-4 for dense models with sizes 151M, 680M, 1.6B, and 6,5B, respectively (Brown) 2020). Regard- ing the insertion of UltraMem and PKM, for UltraMem- 15 M, it's 3:5/6:8/9:11, where 3:5 indicates that UltraMem imput is taken from layer 3 and inserted back into the output of layer 5, and on. For UltraMem-6800 it's For UltraMem-1.66, For
Configuration Key
Value
Weight decay
0.1
81
0.9
82
0.95
LR
6e 4/-. 4/2e 2/1.2e.4
LR schedule
cusime
warmup
Dropout
0.1
Batch size
2048
Sequence length
2048
Training step
238418
1:7/8:12/13:111
it's 1111::///111111:::222223333 PKM-151M, it's 6:6. For PKM-680M, it's 12:12. For PKM-1.6B, it's 16:16. The set- tings for UltraMem and MoE models align with their dense counterparts based on dense parameter size. Table shows the model pa- rameter setting used scaling experiments. What's s more, the common setting for Ultra- Mem is shown 5
Table 4: Training hyper-parameters
Configuration Key
Value
Tucker
Multi-core scoring is
2
Virtual memory expunsion
Aux loss weight
0.001
Aux loss margin
0.15
Table 5: Common UltraMem configuration
Evaludation datasets. We use 10 benchmarks evaluate all kind of models.
1. Knowledge: Massive Multitask Language Understanding (MMLU) (Hendrycks et al. 2020 TriviaQA (Joshi et al. 2017). Graduate-Levv Google-Proof Q&A Benchmark (GPQA) (Rein et al. 2023), A12 Reasoning Challenge (ARC) (Clark et al. 2018).
2. Reasoning: BIG-Bench Hard (BBH) (Suzgun et al. 2022). Boolean Questions
(BoolQ) CCark al., 2019), HellaSwag (Hella) (Zellers et 2019
WinoGrande
(Wino) (Sakaguchi et al. 2021
3. Reading comprehension: Discrete Reasoning Over Paragraphs (DROP) Duua et al.
2019
4. Comprehensive ability: AGIEval (Zhong et al. 2023)
Model
Hidden Dim
Inner Dim
Atin Head
Layer
Top-
Expert
Kdim
Knum
Mem Layer
Parum (B)
FLOPS
4096
Dense 680M
1536
6144
16
24
0.68
1.36
Dense-1.68
2048
1.61
Dense-6.5B
4096
16384
32
32
6.44
12.88
MoE-151MM2hin
1024
2528
16
12
32
2,04
0.35
MoE-680M-2in33
1536
3584
8.95
1.50
MoB-1.68-2in344
2048
4672
16
32
2
34
21.36
3,52
PKM-680M-x12
1536
6144
16
24
35x8
768
2308
8.98
1,50
PKM-1.68-x12
2048
8192
1792
21.44
3.52
UhtraMem-151MP
1024
4096
16
16x2
256
1024
1.71
0.35
UltraMem: 151M x212
1024
4096
16
12
16x2
256
1100
2.03
0.35
UlraMemn-600N-x11
UlraMenn- 1.68-x122
2048
8192
16
32
42x2
448
1792
21.41
3,50
Table 6: Model parameter setting. Top-m means chosen expert number in MoE, means chosen value number times head number PKM and UltraMem. Kdim means the key dimension PKM and UltraMem. Knum means the number of keys, Knum? is the number of values.
F MORE EXPERIMENT RESULTS
17

---

Published as a conference paper at ICLR 2025
w/o query/key LN with query/key LN
Baseline Improved Init.
UltraMem final output
600
Training Perple exity
Baseline
400
Improved Init.
Top1
40
Query conv.
200
20
20
200
400
10a
10 Tokens (Billions)
102
400 (c) UltraMem output standard de- viation
200
Takens (Billions)
Tokens (Billions)
(a) Training perplexity
(b) Topl score
Figure 10: Model training state details. "Topl Score" refers to the highest score among the retrieved keys. "UltraMem Output Std" represents the standard deviation of the outputs from the last layer of UltraMem.
Table 7: All performance metrics of various models
Param FLOPs ARC-CT GPQAT
Trivia MMLU1 QA
BoolQ
Hella
Wino
DROP? Avgt
Model
(B)
(G)
cott
Swagt
Grandef
Evalt
Dense-151M
25.60
26.50
22.57
35.07
26.77
PKM-15IM-x12
42.25
MoB-15IM--insk
2.04
0.35
26.96
26.58
23.24
55.96
31.56
UltraMem
151M-x112
2.03
0.35
25.68
19.42
28.97
25.62
22.65
47.74
43.96
50.83
28.89
Dense-680M
0.68
1.36
24.06
21.09
27.16
24.64
24.65
46,42
48.83
54,93
22.97
30.42
PKM-680M-x12
8,95
1.50
25.51
25.22
26.98
80
57.32
61.72
8.94
25,20
33.97
MoF.6800M2
24.38
62.71
UtraMeem-880M-X1
8.93
1.49
23.72
21.99
55.17
24.97
26.62
48.20
64.15
60.54
25.14
35.88
Dense-1.68
1.61
3.21
26.30
21.76
39.65
26.19
26.41
51.55 .50
58.6
61.72
9.22
22.63
34.81
PKM-L66-x12
3,48
26.71
3.52
67.34
UltraMem- L66-x112
21.41
3.50
25.94
24.66
24.67
30.60
59.8
71.52
66.38
29.99
40.88
Dense 6.5B
6.44
12.88
28.
19.98
57.28
27.68
31.14
68.2
69.73
65.9
9.23
33.12
41.04
Table 8: Independent ablation study of model improvements
Train Loss
Valid. Loss
PKM-15IM-X10
2.604
2,828
softmax
half vdimepmij
+ share query
-0.003
40.002
+ split big mem
00.033
40.05
+ ququryykky LN
-0.002
+0.003
+ IVE
-0.025
-0.023
+ TDQQR
.0.003
-0.007
TDQKR MCS
-0.02
-0.009
value Ir decay
-0.017
-0.007
query
-0.005
-0.001
18

---

Published as a conference paper at ICLR 2025
Dose Md
133N
0.300
0.50
0.275
Mac 631
Accuayy
0.45
-.40
HON 1518 -22
PLO- 20.11
Mot 680M 20:3
Mat 151 20130
Accuracy
0.250
0.225
Done
Dense-1511
Not 60 26-34
0.200
McT-150--1177
MIG 6000 2513
0.175
0.25
0.125
128
256 Consumed tokens (Billon)
1024
128
1024
Consummed takens (Billon)
(a) Average accuracy
(b) BBH-cot-3shot accuracy
wW. Denne-1 mi
Denne-1 AB
1518
15186-15M
Borse 8.50
Burse-550
0.60
FNN 1-01
PON-15114-1
PIN PNN-60315-12
PON-58396:12
0.20
UtraMton 88
UtaaMte
Utaattm 15
UtraMerr-1 15 JM.
U traltem d10R-412
0.45
U traalerr- 600M--12
0.10
0.30
512
1024
1024
Conzumer ed tokens (Billon)
Consumed tokens (Bllion)
(c) DROP accuracy
(d) Hcllaswag accuracy
0.675
Denon FB
Deme-151M
Demso111M
0.650
Dusce92
0.6
Dunce33 30
Derse 600
Mod 3.102334
0.625
Moz 153M 21-32
0.5
2013
6B-42
F8N-4899--12
MiraMom-1 1.5xx22
UniaMtom 15.Mix32
SiraNem-6320-c23
0.2
0.525
0.500
0.1
512
1024
128
256
1024
Consumed tokers (Bilion)
Consumed takens (Bllion)
(c) Winogrande 5shot accuracy
(f) TriviaQA Q 5shot accuracy
Figure 11: The changes in accuracy for all observable evaluation throughout the training.
19


Reviews:
[{"review_id": "DK60WF4q2J", "reviewer": "Reviewer_DbFB", "summary": "Motivated by the high memory access cost bottleneck of Mixtures of Experts (MoEs) for Large Language Model (LLM) inference, the authors propose an alternative methodology, UltraMem, based on Product Ley Memory (PKM) which improves memory access above MoEs, but suffers in performance compared to them. The authors make several observations to improve the vanilla PKM methodology, basing their method on this baseline. The authors then not three problems with the PKM architecture, notably that it does not scale to large value sizes, product key decomposition biases retrieval, and unbalanced GPU computation/communication for large models. The authors propose to decompose the large memory layer of PKM into many smaller memory layers distributed across layers, allowing execution of memory and transformer layers to overlap. The authors propose to use a Tucker decomposition instead of the product key decomposition of PKM. Finally, rather than explicitly maintaining a large memory table for values, the authors propose a virtual memory approach. The authors evaluate UltraMem across several tasks compared to dense and an MoE equivalent model, and compare the performance across these tasks, along with the inference and memory accesses for each, demonstrating significantly faster inference with similar performance to an MoE model.", "questions": "Suggestions:\n* Focus on the high-level motivation before diving into details of methodology and experiments.\n* Move the related work up front in the paper, and put in background required for your method. Include MoEs, PKM, Tucker decompositions, Megatron, any background relevant to IVE and MCS, etc.\n* Fix readability of all the text in the figures.\n* Figure 2 (and to some extent Figure 3) are poorly labelled/captioned, and it's largely left to the reader to figure out what each side represents and how it's related to the method. At a minimum label the subfigures or describe in caption what left/right are. Would also suggest positioning floats at top of page.\n* When including a large table of metrics (i.e. in Table 1), it's helpful to the reader to explain whether lower or higher is better for each, e.g. with an arrow in the header.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* Memory access is perhaps the main bottleneck in inference for contemporary hardware, and so this is a well-motivated and impactful direction to be exploring.\n* The results appear to show impressive improvements over Mixtures of Experts (MoEs) in inference time and memory access while maintaining validation loss/perplexity. \n* The authors methodology is detailed and thoughtful, for example deriving the correct initialization for their method (which would appear to also apply somewhat to PKM).\n* Many of the figures/graphics themselves are good in design and would be helpful to understanding the methodology if it wasn't for the other serious issues with readability and captions (see below).\n* Changes to PKM (before the UltraMem specific changes) are listed in 2.2 as the bag of tricks that improve PKM's performance, and are separate from the author's proposed UltraMem structure. It would be beneficial if more papers took such an approach to be clear about the differences between improved training methodology for baseline methods and a newly proposed method. Ablation in particular is nice about being clear how these changes improve baseline PKM performance.\n* Real-world inference and memory access results on GPU hardware.\n* Large-scale experiments, in particular results are evaluated on various model sizes, from 151M param up to 1.6B. FLOPS for the sparse models.", "weaknesses": "* Overall the paper reads poorly, more as a collection of experimental details and results rather than a cohesive story. I know that's frustrating and perhaps vague feedback to get as an author but I think it's really important to point out as it makes the paper quite hard to read as it is, and reduces the impact of the author's work. To be clear, I do appreciate the experimental and methodological details themselves. However, I believe the authors could do the work much more justice by taking a step back and framing their research story better. In particular, I would recommend focusing first on section 3 as higher-level motivation of the method before diving into building on top of PKM. \n* All the figures containing the majority of results, and some of the methodology figures (Fig 4) are **far** too small. They are so small that the figures are completely unreadable on paper. I measured the font sizes in the figures out of curiousity, and most of the fonts are 2-3pt!\n* Given that the proposed method is largely based on PKM, PKM should be a baseline in the results. The authors mention PKM's \"..effectiveness is significantly inferior to that of MoE.\", without explaining on what measures this is true, or giving the reader to make that judgement in the results. I don't actually doubt the authors on this, but it must be demonstrated in the experiments.\n* While as noted above it's great that the \"bag of tricks\" for PKM is listed, along with the ablation, it also would seem important that the performance of the improved PKM should be evaluated as a baseline over just vanilla PKM, and compared with the improvements from the author's proposed method.\n* The related work is **far** too short given all the work done in this field, much of which is referred to by the authors and built upon in the proposed method. To be honest I'm never a fan of related-work at the end of the paper, but I feel like this paper in particular would have greatly benefited by the related work coming before the methodology rather than at the end of the paper, as instead of having to bring up related work throughout the paper as it is built upon, it could have just been earlier summarized and then referenced throughout, allowing the story to focus on the methodology more clearly.\n* Relies too much on validation loss/perplexity for the evaluation of whether generalization performance is maintained in the main paper/figures. The authors do have six other metrics for tasks in Table 1 and figures in the appendix for these tasks (which they should explain better in the paper/background rather than quickly citing all of them in a list). Given that the authors are in a sense proposing a form of sparsity, and given the results of Jaiswal et al. for pruning (Compressing LLMs: The Truth is Rarely Pure and Never Simple, ICLR 2024) I believe it's more important to focus on the performance in task-specific measures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xFLsJzKFrJ", "reviewer": "Reviewer_jFrZ", "summary": "This paper proposes a novel deep learning architecture called UltraMem, designed to reduce memory access costs during the inference process of large language models, thereby improving inference efficiency. The core innovation of UltraMem lies in its introduction of an ultra-sparse memory layer, which allows the model to activate only a small number of necessary memory units when processing tasks. This approach reduces the number of memory accesses, effectively lowering inference latency. Specifically, , UltraMem surpasses MoE with the same parameters and computation as model capacity increases.", "questions": "see weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper introduces the UltraMem architecture, which demonstrates innovation by significantly reducing inference latency while maintaining computational efficiency. \n2. The paper presents extensive experiments comparing the performance of UltraMem with traditional models (such as MoE and dense models), verifying UltraMem’s advantages in inference speed, memory access costs, and scalability. \n3. The experiments show that UltraMem’s memory access volume grows much more slowly with batch size compared to MoE, enhancing its practical utility.", "weaknesses": "1. The paper lacks references and descriptions of the architecture diagrams within the main text. For example, each step in Figure 4 is not referenced in the text. Additionally, certain terms in the architecture diagrams, such as “fetch values,” are not explained in the text, making the paper difficult to follow.\n2. The paper claims that the proposed UltraMem method has stronger scalability. However, UltraMem was only tested on models with 151M, 680M, and 1.6B parameters, without experiments on larger models, such as the 6.5B model or beyond. This parameter range is insufficient to fully verify UltraMem’s scalability in large-scale models. It is recommended to extend the experimental scale to explore UltraMem’s performance and scalability at extremely large parameter sizes (e.g., 10B and above).\n3. Several modules proposed in the paper, such as Implicit Value Expansion (IVE) and Multi-Core Scoring (MCS), collectively enhance UltraMem’s performance. However, the independent effects of each module are not adequately evaluated. For instance, Table 2 provides ablation experiments but lacks detailed analysis of the independent effects of key modules like IVE and MCS across different tasks or sparsity settings. It is recommended to expand on Table 2 by adding experiments that assess the independent impact of each optimization module on model accuracy.\n4. Figure 4 illustrates the process of Tucker decomposition but lacks analysis of how different decomposition parameters, such as rank r, affect model accuracy. To clarify the impact of different Tucker decomposition configurations on model performance, it is recommended to include more comprehensive ablation experiments to quantify Tucker decomposition’s specific role in UltraMem. Additionally, experiments should be added to analyze the effects of different values of E in the IVE method on experimental results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yUdCqRaIxG", "reviewer": "Reviewer_ch2J", "summary": "In this paper, the authors propose UltraMem, a memory-efficient method designed to replace MLP layers or Mixture of Experts (MoE) in Transformer architectures, particularly for large language models. UltraMem builds on a refined Product Key Memory (PKM), utilizing a 2D grid structure with row-column key decomposition to reduce retrieval costs. Additionally, it employs Tucker Decomposition to further minimize memory usage and computational overhead. Experiments across various language tasks demonstrate UltraMem’s promising effectiveness and efficiency.", "questions": "here are some minor writing issues, such as in Line 094, where it states 'introduce the origin large memory.' These should be easy to address.\n\nMy main questions and observations are:\n\n1. **Comparison with Fine-Grained Mixture of Experts**: How does this method compare to fine-grained Mixture of Experts as in [1], where the hidden dimension is split between experts?\n2. **Training Speed**: How does UltraMem’s training speed compare to that of MoE?\n3. **Placement of UltraMem Blocks**: On what basis were the positions of the UltraMem blocks within the architecture chosen?\n4. **Figure Quality**: Improving the quality of the figures and plots would enhance clarity.\n\n**Suggestions**:\n\n1. It would be beneficial to conduct experiments with popular open-source models like Llama-3.2-1B, Mistral, or others.\n2. Since Section 3, \"Why UltraMem Instead of MoE,\" discusses UltraMem’s advantages over MoE, it would be valuable to include experiments demonstrating that UltraMem consistently outperforms MoE in real-world scenarios, using a comprehensive baseline model for comparison.\n3. Additionally, an ablation study comparing UltraMem's memory efficiency and performance against MoE based on [1] would further strengthen the analysis if applicable.\n\n[1] *Scaling Laws for Fine-Grained Mixture of Experts*.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper addresses an interesting and impactful problem with practical applications in large language model (LLM) research.\n\n- The authors propose an alternative, memory-efficient approach to achieving the performance of Mixture of Experts (MoE) models.\n- They introduce a sparse memory access mechanism and a 2D Product Key Memory structure, which restricts memory access to only the most relevant slots, enhancing efficiency.\n- **Scalability**: The use of Tucker Decomposition improves the scalability of the method, allowing it to handle larger models effectively.\n- **Experiments**: The experiments demonstrate promising results, achieving comparable performance to MoE models of similar scale while maintaining greater memory efficiency.", "weaknesses": "Although the idea presented in this paper is novel, the experimental performance assessment shows some limitations:\n\n1. **Lack of a Proper Baseline**: The authors do not use a well-established MoE baseline to evaluate the performance of their method. Since UltraMem is proposed as a memory-efficient alternative to MoE, it would be beneficial to compare it against real-world MoE implementations from existing open-source models, providing a more comprehensive analysis.\n\n2. **Limited Experiments with Popular Models**: The experiments primarily use custom models, limiting the generalizability of the results. It would be valuable to assess UltraMem’s effectiveness on popular open-source models like Llama, Mistral, or others.\n\n3. **Limited Analysis of Sparsity Levels**: The paper could benefit from a deeper investigation into the effects of different sparsity levels on memory efficiency and performance, as this is a key component of UltraMem’s approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NLyiHz0FHv", "reviewer": "Reviewer_bPzv", "summary": "This paper introduces a novel layer designed to increase model size efficiently. Specifically, the approach refines the existing PKM method and incorporates TDQKR as a sparse method to enhance performance. Additionally, it proposes implicit value expansion to address memory access challenges in large models. A comparative analysis of memory access between the proposed method and MOE is also provided. Experimental results demonstrate significant performance improvements.", "questions": "1. The proposed IVE aims to reduce memory access. Given that it handles the general form of matrix multiplication, could it also be applied to MoE or other methods that involve extensive memory access?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The techniques proposed in this paper are novel. On one hand, the paper thoroughly studies the techniques in PKM, offering valuable empirical insights for future research. On the other hand, it integrates algorithms with memory management to address the memory access issue.   \n2. The paper includes an efficiency analysis focused on memory access.\n3. Experimental results show that the method significantly outperforms existing approaches on large models, being up to six times faster than MOE at the same scale.", "weaknesses": "1. The writing could be improved, particularly since it introduces concepts beyond the domain of algorithms. These should be better explained as background information. For example, the implicit value expansion technique is somewhat confusing when the paper discusses virtual and physical memory, and could benefit from further clarification.\n\n2. The experimental results in the main paper do not fully support the claims made in the methods section. First, while the paper asserts that the method improves PKM, there is no direct comparison except for validation loss. Second, although the paper claims that its memory access method outperforms MOE, this is only demonstrated in the appendix. The main table merely shows FLOPs and model parameters, which do not sufficiently illustrate the method's advantages (including the running time).\n\n### Minor:\n1. The dimensions in Equation (6) are inconsistent.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)