MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers

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.

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. 2. Missing the ablation study about the non-linearity in the FFN module. 3. 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.

Questions

1. What is the training cost of MemoryFormer? Is it fast to train such models? 2. Can MemoryFormer be scaled up to reach billions of parameters like Llama-7B?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

See the weaknesses and questions.

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.

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. • 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. • The experiment section are comprehensive. The results show that MemoryFormer has better performance while consuming much fewer FLOPs. • 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. • 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. • Lack of the experiment for the activation function of Memory Block.

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.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

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.

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.

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. 2 The experiment results demonstrate that the proposed model has better performance than baseline across different sizes, showcasing the scalability of the proposed model. 3 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. 2 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. 3 The formulation of Eq.10 might be a little bit confusing if not reading carefully and could be optimized for readability. [1] Pythia: A suite for analyzing large language models across training and scaling. ICML 2023.

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. 2 This paper reports the FLOPs of different models in Tab.1 and 2. What is the inference latency achieved by MemoryFormer?

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Refer to the weakness.

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.

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. An 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: - 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? - 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? - 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? - 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? Overall, 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.

Questions

See above.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses. I feel the limitations could have been better discussed in the paper, e.g. time complexity.