大規模言語モデル（LLM）の開発、特にデコーダのみのトランスフォーマーモデルからの構築は、自然言語処理（NLP）分野に大きな影響を与え、強化学習、時系列分析、画像処理など、さまざまな深層学習アプリケーションに進出しています。しかし、スケーラビリティと強力なパフォーマンスにもかかわらず、デコーダのみのトランスフォーマーモデルからのLLMはまだ大きな欠点を持っています。トランスフォーマー派生のLLMの注意メカニズムは、トークンの埋め込み間のドット積の類似性の二乗比較を実行し、出力ベクトルへの線形マップを実行するため、計算リソースを大量に必要とします。これには、シーケンスの長さに比例してメモリが必要になり、計算操作（FLOPS）が二乗して増加します。このような大量の計算需要は、モデルのコンテキストの長さを制限し、スケールに応じてオートレグレッシブ生成タスクのコストを高くし、連続データストリームからの学習と無制限のシーケンス処理能力を妨げます。

最近、状態空間モデル（SSM）は、トランスフォーマーアーキテクチャモデルと競合する大規模モデリングベンチマークで優れた能力とパフォーマンスを示しています。また、Mambaなどの最近開発された状態空間モデルは、言語モデリングタスクと長シーケンス処理タスクでトランスフォーマーモデルと同等のパフォーマンスを発揮しています。同時に、Mixture of Expert（MoE）モデルも、推論とトレーニングの計算コストを大幅に削減しながら、デンシモデルと同等の品質を維持することで優れたパフォーマンスを示しています。MambaとMoEモデルを基に、BlackMambaという新しいアーキテクチャが開発され、両者の利点を活かすことができます。BlackMambaの実験結果は、既存のMambaフレームワークとトランスフォーマーベースラインを超えるパフォーマンスを示しています。

この記事では、BlackMambaフレームワークについて詳細に説明します。メカニズム、方法論、そしてアーキテクチャについて、最先端の画像とビデオ生成フレームワークとの比較も含めながら紹介します。はじめましょう。

BlackMamba：状態空間モデル用のMixture of Expertsの紹介

大規模言語モデル（LLM）の進歩、特にデコーダのみのトランスフォーマーアーキテクチャに基づくものは、自然言語処理（NLP）分野に大きな影響を与え、強化学習、時系列分析、画像処理など、さまざまな深層学習アプリケーションに進出しています。しかし、スケーラビリティと強力なパフォーマンスにもかかわらず、デコーダのみのトランスフォーマーモデルからのLLMはまだ大きな欠点を持っています。トランスフォーマー派生のLLMの注意メカニズムは、トークンの埋め込み間のドット積の類似性の二乗比較を実行し、出力ベクトルへの線形マップを実行するため、計算リソースを大量に必要とします。

これまでのところ、大きな努力が払われてきましたが、まだ大きな挑戦があります。状態空間モデル（SSM）とMixture of Expert（MoE）モデルは、最も有望なアーキテクチャ候補です。SSMのアーキテクチャは、トランスフォーマーアーキテクチャではなく、再帰型ニューラルネットワーク（RNN）や畳み込みニューラルネットワーク（CNN）に近いものです。また、連続動的システムからインスパイアされたもので、1次元関数を暗黙の潜在空間を通じてマッピングします。線形動的システムは、連想または畳み込みスキャンを使用して並列計算を効率的に行うことができます。実際のシナリオでは、状態空間モデルの再帰的な性質は、GPUなどの高並列AIハードウェアで採用されにくい原因となっています。しかし、最近開発された状態空間モデル、例えばRWKVやMambaは、並列スキャンカーネルを使用して、再帰操作をGPUに効率的にマッピングすることで、トランスフォーマーモデルと同等の効率でトレーニングを行うことができます。

トランスフォーマーの二乗計算複雑性は、シーケンスの長さに比例して計算コストが増加するため、長いコンテキストでの推論と理解を妨げる大きな制限です。最近の革新は、コンテキストの長さを拡張し、トランスフォーマーを実行可能なスケールでトレーニングし、長いコンテキストでの推論に適用できるようにしました。しかし、推論プロセスはまだ大量の計算リソースとメモリを必要とし、特にキー値（KV）キャッシュの維持に多くのリソースを必要とします。最近の研究は、状態空間モデルの表現能力を高めるために、入力依存のゲーティングメカニズムを導入することに焦点を当てています。これは、トランスフォーマーの注意メカニズムのQKVマトリックスに似ています。

BlackMamba：アーキテクチャと方法論

状態空間モデル

状態空間モデルは、シーケンスの長さに対して線形計算複雑性を持つシーケンスモデルの一種です。状態空間モデルのアーキテクチャは、トランスフォーマーアーキテクチャではなく、再帰型ニューラルネットワーク（RNN）や畳み込みニューラルネットワーク（CNN）に近いものです。また、連続動的システムからインスパイアされたもので、1次元関数を暗黙の潜在空間を通じてマッピングします。線形動的システムは、連想または畳み込みスキャンを使用して並列計算を効率的に行うことができます。

