PaperlayerView exact revision on arXiv

Source-linked AI summary

ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models

Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, Mehrdad Farajtabar

arXiv:2310.04564v1cs.LGcs.AI

TL;DR

LLM inference faces substantial compute and memory demands, and this paper re-evaluates ReLU as an activation for reducing that burden. It studies ReLU relufication and sparsity-based strategies, reporting similar performance with reduced computation, including reductions of up to threefold.

  • Problem

    LLM inference requires substantial computation and memory, while smoother, more complex activations have become favored despite their added computational demands.

    Source: §1 · pp. 1–2

  • Method

    The paper evaluates ReLU in scratch-trained and relufied pretrained models, then uses activation sparsity to devise additional inference strategies.

    Source: §1 · p. 2

  • Results

    Across zero-shot and few-shot tasks, ReLU-based LLMs perform similarly to non-ReLU models at reduced computation; adding ReLU after normalization layers reduces inference FLOPS by up to threefold.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    ReLU can provide activation sparsity and more efficient inference with little performance impact in the evaluated models and tasks.

    Source: §1 · p. 2, §6 · p. 9

    The impact of shifted ReLU in relufication’s second stage remains to be studied, where sparsity is usually not very high.

    Source: §5.3 · p. 8

Abstract

from arXiv · show

Large Language Models (LLMs) with billions of parameters have drastically transformed AI applications. However, their demanding computation during inference has raised significant challenges for deployment on resource-constrained devices. Despite recent trends favoring alternative activation functions such as GELU or SiLU, known for increased computation, this study strongly advocates for reinstating ReLU activation in LLMs. We demonstrate that using the ReLU activation function has a negligible impact on convergence and performance while significantly reducing computation and weight transfer. This reduction is particularly valuable during the memory-bound inference step, where efficiency is paramount. Exploring sparsity patterns in ReLU-based LLMs, we unveil the reutilization of activated neurons for generating new tokens and leveraging these insights, we propose practical strategies to substantially reduce LLM inference computation up to three times, using ReLU activations with minimal performance trade-offs.

1 Introduction

LLM inference is constrained by substantial compute and memory demands, motivating the paper’s re-evaluation of ReLU as a way to exploit activation sparsity. The authors report reduced computation with little performance change and propose further sparsity-based inference strategies.

1 Introduction

  • Activation sparsity offers a favorable accuracy–speedup balance, especially on GPUs, amid LLM inference’s substantial compute and memory demands.Source: §1 · p. 1
  • 32% computation saving reduces OPT inference from 6.6G FLOPS to 4.5G FLOPS per token; activation sparsity exceeds 90% across layers and affects 95% of down-projection rows.The zeroed rows also reduce GPU–CPU weight transfer.Source: §1 · p. 1, Figure 1 · p. 2
  • Fine-tuned pretrained models show no significant within-model performance change across ReLU, GELU, and SiLU, while ReLU requires much less computation.The comparison covers three models and is summarized in Figure 1c.Source: §1 · p. 2, Figure 1 · p. 2
  • Scratch-trained models show no significant performance difference across activation functions, while ReLU is lighter at inference; fine-tuning non-ReLU models with ReLU quickly regains original performance across reasoning and reading-comprehension tasks.The paper proposes adding ReLU after normalization layers to reduce inference FLOPS by up to threefold.Source: §1 · p. 2
  • ReLU models reuse previously activated neurons during token generation, a pattern termed aggregated sparsity; the authors also explore sparsity-guided activation selection, including shifted ReLU.They report speculative-decoding speedup and up to 90% sparsity with performance similar to ReLU for shifted ReLU.Source: §1 · p. 2

2 Related Works

The related-work discussion situates ReLU among smoother activation alternatives and prior sparsity methods, then identifies adapting pretrained non-ReLU models as a gap. It also connects activation sparsity to speculative decoding.

Activation Functions in Transformers

  • Earlier transformer studies reported benefits from alternative activations at smaller scales, while this work trains 1B-parameter models on more than 100B tokens.The paper suggests activation-function effects may be marginal at scale, following scaling laws.Source: §2 · p. 3

Activation Sparsity

  • Prior sparsity methods often assume pretrained ReLU models or focus on one architecture; this work targets non-ReLU pretrained models across architectures without a separate prediction module.Source: §2 · p. 3

Speculative Decoding and Sparsity

  • Activation sparsity has a temporal pattern that can enhance speculative decoding, and the paper provides guidelines for selecting sparsity parameters.Source: §2 · p. 3

Related work scope

  • The discussion defers orthogonal work on model compression, sparse attention, and Mixture of Experts to Appendix A.Source: §2 · p. 3

3 Does the Activation Function Impact Performance?

The section tests whether activation choice affects model performance and inference efficiency, using scratch training to compare activation functions. It reports similar performance but differing sparsity, with ReLU yielding greater sparsity.

3 Does the Activation Function Impact Performance?

  • The section’s central comparison is that changing activation functions minimally impacts performance, while inference efficiency changes substantially.Source: §3 · p. 3

