Learn To be Efficient: Build Structured Sparsity in Large Language Models

📄 arXiv: 2402.06126v4 📥 PDF

作者: Haizhong Zheng, Xiaoyan Bai, Xueshen Liu, Z. Morley Mao, Beidi Chen, Fan Lai, Atul Prakash

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-02-09 (更新: 2024-12-12)

期刊: Advances in Neural Information Processing Systems (NeurIPS) 2024

💡 一句话要点

提出LTE算法以提升大语言模型的结构化稀疏性

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 激活稀疏性 推理效率 结构化稀疏性 机器学习 自然语言处理 算法优化

📋 核心要点

  1. 现有方法仅在训练后利用激活稀疏性，未能充分挖掘其潜力，导致推理效率低下。
  2. 提出的LTE算法通过结构化激活稀疏性训练LLMs，使其在推理时激活更少的神经元，优化稀疏性与性能的平衡。
  3. 实验结果显示，LTE在语言理解和生成任务上均超越了现有基线，且在50%稀疏性下显著降低了推理延迟。

📝 摘要（中文）

大语言模型（LLMs）在参数规模上取得了显著成功，但其推理开销较高。激活稀疏性为降低这一成本提供了自然的解决方案。然而，现有方法仅在训练后利用这种稀疏性，未能进一步放大其潜力。本文假设LLMs可以通过实现更结构化的激活稀疏性来学习高效性。为此，我们提出了一种新颖的训练算法Learn-To-be-Efficient（LTE），旨在训练效率感知的LLMs，使其激活更少的神经元，从而在稀疏性与性能之间实现更好的平衡。与主要针对ReLU模型的SOTA MoEfication方法不同，LTE也适用于使用非ReLU激活的LLMs，如LLaMA。广泛的评估表明，LTE在语言理解、生成和指令调优任务上均优于SOTA基线，并且通过硬件感知的自定义内核实现，LTE在50%稀疏性下将LLaMA2-7B的推理延迟减少了25%。

🔬 方法详解

问题定义：本文旨在解决大语言模型在推理过程中高开销的问题。现有方法主要在训练后利用激活稀疏性，未能在训练阶段进一步优化稀疏性，导致效率不足。

核心思路：论文提出的LTE算法通过在训练过程中引导模型学习结构化的激活稀疏性，使模型能够在推理时激活更少的神经元，从而提高推理效率。

技术框架：LTE算法的整体架构包括数据预处理、模型训练和推理优化三个主要阶段。在训练阶段，模型通过特定的损失函数学习激活稀疏性；在推理阶段，利用硬件感知的自定义内核实现高效推理。

关键创新：LTE的核心创新在于其能够在非ReLU激活的LLMs中应用，而不仅限于传统的ReLU模型。这一特性使得LTE在更广泛的模型架构中具有适用性，显著提升了稀疏性利用的灵活性。

关键设计：在参数设置上，LTE采用了特定的损失函数来平衡稀疏性与性能，同时设计了适应不同激活函数的网络结构，以确保模型在不同任务中的有效性。通过这些设计，LTE能够在保持模型性能的同时，显著提高推理效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LTE在语言理解、生成和指令调优任务上均优于现有的SOTA基线，尤其是在50%稀疏性下，LLaMA2-7B的推理延迟减少了25%。这一显著提升展示了LTE在实际应用中的潜力和价值。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和智能助手等。通过提高大语言模型的推理效率，LTE算法能够在资源受限的环境中实现更快速的响应，提升用户体验。此外，随着模型规模的不断扩大，优化推理效率将对实际应用产生深远影响。

📄 摘要（原文）

Large Language Models (LLMs) have achieved remarkable success with their billion-level parameters, yet they incur high inference overheads. The emergence of activation sparsity in LLMs provides a natural approach to reduce this cost by involving only parts of the parameters for inference. However, existing methods only focus on utilizing this naturally formed activation sparsity in a post-training setting, overlooking the potential for further amplifying this inherent sparsity. In this paper, we hypothesize that LLMs can learn to be efficient by achieving more structured activation sparsity. To achieve this, we introduce a novel training algorithm, Learn-To-be-Efficient (LTE), designed to train efficiency-aware LLMs to learn to activate fewer neurons and achieve a better trade-off between sparsity and performance. Furthermore, unlike SOTA MoEfication methods, which mainly focus on ReLU-based models, LTE can also be applied to LLMs like LLaMA using non-ReLU activations. Extensive evaluation on language understanding, language generation, and instruction tuning tasks show that LTE consistently outperforms SOTA baselines. Along with our hardware-aware custom kernel implementation, LTE reduces LLaMA2-7B inference latency by 25% at 50% sparsity.