Paper: ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models (ICLR-2024)

Reviewer comment (Novelty):
The observation that replacing activation functions like GeLU, SiLU with ReLU only marginally influences performance is not new. It is also mentioned and discussed in [1][2].

Author response:
We appreciate the reviewer's observation regarding the impact of replacing activation functions like GeLU and SiLU with ReLU. While the marginal influence of such replacements on performance has indeed been noted in previous works, our study diverges in both scope and focus:

* **Scale and Scope of Study**: Our research extends beyond the replacement of activation functions in existing models, conducting a comprehensive investigation across various activation functions, models, and evaluation tasks. This significantly larger scale study is particularly noteworthy, as previous works (e.g., [1], [2]) have primarily focused on smaller-scale experiments, such as replacing GELU with ReLU in ViT or BERT. In contrast, we train 1 billion parameter models on over 200 billion tokens, thus providing a more robust and extensive evaluation on various activation functions, architectures and tasks.
* **Inference Efficiency of LLMs**: The central motivation of our work is to analyze the implications of different activation functions on the inference efficiency of large language models (LLMs). Our findings demonstrate that commonly used activation functions do not offer a favorable balance between performance gains and efficiency or latency losses during inference. We argue in the paper that with additional one-time cost of longer training on relufied model one can get speed up during inference.
* **Exploration of Activation Function Consequences**: Our work delves into understanding the broader consequences of using different activation functions in LLMs. For instance, we investigate the pre-activation distribution, as shown in Figure 5, which leads us to explore a rarely used activation function for LLMs: the shifted ReLU. This exploration reveals that shifted ReLU can significantly increase sparsity while maintaining competitive performance, a finding detailed in Section 5.3 and Appendix F.
* **Practicality in Pretrained LLMs**: We emphasize the practicality of modifying activation functions in already pretrained LLMs. Our study illustrates that fine-tuning these models with ReLU does not result in a significant drop in accuracy, highlighting a trade-off that has been underexplored in prior research.  


In summary, while acknowledging the contributions of previous works, our study offers a unique perspective by examining the consequences of activation function selection in LLMs from an “inference efficiency” standpoint. We believe that our comprehensive approach and the novel insights generated, particularly in the context of model efficiency and practical adaptability, represent a significant advancement in the field.

Reviewer comment (Evaluation):
The evaluation (Table 1 & 2) majorly focuses on zero-shot learning and ICL scenarios. Although I understand that zero-shot learning and ICL are common settings to compare LLM performance, it can be helpful to compare the model performance on generation tasks to better understand how different activations influence model performance.

Author response:
Thank you for your valuable comment. In Appendix F.2, we have included new results that compare the generation quality of several relufied models. Overall, similar to the zero-shot scenario, the relufication process does not significantly impact perplexity. More interestingly, the TruthfulQA scores of the models seem to improve in the majority of cases, which encourages us to further investigate the impact of relufication on generation quality for other scenarios and additional models. In addition, we have included additional perplexity scores throughout Appendix F (e.g., Tab.4, Fig.14b) that are in line with our conclusions on zero-shot tasks.  

We thank the reviewer again and hope that these enhancements to our empirical evaluations, inspired by your suggestion, will strengthen the paper.

Reviewer comment (Experiments):
I am confused by the applications studied in Section 5.1, what does “loading new weights” mean here? Shouldn’t all weights be pre-loaded to the GPU HBM in common inference frameworks like vLLM [3]?

Author response:
We apologize for the confusion. The reviewer is indeed correct in their assumption that weights are pre-loaded into GPU memory. However, our research highlights an additional I/O cost incurred during the token generation phase. This cost arises from transferring weights between GPU memory and computational caches and registers.
Our findings on aggregated sparsity, as shown in Figure 7, reveal that this cost can be significantly reduced through neuron reuse. By reusing neurons, the total weight transfer volume during multiple token generations is substantially lower compared to scenarios without reuse.  
Furthermore, our findings have implications beyond this initial scenario. In cases where the entire model cannot fit in the GPU memory, but only a smaller subset can (i.e., by sharding or offloading), aggregated sparsity can be beneficial. For example, we can utilize a portion of the down projection layer's loaded weights for several subsequent generation steps without a significant loss in the quality of generation, as detailed in Figure 7(c).  
We have updated our manuscript for clearer articulation of these points. We believe these findings open up intriguing optimization opportunities for exploration.

Reviewer comment (Evaluation):
Although the paper claims that replacing other activations with ReLU only has negligible impacts on the performance, the accuracy drop seems not to be so marginal. However, I agree that this replacement can also be a potential good trade-off between model performance and model efficiency.

