MemoryFormer : Minimize Transformer Computation by Removing Fully-Connected Layers

1. Data regarding the number of parameters.

We report the number of parameters for each model in the row "Model Size" in table below.

It is worth noting that, in a traditional transformer that is based on fully-connected layers, all parameters are involved in the matrix-multiplication. However, 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.

2. Ablation study about the non-linearity in the Memory Block.

Following 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.

As 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.

3. Training cost of the MemoryFormer.

We report the training speed for MemoryFormers:

Given 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.

4. Scalability of the MemoryFormer.

A 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.

1. Combining MemoryFormer with other efficient attention method.

We incorporate the efficient attention modules proposed by Linformer and cosFormer with MemoryFormer-tiny, and train each model for 8000 steps on PILE dataset.

The 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. While 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.

2. The parameter size of the proposed Memory Layer and MemoryFormer against fully-connected layer baseline.

We 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.

The number of parameters of the MemoryFormer models are reported in the row "Model Size" of the table below:

It is worth noting that, in the traditional transformer that is based on fully-connected layers, all parameters are involved in the matrix-multiplication. However, 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.

3. Experiment for the activation function.

Following 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.

As 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.

4. Scalability of the MemoryFormer.

The 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.

1. Data regarding the inference latency.

We 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.

As shown in the table, the inference of MemoryFormer is faster than Pythia baseline on CPU while MemoryFormer has much less FLOPs than Pythia baseline.

2. Evaluation on more tasks.

We further evaluate the models on Lambada and SciQ datasets, and report the complete results as follows.

After 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.

3. Improving readability for Eq.10.

Thanks 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.

4. Which model size is used in the ablation study?

We feel sorry we forgot to mention such important information. We use MemoryFormer-tiny to conduct the experiments for ablation study throughout Sec.3.3.

1.1. What is the experimental setup used to trained Pythia models? How many steps are used during training?

All 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. We will refine the information regarding experimental setup in Sec.3.

1.2. How are the loss curves for both the Pythia baseline and the memoryFormer?

We 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.

1.3. Which hardware was used?

We use NVIDIA-A100 server with 8 GPUs for training and testing.

2. The comparison in terms of number of parameters. What scores would you get using a Pythia model with a similar number of parameters?

The 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. In 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.

It is worth noting that, in a traditional transformer that is based on fully-connected layers, all parameters are involved in the matrix-multiplication. However, 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.

3. Discussion on time complexity.

The speed of the proposed method is indeed an important aspect. We report the average training speed for models of similar sizes:

The 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.

According 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.

4.1 Accuracies obtained when answering at random?

We test the accuracies of randomly initialized models:

PIQA, 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%. The average score is a reasonable indicator for the performance of the model.

4.2 What are the perplexities reached by the different models?

Perplexities reached by the different models:

MemoryFormer has lower ppl than corresponding baseline.

Reference:
[1] A survey on processing-in-memory techniques: Advances and challenges. Elsevier 2023.
[2] Computing Utilization Enhancement for Chiplet-based Homogeneous Processing-in-Memory Deep Learning Processors. ACM GLSVLSI 2021.
[3] Dojo: The microarchitecture of tesla’s exa-scale computer. IEEE HCS 2022.
[4] The Groq Software-defined Scale-out Tensor Streaming Multiprocessor: From chips-to-systems architectural overview. IEEE HCS 2022.
[5] Chiplet Cloud: Building AI Supercomputers for Serving Large Generative Language Models. arxiv/2307.02666.

Greetings, Anastasiia

This work differs from LookupFFN in the following three aspects:

(1) Operator design level difference. (suppose $x$ is input and $y$ is output)

Forward of LookupFFN still contains a fully-connected layer, which increases FLOPs:

$z_k=x R_k , R_k \in \mathcal{R}^{~ d \times \tau} , y = \sum_{k=1}^K [T_k] _{z_k}$.

Forward of MemoryLayer does NOT include fully-connected layer, which greatly reduces FLOPs:

$z_k$= split($x$, num_chunk = K) = $x$.view(K, x.size(-1)/K), $y=\sum_{k=1}^K [T_k] _{z_k}$.

(2) Architecture design level difference.

LookupFFN is the replacement for the whole FFN module (two fully-connected layers). There are multiple FC layers in a block (Q/K/V/output_projection/$R_k$). LookupFFN does not change the attention module.

MemoryLayer is the replacement for a single fully-connected layer. There are NO FC layer in a block. The MemoryFormer change the architecture of the attention module, where the Q/K/V projection are replaced with MemoryLayers and the output_projection is removed.

(3) Experimental setting difference.

LookupFFN is tested with encoder-only model , while MemoryFormer is tested with decoder-only model.

---

---

This work expands the great idea of ​​LookupFFN, who accelerate the inference speed by utilizing LookupFFN (FC layer + LSH + learnable tables) to reduce the FLOPs of FFN module (two FC layers).

MemoryFormer try to reduce the FLOPs of the whole transformer much further. Our practice is that we replaced all FC layers with pure (LSH + learnable tables).

---

---

I hope my reponse could address your concern.

Ning