ReLU$^2$ Wins: Discovering Efficient Activation Functions for Sparse LLMs (2402.03804v1)

Abstract: Sparse computation offers a compelling solution for the inference of LLMs in low-resource scenarios by dynamically skipping the computation of inactive neurons. While traditional approaches focus on ReLU-based LLMs, leveraging zeros in activation values, we broaden the scope of sparse LLMs beyond zero activation values. We introduce a general method that defines neuron activation through neuron output magnitudes and a tailored magnitude threshold, demonstrating that non-ReLU LLMs also exhibit sparse activation. To find the most efficient activation function for sparse computation, we propose a systematic framework to examine the sparsity of LLMs from three aspects: the trade-off between sparsity and performance, the predictivity of sparsity, and the hardware affinity. We conduct thorough experiments on LLMs utilizing different activation functions, including ReLU, SwiGLU, ReGLU, and ReLU$^2$. The results indicate that models employing ReLU$^2$ excel across all three evaluation aspects, highlighting its potential as an efficient activation function for sparse LLMs. We will release the code to facilitate future research.

ReLU Wins: Discovering Efficient Activation Functions for Sparse LLMs

The paper "ReLU Wins: Discovering Efficient Activation Functions for Sparse LLMs" by Zhengyan Zhang et al. investigates the efficacy of different activation functions in the context of deploying LLMs using sparse computation. It posits that ReLU, despite being the conventional choice, is exceptionally well-suited for achieving high sparsity ratios, efficient prediction, and hardware-friendly implementation. This perspective challenges the focus on more novel activation functions like SwiGLU and ReGLU, especially in resource-constrained environments.

Summary of Key Contributions

The paper presents several key contributions:

1. Generalized Activation Definition: The authors expand the traditional approach of zero-activation-based sparsity to a more generalized definition that includes non-zero activation values. This is achieved by defining neuron activation based on output magnitudes exceeding a magnitude threshold.
2. Framework for Evaluation: A systematic framework is proposed to assess activation functions from three perspectives: the trade-off between sparsity and performance, predictivity of sparsity, and hardware affinity.
3. Experimental Validation: Comprehensive experiments are conducted on LLMs with different activation functions, including ReLU, SwiGLU, ReGLU, and ReLU. ReLU consistently demonstrates superior performance across all metrics.

Implications of Key Findings

Generalized Activation and Sparsity

By redefining activation to consider neurons with small output magnitudes as inactive, the paper reveals that non-ReLU LLMs can also exhibit sparse activation. This is a significant shift in understanding, as it challenges the prevalent notion that sparse activation is predominantly seen in ReLU-based LLMs. It opens new avenues for efficient sparse computation by broadening the scope of activation functions that can be considered.

Performance vs. Sparsity

The experiments demonstrate that ReLU provides the best trade-off between sparsity and performance. Specifically, the performance degradation is less than 0.1% at a sparsity ratio of approximately 90%. This is critical for practical applications where minor performance sacrifices are acceptable for significant computational savings.

Predictivity and Hardware Affinity

Predictivity of neuron activation is crucial for efficient deployment, and models using ReLU and ReLU exhibit the highest neuron activation predictivity. These models also show better performance in terms of hardware affinity. For instance, the I/O overhead of feed-forward networks using ReLU can be reduced by 92%. This has profound implications for optimizing LLM inference on real-world hardware, particularly in scenarios with stringent resource constraints.

Theoretical and Practical Implications

From a theoretical standpoint, this paper underscores the importance of revisiting and rigorously evaluating conventional components like ReLU in the context of new applications such as sparse computation. Practically, the findings advocate for a more informed selection of activation functions when designing LLMs intended for deployment in low-resource environments. The results suggest that ReLU, while perhaps seen as an older activation function, may still hold significant advantages that warrant its continued use and further exploration.

Speculations on Future Developments

Future research could focus on further refining the generalized activation definition to account for dynamic thresholds that adapt to different LLM architectures and tasks. Additionally, exploring hybrid models that combine the efficiency of ReLU with the non-linearity properties of newer activation functions might yield even more efficient sparse LLMs. There is also scope for developing hardware accelerators specifically optimized for the sparse computational patterns identified through this paper.

In conclusion, the paper by Zhang et al. provides a robust argument for Reevaluating ReLU within the field of sparse LLMs. It demonstrates that, despite the availability of newer activation functions, ReLU remains a highly effective choice for efficient and scalable deployment of LLMs. This finding could significantly influence future research and development in the area of efficient AI model deployment.