Author response:
We agree with your assessment of the favorable trade-off between model performance and efficiency. The performance gap with ReLU activation is relatively minor and can be offset by slightly prolonged training. Considering the widespread application of current LLMs, this additional training and fine-tuning is a justified investment for the efficiency gains achieved, as illustrated in Appendix B. These gains are particularly valuable in practical scenarios, highlighting the relevance of our approach in real-world deployments.

Reviewer comment (Experiments):
Authors should provide more empirical comparisons to other size-reduction methods to validate thesucess of their strategy. The approach they develop is not compared to any pruning methods such as [https://openreview.net/forum?id=0GRBKLBjJE, https://arxiv.org/abs/2003.03033] which could be seen as competition.

Author response:
Thank you for suggesting this interesting experiment. In Appendix F.3, we have compared relufied models with several pruning methods. To have a fair comparison, we fixed the pruned mask and finetuned these models on the same dataset as our relufied models. Furthermore, for hardware compatibility, we utilized the 2:4 pruning method.   

Our results show that the weight pruning methods exhibit significantly lower zero-shot accuracy across several tasks. We partially attribute this to the inherent differences between activation sparsity and weight pruning, with the latter permanently removing parts of the model and consequently reducing its overall computational capacity and knowledge.   

This additional result, we believe, has significantly enhanced the comprehensiveness of our study, and we thank the reviewer for their valuable suggestion.

Reviewer comment (Evaluation):
What are the latency speedups of this approach on the hardware it was tested on?

Author response:
We acknowledge that our initial manuscript did not delve deeply into hardware support, primarily to avoid redundancy with similar discussions in the literature [1].   

However, recognizing the importance of this aspect, we have been working on a custom GPU kernel for a while, and we have reported latency improvements on commonly accessible hardware, such as the MacBook Pro. While we are continuing to improve our implementation, details of our efficient kernel for developers familiar with GPU programming, are provided in Appendix B.1. Additionally, Figure 10(b) demonstrates how our sparse-vector dense-matrix multiplication scales with increasing levels of sparsity. This kernel is adaptable for use in other frameworks supporting Metal Performance Shaders (MPS), such as PyTorch.

Reviewer comment (Reproducibility):
What are prerequisites to make sure a user can realize the full benefits from this approach? I understand one needs a specific implementation of the NN code to take advantage of the sparsity.

Author response:
We acknowledge that our initial manuscript did not delve deeply into hardware support, primarily to avoid redundancy with similar discussions in the literature [1].   

However, recognizing the importance of this aspect, we have been working on a custom GPU kernel for a while, and we have reported latency improvements on commonly accessible hardware, such as the MacBook Pro. While we are continuing to improve our implementation, details of our efficient kernel for developers familiar with GPU programming, are provided in Appendix B.1. Additionally, Figure 10(b) demonstrates how our sparse-vector dense-matrix multiplication scales with increasing levels of sparsity. This kernel is adaptable for use in other frameworks supporting Metal Performance Shaders (MPS), such as PyTorch.

Reviewer comment (Experiments):
We observe different aggregated sparsity ratios at different layers: the deep layer seems less sparse, according to Figure 7(a). Then, for the perplexity experiment with aggregated sparsity, did the authors use the same γ for all layers? If not, could this help recover more performance?

Author response:
In our initial approach, we used a fixed gamma across all layers. However, based on your suggestion, we experimented with a mixed reuse strategy in Fig. 15 Appendix F.4. This strategy employs varying gamma values for different layers, with promising initial results. As gamma increases, we observe that shorter reuse periods in later layers do not entirely compensate the extended reuse duration in earlier layers. This insight has led us to believe that more sophisticated mixing strategies could be a fruitful area for future research. We appreciate your suggestion and believe it has significantly contributed to the depth of our study.

Reviewer comment (Experiments):
Could the authors elaborate on how they would imagine optimizing the finetuning process in relufication further to recover the full performance?

Author response:
This is an interesting question. There are several directions to further optimize the finetuning process:

* Increasing the number of tokens used for finetuning could further minimize the performance gap: This extension would allow the model more time to adapt to the new activation function while retaining the learned knowledge. To make this more practical, we can explore efficient finetuning methods such as low-rank methods (e.g., LoRA).
* More tailored training regimes: Currently, we follow the pretraining recipe of our base models. However, we believe that for short finetuning periods, we can use larger learning rates.
* Exploring the impact of finetuning data distribution: One interesting hypothesis is that finetuning on data with a similar distribution to the training data can help stabilize the training, reduce the impact of catastrophic forgetting, and help the model recover performance faster at larger learning rates.
* Another option may be to relufy layer by layer instead of doing it for all layer at once. This will lead to smaller changes which give the network easier time to adapt and recover. 

We hope this answers the reviewer's question. We are eager to know the reviewer's thoughts on these future directions