Spark Transformer：重新激活FFN和注意力中的稀疏性
机器学习
2025-10-24 v2 机器学习
摘要
发现训练 Transformer 中的惰性神经元现象后，引发了对激活稀疏性以提高大模型效率的浓厚兴趣。虽然已取得显著进展将稀疏性转化为实际性能优势，但现代 Transformer 已远离关键此现象的 ReLU 激活函数。现有重新引入激活稀疏性的尝试往往会降低模型质量、增加参数量、增加训练复杂度或减慢训练速度。稀疏注意力（将稀疏激活应用于注意力机制）也面临类似挑战。本文提出 Spark Transformer，一种新型架构，在保持模型质量、参数量和标准训练流程的同时，在FFN和注意力机制中实现高度激活稀疏性。我们通过 top-k 掩码实现稀疏性，显式控制稀疏程度。关键在于引入统计 top-k 方法，这是一种面向硬件加速器的、线性时间的近似算法，避免了昂贵的排序操作，并缓解了标准 top- 操作带来的显著训练延迟。此外，Spark Transformer 重新分配现有FFN 参数和注意力键嵌入，以构建低成本预测器用于识别激活条目。该设计不仅缓解了强制稀疏性导致的质量损失，还增强了实际性能优势。使用 Gemma-2 配方预训练，Spark Transformer 在标准基准测试中表现出竞争力，同时实现显著稀疏性：仅激活8%的FFN神经元，每个token最多注意256个token。这种稀疏性导致FLOPs减少2.5倍，CPU上解码性能提升最高可达1.79倍，GPU上提升最高可达1.40倍。
引用
@article{arxiv.2506.06644,
title = {Spark Transformer: Reactivating Sparsity in FFN and Attention},
author = {Chong You and Kan Wu and Zhipeng Jia and Lin Chen and Srinadh Bhojanapalli and Jiaxian Guo and Utku Evci and Jan Wassenberg and Praneeth Netrapalli and Jeremiah J. Willcock and Suvinay Subramanian and Felix Chern and Alek Andreev and Shreya Pathak and Felix Yu and Prateek Jain and David E. Culler and Henry M. Levy and Sanjiv Kumar},
journal= {arXiv preprint arXiv:2506.06644},
year = {2025}
}
备注
NeurIPS 2025