3.1 Experimental setup

  • The experiments use OPT, Llama, and Falcon, RefinedWeb for pretraining and fine-tuning, and few-shot tasks spanning abilities such as reasoning and reading comprehension.WikiText validation is used to measure sparsity and record preactivation distributions.Source: §3.1 · p. 3

3.2 Training from scratch: performance and sparsity

  • For OPT 1.3B models trained from scratch, performance is similar across activation functions, while sparsity increases from SiLU toward ReLU as β rises; ReLU drops significantly more values than GELU and SiLU.Source: Figure 2 · p. 4, §3.2 · p. 4

4 Relufication

Relufication replaces non-ReLU activations in pretrained models and can add ReLUs after normalization layers. The modified models gain activation sparsity while retaining comparable performance, with finetuning helping them recover performance after the architectural changes.

4.1 Stage 1: replacing non-ReLU activations

  • Stage 1 replaces Falcon’s GELU and Llama’s SiLU with ReLU, then finetunes the models on 30 billion RefinedWeb tokens; OPT’s existing ReLUs remain unchanged.Activation sparsity increases, while preactivation distributions change little during short finetuning and models quickly recover performance lost after the surgery.Source: §4.1 · p. 5, Figure 5 · p. 6

4.2 Stage 2: Pushing for more sparsity

  • Across several zero-shot tasks, relufied models have significantly higher activation sparsity and FLOPS efficiency, with comparable performance within each model group.The reported small performance gap may partly reflect the limited finetuning duration: 30B tokens for stage 1 and 50B for stage 2.Source: Table 1 · p. 6, §4.2 · p. 6
  • Stage 2 inserts ReLUs after normalization layers to sparsify inputs to additional attention and FFN projections beyond the down projection.The addressed matrix-vector multiplications account for about 55% of total computation; the down projection accounts for roughly 30%.Source: §4.2 · p. 6
  • Relufication’s stages do not significantly reduce zero-shot accuracy while using less compute, and five-shot MMLU performance also remains similar across activation choices.The MMLU comparison is reported in Table 2.Source: §4.2 · pp. 6–7, Table 2 · p. 7

5 Applications

The section explores how ReLU activation sparsity can support neuron reuse, faster speculative decoding, and sparser shifted activations. Across these applications, it reports reuse patterns and efficiency gains while examining performance trade-offs.

5.1 Aggregated Sparsity: reusing previously activated neurons

  • For OPT-6.7B on WikiText prompts, about 50% of neurons remain unused across the first 150 tokens, and observed reuse exceeds the random pattern.The authors report that the pattern also holds for other ReLU models and datasets.Source: §5.1 · p. 7
  • Reusing previously loaded down-projection weights across tokens causes only a slight perplexity increase, whereas random selection notably increases perplexity.In the example, generation alternates γ-token blocks of conventional computation and weight reuse after an initial 128 tokens.Source: §5.1 · p. 7

5.2 Activation sparsity and speculative decoding

  • At γ = 16, sparse speculative decoding is 1.27x faster than standard decoding, versus 1.20x for random sparsity.At γ = 64, aggregated sparsity yields around 1.14x speedup, while random sparsity's speedup is almost negligible.Source: Figure 7 · p. 8, §5.2 · p. 8

5.3 The shifted ReLU activation

  • Choosing b from the preactivation distribution enables ReLU(x − b) to increase sparsity without changing the loss function or training regime.For the cited distribution, b = 1 drops 95% of preactivations.Source: §5.3 · p. 8
  • The authors identify further study of sparse ReLU variants and shifted ReLU in stage 2 as future directions.They note that sparsity is usually not very high in stage 2.Source: §5.3 · p. 8
  • Shifted ReLU has on-par accuracy with ReLU and quickly recovers lost performance during fine-tuning while maintaining high activation sparsity.The gap is wider early in training and narrows as more tokens are seen.Source: §5.3 · p. 8, Figure 8 · p. 9

6 Conclusion

The conclusion reports that ReLU-based LLMs perform similarly to non-ReLU models at reduced computation, and points to further investigation of structured activation sparsity.

6 Conclusion

  • Across several zero-shot and few-shot tasks, ReLU-based LLMs perform similarly to their non-ReLU counterparts at significantly reduced computation.The study also reports that activation-function choice during pretraining and fine-tuning does not significantly affect performance.Source: §6 · p. 9
  • The authors explore aggregated sparsity and shifted ReLU as directions for improving token-generation speed and efficiency.These directions follow their observations of sparsity patterns in ReLU LLMs.Source: §6 · p. 9
  • The authors hope the findings motivate further study of well-structured activation sparsity to enhance model efficiency.Source: §6 · p. 9

A Extended Related Works

The related work contrasts this study's large-scale investigation of activation functions with prior work on activations, sparsity, compression, and speculative decoding. It positions relufication as a distinct approach that can complement other efficiency techniques.

