Q-Sparse: All Large Language Models can be Fully Sparsely-Activated
作者: Hongyu Wang, Shuming Ma, Ruiping Wang, Furu Wei
分类: cs.CL, cs.LG
发布日期: 2024-07-15 (更新: 2024-07-24)
备注: Work in progress
💡 一句话要点
Q-Sparse：实现大语言模型全稀疏激活，提升推理效率
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 稀疏激活 大语言模型 推理效率 top-K 稀疏化 straight-through-estimator
📋 核心要点
- 现有大语言模型推理效率低，激活稠密是主要瓶颈之一。
- Q-Sparse 通过 top-K 稀疏化激活，结合 straight-through-estimator 训练，实现全稀疏激活。
- 实验表明，Q-Sparse 在多种场景下均有效，且与 BitNet 结合能显著提升效率。
📝 摘要（中文）
我们提出了Q-Sparse，一种简单而有效的训练稀疏激活大语言模型（LLM）的方法。Q-Sparse 实现了 LLM 中激活的全稀疏性，从而在推理时带来显著的效率提升。这是通过对激活应用 top-K 稀疏化以及使用 straight-through-estimator 进行训练来实现的。我们还引入了 Block Q-Sparse 用于批量训练和推理。这项工作的主要成果是：（1）Q-Sparse 可以实现与基线 LLM 相当的结果，同时在推理时效率更高；（2）我们提出了稀疏激活 LLM 的推理最优缩放定律；（3）Q-Sparse 在不同的设置中都有效，包括从头开始训练、对现有 LLM 进行持续训练以及微调；（4）Q-Sparse 适用于全精度和 1-bit LLM（例如，BitNet b1.58）。特别是，BitNet b1.58 和 Q-Sparse（可以配备 MoE）的协同作用为未来 LLM 的效率（包括成本和能源消耗）的革命性发展奠定了基石，并提供了一条清晰的道路。
🔬 方法详解
问题定义：现有的大语言模型在推理时计算量巨大，其中激活的稠密性是导致效率低下的一个重要原因。传统的模型激活通常是稠密的，这意味着每一层的所有神经元都会参与计算，这带来了大量的冗余计算，尤其是在大规模模型中。因此，如何降低激活的计算量，提高推理效率，是当前大语言模型研究的一个重要挑战。
核心思路：Q-Sparse 的核心思路是通过稀疏化激活来减少计算量。具体来说，它只保留每一层中激活值最大的前 K 个神经元，而将其他神经元的激活值置为零。这样，在推理时只需要计算这些被激活的神经元，从而大大减少了计算量。这种方法的核心在于如何在训练过程中有效地学习到哪些神经元应该被激活，以及如何保证稀疏化后的模型性能不下降。
技术框架：Q-Sparse 的整体框架包括以下几个主要步骤：1) 前向传播：输入数据经过模型，计算每一层的激活值。2) 稀疏化：对每一层的激活值进行 top-K 稀疏化，只保留激活值最大的前 K 个神经元。3) 反向传播：使用 straight-through-estimator (STE) 来近似梯度，从而可以训练稀疏化的模型。4) 更新参数：使用梯度更新模型的参数。此外，论文还提出了 Block Q-Sparse，用于支持批量训练和推理，进一步提高效率。
关键创新：Q-Sparse 的关键创新在于它实现了一种简单而有效的全稀疏激活方法，并且能够与现有的 LLM 架构和训练方法兼容。与以往的稀疏化方法相比，Q-Sparse 更加简单易用，并且能够实现更高的稀疏度。此外，Q-Sparse 还提出了针对稀疏激活 LLM 的推理最优缩放定律，为未来的模型设计提供了理论指导。
关键设计：Q-Sparse 的关键设计包括以下几个方面：1) Top-K 稀疏化：选择合适的 K 值是关键，需要根据模型的规模和任务的复杂度进行调整。2) Straight-through-estimator：STE 的选择会影响训练的稳定性和收敛速度。3) Block Q-Sparse：通过将激活分成多个块，可以更好地利用硬件资源，提高批量训练和推理的效率。4) 损失函数：可以使用标准的交叉熵损失函数，也可以根据具体的任务进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Q-Sparse 能够在保持与基线 LLM 相当的性能的同时，显著提高推理效率。例如，在某些任务上，Q-Sparse 能够实现高达 2 倍的推理速度提升，并且能够与 BitNet b1.58 等 1-bit LLM 协同工作，进一步提高效率。此外，论文还验证了 Q-Sparse 在不同设置下的有效性，包括从头开始训练、对现有 LLM 进行持续训练以及微调。
🎯 应用场景
Q-Sparse 有潜力广泛应用于各种需要高效推理的大语言模型应用场景，例如移动设备上的自然语言处理、边缘计算环境下的智能助手、以及对延迟敏感的在线服务。通过降低计算成本和能源消耗，Q-Sparse 有助于推动大语言模型在资源受限环境下的普及和应用。
📄 摘要（原文）
We introduce, Q-Sparse, a simple yet effective approach to training sparsely-activated large language models (LLMs). Q-Sparse enables full sparsity of activations in LLMs which can bring significant efficiency gains in inference. This is achieved by applying top-K sparsification to the activations and the straight-through-estimator to the training. We also introduce Block Q-Sparse for batch training and inference. The key results from this work are, (1) Q-Sparse can achieve results comparable to those of baseline LLMs while being much more efficient at inference time; (2) We present an inference-optimal scaling law for sparsely-activated LLMs; (3) Q-Sparse is effective in different settings, including training-from-scratch, continue-training of off-the-shelf LLMs, and finetuning; (4) Q-Sparse works for both full-precision and 1-bit LLMs (e.g., BitNet b1.58). Particularly, the synergy of BitNet b1.58 and Q-Sparse (can be equipped with MoE) provides the cornerstone and a clear path to revolutionize the efficiency, including cost and energy consumption, of future LLMs.