MITと英伟达が開発、長文処理速度14.7倍に飛躍!FlashMoBAがMoBAのハードルを突破
大規模言語モデル(LLM)の長文処理性能を14.7倍向上させる画期的な技術が、マサチューセッツ工科大学(MIT)の韓松教授率いるH.A.N.実験室と英伟达(NVIDIA)の共同研究チームによって開発された。この成果は、長文処理における計算コストの爆発的増加という根本的な課題を、ハードウェアに最適化されたソフトウェア設計によって解決した。 長文処理の最大の障壁は、従来の自己注意機構(self-attention)が序列長の2乗に比例して計算量が増加する点にある。2024年2月に月之暗面(Moonshot AI)が提案した「ブロック注意混合」(MoBA)は、この問題を理論的に解決する画期的なアーキテクチャとして注目された。MoBAは、長文を複数のデータブロックに分割し、各クエリに対して関連性の高い少数のブロックのみを動的に選択する「ルーター」を導入。これにより、計算量を二次から近似線形に削減する可能性を示した。 しかし、MoBAの実用化には大きな壁があった。特に、小サイズのブロックを扱う場合、ルーターの評価やメモリ管理のオーバーヘッドが計算効率を逆に低下させ、理論的な利点が実際の性能に反映されにくかった。MITと英伟达の研究チームは、この課題を解決するため、GPUに特化した新アーキテクチャ「FlashMoBA」を設計。その鍵は、ハードウェアの特性を意識した完全なCUDAカーネルの再設計にある。 FlashMoBAの主な革新は2点。第一に「FlashTopK」は、質心計算とtop-k選択を1つのパイプラインで行い、従来の巨大なスコア行列を生成せずに処理。これにより、メモリ使用量の削減と計算の高速化を両立。第二に「収集・致密化(Gather-and-Densify)」戦略では、分散した小ブロックをGPUの高速キャッシュに一括収集し、その後、GPUに最適な密集型演算に変換。これにより、不連続なメモリアクセスやスケジューリングの負荷を大幅に削減。 実験結果では、序列長64Kで原始MoBA実装比7.4倍の速度向上、メモリ使用量は6.1倍削減。特に、128K以上の長序列では原始MoBAがメモリオーバーフロー(OOM)を起こすが、FlashMoBAは512Kまで安定して処理可能。さらに、業界標準のFlashAttention-2と比較しても、長序列では最大14.7倍の加速を達成した。 性能向上に加え、モデルの精度も損なわれていない。小ブロック化により「注意の希薄化」問題を緩和し、長文における情報の焦点が明確になる。実験では、MoBAベースモデルがより高い言語モデリング精度と長文検索精度を達成し、従来の密集型注意機構と同等、あるいはそれを上回る結果を示した。 この成果は、AIモデルの実用化における長文処理の限界を大きく押し広げ、将来の超長文理解や知識ベースの生成システムに大きな影響を与えると期待されている。
