SpikingBrain Technical Report: Spiking Brain-inspired Large Models

主流的基于Transformer的大语言模型面临着严重的效率瓶颈：训练计算量随着序列长度呈平方级增长，而推理时内存则线性增长，这限制了对长上下文的处理能力。在非NVIDIA平台上构建大模型也对稳定且高效的训练提出了挑战。为此，我们提出了SpikingBrain，这是一系列受大脑启发的模型，专为高效进行长上下文训练和推理而设计。SpikingBrain依托MetaX GPU集群，从三个方面展开工作：（1）模型架构：采用线性注意力和混合线性注意力架构，并引入自适应脉冲神经元；（2）算法优化：构建高效的基于转换的训练流程，并设计专用的脉冲编码框架；（3）系统工程：开发了针对MetaX硬件定制的训练框架、算子库以及并行策略。 通过这些技术，我们构建了两个模型：SpikingBrain-7B，一个线性注意力的大语言模型，以及SpikingBrain-76B，一个混合线性注意力的专家混合（MoE）大语言模型。这些模型验证了在非NVIDIA平台上进行大规模大语言模型开发的可行性。SpikingBrain在仅使用约1500亿token进行持续预训练的情况下，即可实现与开源Transformer基线模型相当的性能。我们的模型显著提升了长序列训练的效率，并实现了（部分）恒定内存占用和事件驱动的脉冲推理行为。例如，对于400万token的长序列，SpikingBrain-7B在“首字生成时间”（Time to First Token）上实现了超过100倍的加速。在数百块MetaX C550 GPU上，训练可以稳定运行数周，其中7B模型达到了23.4%的模型FLOPs利用率。所提出的脉冲机制实现了69.15%的稀疏性，从而支持低功耗运行。总体而言，本研究展示了受大脑启发的机制在推动下一代高效、可扩展的大模型设计方面的巨大潜力。

主流基于Transformer的大语言模型在训练和推理效率上面临瓶颈，训练计算随序列长度平方增长，推理内存线性增长，限制了长上下文处理。此外，在非NVIDIA平台上构建大模型也面临稳定性和效率挑战。

论文提出SpikingBrain，一种受脑科学启发的模型架构，结合线性与混合线性注意力机制、转换式训练流程、专用脉冲编码框架，以及针对MetaX硬件的系统级优化，实现高效的长序列训练与推理。

1. 开发了两个模型：SpikingBrain-7B（线性LLM）和SpikingBrain-76B（混合线性MoE LLM） 2. SpikingBrain-7B在4M token序列上实现了超过100倍的首字生成速度提升 3. 训练在数百块MetaX C550 GPU上稳定运行数周，7B模型达到23.4%的FLOPs利用率 4. 提出脉冲机制，达到69.15%稀疏性，支持低功耗运行 5. 仅使用约1500亿token进行持续预训练即可达到与开源Transformer基线相当的性能

1. Efficient Training of Very Long Sequence Models Using Sparsity Pattern Locking 2. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Aware Algorithms 3. LongNet: Long Sequence Modeling with Sparse Attention and Homogeneous Sequence Operations 4. Mixture-of-Experts (MoE) Architectures for Large Language Models 5. Sparse Transformers and Linear Transformers for Efficient Sequence Modeling