Source-linked AI summary
ReLU$^2$ Wins: Discovering Efficient Activation Functions for Sparse LLMs
Zhengyan Zhang, Yixin Song, Guanghui Yu, Xu Han, Yankai Lin, Chaojun Xiao, Chenyang Song, Zhiyuan Liu, Zeyu Mi, Maosong Sun
TL;DR
Sparse computation may reduce the substantial inference resources LLMs require, but prior approaches focus on ReLU and exact-zero activations. This paper defines activation using neuron output magnitudes and a threshold, then evaluates activation functions across sparsity–performance trade-offs, predictivity, and hardware affinity. ReLU2 excels across all three aspects.
Problem
Sparse LLM inference for low-resource scenarios is resource-intensive, and prior approaches focus on ReLU-based sparsity from zero-valued activations.
Source: §1 · p. 1Method
The paper defines activation using neuron output magnitude and a threshold, then evaluates ReLU, SwiGLU, ReGLU, and ReLU2 across trade-offs, predictivity, and hardware affinity.
Source: §1 · pp. 1–2Results
Takeaways & Limitations
ReLU2 is a promising activation function for sparse LLMs, and the broader magnitude-based definition finds sparse activation beyond ReLU.
Source: §6 · pp. 8–9
Abstract
from arXiv · showhide
Sparse computation offers a compelling solution for the inference of Large Language Models (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.
1. Introduction
The paper broadens sparse activation beyond exact zeros by using neuron output magnitudes and a threshold, and asks which activation function best supports sparse LLMs. It evaluates sparsity and performance trade-offs, predictivity, and hardware affinity, finding ReLU2 strongest across these aspects.
1. Introduction
- Sparse computation targets low-resource inference by skipping weakly contributing parameters, while prior sparse-deployment efforts primarily use ReLU and exact-zero activations.The paper asks which activation function is optimal for sparse LLMs.Source: §1 · p. 1
- Defining activation by neuron output magnitude and a threshold captures small, nonzero outputs; this reveals sparse activation in non-ReLU LLMs and greater sparsity in ReLU LLMs than zero-based counting.The paper reports a long-tailed distribution of neuron output magnitudes, with many small values.Source: §1 · pp. 1–2
- The evaluation framework examines sparsity–performance trade-offs, predicts inactive neurons before computation, and studies hardware affinity, including token and neuron relationships that can improve memory reuse and contiguous access.Threshold adjustment explores the trade-off, while a lightweight predictor assesses predictivity.Source: §1 · p. 2
- ReLU2 excels across all three aspects; performance degradation is less than 0.1% near 90% sparsity, and FFN I/O overhead can be reduced by 92% using hardware features.The experiments compare ReLU, SwiGLU, ReGLU, and ReLU2.Source: §1 · p. 2
2. Related Work
Related work frames efficient LLM inference as an algorithm-and-systems challenge and describes sparse activation as one approach to that challenge. It identifies open questions about sparse activation beyond ReLU and relates neuron-level sparsity to conditional computation.
Efficient Inference of LLMs
- LLM inference work reduces time and memory through parameter compression, structural changes, and decoding methods, while systems work integrates computation with hardware features.Sparse activation combines algorithmic choices, such as activation selection and prediction, with systems work that exploits sparsity on specific hardware.Source: §2 · p. 2
Sparse Activation of LLMs
- Prior work observes sparse activation in ReLU-based LLMs and explains it through training, but whether it extends to other activation functions and how ReLU variants differ remained unclear.The paper explores ReLU variants and non-ReLU functions to address these questions.Source: §2 · pp. 2–3
Conditional Computation of LLMs
- Unlike MoE's conditional computation introduced before training, naturally emerging sparse activation offers a post-training opportunity, with each neuron treated as an expert.Prior work describes transforming sparsely activated feed-forward networks, and later attention networks, into MoE-like computation.Source: §2 · p. 3
3. Is Non-ReLU LLM Sparsely Activated?
Because non-ReLU activations are not exactly zero for negative inputs, the authors investigate whether sparse activation also appears in non-ReLU LLMs. They analyze neuron-output magnitudes in LLaMA-2 and relate tail truncation to sparsity and task performance.
3. Is Non-ReLU LLM Sparsely Activated?
- Because non-ReLU activations are not exactly zero for negative inputs, the study tests LLaMA-2 as a representative non-ReLU model for sparse activation.The question is whether sparse activation also exists in non-ReLU LLMs.Source: §3 · p. 3
- Normalized neuron-output magnitudes have similar long-tailed distributions across layers, while average output magnitude increases with layer index.The analysis records output magnitudes for LLaMA-2 7B on a mixed corpus.Source: §3 · p. 4
- At a sparsity ratio of 0.7, average CETT is 0.2, and CETT increases much more slowly than sparsity.CETT measures the impact of truncating tail neuron outputs on FFN computation.Source: §3 · p. 4
- At a sparsity ratio of 0.7, LLaMA performance drops by about 1%; at 0.95, it drops by over 5%.Performance is not sensitive to tail truncation below a sparsity ratio of 0.7.Source: §3 · p. 4
- The proposed adaptive threshold is the largest output-magnitude threshold whose CETT remains below a target upper bound, found by binary search over candidate quantiles.The threshold is designed to account for magnitude variation across layers and models.Source: §3 · pp. 4–5
4. Key Factors for Sparse LLM Deployment
The section assesses sparse LLM deployment through performance–sparsity trade-offs, predictivity, and hardware affinity. It evaluates threshold-based sparsity alongside token and neuron activation relationships.
4. Key Factors for Sparse LLM Deployment
- The sparsity ratio is the share of neurons whose output magnitudes fall below a threshold, and performance is evaluated across sparsity ratios to identify the trade-off.A higher ratio offers greater optimization potential when performance is not significantly sacrificed.Source: §4 · p. 5
- Predictivity measures whether neuron activation can be forecast before FFN computation; a two-layer neural-network predictor is evaluated by activated-neuron recall and prediction sparsity.Higher recall and prediction sparsity indicate better predictivity for predictors of the same size.Source: §4 · p. 5
- Token-level overlap in activated neurons is measured as the reuse ratio, indicating potential to cache prior-token parameters and reduce parameter I/O.The reuse ratio averages overlap between the current token and the previous k tokens.Source: §4 · p. 5
- The top-average co-activation gap compares each neuron's strongest co-activation with the matrix-wide average; a high gap favors placing co-activated neuron pairs at contiguous addresses.The matrix records pairwise co-activation frequency, normalized by the activation count of the first neuron.Source: §4 · pp. 5–6
- At 100 billion training tokens, the 1B models using SwiGLU, ReGLU, and ReLU2 have very close evaluation performance.Figure 2 also reports training loss and performance dynamics across evaluation datasets.Source: Figure 2 · p. 6
- ReLU2 achieves the best trade-off between performance and sparsity, and has higher sparsity than other functions at the same cumulative error in most cases.A cumulative error of 0.2 is an inflection point for model performance across all activation functions.Source: Figure 3 · p. 6
5. Finding Best Function for LLM Deploying
The experiments compare activation functions across training performance, sparsity trade-offs, predictivity, and hardware affinity. ReLU2 performs strongly across these dimensions, with results also showing how sparsity and reuse vary by model scale and activation function.
5. Finding Best Function for LLM Deploying
- ReGLU and ReLU2 ultimately match SwiGLU performance, while ReLU consistently underperforms.ReLU2 lags slightly during intermediate training but improves toward the end.Source: §5 · p. 6
5.1. Trade-off between Performance and Sparsity
- At 0.95 sparsity, ReLU2’s CETT is 0.2, nearly half that of ReLU and ReGLU; CETT 0.2 also marks the empirical upper bound before performance drops significantly.ReLU2 offers the best performance–sparsity trade-off across the tested sparsity ratios.Source: §5.1 · pp. 6–7
- SwiGLU’s CETT–sparsity tendency is similar across model scales, whereas ReGLU becomes sparser as scale increases.For ReGLU, larger-model curves indicate less impact from removing tail neuron outputs.Source: Figure 4 · p. 6, §5.1 · p. 7
5.2. Predictivity
- With top-k prediction fixed at 0.2 sparsity, ReLU2 and ReLU have similar recall, both higher than SwiGLU and ReGLU.Recall increases with sparsity for all four activation functions.Source: §5.2 · p. 7, Figure 5 · p. 7
- With threshold-based prediction, ReLU2 leads in recall and prediction sparsity; from 0.85 to 0.95 sparsity, prediction sparsity rises by over 0.1 while recall changes by under 1% for SwiGLU, ReLU, and ReLU2.Neurons are predicted active when their prediction score exceeds 0.5.Source: §5.2 · p. 7, Figure 5 · p. 7
5.3. Hardware Affinity
- For token reuse, a sliding window of five yields twice the overall reuse ratio of a window of one, and ReLU2 has higher optimization potential at high activation ratios.The comparison sets model sparsity ratios close to 0.95.Source: §5.3 · pp. 7–8, Figure 6 · p. 8
- ReLU variants have a top-average co-activation gap of about 0.2 versus about 0.1 for SwiGLU; after layer eight, ReLU2’s gap exceeds those of ReLU and ReGLU.A larger gap makes reallocating highly co-activated neuron pairs to contiguous memory addresses more beneficial, as described for these models.Source: §5.3 · p. 8, Figure 7 · p. 8
5.4. ReLU2: Best Function for LLM Deploying
- ReLU2 scores highly across all four evaluation aspects; estimates indicate 56% lower FLOPs with less than 1% performance loss and 92% lower FFN I/O overhead.The I/O estimate uses both sparsity and reuse ratios.Source: §5.4 · p. 8, Table 1 · p. 8
6. Conclusion
The conclusion presents magnitude-threshold activation as a general way to identify sparse activation beyond ReLU’s exact-zero values. The appendices specify evaluation datasets and the predictor’s training setup.
6. Conclusion
- A magnitude threshold on neuron output magnitudes reveals sparse activation in non-ReLU LLMs and sparser activation in ReLU-based LLMs than the traditional definition.Across the three evaluation aspects, the conclusion identifies ReLU2 as excelling and highlights its potential for sparse LLMs.Source: §6 · pp. 8–9
A.1. Evaluation Datasets
- LLaMA-2 7B evaluation uses MMLU, ARC challenge, Winogrande, HellaSwag, TruthfulQA, and GSM8k; 1B evaluation excludes MMLU and GSM8k for insufficient model capacity.TruthfulQA is inverted when computing the 7B average accuracy; the 1B evaluation also includes ARC easy, LAMBADA, PIQA, and OpenBookQA.Source: §A.1 · p. 13
A.2. Details of Predictor
- The FFN predictor is trained and validated on inputs paired with activation records, minimizing binary cross-entropy between predicted scores and those records.A neuron is recorded active when its output magnitude exceeds the threshold; a prediction score above 0.5 predicts activation.Source: §A.2 · p. 13
A.3. Training of 1.3B Model
This appendix section introduces the 1.3B model’s architecture, data, and hyperparameters, and points to its final evaluation results.
A.3. Training of 1.3B Model
- The subsection introduces training details for the 1.3B model, including its architecture, data types, and hyperparameters, and states that evaluation results for the final 1.3B models appear in Table 2.Source: §A.3 · p. 13
- Table 2 reports accuracy percentages for 1B models across evaluation datasets and an average.The table caption distinguishes ARC easy and challenging sets as ARC:E and ARC:C.Source: Table 2 · p. 13, §A.3 · p. 13
- ARC:E denotes the easy set of ARC, while ARC:C denotes its challenging set.Source: Table 2 · p. 13
A.3.1. ARCHITECTURE
The architecture section describes a LLaMA-2-like model and specifies its activation functions, intermediate hidden sizes, and attention components.
A.3.1. ARCHITECTURE
- The model adopts an architecture similar to LLaMA-2.Table 3 is identified as containing architecture details.Source: §A.3.1 · p. 13, Table 3 · p. 14
- The experiments use ReLU, ReLU2, ReGLU, and SwiGLU, with intermediate hidden sizes 8192 for ReLU variants and 5376 for ReGLU and SwiGLU.The different sizes are set to ensure equal parameter counts.Source: §A.3.1 · pp. 13–14
- The attention block follows LLaMA2-7B’s architecture, with RMSNorm pre-normalization and RoPE positional embeddings.Source: §A.3.1 · p. 14
A.3.2. PRE-TRAINING DATA
The 1B models were pre-trained for one epoch on 100 billion tokens sampled from SlimPajama and refinedweb.
A.3.2. PRE-TRAINING DATA
- The 1B models were trained for one epoch on 100 billion tokens sampled from SlimPajama and refinedweb.The two natural-language corpora were sampled and combined for pre-training.Source: §A.3.2 · p. 14
A.3.3. TRAINING HYPERPARAMETERS
Training used AdamW with specified β values and a cosine learning-rate schedule; other listed settings used defaults.
A.3.3. TRAINING HYPERPARAMETERS
- Training used AdamW with β1 = 0.9 and β2 = 0.95.Source: §A.3.3 · p. 14
- The learning rate followed a cosine schedule, while weight decay and gradient clipping used default hyper-parameters.Source: §A.3.3 · p. 14
A.4. Further Pre-training of LLaMA using ReGLU
ReGLU-based LLaMA models were comparable to original LLaMA-2 at 7B and 13B, but trailed at 70B; the section also describes decomposing FFNs into neuron contributions.
B. Neuron Decomposition of FFNs
- For the two-layer FFN used by the original Transformer and Falcon, d_ff neurons have vector outputs that sum to form the FFN output.Bias terms are omitted for simplicity in the discussion.Source: B. Neuron Decomposition of FFNs · pp. 14–15
A.4. Further Pre-training of LLaMA using ReGLU
- At 7B and 13B, ReGLU-based LLaMA models were comparable to original LLaMA-2, but at 70B they lagged by about 6% on average.The authors conjecture insufficient pre-training data caused the 70B gap and expect more data could improve performance.Source: §A.4 · p. 14, Table 5 · p. 15
B. Neuron Decomposition of FFNs
- Each neuron's activation value is scalar, while its output representation is a vector.Source: B. Neuron Decomposition of FFNs · p. 15
C. Threshold-finding Method
The threshold-search procedure selects a threshold using neuron-output magnitudes and a CETT upper bound; the section also reports hot-neuron patterns across model sizes.
C. Threshold-finding Method
- The threshold search uses neuron outputs and a CETT upper bound to select a threshold, calculating output-magnitude quantiles and CETT across candidate bounds.The tested CETT upper bounds are 0.01, 0.02, 0.04, 0.1, 0.2, 0.3, 0.4, and 0.5.Source: C. Threshold-finding Method · p. 15
- Figure 8 plots activation CDFs for 1B models, with neuron proportion sorted descending on the x-axis; Figure 9 covers models around 7B parameters.OPT 6B uses ReLU, while Persimmon 8B uses ReLU2.Source: Figure 8 · p. 16, Figure 9 · p. 16, C. Threshold-finding Method · p. 16
D. Hot-activated Neurons
- The 1B models show no significant hot-neuron presence, whereas some layers in all examined ~7B models do; OPT 6B's lower-layer top 20% are active over 80% of the time.The authors say the size-related difference is not yet fully understood and suggest 1B models may be too small to show hot-activated neurons.Source: D. Hot-activated Neurons · p. 16