Spark Transformer: Reactivating Sparsity in FFN and Attention

2025年06月07日
  • 简介
    在训练好的 Transformer 中发现的“懒惰神经元”现象表明,对于每个 token,其前馈网络(FFN)中的绝大多数神经元都处于非激活状态。这一发现引发了对激活稀疏性研究的极大兴趣,以提升大规模模型的效率。尽管在将这种稀疏性转化为实际运行时间收益方面已取得显著进展,但现代 Transformer 已逐渐远离了对该现象至关重要的 ReLU 激活函数。重新引入激活稀疏性的现有方法往往会导致模型质量下降、参数数量增加,或者使训练过程复杂化甚至变慢。而稀疏注意力(即将稀疏激活应用于注意力机制)通常也面临类似的挑战。 本文提出了 Spark Transformer,这是一种新颖的架构,能够在 FFN 和注意力机制中实现高水平的激活稀疏性,同时保持模型质量、参数数量和标准训练流程不变。我们的方法通过 top-k 掩码实现稀疏性,从而对稀疏性水平进行显式控制。关键在于,我们引入了统计 top-k,这是一种硬件加速器友好的、线性时间近似算法,能够避免昂贵的排序操作,并缓解标准 top-k 操作带来的显著训练减速问题。此外,Spark Transformer 重新分配现有的 FFN 参数和注意力键嵌入,构建了一个低成本的预测器,用于识别激活项。这一设计不仅减轻了强制稀疏性带来的质量损失,还进一步提升了运行时间收益。通过 Gemma-2 预训练方案,Spark Transformer 在标准基准测试中表现出竞争力,同时展现出显著的稀疏性:仅 8% 的 FFN 神经元被激活,且每个 token 最多只关注 256 个 token。这种稀疏性使得浮点运算量(FLOPs)减少了 2.5 倍,从而在 CPU 上实现了高达 1.79 倍的解码运行时间加速,在 GPU 上则达到了 1.40 倍的加速效果。
  • 图表
  • 解决问题
    本论文试图解决如何在现代Transformer架构中实现高激活稀疏性(activation sparsity),从而显著减少计算量并加速推理过程,同时不牺牲模型质量、参数数量或标准训练流程。这是一个重要但尚未完全解决的问题,尤其是在ReLU激活函数逐渐被替代的背景下。
  • 关键思路
    论文提出了一种名为Spark Transformer的新架构,通过引入统计top-k算法来近似实现线性时间复杂度的稀疏性控制,避免了传统top-k操作带来的排序开销。此外,该方法重新分配FFN和注意力机制中的参数以构建低成本预测器,用于选择激活项,从而缓解稀疏性对模型性能的影响。相比现有方法,Spark Transformer不仅实现了更高的稀疏性,还保持了模型质量和效率。
  • 其它亮点
    1. Spark Transformer在前馈网络(FFN)中仅激活8%的神经元,并将每个token的注意力范围限制为256个token,显著减少了FLOPs(2.5倍降低)。2. 在CPU和GPU上分别实现了1.79倍和1.40倍的解码加速。3. 使用Gemma-2预训练方案,在多个基准测试中表现出与密集模型相当的性能。4. 提出了硬件友好的统计top-k算法,解决了稀疏化训练速度慢的问题。5. 论文未提及是否开源代码,但其设计思路值得进一步研究,例如探索更高效的稀疏预测器或扩展到多模态任务。
  • 相关研究
    近期相关工作包括:1.《Sparse Transformers》提出了稀疏注意力机制,但在大规模数据上的表现有限;2.《Routing Transformers》引入动态路由策略以减少计算量,但可能增加训练复杂度;3.《Prime: Prime Number Theorem Inspired Sparsity for Speeding-up Transformers》利用素数理论实现稀疏性优化;4.《Efficient Sparse Transformer Architectures for Natural Language Understanding》探讨了不同稀疏模式对语言理解任务的影响。这些研究均旨在提高Transformer的效率,但Spark Transformer通过参数重分配和统计top-k算法提供了一种新的解决方案。
PDF
原文
点赞 收藏 评论 分享到Link

沙发等你来抢

去评论