提出Spark Transformer,让模型在保持性能的同时实现高效稀疏计算。
Spark Transformer: Reactivating Sparsity in FFN and Attention
- 用top-k掩码控制稀疏度,显式调节激活比例。
- 仅8%前馈神经元激活,每令牌最多关注256个,FLOPs降低2.5倍。
- 无需额外参数,训练不慢,适合部署在CPU/GPU上加速推理。
训练后的Transformer中,前馈网络(FFN)绝大多数神经元对每个标记都处于非激活状态,这一‘懒惰神经元’现象激发了对激活稀疏性的研究,以提升大模型效率。尽管已有进展将稀疏性转化为实际运行时间优势,但现代Transformer已不再使用对稀疏性至关重要的ReLU激活函数。现有重引入稀疏性的方法常导致模型质量下降、参数量增加或训练变慢。将稀疏性应用于注意力机制(稀疏注意力)也面临类似挑战。本文提出Spark Transformer,一种新架构,在保留模型质量、参数量和标准训练流程的前提下,同时在FFN和注意力机制中实现高水平激活稀疏性。通过top-k掩码实现稀疏度的显式控制;关键创新是提出统计top-k算法——一种硬件友好的线性时间近似方法,避免昂贵排序操作,缓解传统top-k算子带来的显著训练延迟。此外,该模型重新分配现有的FFN参数和注意力键嵌入,构建低成本预测器以识别应激活的输出项。这一设计不仅缓解了强制稀疏带来的性能损失,还进一步提升了实际运行时间收益。基于Gemma-2预训练方案,Spark Transformer在标准基准测试中表现优异,其中只有8%的FFN神经元被激活,每令牌最多关注256个输入项。这种稀疏性使计算量减少2.5倍,推理时延在CPU上最高提速1.79倍,GPU上最高提速1.40倍。
原文摘要 · Abstract (English)
The discovery of the lazy neuron phenomenon in trained Transformers, where the vast majority of neurons in their feed-forward networks (FFN) are inactive for each token, has spurred tremendous interests in activation sparsity for enhancing large model efficiency. While notable progress has been made in translating such sparsity to wall-time benefits, modern Transformers have moved away from the ReLU activation function crucial to this phenomenon. Existing efforts on re-introducing activation sparsity often degrade model quality, increase parameter count, complicate or slow down training. Sparse attention, the application of sparse activation to the attention mechanism, often faces similar challenges. This paper introduces the Spark Transformer, a novel architecture that achieves a high level of activation sparsity in both FFN and the attention mechanism while maintaining model quality, parameter count, and standard training procedures. Our method realizes sparsity via top-k masking for explicit control over sparsity level. Crucially, we introduce statistical top-k, a hardware-accelerator-friendly, linear-time approximate algorithm that avoids costly sorting and mitigates significant training slowdown from standard top-$k$ operators. Furthermore, Spark Transformer reallocates existing FFN parameters and attention key embeddings to form a low-cost predictor for identifying activated entries. This design not only mitigates quality loss from enforced sparsity, but also enhances wall-time benefit. Pretrained with the Gemma-2 recipe, Spark Transformer demonstrates competitive performance on standard benchmarks while exhibiting significant sparsity: only 8% of FFN neurons are activated, and each token attends to a maximum of 256 tokens. This sparsity translates to a 2.5x reduction in FLOPs, leading to decoding wall-time speedups of up to 1.79x on CPU and 1.40x on GPU.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。