MemoryFormer : Minimize Transformer Computation by Removing Fully-Connected Layers
Authors: Ning Ding
Yehui Tang
Haochen Qin
Zhenli Zhou
Chao Xu
Lin Li
Kai Han
Liao Heng
Yunhe Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: transformer architecture
large language model
computational complexity reduction
Primary area: deep_learning_architectures
TL;DR: Use "retrieving-from-lookup-tables" operation which has near-zero FLOPs to replace the matrix multiplication used in fully-connected layer to greatly reduce the computational complexity of transformer models.

Abstract: In order to reduce the computational complexity of large language models, great efforts have been made to to improve the efficiency of transformer models such as linear attention and flash-attention. However, the model size and corresponding computational complexity are constantly scaled up in pursuit of higher performance. In this work, we present MemoryFormer, a novel transformer architecture which significantly reduces the computational complexity (FLOPs) from a new perspective.  We eliminate nearly all the computations of the transformer model except for the necessary computation required by the multi-head attention operation. This is made possible by utilizing an alternative method for feature transformation to replace the linear projection of fully-connected layers. Specifically, we first construct a group of in-memory lookup tables that store a large amount of discrete vectors to replace the weight matrix used in linear projection. We then use a hash algorithm to retrieve a correlated subset of vectors dynamically based on the input embedding. The retrieved vectors combined together will form the output embedding, which provides an estimation of the result of matrix multiplication operation in a fully-connected layer. Compared to conducting matrix multiplication, retrieving data blocks from memory is a much cheaper operation which requires little computations. We train MemoryFormer from scratch and conduct extensive experiments on various benchmarks to demonstrate the effectiveness of the proposed model.

MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers
Ning Ding"", Yehui Tang"* Haochen Qin Zhenli Zhou', Chao Xu', Lin Li
Kai Han Heng Liao², Yunhe Wang
State Key Lab of General AI, School of Intelligence Science and Technology, Peking University Huawei HiSilicon
Huawei Noah's 's Ark Lab.
dingningdatu.pka pku.ppuuuuuuu edu.cndu.c cn (yehui, ayehui.tang, yunhe. (awng))euwwwe.ccommm
xuchaoRcis, ..u. edu.
Abstract
In order reduce the computational complexity of large language models, great efforts have been made to to improve the efficiency of transformer models such as linear attention and flash-attention. However, the model size and corresponding computational complexity are constantly scaled up in pursuit of higher perfor- mance. In this work, we present MemoryFormer, novel transformer architecture which significantly reduces the computational complexity (FLOPs) from new perspective. We eliminate nearly all the computations of the transformer model except for the necessary computation required by the multi-head attention operation. This is made possible by utilizing an alternative method for feature transformation replace the linear projection of fully-commn connectdd layers. Specifically, first construct a group of in- mn-merrry lookup tables that store large amount of discrete vectors replace the weight matrix used in linear projection. We then use hash algorithm to retrieve correlated subset of vectors dynamically based on the input embedding. The retrieved vectors combined together will form the output embedding, which provides an estimation of the result of matrix multiplication operation fully-connected layer. Compared conducting matrix multipli- cation, retrieving data blocks from memory is much cheaper operation which requires little computations. We train Memory Former from scratch and conduct extensive experiments on various benchmarks to demonstrate the effectiveness of the proposed model.
Introduction
The Transformer model has made magnificent achievement in deep learning community since it was made public. The Transformer not only successfully leads the revolution the field of natural language processing due to its excellent performance, but also motivates the innovation in model architecture in other fields such as computer vision and speech recognition. Recently, large language models (LLMs), transformers that are extremely scaled up size, have drawn remarkable attention of both researchers and non-researchers across the globe. The unprecedented emergent abilities that LLMs demonstrate attract an increasing number of investments and researches, which point out potential pathway for the artificial general intelligence.
However, what comes along with the scaling lay is only more intelligence, but also greater consumption of computing resources. The ever-increasing computational complexity is currently the main obstacle hindering the application and popularization of L.MSs. In response to this situation,
Equal Contribution. *Correpooniigg Author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

