[论文解读] ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models

Iman Mirzadeh, Keivan Alizadeh+6 more|arXiv (Cornell University)|Oct 6, 2023

Topic Modeling被引用 7

一句话总结

本文主张在大型语言模型（LLMs）中重新引入ReLU激活函数，以利用其固有的激活稀疏性，实现高达3倍的推理FLOPS减少，同时性能损失极小。通过利用ReLU带来的结构化稀疏性，并在归一化层后引入ReLU层，该方法实现了高效的推测解码与模型压缩，表明在计算资源受限的条件下，基于ReLU的LLM在效率和准确性方面可超越更小的密集模型。

ABSTRACT

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.

研究动机与目标

• 重新评估ReLU激活在LLMs中的应用，尽管当前趋势是采用更平滑、更复杂的替代方案（如GELU和SiLU）。
• 探究基于ReLU的LLM是否能通过结构化激活稀疏性实现显著的推理效率提升。
• 探索对现有非ReLU LLM（如Llama、Falcon）进行微调与架构修改，以实现FLOPS显著降低的实用方法。
• 证明在推理FLOPS预算受限的情况下，基于ReLU的模型可在准确性上超越更小的密集模型。
• 揭示激活稀疏性的新应用，例如用于推测解码的聚合稀疏性，以及用于替代激活函数设计的预激活分析。

提出的方法

• 对现有非ReLU LLM（如Llama、Falcon）进行微调，引入ReLU激活，以恢复原始性能的同时提升推理效率。
• 在归一化层后插入额外的ReLU层，以进一步增强激活稀疏性并降低FLOPS。
• 提出一种稀疏推测解码框架，利用聚合稀疏性加速自回归生成。
• 利用ReLU训练模型的预激活分布，指导设计替代的稀疏激活函数（如移位ReLU）。
• 测量并利用“聚合稀疏性”——即在生成过程中跨token步骤重复激活的神经元——以减少冗余计算。
• 应用结构化稀疏模式（例如，将权重矩阵中的整行置零）以加速现代硬件上的推理，尤其是在内存受限的场景下。

(a) Sparsity of different models

实验结果

研究问题

• RQ1 在LLM中使用ReLU是否能在不牺牲模型性能的前提下带来显著的推理效率提升？
• RQ2 非ReLU LLM（如Llama、Falcon）能否通过引入ReLU激活实现有效微调，以恢复原始性能并降低FLOPS？
• RQ3 在自回归生成中，激活稀疏性在多大程度上可被利用以加速推测解码？
• RQ4 预激活分布能否用于设计保持性能的同时最大化稀疏性的替代稀疏激活函数？
• RQ5 在相同推理FLOPS预算下，经ReLU化的大型模型是否能在准确性上超越更小的密集模型？

主要发现

• 基于ReLU的LLM在OPT模型上实现了高达32%的推理FLOPS减少（从6.6G降至4.5G FLOPS/ token），归因于超过90%的激活稀疏性。
• 对非ReLU LLM进行ReLU激活微调后，可在推理和阅读理解任务上恢复原始性能，精度损失极小。
• 通过在归一化层后插入ReLU层并利用聚合稀疏性，该方法将推理FLOPS降低了三倍。
• 基于ReLU的稀疏推测解码在较高生成速度（γ）下相比标准推测解码实现了更高的加速比，得益于持续的稀疏性。
• 经ReLU化的大型模型（如OPT 6.7B）在等效FLOPS下超越了更小的密集模型，位于原始模型缩放曲线之上，零 shot准确率最高提升2%。
• 预激活分析表明，ReLU类分布自然出现在训练过程中，从而可设计出移位ReLU函数，实现高达90%的稀疏性，且性能与标准ReLU相当。

(b) Sparsity for Efficiency