A Extended Related Works

  • Unlike prior smaller-scale activation-function comparisons, this work trains billion-scale models and data, finding that extended training can diminish performance gaps between ReLU and other functions while saving computation.Source: A Extended Related Works · p. 15
  • Unlike sparsity-prediction studies that assume pretrained ReLU activations, this approach replaces other activations with ReLU and fine-tunes for MLP activation sparsity.It also reports that adding ReLU before QKV and feedforward layers can reduce FLOPS at a minor accuracy cost, without a sparsity predictor.Source: A Extended Related Works · p. 15
  • Relufication focuses on activations before weight multiplications, unlike prior work replacing attention softmax with ReLU [68] [32].Source: A Extended Related Works · p. 15
  • Relufication is orthogonal to quantization and distillation, and its activation sparsity can combine with weight-sparse models to reduce compute further.The passage distinguishes activation sparsity, which supports nonzero-row matrix multiplication, from typically unstructured weight sparsity.Source: A Extended Related Works · p. 15
  • Relufication can be combined with Mixture of Experts, whose expert routing is a group-form subset of activation sparsity, and can support MoEfication by increasing FFN sparsity.Source: A Extended Related Works · p. 15
  • The study examines incorporating sparsity into speculative decoding and provides parameter-selection guidelines for that setting.It reports that adding sparsity can improve speculative-decoding performance.Source: A Extended Related Works · p. 15

B Discussion on Activation Sparsity and Inference Efficiency

The discussion argues that FLOPS can serve as an efficiency proxy when activations are sparse, because sparsity affects both computation and weight transfer. It reports a correlation between FLOPS and token-generation time for OPT on an NVIDIA A100 node.

B Discussion on Activation Sparsity and Inference Efficiency

  • Latency is the most practical efficiency measure for a target device, but device-specific properties motivate using a more ubiquitous proxy.The discussion identifies computation and I/O transfer as major efficiency factors and says I/O is the major inference bottleneck for today's large models.Source: B Discussion on Activation Sparsity… · pp. 15–16
  • For OPT on an NVIDIA A100 node, FLOPS provide a reasonable approximation of token-generation time, with high correlation, especially for LLMs with activation sparsity.Figure 9 presents the comparison of FLOPS and time needs to generate tokens.Source: Figure 9 · p. 16, B Discussion on Activation Sparsity… · p. 16
  • Sparse activations reduce computation through sparse vector-dense matrix multiplication at inference and sparse-matrix-dense matrix multiplication during training.The authors describe this as semi-structured sparsity that can exploit transferring weights in row-sized chunks and existing accelerator support.Source: B Discussion on Activation Sparsity… · p. 16
  • During inference, row-major matrix storage permits skipping weight rows for zero outputs, reducing I/O transfer.The passage identifies weight transfer to device cache as the main bottleneck during token generation and cites approximately 99.3% of total latency as I/O [51].Source: B Discussion on Activation Sparsity… · p. 16

C Activation Sparsity and Speculative Decoding

Sparse speculative decoding uses aggregated activation sparsity to evaluate only the nonsparse portion of the larger model during verification. Under an I/O-bound latency model, selecting the speculation length using sparsity can improve speedup over standard speculative decoding and autoregressive decoding.

C Activation Sparsity and Speculative Decoding

  • The larger model verifies γ tokens proposed by a smaller model, avoiding sequential runs of the larger model for each token.Speculative decoding selects matching tokens from the larger model's verification.Source: C Activation Sparsity and Speculative… · p. 16
  • In the I/O-bound latency model, sparse verification runs the smaller model plus only the larger model's nonsparse portion, assuming i.i.d. token acceptance.The larger model's latency for γ tokens is approximated by its runtime multiplied by one minus average aggregated sparsity.Source: C Activation Sparsity and Speculative… · pp. 16–17
  • Theorem 1 gives sparse speculative decoding's expected latency improvement over standard speculative decoding as standard cost divided by sparse cost.Both costs include the smaller-model run; sparse decoding replaces the full larger-model run with its nonsparse portion.Source: C Activation Sparsity and Speculative… · p. 17
  • Theorem 2 gives the expected latency improvement of sparse speculative decoding over autoregressive decoding using the larger model.The comparison accounts for token acceptance probability and aggregated sparsity.Source: C Activation Sparsity and Speculative… · p. 17
  • At α = 0.8 and c = 0.02, sparse speculative decoding peaks at γ = 10, beating standard decoding at γ = 12 and sparse decoding at γ = 12.The chosen sparse-decoding γ is smaller than standard decoding's optimal γ, with a gap always under 20%.Source: C Activation Sparsity and Speculative… · p. 17, Figure 10 · p. 17

D Pre-activation Distribution of OPT models trained from scratch

OPT 1.3B models trained from scratch with four activation variants show pre-activation distributions that diverge over training. Moving from SiLU toward ReLU concentrates the distributions around zero.

D Pre-activation Distribution of OPT models trained from scratch

E Is a sparsified large model better than dense smaller ones?

Relufied OPT models lie above the original dense-model scaling curve in performance versus efficiency comparisons. The examples indicate that, at equivalent FLOPS, dense models perform lower in zero-shot performance.

E Is a sparsified large model better than dense smaller ones?

Loading 2310.04564v1…