great efforts have been made towards optimizing the architecture of transformer model by the research community. Some works using traditional methods such as model pruming and weight quantization are able to lower the computational complexity of L.LMs to some degree. Another line of works that are specialized for transformer re-design the self-attention mechanism, which the key to sequence modeling. They use sliding-windows windows or kernel function to reduce the complexity from quadratic to sub-quadratic or even linear with respect to the sequence length, while maintaining a comparable performance.
According to our observation, in most application scenarios, only a small proportion of the computa- tional complexity comes from the multi-head attention (MHA) operation, while the majority of the computation comes from the fully-connected (FC) layers in the transformer model. Specifically, given standard transformer model with the hidden size being d and the length of input sequence being the amount of floating-point computation of the MHA operation is 2s and the computation in all the FC layers is 12sd The computation required by the MHA becomes dominant only when & >> 6d. That's say, for an LLM with hidden size 4096, the sequence length needs be larger than 24K.
Another observation we have is that, at present, the inference stage of deep neural network relies on the parallel-computing cores of the graphics processing unit (GPU), while the CPU and the random- access memory (RAM) resources in the computer system are left almost unused, despite the fact that the size of the RAM easily reaches terabytes and the CPU has hundreds of cores (e.g. NVIDIA DGX A100 has 2TB RAM and 128 CPU cores (2I)). What's more, CPU manufacturers have started developing tensor core to accelerate parallel-computing which might make the low-latency CPU inference feasible in the future.
Based on the above observations, we present novel transformer architecture in this work, named MemoryFormen minimize the required computational complexity from new perspective. Instead of the fully-connected layers used in a standard transformer, MemoryFormer uses the Memory Layer process the feature vector (the embedding of token). Specifically, the Memory Layer contains group of in memory hash tables which store a large amount of discrete vectors. It uses locality- sensitive hashing (LSH) algorithm to retrieve subset of vectors from the hash tables that are correlated with the imput token embedding. These retrieved vectors are then aggregated with different weights to form the output of the Memory Layer. The result of hashing- aorhing-and-aan and-agggegg operation performed by the Memory Layer provides an estimation of the matrix multiplication of the fully- connected layer. We also design a method to make all the vectors stored in hash tables learnable via the back-prod poopppropagation of gradients. Therefore, the MemoryFormercan be trained from scratch in an end-to-end manner.
Since the amount of computation produced by the hash operation is negligible, the absolute majority of the computations now results from the matrix- multiplications within the multi-hear attention. Fig-
300
1en:
Transformer
seg 10r-1024 MemoryFormer
250
204- Transformer
Inference blokk
seq,
2e4- MemoryFormet
ure [ demonstrates the comparison of the compu- tational complexity of one block between the pro-
200
150
posed MemoryFormer (red) and the baseline trans- former (blue). The MemoryFormer block only re-
100
quires 199 of the FLOPs compared with the
baseline transformer block when sequence length
50
2048 and hidden size d 2048. The effect of FeOPstrduuction is going to be more significant as the model size is scaled up. Replacing all fully- connected layers with Memory Layers allows us to trade the memory resources for less computational complexity.
512
1024
1536
2048
25.00
3072
model hidden size
1
Figure 1: FLOPs with different model hidden size and sequence lengths.
This work not only proposes a new FLOPs-reduction strategy that is different from existing ap- proaches, but also provides guiding significance for the hardware design (e.g. bigger bus width and higher cache hit rate) of the next-generation parallel-computngg platform. We train series of MemoryFormers different sizes and validate the effectiveness of these models on multiple public benchmarks. The experiment results show that method can achieve performance comparable to the bascline Transformer with significantly less computation.

---

2 MemoryFormer
2.1 Background
In the standard Transformer model, there are two main types of operations to perform feature transformation on the sequence of token embeddings. One is the multi-head attention (MHA), the most important operation in transformer block which captures the long-range inter-ellatonnshiss among different tokens within the sequence. The other is the ubiquitous fully- that performs linear projection on each token in the sequence separately. In addition to the projections of (We, WK, } prior to the MHA, the feed-forward network (FFN) is also composed of multiple FC layers.
(FC) layer
Let x E Rd be a row vector representing any token embedding, fully-connected layer parameterized by weight matrix W Rdxh applies a linear projection to formulated as y xW, where y g the output token embedding. For a sequence composed of tokens, this would become a matrix multiplication in Eq. [ with computational complexity of (sdhh).
Y = X y R*xh.
In finiteddimensional vector space, the fully-connected layer is a continuous linear operator, which means, for two adjacent input feature vector X1 and Xg, given the same weight matrix W, the projected vectors y1 W and y2 = X W are most likely to be similar as well. In this work, we want to find an alternative mapping function which has much less computational complexity than (ssh) yet generally in accord with the properties of the linear projection. If this is achieved, we can use this alternative method replace all the FC layers to perform feature transformation while reducing computation.
In this paper, we use normal-font lowercase letters to denote scalars (e.g. and h), bold-font lowercase letters to denote vectors (e.g. x and y), and bold-font uppercase letters to denote matrices (e.g. W). We use the notation [-]x to represent the i-th row of a matrix, and also use [:]: to represent the i-th entry of a vector.
2.2 Compute-less Locality- Hashing
Unlike ordinary hash functions that are designed to avoid collisions for different encoded items, the aim of senssity-sessitive hashing (LSH) is map similar items to the same hash bucket (memory location) in the hash table. For example, in the text-to-image search system, several different descriptions for the same image expected to have an identical hash code after encoded by the LSH function, and thus retrieve the same image.
We decide to apply LSH function in the embedding space to encode any input feature vectors. Assuming x hashed to a specific bucket that stores a vector ý by the LSH function, then the hashing result for some adjacent neighbor vectors of x will correspond to the same hash bucket and retrieve as well. If the value y in the hash table an approximation of the result of linear operation y = xW for the input vector we can use this method, that is to find an estimated result in hash table, to replace the fully-connected layer, while only using the FLOPs of the hashing operation.
Firstly, we construct an in-memory hash table parameterized by the matrix T E = which stores 2d vectors [T) R Traditional LSH functions, such Hyperplane-Lske [5), incorporate multiple linear projections to generate the hash code for table indexing. To avoid any unnecessary computations, utilize much simpler LSH function to generate the hash code h(x). Specifically, the process of hashing and retrieving the result from table T is formulated as:
(x) = integer(sign( x)),
(2)
- 1, if
[x]
< 0.
sign([>}+)
(3)
1, if
x
0,
d-1
[s] +1
integer(s)
2
(4)
2
1=0
y = T)
(5)
3

---

where x E R is the input vector, s = sign(x) E f-1,1}4 is the corresponding binary representation (hash code), integer(-) function converts the binary representation to the corresponding non-negative integer used as the index number of the hash table, h(x) E 20,1,2..... 2 1]. Notably, the space complexity of such hash table is (2) The required memory space would be 10 terabytes (TB) when using float16 datatype with d = h = 512, which is impractical for any modern computer system.
To tackle this problem, we propose to evenly split x E Ie into K non-overlapping chunks and handle them separately:
2 sxlit(x, num_chunk = K), = 1,2,
(6)
where 2g E GR. d each sector respectively. Therefore, the output result is
R and d is evenly divisible by K. We then set up hash table Ta IR22 xh for
K
T Tkln(n=)
Since Z has a smaller bit width after binarization and thus the corresponding hash table Ts would comsume less memory space for storage. The space complexity of Eq. 7 is O(K2Th). h). When = h = 512, = 8, K = 64 and data type is float16, the storage required by all K hash tables is
16 MMg) Figure 2 is simple demonstration of the proposed locality-sensitive hashing.
2.3 Memory Layer
So far, the mentione formulation is able to simulate the forward pass of fully-connected layer. And the values store in the hash tables can be updated via back-propagationn derivative of the loss function with respect to the hash ta- ble is (.(e))
sign(22)=(-..1)
bucket
bucket3
The
Zz=
by m However, the input vector is unable to have gradient since it's hashed to multiple inte- gers hu(z)) used as the index number for retrieval, which is a non-differentiable operation. If we can reformulate Eq. (7) y = > p(zx) (izx)-((TTT -(Tx)a(a)) ) add coefficient p(zx) weight each retrieved item, where p(zz) is function of the variable ag, the gradients can be back-propagated to the input x via [Tи))(es)=
0f.
Zi Z1-(0.7,0.1
Zs=(0.4.-0.6))
bucketO
-1 buckel
Figure 2: A demonstration with
T 2 and K = 3, where 21 is hashed to the bucket2 of T2, Z2 is hashed to the bucketl of T2, Zz is hashed to the bucket2 of T3.
As we can observe from Figure 2 many sub-vectors with various directions and amplitudes can still be hashed to the same bucket as long their signs are identical, but the angle between the bucket's s representative binary vector (each entry is either or -1) and these sub-vector. scaled cosine similarity, which takes into account both the direction and amplitude of Zk. to measure the relevance between Zµ and its corresponding hash bucket h(z;)::
different, which is defined by the cosine value cos(z;, sign sign(zz)). We use
sim(ze, l((z)) )} = |2 ||sign (zx) ||2 cos(2), sign (2x)) = (Zx, sign(ze)),
(8)
where (.) computes the inner- product of two vectors. Considering all the 2" buckets that z is possibly hashed to in the lookup table Tr simultaneously, we define the probability that N is specifically mapped to the h(2z)-th hash bucket:
p(zz) explsim(zz hzzx))/t]
ezp sign(z
Σ explsim 1)/t]
expl (
where t is the temperature paree--arammttr, integer;" (1) 1-1, is function that maps an non-negative integer 0 < i 2" to the corresponding 7-bit binary representation. Note that (:, integer;" (i)) operator takes the summation after elementwise selective sign-flippigg over any 7-dimensional vector, therefore we have

---

odppl
Winghted Surt
Memory Block
x.V
a (QK y
PA
Integer
K
Mamory Jayar
Layer
Input
Figure 3: Left: The schematic diagram of the Memory Layer Right: One building block of the MemoryFormer.
(Zg, sign (2))
[zalil,
27
expl Zk. integer, (1)) II (erp((axl)) + exp(-(zzpl ==0
(10)
1=0
erp( El [
1
p(z))=
capp [x]] + expp - [2n]:/1)
I expl (arl-2212z11) exp(-21(2u).//1)
This way, we can formulated the Memory Layer as
K
K
T TK)a())) I + exp( 2 [
y
p(z)) [Ta]a(a))
(1I)
The left part of Figure 3 illustrates the schematic the Memory Layer. For a sequence composed of s dddimensional tokens. and the dimensionality of output embeddings is h, the computation complexity of a Memory Layer in Eq. 11 is ((s(T + 1)) s O sdh This one order of magnitude smaller than the cunne-coneected layer which is O(sdh) when 10.
According to Eq. [I], we can compute the derivative of the loss function with respect to both the hash tables and the input vector follows:
aL
p(zx)
OL by
if h(zz) = i 10. ,7 -1,,
T
0,
if hf((zz) ti,
(12)
0L
OL
Op(zk)
concat( [... TR]
for k in range(1, K + 1)]).
πx
Oy
Ozk
2.4 Architecture of MemoryFormer
follow the generic design paradigm of the standard transformer architecture 25 using N stacked blocks to build the MemoryFormer. The right part of Figure 3 depicts one building block.
Multi-Head Attention Given an input sequence X = (x1, X2, Xx) T E Raxdd _ Norm(-) layer
first normalizes the input. There Memory Layers transform the normalized into Q, K, V E Raxd respectively. The tokens in Q. K, V are then evenly split into multiple sub-vectors for multi-head purpose. The calculation of multi-head attention remains untouched as in 225 Therefore, any other efficiet self-attention techniques such Flash Attention [10), Linear Attention [16] and KV-Cache can be seamlessly incorporated into MemoryFormer to further increase the forward-time efficiency.
Feed-Forward Network used in standard transformers. The Memory Block is composed of 2 consecutive Memory Layers, each of which is preceded by a Norm(.) layer. The norm layer vital by setting zero-mean distribution for the input embedding before the hashing operation, and thus the sign function in Eq. (3)
Memory Block In MemoryFormer we use the Memory Block to replace

---

can generate -1 and +1 evenly. Therefore, the output of Eq. 4 can have an uniform distribution so that every bucket in the hash table will be retrieved with equiprobability.
Another detail about the Memory Block that we omit the intermediate activation function (e.g. ReL.U, GELU). The hashing operation is a non-linear operation itself. An extra nonlinear function is thus redundant. We have verified through experiments that discarding the non-linear function between the two Memory Layers has no effect on the performance.
In a traditional FFN module, in order to increase the model capacity and performance, the dimen- sionality of the output token embeddings of the first FC layer is expanded by 4 times, and then restored to hidden size d by the second FC layer. To keep aligned with this design pattern, we set the output dimensionality of the first Memory Layer to be (T + 2) K. Remeber that d TK, that is, the size of each hash table is T E TR2 x(r-2)-K in the first Memory Layer of the Memory Block. Therefore, the bit width of K sub-vectors 2μ in the second Memory Layer is bits larger than that of the sub-vectors in the first layer. The size of hash tables in the second layer is T E R2( xd which leads to a capacity times larger than the first layer while restores the dimensionality of the output embeddings back to d.
Computational Complexity. So far, the above computing process of one MemoryFormer block is formulated as follows:
X NXrm(X).
Q = MemoryLayerg (X), K MemoryLaayyy (X),
MemorrLayery (X),
(13)
Z X + MultiHeadAAtteniion( K, V),
Y = Z + MemoryLaycr (ooommMMmrryLayrr (Norm(X))).
The amount of peating-ppii computation of a standard transformer block is 2s 12sdf, while the amount of computation of a MemoryFormer block is only about 2sdd sd 22 + 6Ksd. The computations originating from FC layers in standard transformer are eliminated by an order of magni- tude. The absolute majority of the computational workload now comes from the MultiHeadAttentinn
3 Experiment
In this section, we conduct thorough experiments on multiple NLP benchmarks to validate the efficiency and effectiveness of the MemoryFormer across different scales. We also compare our model with existing efficient transformer methods.
Given the fact that most large language models only open-source the checkpoint without providing detailed training information, reproducing them is unachievable. Therefore, we employ Pythia []) well-developed LLM training framework with completely available dataset and detailed model hyper-parameters, to implement our method. As for training data, the Pile [12] dataset contains 825 GiB corpus with 22 diverse high-quality subsets, which either pre-exists is constructed from professional and academic fields. We use exactly the same optimizer, scheduler and other hyper-parameters following the setting Pythia to conduct fair comparisons.
We choose six widely-used evaluation task for our approach: PIQA [4, WinoGrande 23 WSC [24],
ARC-E, ARC-C 9. and LogiQA[T9], These tasks range from knowledge to reasoning, forming comprehensive benchmark for evaluating the all-round capability of large language models.
3.1 Evaluation Across Different Scales
We choose P0th Pythia 160M, Pythia-410M as the baseline models upon which we build our MemoryFormers, Specifically, MemoryFormer-ting has the same hidden size and number of layers as Phythia-70M. MemoryFormer-smaal has the same hidden size and number of layers as Phythia- 16M and MemoryFormer-bas has the same hidden size and number of layers as Phythia-410M As for the hyper-parameter of Memory Layer, fix the value of to be 8, while the number of hash tables K is 64, 96 and 128 respectively for MemoryFormer-ting -small and base model. Notably, considering the sparsity of gradients of the hash tables, set the learning rate to be 3 times of the baseline learning rate used by the corresponding Pythia model. This is ablated in Sec. Sec. It's worth noting that, the only one fully-connected layer in the MemoryFormer is the classifier head.
Tab. I reports both the computational complexity and the evaluation results of our models compared to bascline. The FLOPs is calculated for one transformer block with the input sequence length of

---

Table 1: Zero-shot evaluation results on public NLP benchmarks. We use "MF" as the abbreviation for MemoryFormer. "Attn." refers to the computation of o (QK V Inference FLOPs are measured for one block with sequence length of 2048.
Model
Pythia-70M
MF-tiny
Pythia- 160i
MF-small
Pythia-410M
MF-base
Layers
6
6
12
12
24 1024
24 1024
Hidden Size
512
512
768
768
FLOPs w/o Attn.
6,4 G
0.4 G
14.5 G
1.0 G
25.8 G
1.6 G
Total FLOPs
10.7 G
4.7
20.9 G
7.4
34.4
10.2 G
PIQA
0.585
0.602
0.618
0.642
0.675
0.698
WinoGrande
0.511
0.522
0.497
0,523
0.534
0.546
WSC
0.365
0.375
0.365
0.394
0.471
0.385
ARC-E
0.380
0.437
0,440
0,461
0.517
0.585
ARC-C
0.177
0.228
0.201
0.247
0.202
0.259
LogiQA
0.232
0.260
0.210
0.272
0.209
0.272
Avg.
0.375
0.404
0.389
0.423
0.435
0.458
Table 2: Comparison different efficient transformer methods based on Pythia-410M. Inference FLOPs measured for one block with sequence length of 2048.
Model
Pythia-410M
Linformer
cosFormer
Performer
MemoryFormer--ase
FL.PPs
34.4
26.1G
30.0 G
26.7
10.2 G
PIQA
0.675
0.527
0.522
0.643
0.698
WinoGrande
0.534
0.511
0.506
0.496
0.546
WSC
0.471
0.635
0.605
0.433
0.385
ARC-E
0.517
0.265
0.267
0.470
0.585
ARC-C
0.202
0.244
0.263
0.231
0.259
LogiQA
0.209
0.207
0.264
0.236
0.272
Avg.
0.435
0.398
0.405
0.418
0.458
2048. The experiment results show that MemoryFormer has the minimum computation except for the necessary computation of self-attention. Across all three different model sizes, we achieve better average accuracy on the benchmark than the Pythia baseline. This suggests that the proposed method able to greatly reduce the computational complexity without compromising the performance.
3.2 Comparison with Efficient Transformers
We also compare our models with existing efficient transformer methods to show the superiority of MemoryFormer in both performance and efficiency. We choose Pythia-410M the baseline, and replace the multi-head attention module of Pythia with the ones proposed by Linformer [26), Cosformer [22), and Performer { respectively. Tab. (abb22oo demonstratse the experiment results the benchmark and the inference FLOPs of each model. We measure the FLOPs using one transformer block with the sequence length of 2048. As shown in Tab, 2 these efficient attention methods can obtain FLOPs-reduction but with considerable performance degradation, while MemoryFormer sig- nificantly eliminates the computations and gains better performance. On the other hand, even though existing efficient transformer methods can reduce the computation cost of the self-attention operation, yet can observe from Tab. that the majority of the computation originates from the fully- connected layers in both the MHA and FFN module as discussed previously. The linear projection operation accounts for the biggest part of the workload when the sequence length extremely large in most of the practical scenarios. Utilizing Memory Layer in the embedding space to replace FC layers does provide new solution minimize the FLOPs of LLMs.
3.3 Ablation Study
Tradeoff between and K. In Eq. 6 the explosive growth of memory usage of the hash tables. We study the model performance with different (r, K) combination by controlling the model hidden size to be the same. We report the
we split an input embedding in to K sub-vectors to avoid

---

Table 3: Ablation study on different T and K. Memory Size refer to the storage space required by the Memory Layer Q.
Table 4: Val. PPL at 8000
training steps with various LR.
d
T
K
Val. PPI.
FLOPss
Memory Size
LR
Val. PPL
le-3
19.86
512
4
128
19.01
0.14 G
2.1 MB
2e-3
19.07
512
8
64
18.82
0.07 G
16.8 MB
3e-3
18.82
510
10
51
18.67
0.06 G
53.5 MB
4e-3
18.84
Table 5: Different expanding bits of Memory Block. #Expanding Bit= 7' T denotes the number of
extra bit of 22 after expansion. Memory Size denotes the storage space required by Memory Block.
#Expanding Bit
r'
Val. PPL +
Size of Hash Tables
Memory Size
0
19.89
TMI E R256x112
TM R 2566x122
33.6 MB
1
19.26
T
E R256x576
Ty R 51225112
52.4 MB
2
10
18.82
TMI E R256x610
T T
R 1024xx612
88.1 MB
3 11
18.54
The E R256x701
220455522
157.3 MB
experiment results in Tab. 3 along with the computation in FLOPs of the corresponding Memory Layer with sequence length of 2048. As shown in Tab. B as the bit width of z increases, the model performance continues to grow due to the exponentially enlarged capacity of hash tables. However, the memory consumption of the Memory Layer also increases drastically and its computational complexity soon reaches the lower bound. To trade off between the efficiency and memory usage, we conjecture that 7 8 is good option for MemoryFormer.
Larger learning rate. From Eq. (12) we can observe that the gradients of the hash table are sparse during backward propagation. Some buckets in the hash table might not get updated in one training step. We conjecture that larger learning rate will help remedy the back-of-gradients situation. We train MemoryFormer-tin for 8000 steps with different LR and report the PPL of validation set in Tab. 4
We use initial LR=1e-3 as the baseline learning rate following the settings of Pythia-70M. The best performance is achieved with the learning rate times of the baseline.
Expanding bit in the Memory Block. As mentioned in Sec. 2.4 in order to increase the model capacity, we enlarge the dimensionality of the output embedding of the first layer of the Memory Block, which consequently expands the bit-width of the the sub-vector Memory Block. We use MemoryFormer-ting with hidden size d = 512, bit-width 7 8 and number of hash tables K- = 64 as the baseline model, and report the perplexity results of using different number of expanding bits after 8000 training steps in Tab. 5
in the second layer of the
We use T to denote the hash tables of the 1st layer of the Memory Block, T to denote the hash tables of the 2nd layer, and use 7' to denote the bit-width of 2₁ in the second layer after expansion. As shown in Tab. 5 the validation PPL continuously decreases as 7' gets larger. However, the memory space consumed by the hash tables keeps increasing exponentially. Therefore, we choose 2 as the number of expanding bit in the Memory Block for the trade-off between the space complexity and performance.
Removing non- non-lincarity in the Memory Block. We also mentioned in Sec. 2.4 that all activation functions are discarded in the MemoryFormer. We trained a MemoryFormer-inn where we insert a GeLU layer between the two consecutive Memory Layers of the Memory Block for each building block, and report the testing scores on multiple tasks. As shown in Tab, 6 adding an extra GeLU into the Memory Block leads to almost identical result to the baseline MemoryFormer.
on PILE dataset,
3.4 Visualization
Distribution of Hash Bucket. In Memory Layer, we do not want most of the sub-vector z hashed to a few popular buckets while the rest of buckets in the table are rarely selected. We expect 2u to be hashed to all the buckets with the same probability. so, the output space of the Memory Layer would be diversified enough to enlarge the capacity of MemoryFormer. We use 2048 sequences of 1024 token length to visualize the distribution of the frequency that each bucket retrieved within hash table in Figure 4 Specially, we choose the first table T1 and the last table Toi of the Q.
8

---

Table 6: Ablation study on whether to use the non-linearity in the Memory Block.
Model
PIQA
WinoGrunde
WSC
ARC-E
ARC-C
LogiQA
Avg.
FemoryFormer-tiny MemoryFormer- tiny + GeLU
0.602
0.522 0.522
0.375
0.437
0.228
0.26
0.404
0.595
0.375
0,441
0.211
0.265
0.402
Lot tablc offonnrry Layer
Lct tabll of Romcry Layer
3st table of Momoary LayerV
bot of Nenory Blck Layeri
isttabre of Humory Book Liye
15800
2000
1:700
17500
Froqurecy
3230
100
208
last
130
200
3100
150
Figure 4: The frequency at which each bucket in the hash table is retrieved.
K. V projection layers and the two layers in the FFN module from the first building block of the MemoryFormer-ting model. As shown in Figure the number of times each bucket is hashed to by is generally uniform.
4 Related Works
(LSHH) a special kind of hash function which is designed to maximize the collision probability for similar input items. It's wildly adopted in deep learning-based applications. Large scale image retrieval system [18 30) uses LSH to locate the similar images. Reformer[17] and YOSO (28) both use LSH algorithm to reduce the memory consumption and improve the computational efficiency of the self-attention module. LookupFEN [29] adopt the LSH to accelerate the inference speed of the
Locality Sensitive Hashing. Locality sensitive hashing 11
Network. SLIDE 7 and MONGOOSE (6] improve the converging speed of neural network training process with the help of locality sensitive hashing algorithms.
Efficient Transformers. Minimizing the computational complexity of the transformer model has always been a center task for the deep learning community. Many previous works are dedicated to reducing the complexity of multi-head attention module to sub-quadratic such as CosFormer [22], PerFormer (8), LinFormer [26] and so on. [14] uses sliding window constrain the attention map within a local range. Besides, there are some researches [13 20 27 exploiting the sparsity of the intermediate activation in the FFN (MLP) module to reduce the computation. Recently, with the development of large language models, practical engineering method like FlashAttention [T0 brings substantial optimization for the self-attention mechanism.
5 Conclusion
In this work, we propose MemoryFormer, a novel transformer architecture that significantly reduces the computational complexity (FLOPs) of transformer model from a new perspective. Unlike existing methods that opt to optimize the computation of the multi-head attention operation, this work provides new solution from a new perspective. We observe that, in most scenarios the vast majority of computations originates from the fully-connected layers in the transformer model. Thus we focus on removing them. The MemoryFormerte algorithm to perform feature transformation in the embedding space, to replace all the FC layers. It achieves a similar function as the computation-heat less FL.OPs. We successfully eliminate nearly all the computations of the transformer model except for the necessary ones required by the self-attention operation. We validate the efficiency and the effectiveness of MemoryFormer via extensive experiments public NLP benchmarks.
Memory Layer, which uses locality-sensitiv hashing
matrix multiplication operation but with much

---

Acknowledgement
This work is supported by the National Key R&D Program of China under Grant No. N0.20222D01600 and the National Natural Science Foundation of China under Grant No.62276007. We gratefully acknowledge the support of MindSpore, CANN and Ascend AI Processor used in this research.
References
[1] Alexandr Andoni, Piotr Indyk, Huy L Nguyen, and Ilya Razenshteyn. Beyond locality-sensitive hashing- In Proceedings of the twenty-fifhh annual ACM-SIAM symposium on Discrete algorithms, pages 1018-1028.
SIAM, 2014.
[2] Prabavathy Balasundaram, Sriram Muralidharan, and Sruthi Bijoy. An improved content based image retrieval system using unsupervised deep neural network and locality sensitive hashing. 2021 5th International Conference on Computer, Communication and Signal Processing (ICCCSP), pages 1-7.
IEEE, 2021.
[3] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Fric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[4] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical common- sense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages
2020.
[5] Moses S Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium Theory of computing, pages 380-388, 2002.
[6] Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshu- mali Shrivastava, and Christopher Re. Mongoose: A learnable Ish framework for efficient neural network training. In International Conference on Learning Representations, 2020.
[7] Beidi Chen, Tharun Medini, James Farwell, Charlic Tai, Anshumali Shrivastava, et al. Slide: defense of smart algorithms over hardware acceleration for large-scale deep learning systems. Proceedings of Machine Learning and Systems, 2:291-306, 2020.
[8] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint ar Xiv: /4794, 2020.
[9] 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 arXiv: 1803.084575 2018.
[10] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 16344-16359, 2022.
[11] Anirban Dasgupta, Ravi Kumar, and Tamás Sarlós. Fast 17th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 1081, 2011.
hashing. In Proceedings of the
[12] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace Hc, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling, arXiv preprint arXiv: 00027, 2020.
[13] Matteo Grimaldi, Darshan C Ganji, Ivan Lazarevich, and Sudhakar Sah, Accelerating deep neural networks via semi-structured activation sparsity. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1179- 1188, 2023.
[14] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel. Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint 2.3
2023.
[15] Ke Jiang, Qichao Que, and Brian Kulis. Revisiting kernelized -ocality-sensitive sensitive hashing for improved large- scale image retrieval. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 4933 4941, 2015.
10

---

[16] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are runs: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.
[17] Nikita Kitaev, Lukasz Kaiser, and Auselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001. 04451, 2020.
[18] Brian Kulis and Kristen Grauman. Kemeliiedd lecality-sensitive hashing for scalable image search. In 2009 IEEE 12th international conference on computer vision, pages 2130 2137. IEEE, 2009.
[19] Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint urXiv:2007.08222 2020.
[20] Zichang Liu, Wang. Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang. Yuandong Tian, Christopher Re, et Deja vu: Contextual sparsity for efficient Ilms at inference time. In International Conference Machine Learning, pages 22137 22176. PMLR, 2023.
[21] NVIDIA. Nvidia dgx a100 user guide, 2023.
URL. https://doca.na nvidia. cow/dgx/ htggx//dd////// pdf
dgxa100-usee uueo-user-ggn seneegguie.ppee
[22] Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong. and Yiran Zhong, cosformer: Rethinking softmax in attention. arXiv preprint arXiv:2202.08793 087799 2022.
[23] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial
winograd schema challenge at scale. arXiv preprint 19Xivv1907.10641, 2019.
[24] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 6409)99-106, 2021.
[25] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[26] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 2020.
[27] Varun Yerram, Chong You, Srinadh Bhojanapalli, Sanjiv Kumar, Prateek Jain, Praneeth Netrapalli, et al. Hire: High recall approximate top-k estimation for efficient I1m inference, arXiv preprini 200.00040..093600 2024.
[28] Zhanpeng Zeng. Yunyang Xiong, Sathya Ravi, Shailesh Acharya. Glenn M Fung. and Vikas Singh. You only sample (almost) once: Linear cost self-attention via bernoulli sampling. In International conference on machine learning, pages 12321-12332. PMLR, 2021.
[29] Zhanpeng Zeng, Michael Davies, Pranav Pulijala, Karthikeyan Sankaralingam, and Vikas Singh. Lookupffn: making transformers compute- lite for cpu inference. In International Conference on Machine Learning, pages 40707 40718 18. PMLR, 2023.
[30] Zheng Zhang, Jianning Wang, Lei Zhu, Yadan Luo, and Guangming Lu. Deep collaborative graph hashing for discriminative image retrieval. Pattern Recognition, 139:109462, 2023.

---

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

---

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

---

Answer:
[No]
Justification:
The code and data will not be provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm
Bttie//guiide//odeSubbmmn cesssnnnooccyyl for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ //nips. /pipp/.ccppullc/guuiii httpp///iiiicuiiceoo/ooeeeaaaaia CoodeSubeeiieeeeeeeeee ssi asiinnolle onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The details are provided the Sec.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification:
The error bar is not provided.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
14

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The details provided in the Sec.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: https::/ //terrips///r/ps..... http://ss...cuuoooocaaacoooooooa
Answer: [Yes]
Justification: is conformed.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
15

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
There is risk in the paper.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
CC-BY 4.0 is included for each asset
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwppppaipptheeee throdee.wtecce...... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
16

---

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


Reviews:
[{"review_id": "QUkUGlrgI0", "reviewer": "Reviewer_omZU", "summary": "This work proposes to replace most linear layers of transformers by trainable hash-tables. The new modules---called memory layers---rely on locality sensitive hashing to obtain relevant indices within several hash tables, and returns a linear combination of the associated vectors. To overcome the non-differentiability of the hashing operation, the weights of the linear combination are computed from the inner products between each input vector and its hashed representation. They use memory layers to replace key, query, and value projections, as well as to replace the down and up-projection matrices of the feedforward blocks. Playing on the coarseness of the hashing, this new module can reduce the FLOPs required when compared to a traditional matrix multiplication. They train modified Pythia models and show their approach is outperforming baseline models on multiple reasoning tasks, while using significantly less FLOPs.", "questions": "See above.", "limitations": "See weaknesses. I feel the limitations could have been better discussed in the paper, e.g. time complexity.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I find the paper well motivated. Tackling the often dominating cost of the FFW operations is an important research direction. The idea developed in this work is novel and the results are surprisingly good. Improving upon the transformer architecture is not an easy feat and the results suggests that the memoryFormer is better on reasoning tasks while requiring less FLOPs. I find it interesting that the non-linear nature of the hashing operation allows to omit the use of activation functions. \n\nAn interesting work overall, but I believe some results/details are missing (see weaknesses).", "weaknesses": "To fully evaluate the impact of this work, I am missing some informations:\n- What is, in more details, the experimental setup used to trained Pythia models? How many steps are used during training? How are the loss curves for both the Pythia baseline and the memoryFormer? Which hardware was used? \n- You should do a comparison in terms of number of parameters. The memoryFormer likely uses an order of magnitude more trainable parameters compared to Pythia models. This makes the comparison between e.g. Pythia-70M and MF-tiny questionable. What scores would you get using a Pythia model with a similar number of parameters?\n- Discussion on time complexity are missing from the analysis. I am assuming that the memoryFormer is slower. Speed depends on hardware and I understand that it seems unfair to compare your approach to very optimized matrix-multiplication kernels, but this is still an important limitation to discuss. How many iterations per second during training for a memoryFormer and a Pythia model with the same number of parameters? What is the inference speed? This work seems clearly oriented at proposing a cheaper transformer, and successfully reduces the FLOPs required. Knowing if this translates into speed gains today---or giving ideas on how easy it would be to leverage this gain tomorrow---seems important to evaluate the impact of this work. Given a time budget and a specific machine (cpu or gpu), which size of models would fit the time budget? For those model sizes, would a memoryFormer provide better scores than a baseline Pythia model?\n- At the small models scale, I am not sure how trustworthy the scores on reasoning tasks are. What would be the accuracies obtained when answering at random? For instance, WinoGrande is a binary task, and scores are often close to 50%. What are the perplexities reached by the different models?\n\nOverall, you propose an interesting architecture, but I am not entirely convinced by the evaluation. A more faithful account of the implications of using memory blocks on time complexity and on the number of parameters would help.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **1.1. What is the experimental setup used to trained Pythia models?  How many steps are used during training?**\nAll models are trained for 143000 steps. The total batch size is 1024 samples and the sequence length of each sample is 2048. We use exactly the same experimental setups, such as Adam optimizer, cosine lr decay scheduler, rotary positional embeddings, tokenizer, etc., following the routine done by original Pythia paper. The only difference is that, the learning rate for MemoryFormer is set to be 3 times of the LR used by its corresponding Pythia baseline, which is mentioned in Line.230. \nWe will refine the information regarding experimental setup in Sec.3.\n\n### **1.2. How are the loss curves for both the Pythia baseline and the memoryFormer?**\nWe plot the loss curves of training stage for Pythia-410M and MemoryFormer-base model in the rebuttal PDF file. We'll add a more comprehensive figure to our paper in final version.\n### **1.3. Which hardware was used?**\nWe use NVIDIA-A100 server with 8 GPUs for training and testing.\n### **2. The comparison in terms of number of parameters. What scores would you get using a Pythia model with a similar number of parameters?**\nThe aim of MemoryFormer is to reduce the FLOPs of transformer, which is achieved by increasing the number of parameters in exchange for much less computational complexity.\nIn the table below, we report the number of parameters in the row of \"Model Size\". We also show the FLOPs and scores for Pythia models of comparable model sizes.\n\n||Pythia-70M|MF-tiny|Pythia-160M|MF-small|Pythia-410M|MF-base|pythia-1B|pythia-2.8B|pythia-6.9B|\n|-|-|-|-|-|-|-|-|-|-|\n|Layers/Hidden_dim|6/512|6/512|12/768|12/768|24/1024|24/1024|16/2048|32/2560|32/4096|\n|__Model Size__|70M|460M|160M|1.9B|410M|6.7B|1B|2.8B|6.9B|\n|FLOPs w/o Attn.|6.4G|0.4G|14.5G|1.0G|25.8G|1.6G|103.1G|161.1G|412.3G|\n|Total FLOPs|10.7G|4.7G|20.9G|7.4G|34.4G|10.2G|120.3G|182.5G|446.7G|\n|PIQA|0.585|0.602|0.618|0.642|0.675|0.698|0.700|0.741|0.760|\n|WinoGrande|0.511|0.522|0.497|0.523|0.534|0.546|0.529|0.582|0.631|\n|WSC|0.365|0.375|0.365|0.394|0.471|0.385|0.365|0.385|0.442|\n|ARC-E|0.380|0.437|0.440|0.461|0.517|0.585|0.585|0.635|0.686|\n|ARC-C|0.177|0.228|0.201|0.247|0.202|0.259|0.245|0.301|0.331|\n|LogiQA|0.232|0.260|0.210|0.272|0.209|0.272|0.212|0.214|0.215|\n|Avg.|0.375|0.404|0.389|0.423|0.435|0.458|0.439|0.476|0.511|\n\n**It is worth noting that**, in a traditional transformer that is based on fully-connected layers, all parameters are involved in the matrix-multiplication. \nHowever, the Memory Layer in the MemoryFormer is a sparsely activated module. Only a small fraction ($\\frac{1}{2^\\tau}$, usually $\\frac{1}{256}$) of vectors stored in the hash tables are retrieved during inference while most of the vectors are not used. Therefore, the model size of a MemoryFormer is not the real number of parameters involved in the forward pass. \n### **3. Discussion on time complexity.**\nThe speed of the proposed method is indeed an important aspect. We report the average training speed for models of similar sizes:\n\n||MF-tiny|Pythia-410M|MF-base|pythia-6.9B|\n|-|-|-|-|-|\n|Model Size|460M|410M|6.7B|6.9B|\n|Training speed (sec/iter)|1.98|1.96|6.60|6.57|\n\nThe inference latency is related to specific hardwares. We show the latency tested on Intel Xeon 8377C CPU and A100 GPU using custom kernel in the table below. We take the average time of feeding a sequence of 2048 tokens into the model for 100 times.\n\n||MF-tiny|Pythia-410M|MF-base|pythia-6.9B|\n|-|-|-|-|-|\n|Model Size|460M|410M|6.7B|6.9B|\n|CPU Inference (ms)|200.4|1586.3|1581.6|9788.8|\n|GPU Inference (ms)|12.1|33.5|64.8|163.2|\n\nAccording to the table, the MemoryFormer achieves lower inference latency on CPU and GPU than pythia model of similar number of parameters. This result could be bettered via optimizing the hash kernel. Currently, MemoryFormer's latency on CPU/GPU platform is primarily bounded by memory bandwidth. GPU hardware is a computationally intensive chip, which gives rise to MatMul-based models with large FLOPs. Yet there are other types of computing chips being broadly developed. For example, Processing-in-Memory (PIM) architecture[1] is a frontier ASIC. Such hardwares[2,3,4,5] incorporate the arithmetic units into the design of RAM circuit to provides extremely large bandwidth with massive memory space, but have relatively lower compute capability. The inference paradigm of the proposed MemoryFormer is different from traditional transformer, which indicates that LLMs can be deployed on low-compute-capability platforms while achieving ideal performance and latency. \n### **4.1 Accuracies obtained when answering at random?** \nWe test the accuracies of randomly initialized models:\n||PIQA|WinoGrande|WSC|ARC-E|ARC-C|LogiQA|Avg.| \n|-|-|-|-|-|-|-|-|\n|Random Pythia|0.509|0.489|0.516|0.251|0.255|0.249|0.378|\n|Random MF|0.501|0.488|0.519|0.263|0.235|0.277|0.381|\n\nPIQA, WinoGrande, WSC are binary tasks, so their results are around 50%. The answers of ARC-E, ARC-C and LogiQA include 4 choices, so their results are around 25%. \nThe average score is a reasonable indicator for the performance of the model.\n### **4.2 What are the perplexities reached by the different models?** \nPerplexities reached by the different models:\n||Pythia-70M|MF-tiny|Pythia-160M|MF-small|Pythia-410M|MF-base|\n|-|-|-|-|-|-|-|\n|ppl|24.04|19.88|12.68|11.02|9.12|8.94|\n\nMemoryFormer has lower ppl than corresponding baseline.\n\nReference:\\\n[1] A survey on processing-in-memory techniques: Advances and challenges. Elsevier 2023.\\\n[2] Computing Utilization Enhancement for Chiplet-based Homogeneous Processing-in-Memory Deep Learning Processors. ACM GLSVLSI 2021.\\\n[3] Dojo: The microarchitecture of tesla’s exa-scale computer. IEEE HCS 2022.\\\n[4] The Groq Software-defined Scale-out Tensor Streaming Multiprocessor: From chips-to-systems architectural overview. IEEE HCS 2022.\\\n[5] Chiplet Cloud: Building AI Supercomputers for Serving Large Generative Language Models. arxiv/2307.02666.", "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": "W6ZH8bj5Af", "reviewer": "Reviewer_vj2X", "summary": "This paper introduces a novel neural network call MemoryFormer. This is a modified version for transformer that eliminates the dense layer (FC layer). The proposed MemoryFormer tries to address the problem of high computational complexity of decoder-style generative models. Concretely, the author proposes the Memory Layer which uses locality sensitive hashing (LSH) to achieve the embedding projection which is done by the fully-connected layer in normal transformer. This way, the MemoryFormer eliminates the computational cost of linear projections and keeps only the computational cost of multi-head attention. The author conducts sufficient experiments at different model scale on different NLP benchmarks. The experiment results show that MemoryFormer has better performance than other transformer models with the same #layers and hidden dimension while obviously reduce the FLOPs. The method in this paper shows a different but possible way of building future large language models.", "questions": "1 Which model size is being used in the ablation study in Section 3.3? Is it MemoryFormer-tiny, or -small, or base? This paper seems to forget to mention about this information.\n2 This paper reports the FLOPs of different models in Tab.1 and 2. What is the inference latency achieved by MemoryFormer?", "limitations": "Refer to the weakness.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1 This paper proposes an interesting way to reduce the computational complexity of transformer neural network. The proposed Memory Layer uses LSH algorithm as a replacement for the linear projections of linear layer. It’s novel to use the memory space to store feature vectors instead of compute them on the fly.\n2 The experiment results demonstrate that the proposed model has better performance than baseline across different sizes, showcasing the scalability of the proposed model.\n3 The alternative plan for FC layer in transformer is underexplored as it is the most basic building component in neural networks.", "weaknesses": "1 This paper doesn’t have any data regarding the inference latency of the proposed MemoryFormer neither on CPU or GPU. The purpose of this paper is reducing FLOPs of transformer, I think such experiment is important.\n2 The baseline model Pythia[1] uses 8 benchmarks to evaluate Pythia models while this paper uses 6 benchmarks. This is not a big deal but I hope the author can report the results on the other two.\n3 The formulation of Eq.10 might be a little bit confusing if not reading carefully and could be optimized for readability.\n\n[1] Pythia: A suite for analyzing large language models across training and scaling. ICML 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **1. Data regarding the inference latency.**\nWe measure the inference latency for different models on Intel Xeon 8377C CPU and A100 GPU using custom kernel, and show the results in the table below.\n\n||Pythia-70M|MF-tiny|Pythia-160M|MF-small|Pythia-410M|MF-base|\n|-|-|-|-|-|-|-|\n|Layers|6|6|12|12|24|24|\n|Hidden Size|512|512|768|768|1024|1024|\n|FLOPs w/o Attn.|6.4G|0.4G|14.5G|1.0G|25.8G|1.6G|\n|Total FLOPs|10.7G|4.7G|20.9G|7.4G|34.4G|10.2G|\n|**CPU latency (ms)**|250.7|200.4|597.6|520.8|1586.3|1581.6|\n|**GPU latency (ms)**|7.2|12.1|15.6|27.6|33.5|64.8|\n\nAs shown in the table, the inference of MemoryFormer is faster than Pythia baseline on CPU while MemoryFormer has much less FLOPs than Pythia baseline.\n\n### **2. Evaluation on more tasks.**\nWe further evaluate the models on Lambada and SciQ datasets, and report the complete results as follows.\n||Pythia-70M|MF-tiny|Pythia-160M|MF-small|Pythia-410M|MF-base|\n|-|-|-|-|-|-|-|\n|Lambada|0.23|0.326|0.398|0.408|0.529|0.524|\n|PIQA|0.585|0.602|0.618|0.649|0.675|0.698|\n|WinoGrande|0.511|0.522|0.497|0.511|0.534|0.546|\n|WSC|0.365|0.375|0.365|0.365|0.471|0.385|\n|ARC-E|0.38|0.437|0.44|0.472|0.517|0.585|\n|ARC-C|0.177|0.228|0.201|0.238|0.202|0.259|\n|SciQ|0.654|0.692|0.774|0.788|0.824|0.811|\n|LogiQA|0.232|0.260|0.210|0.295|0.209|0.272|\n|Avg.|0.392|0.431|0.438|0.466|0.495|0.510|\n\nAfter evalutaed on the remaining two tasks, the MemoryFormer model still achieves higher average scores than its Pythia baseline which has the same hidden size and number of layers.\n\n### **3. Improving readability for Eq.10.**\nThanks for the suggestion. We are considering separating the third line of Eq.10 from the equation, to make the transition of $p(z_k)$ more smooth. We'll revise it in the final version.\n\n### **4. Which model size is used in the ablation study?**\nWe feel sorry we forgot to mention such important information. We use MemoryFormer-tiny to conduct the experiments for ablation study throughout Sec.3.3.", "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": "Qtxtw70eLW", "reviewer": "Reviewer_vbej", "summary": "A FLOPs-reduction strategy designed for transformer is proposed by this paper. The author introduces MemoryFormer, a transformer model that is built using the Memory Layer instead of fully-connected layer. According to the paper, the author claims that the MemoryFormer has the minimum computation because the Memory Layer uses a locality-sensitive hashing algorithm with low computational complexity to compute feature projection instead of using matrix multiplication performed by Fully-Connected layer. The author also designs a differentiable mechanism for the hash tables to be updated, therefore making the MemoryFormer end-to-end trainable like baseline transformer. Besides, the author keeps the necessary computation of self-attention operation untouched, which makes this work orthogonal to other existing efficient transformer designs. Extensive experiments show that the proposed MemoryFormer performs better than other transformer baselines on 6 tasks when built with the same hyper-parameters, such as number of layers and hidden sizes.", "questions": "Since current LLMs require scaling up in model size for more intelligence, I wonder if the author experimented on a larger size of the proposed model? Scalability is crucial.", "limitations": "This paper proposed an efficient operation for transformer model, which is experimented via generative decoder model. I’m concerned about the scalability of the proposed method.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "• The heavy computations of decoder model hinders the further development for LLMs. This paper aims at solving this main issue that the LLMs are currently facing. Exploiting storage to reduce inference computation is a fancy idea.\n• The proposed hashing-based Memory Layer requires much less computation than Fully-Connected layer in forward pass. The back-propagation mechanism of the hash tables proposed in Sec.2.2 sounds reasonable.\n• The experiment section are comprehensive. The results show that MemoryFormer has better performance while consuming much fewer FLOPs.\n• The content of paper is well-written and well-organized. The equations are clear and easy to understand.", "weaknesses": "• The proposed method could be further tested with other attention method, for example, incorporating linear attention mechanism into MemoryFormer. However, such experiment is missing.\n• The author didn't report the parameter size of the proposed Memory Layer against the corresponding Fully-Connected layer, and did not report the number of parameters of the MemoryFormer model.\n• Lack of the experiment for the activation function of Memory Block.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **1. Combining MemoryFormer with other efficient attention method.**\nWe incorporate the efficient attention modules proposed by Linformer and cosFormer with MemoryFormer-tiny, and train each model for 8000 steps on PILE dataset.\n|Attn. Type|baseline MHA|Linformer|cosFormer|\n|-|-|-|-|\n|**Val ppl.**|18.82|18.91|18.86|\n\nThe loss curves of 3 models before 8000 steps are similar. This confirms the fact that the proposed MemoryFormer is orthogonal to the type of attention module and can works well with them.\nWhile the validation perplexity of the baseline attention model is indeed lower after 8000 training steps, the impact of linear attention on the MemoryFormer should be reflected by the testing scores on multiple tasks after full training. Due to the limited time for rebuttal, we're unable to fully train these models. We'll add this experiment to the final version.\n\n\n\n### **2. The parameter size of the proposed Memory Layer and MemoryFormer against fully-connected layer baseline.**\n\nWe report the parameter size for Memory Layer and the corresponding Fully-Connected layer of the same hidden dimension. We assume the output dimension is the same as input. We also report the FLOPs required by the two kinds of layers to process a sequence of 2048 tokens.\n\n|Hidden dim|512|768|1024|\n|-|-|-|-|\n|Parameter size of Memory Layer|8.39M|18.87M|33.55M|\n|Parameter size of FC Layer|0.26M|0.59M|1.05M|\n|FLOPs of Memory Layer|0.07G|0.15G|0.27G|\n|FLOPs of FC Layer|0.54G|1.21G|2.15G|\n\nThe number of parameters of the MemoryFormer models are reported in the row \"Model Size\" of the table below:\n||Pythia-70M|MF-tiny|Pythia-160M|MF-small|Pythia-410M|MF-base|pythia-1B|pythia-2.8B|pythia-6.9B|\n|-|-|-|-|-|-|-|-|-|-|\n|Layers/Hidden_dim|6/512|6/512|12/768|12/768|24/1024|24/1024|16/2048|32/2560|32/4096|\n|__Model Size__|70M|460M|160M|1.9B|410M|6.7B|1B|2.8B|6.9B|\n|FLOPs w/o Attn.|6.4G|0.4G|14.5G|1.0G|25.8G|1.6G|103.1G|161.1G|412.3G|\n|Total FLOPs|10.7G|4.7G|20.9G|7.4G|34.4G|10.2G|120.3G|182.5G|446.7G|\n\n**It is worth noting that**, in the traditional transformer that is based on fully-connected layers, all parameters are involved in the matrix-multiplication. \nHowever, the Memory Layer in the MemoryFormer is a sparsely activated module. Only a small fraction ($\\frac{1}{2^\\tau}$, usually $\\frac{1}{256}$) of vectors stored in the hash tables are retrieved for computation during inference while most of the vectors are not used. Therefore, the model size of a MemoryFormer is not the real number of parameters involved in the forward pass. The proposed FLOPs-reduction method is achieved via using more memory space.\n\n\n### **3. Experiment for the activation function.**\nFollowing the experiment setup of training the MemoryFormer-tiny model, we insert a GeLU layer between the two consecutive Memory Layers of the Memory Block for each building block of the MemoryFormer-tiny. We report the testing scores on multiple tasks after training this model on PILE dataset.\n\n|Model|PIQA|WinoGrande|WSC|ARC-E|ARC-C|LogiQA|Avg.|\n|-|-|-|-|-|-|-|-|\n|MemoryFormer-tiny|0.602|0.522|0.375|0.437|0.228|0.26|0.404|\n|MemoryFormer-tiny + GeLU|0.595|0.522|0.375|0.441|0.211|0.265|0.402|\n\nAs shown in the table, adding an extra GeLU into the Memory Block leads to almost identical result to (or even slightly worse than) the baseline MemoryFormer. We will add this experiment into the ablation study.\n\n### **4. Scalability of the MemoryFormer.**\nThe scalability is important for large models to develop emergent abilities. A series of experiments on different model sizes (-tiny/-small/-base) proved the basic scalability of MemoryFormer. As for larger-sized model, We are currently training a MemoryFormer-large based on Pythia-1.4B baseline, which has 24 layers and the hidden dimension of 2048. The trend of the training loss curve is decreasing normally, which further verifies the scalability.", "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": "qoImygWW8M", "reviewer": "Reviewer_eumW", "summary": "In this paper, the authors present a new transformer model named MemoryFormer, which utilizes locality-sensitive hashing to replace the matrix multiplication operation of the fully-connected layer. This paper can minimize the computational complexity of transformer by removing all the computations except for self-attention. The idea is to use hash tables to store pre-computed features in the memory and retrieve them during inference. The authors evaluate the MemoryFormer on six tasks and compare with other efficient transformer models to show the superiority in saving computations.", "questions": "1. What is the training cost of MemoryFormer? Is it fast to train such models?\n2. Can MemoryFormer be scaled up to reach billions of parameters like Llama-7B?", "limitations": "See the weaknesses and questions.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper is well-written and easy to follow. Existing works that try to make the transformer model computationally efficient mainly focus on the self-attention module and FFN module, while this paper tries to address this issue by focusing on the most fundamental part of transformer, the fully-connected layer. This makes the paper stand out for its novelty. The authors use locality-sensitive hashing to simulate the linear projection and use gradient descent via backpropagation to learn the hash tables where the embeddings are stored. The leverage of the memory resources rather than massive parallel matrix multiplication makes the presented model potentially feasible for CPU deployment and inference.", "weaknesses": "1. Although the authors experiment on three different model sizes, but they do not provide the exact number of parameters of these models. It would be better if the authors could provide the parameter number.\n2. Missing the ablation study about the non-linearity in the FFN module.\n3. It would be better to provide some discussion about the training cost given that frequently hashing and retrieving would cost much VRAM I/O of GPU for training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **1. Data regarding the number of parameters.**\nWe report the number of parameters for each model in the row \"Model Size\" in table below.\n||Pythia-70M|MF-tiny|Pythia-160M|MF-small|Pythia-410M|MF-base|pythia-1B|pythia-2.8B|pythia-6.9B|\n|-|-|-|-|-|-|-|-|-|-|\n|Layers/Hidden_dim|6/512|6/512|12/768|12/768|24/1024|24/1024|16/2048|32/2560|32/4096|\n|__Model Size__|70M|460M|160M|1.9B|410M|6.7B|1B|2.8B|6.9B|\n|FLOPs w/o Attn.|6.4G|0.4G|14.5G|1.0G|25.8G|1.6G|103.1G|161.1G|412.3G|\n|Total FLOPs|10.7G|4.7G|20.9G|7.4G|34.4G|10.2G|120.3G|182.5G|446.7G|\n\n**It is worth noting that**, in a traditional transformer that is based on fully-connected layers, all parameters are involved in the matrix-multiplication. \nHowever, the Memory Layer in the MemoryFormer is a sparsely activated module. Only a small fraction ($\\frac{1}{2^\\tau}$， usually $\\frac{1}{256}$) of parameters stored in the hash tables are retrieved for computation during inference while most of the parameters are not used. Therefore, the model size of a MemoryFormer is not the real number of parameters involved in the forward pass. The proposed FLOPs-reduction method is achieved via using more memory space.\n\n### **2. Ablation study about the non-linearity in the Memory Block.**\nFollowing the experiment setup of training the MemoryFormer-tiny model, we insert a GeLU layer between the two consecutive Memory Layers of the Memory Block for each building block of the MemoryFormer-tiny. We report the testing scores on multiple tasks after training this model on PILE dataset.\n\n|Model|PIQA|WinoGrande|WSC|ARC-E|ARC-C|LogiQA|Avg.|\n|-|-|-|-|-|-|-|-|\n|MemoryFormer-tiny|0.602|0.522|0.375|0.437|0.228|0.26|0.404|\n|MemoryFormer-tiny + GeLU|0.595|0.522|0.375|0.441|0.211|0.265|0.402|\n\nAs shown in the table, adding an extra GeLU into the Memory Block leads to almost identical result to (or even slightly worse than) the baseline MemoryFormer. We will add this experiment into the ablation study.\n\n### **3. Training cost of the MemoryFormer.**\n\nWe report the training speed for MemoryFormers:\n||MemoryFormer-tiny|MemoryFormer-small|MF-base|\n|-|-|-|-|\n|Training speed (sec/iter)|1.98|4.11|6.60|\n\nGiven the fact that GPU is a computationally intensive device with moderate memory i/o bandwidth, and we use custom CUDA kernel that is unable to fully exploit the GPU performance for hashing operation, the training throughput of MemoryFormer is still acceptable. For example, using 2 servers with 8 A100 GPUs, MemoryFormer-small finishs training for 143000 steps on PILE dataset in 6.8 days. The training speed could be faster if our custom CUDA kernel is professionally tuned and optimized.\n\n### **4. Scalability of the MemoryFormer.**\nA series of experiments on different model sizes (-tiny/-small/-base) proved the scalability of MemoryFormer. As for larger-sized model, We are currently training a MemoryFormer-large based on Pythia-1.4B baseline, which has 24 layers and the hidden dimension of 2048. The trend of the training loss curve of such a model is decreasing normally, which further verifies the scalability.", "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:
We provide the training loss curves for Pythia-410M and MemoryFormer-base in the rebuttal PDF file.

Meta-review: This paper introduces a novel transformer architecture that replaces fully-connected layers with memory layers using a locality-sensitive hashing (LSH) approach. This significantly reduces computational complexity (FLOPs) while maintaining model performance. 

Strengths:
* Novelty: The paper presents a novel method of reducing FLOPs by replacing linear projections with memory layers, offering an efficient alternative to fully-connected layers, a fundamental part of transformers. The approach is both unique and timely, given the push to make large language models (LLMs) more efficient.
* Performances: MemoryFormer outperforms baseline models like Pythia on several reasoning tasks while requiring significantly fewer FLOPs. The experiments are comprehensive across different model sizes, and the paper demonstrates better scalability than existing methods.

Weaknesses and Limitations:
* Training Complexity: Several reviewers raised concerns about the training speed and hardware dependency. The method, while reducing FLOPs, does not necessarily translate into faster training, particularly when compared to heavily optimized matrix-multiplication-based models.
* Parameter Size: While the paper reduces computational complexity, it does so at the expense of increasing the model size, raising questions about memory efficiency and scalability.
* Limited Evaluation: Some reviewers pointed out the need for additional tasks and metrics in evaluation, including a clearer comparison of parameter sizes and in-depth latency studies, which were partially addressed in the rebuttal but could have been more thorough.

Overall, the paper is technically solid, presents a high-impact novel idea, and provides compelling experimental results. However, the evaluation could benefit from deeper exploration of limitations around speed and scalability. Reviewers generally support the paper’s acceptance, but some remain cautious about the practicality of the proposed method on large-scale models. I recommend acceptance as the new idea may spark further discussions.
Decision: Accept (poster)