Yuqi Pan Yupeng Feng Jinghao Zhuang Siyu Ding et al

要約

主流のTransformerベースの大規模言語モデルは、大きな効率的ボトルネックに直面している。具体的には、学習時の計算量がシーケンス長に二次的に増加し、推論時のメモリ使用量が線形に増加するため、長文脈処理が制限されている。また、NVIDIA以外のプラットフォーム上で大規模モデルを構築する際も、安定かつ効率的な学習が困難である。こうした課題に対応するため、本研究では、効率的な長文脈学習および推論を実現するための脳にインスパイアされたモデル群「SpikingBrain」を提案する。SpikingBrainはMetaX GPUクラスタを活用し、以下の3つの側面に焦点を当てている。（1）モデルアーキテクチャ：適応型スパイクニューロンを備えた線形およびハイブリッド線形アテンション構造；（2）アルゴリズム的最適化：効率的な変換ベースの学習パイプラインと専用のスパイク符号化フレームワーク；（3）システム工学：MetaXハードウェアに最適化されたカスタム学習フレームワーク、オペレータライブラリ、並列化戦略。これらの技術を用いて、2つのモデルを構築した。SpikingBrain-7B（線形LLM）とSpikingBrain-76B（ハイブリッド線形MoE LLM）である。これらのモデルは、NVIDIA以外のプラットフォーム上での大規模LLM開発の実現可能性を示している。SpikingBrainは、オープンソースのTransformerベースのモデルと同等の性能を発揮しつつ、継続的プレトレーニングに約1500億トークンしか使用していない。本モデルは長シーケンス学習の効率を大幅に向上させ、（部分的に）定常的なメモリ使用とイベント駆動型のスパイク行動による推論を実現している。例えば、SpikingBrain-7Bは400万トークンのシーケンスに対して、最初のトークン生成までの時間（Time to First Token）で100倍以上の高速化を達成した。数百台のMetaX C550 GPU上で数週間にわたり安定した学習が可能であり、7BモデルではモデルFLOPs利用率が23.4％に達した。提案するスパイク機構は69.15％のスパース性を達成し、低消費電力動作を可能にしている。総じて、本研究は、脳にインスパイアされたメカニズムが、次世代の効率的かつスケーラブルな大規模モデル設計を推進する可能性を示している。