Mambaは、深層学習に焦点を当てたシーケンスモデリングのためのアーキテクチャであり、カーネギーメロン大学のAlbert Guとプリンストン大学のTri Daoという研究者によって開発された。これは、トランスフォーマーモデルの限界、特に長いシーケンスの処理における限界に対処するために導入され、Structured State Space sequence（S4）モデルに基づいて構築されている。このアーキテクチャは、人工知能コミュニティ内で、言語、音声、ゲノミクスに関わるタスクにおけるトランスフォーマーベースの設計の代替として注目を集めている。
Mambaの背後にある中心的な動機は、拡張されたコンテキストを扱う際のトランスフォーマーの計算非効率性に由来する。従来のトランスフォーマーは、シーケンス長に応じて二次的にスケールするアテンションメカニズムに依存しており、長距離依存関係に対してリソース集約的である。Mambaは、競争力のあるパフォーマンスを維持しながら、より効率的なソリューションを提供することを目指しており、機械学習研究における重要な発展として位置づけられている。
アーキテクチャ
Mambaは、Structured State Space sequenceモデル（S4）を組み込むことで、長いデータシーケンスの効果的な処理を可能にしている。S4は、連続時間、リカレント、畳み込みモデルの強みを組み合わせることで長い依存関係をモデル化し、不規則にサンプリングされたデータの処理、無制限のコンテキストの維持、トレーニングおよびテスト段階の両方での計算効率の維持を可能にする。
S4に基づき、Mambaは、入力に基づいて構造化状態空間モデル（SSM）パラメータを適応させる独自の選択メカニズムを導入している。このメカニズムにより、モデルはシーケンス内の関連情報に選択的に焦点を当て、関連性の低いデータを効果的にフィルタリングできる。モデルは、時間不変のフレームワークから時間変動のフレームワークへと移行し、これが計算と効率の両方に影響を与える。
この時間変動性から生じる計算上の課題に対処するため、Mambaは、GPUなどの現代のハードウェア上で効率的な計算を可能にするハードウェア認識アルゴリズムを採用している。このアルゴリズムは、カーネル融合、並列スキャン、再計算技術を使用し、メモリ集約型レイヤーでの拡張状態の実体化を回避する。これにより、パフォーマンスとメモリ使用量が最適化され、以前の方法と比較して長いシーケンスの処理において大幅に効率的なアーキテクチャが実現される。
さらに、Mambaは、SSM設計をMLPブロックと統合することでアーキテクチャを簡素化し、均質で合理化された構造を作り出している。この設計は、言語、音声、ゲノミクスを含むさまざまなデータタイプにわたる一般的なシーケンスモデリングをサポートし、トレーニングと推論の両方で効率を維持する。
バリアント
MoE-Mambaは、Mambaアーキテクチャと混合エキスパート（MoE）レイヤーを統合している。この組み合わせにより、より効率的な実装が可能になり、モデルはMambaと同等のパフォーマンスを2.2倍少ないトレーニングステップで達成し、トランスフォーマーに対するMambaの推論パフォーマンスの利点を維持できる。モデルの設計は、MambaとMoEレイヤーを交互に配置し、シーケンス全体のコンテキストを効率的に統合し、各トークンに最も関連するエキスパートを適用できるようにする。
応用と影響
Mambaの設計は、大規模言語モデルやその他のシーケンスベースのアプリケーションに影響を与えている。計算コストが線形にスケールする長いシーケンスを処理する能力は、文書分析、音声処理、ゲノムシーケンスモデリングなどのタスクにとって魅力的である。研究者は、生成AIシステムを含むさまざまな領域で、トランスフォーマーアーキテクチャの潜在的な代替としてMambaを探求してきた。
このアーキテクチャの選択的スキャンメカニズムとハードウェア認識実装は、状態空間モデルにおけるその後の研究に影響を与えている。これは、効率とパフォーマンスの観点からトランスフォーマーベースのアプローチと比較されており、長いコンテキストのシナリオでの利点を強調する研究がある。最近の開発として、Mambaは、その強みを他の技術と組み合わせたさらなるバリアントやハイブリッドモデルに影響を与えている。
トランスフォーマーとの比較
Mambaは、シーケンスモデリングへのアプローチにおいてトランスフォーマーと根本的に異なる。トランスフォーマーは、すべての位置を同時に処理するマルチヘッドアテンションメカニズムを使用する一方、Mambaは、シーケンスを順次処理するリカレントスタイルの状態空間定式化を採用している。この違いは、明確なトレードオフをもたらす。Mambaは、長いシーケンスに対してより高速な推論と一定のメモリ使用量を提供する一方、トランスフォーマーは、並列化可能なトレーニングと確立されたインフラストラクチャを提供する。
Mambaの選択メカニズムは、アテンションの重みがコンテキストに応じて変化するのと同様に、入力コンテンツに基づいて状態空間パラメータを動的に調整できる。しかし、Mambaは、アテンションの二次的な複雑さなしにこれを達成しており、リアルタイム処理や非常に長い入力を必要とするアプリケーションに特に適している。これらの特性により、Mambaは、ニューラルネットワークアーキテクチャの継続的な進化において注目すべき代替として位置づけられている。