トランスフォーマーの二乗計算複雑性は、シーケンスの長さに比例して計算コストが増加するため、長いコンテキストでの推論と理解を妨げる大きな制限です。最近の革新は、コンテキストの長さを拡張し、トランスフォーマーを実行可能なスケールでトレーニングし、長いコンテキストでの推論に適用できるようにしました。しかし、推論プロセスはまだ大量の計算リソースとメモリを必要とし、特にキー値（KV）キャッシュの維持に多くのリソースを必要とします。最近の研究は、状態空間モデルの表現能力を高めるために、入力依存のゲーティングメカニズムを導入することに焦点を当てています。これは、トランスフォーマーの注意メカニズムのQKVマトリックスに似ています。

これらの努力は、状態空間モデルの表現能力を高め、トランスフォーマーと同等のパフォーマンスを実現することを目的としています。Mambaは、状態空間モデルの一種で、トランスフォーマーと同等のパフォーマンスを実現することで注目されています。Mambaは、入力依存のゲーティングを状態空間モデルの再帰に適用し、並列スキャンカーネルを使用して効率的な計算を実現しています。

Mixture of Expertモデル

Mixture of Expert（MoE）モデルは、推論コストとパラメータ数を分離することで、計算コストを削減します。MoEモデルは、トークンを特定のマルチレイヤーパーセプトロン（MLP）エキスパートにルーティングします。理想的には、各エキスパートは特定の入力タイプを処理するように設計されています。ルーティングメカニズムは、コンパクトなニューラルネットワークを使用して、各トークンに最も適したエキスパートを決定します。このアプローチは、デンシモデルと同等のパラメータ数を持つモデルと同等の表現能力を維持しながら、計算コストを大幅に削減することを目的としています。

アーキテクチャ

BlackMambaのアーキテクチャは、標準のトランスフォーマーモデルをベースにしています。BlackMambaは、トランスフォーマーのマルチレイヤーパーセプトロンブロックをルーティングされたエキスパートレイヤーに置き換え、さらにトランスフォーマーの注意ブロックをMamba状態空間モデルのレイヤーに置き換えています。BlackMambaのアーキテクチャは、以下の図に示すように、Mamba状態空間モデルのレイヤーとルーティングされたエキスパートレイヤーを組み合わせたものです。

トレーニングとデータセット

BlackMambaモデルは、300億トークン以上のカスタムデータセットでトレーニングされています。BlackMambaは、エキスパートのマルチレイヤーパーセプトロンにSwiGLUアクティベーション関数を使用しています。BlackMambaのトレーニングでは、8つのエキスパートが使用されています。これは、メモリフットプリントと推論コストのバランスをとるために選択された数です。BlackMambaのカスタムデータセットは、Starcoder、SlimPajama、Pileなどの既存のオープンソースデータセットの混合です。以下の表は、BlackMambaのトレーニングに使用されたデータセットの重みを示しています。合計で1.8兆トークンが含まれています。

BlackMamba：結果

BlackMambaとMambaの比較のため、両者は同じトレーニングパラメータとトレーニングデータでトレーニングされています。BlackMambaは、Mambaとトランスフォーマーモデルを超えるパフォーマンスを示しています。以下の図は、BlackMambaがシーケンスの長さに応じてオートレグレッシブに生成する時間を示しています。

さらに、BlackMambaは、Mixture of ExpertモデルとMambaモデルの推論コストの利点を組み合わせて、トランスフォーマーモデル、Mambaモデル、MoEモデルと比較して推論時間が大幅に短縮されています。BlackMambaの推論時間の利点は、シーケンスの長さに比例して増加します。以下の図は、BlackMambaモデルの340百万パラメータと640百万パラメータのエキスパートへのトークンの割り当てを示しています。多くの層でエキスパートのバランスが高いことがわかります。これは、BlackMambaモデルで改良されたシンクホーンアルゴリズムの結果です。

以下の表は、BlackMambaフレームワークの評価スコアを、さまざまなオープンソースの事前トレーニング済み言語モデルと比較しています。BlackMambaは、ほとんどのフレームワークを超える、または同等のパフォーマンスを示しています。さらに、BlackMambaを超えるモデルは、BlackMambaよりもはるかに多くのパラメータを持っています。パフォーマンスの差はわずかであり、BlackMambaが少ないパラメータで優れたパフォーマンスを示す能力があることを示しています。

最終的な考え

この記事では、Mamba状態空間モデルとMixture of Expertモデルを組み合わせたBlackMambaアーキテクチャについて説明しました。BlackMambaの実験結果は、既存のMambaフレームワークとトランスフォーマーベースラインを超えるパフォーマンスを示しています。BlackMambaの優れたパフォーマンスは、MambaとMoEの両者の能力を組み合わせることができ、MoEからの安価で高速な推論とMambaからの線形計算複雑性を実現していることを示しています。BlackMambaは、強力にトレーニングされた大規模言語モデル、既存のMambaフレームワーク、Mixture of Expertモデルを超えるトレーニングFLOPSと推論コストで優れたパフォーマンスを示しています。さらに、BlackMambaは、Mixture of ExpertモデルとMambaフレームワークの両者から、生成FLOPSと削減されたトレーニングを継承しています。