Artificial Intelligence

BlackMamba: 状態空間モデルの専門家の混合

公開済み

1ヶ月前

2024 年 3 月 26 日

デコーダーのみのトランスフォーマーモデルから構築された大規模言語モデル (LLM) の開発は、自然言語処理 (NLP) ドメインの変革において重要な役割を果たし、また、強化学習、時系列分析、画像処理など。ただし、スケーラビリティと強力なパフォーマンスにもかかわらず、デコーダのみのトランスフォーマモデルから構築された LLM は依然として重大な欠点に直面しています。表現力豊かではありますが、トランスフォーマー由来の LLM のアテンションメカニズムは、推論とトレーニングの両方で大量の計算リソースを必要とし、シーケンス長と二次 FLOP 用に大量のメモリを必要とします。この高い計算要件により、トランスフォーマーモデルのコンテキストの長さが制限され、自己回帰生成タスクが規模に比例して高価になり、連続データストリームからの学習や真に無制限のシーケンス処理の機能が妨げられます。

最近では、状態空間モデル (SSM) は、シーケンス長と線形時間の関数としてメモリの複雑さを達成しながら、大規模モデリングベンチマークでトランスアーキテクチャモデルと競合する、驚くべき機能とパフォーマンスを実証しました。さらに、最近リリースされた状態空間モデルである Mamba は、さまざまな言語モデリングや長時間シーケンスの処理タスクにおいて優れたパフォーマンスを示しています。同時に、Mixture of Expert (MoE) モデルも、より大きなメモリフットプリントを犠牲にしたにもかかわらず、推論のレイテンシと計算コストを大幅に削減しながら、優れたパフォーマンスを示しました。この記事では、Mamba モデルと MoE モデルに基づいて、Mamba 状態宇宙モデルと MoE モデルを組み合わせて、両方のフレームワークが提供する利点を活用する新しいアーキテクチャである BlackMamba について説明します。 BlackMamba の実験では、トレーニング FLOP と推論の両方において、既存の Mamba フレームワークとトランスフォーマーベースラインを上回るパフォーマンスを発揮する能力が実証されました。 BlackMamba フレームワークの優れたパフォーマンスは、Mamba フレームワークと MoE フレームワークの機能を効果的に組み合わせて、MoE からの高速かつコスト効率の高い推論と、Mamba からの線形複雑さの生成を提供できることを示しています。

この記事は、BlackMamba フレームワークについて詳しく説明することを目的としています。このフレームワークのメカニズム、方法論、アーキテクチャを、最先端の画像およびビデオ生成フレームワークとの比較とともに探ります。始めましょう。

BlackMamba : 状態空間モデルの MoE の概要

大規模言語モデル (LLM) の進歩、特にデコーダのみのトランスフォーマーアーキテクチャに基づくものは、顕著な影響を与えています。自然言語処理（NLP）強化学習、時系列分析、画像処理などを含むさまざまな深層学習アプリケーションに拡張されました。それにもかかわらず、これらのデコーダのみのトランスベースの LLM は、そのスケーラビリティと堅牢なパフォーマンスにもかかわらず、顕著な課題に直面しています。トランスベースの重要な機能であるアテンションメカニズム LLMs、推論とトレーニングの両方に大量の計算リソースを必要とします。これには、シーケンスの長さに応じて増加するメモリの必要性と、二次関数的に増加する計算演算 (FLOP) が必要になります。このような集中的な計算の必要性により、モデルのコンテキストの長さが制限され、モデルの規模が拡大するにつれて自己回帰生成タスクのコストが上昇し、連続データストリームから学習したり、無制限の長さのシーケンスを効率的に処理したりするモデルの能力が妨げられます。

これらの制限を克服するためにここ数年で多大な努力が払われ、標準的なデンスアテンショントランスフォーマーモデルに代わるアーキテクチャの考案に注目が移されており、SSM と MoE モデルが最も有望な候補アーキテクチャとなっています。トランスフォーマーアーキテクチャモデルよりも状態空間モデルを優先することで得られる主な利点は、トランスフォーマーによって提供される 2 次の複雑さとは対照的に、SSM によって提供される入力シーケンスの長さに関する線形の計算複雑さです。理論的には、入力シーケンスの長さに対する線形の計算複雑性により、状態空間モデルは、指定された FLOPS または 1 秒あたりの浮動小数点演算のバジェットについて、トランスフォーマーアーキテクチャモデルよりも大きなシーケンスを処理し、KV キャッシュを使用せずに計算で自己回帰生成定数をレンダリングすることができます。 Mamba、RetNet、その他いくつかを含む最近開発された状態空間モデルは、同様のスケーリング特性を持つトランスフォーマーに対する競合言語モデリングタスクのパフォーマンスとともに、効率的な長期シーケンス推論とトレーニングを実証しました。一方、エキスパートモデルの混合アーキテクチャは、密なモデルと同等の品質を達成するために不可欠な推論とトレーニングの FLOP を大幅に削減できるため、密なトランスフォーマーの代替として人気が高まっています。 MoE (Mixture of Experts) モデルは、単一の順方向パス中に全パラメーターのまばらな選択のみをアクティブにすることによって動作します。彼らはルーティング機能を利用して、与えられたコンテキストに基づいてどの「専門家」が行動を起こすよう求められるかを決定します。このアプローチにより、推論の計算コストとパラメータの総数が分離され、パラメータの数とメモリ要件が増加しても、固定の推論予算内でパフォーマンスの向上が可能になります。

このアーキテクチャの進歩は、従来の変圧器に比べて顕著な利点をもたらし、さらなる開発へのエキサイティングな方向性を示しています。これらの機能強化を結合された Mamba-MoE モデルに統合すると、言語モデリング機能と効率が標準のトランスフォーマーモデルを超えて大幅に加速される可能性があると考えられます。従来の高密度トランスフォーマーモデルと比較した Mamba-MoE アーキテクチャの予想される利点は次のとおりです。

マンバ： トレーニングフェーズと推論フェーズの両方で、入力シーケンスの長さに応じて線形の計算複雑さを実現します。これにより、一定の時間枠および一定のメモリ使用量で自己回帰生成を実行できます。

教育省: より高密度のバージョンと同等のパラメーター数を持つモデルに匹敵するモデル品質レベルを維持しながら、より小規模で高密度のベースラインモデルと同等の推論速度とトレーニング計算効率を提供します。

そうは言っても、トランスフォーマーアーキテクチャモデルは依然として最先端であり、言語モデリングタスクとシーケンス処理タスクで一貫した驚くべき強力なパフォーマンスを実証していることを述べておくことが重要です。トランスフォーマーアーキテクチャの中核では、シーケンス内の異なるトークンの埋め込み間の内積類似性の二次的な全対全比較を実行し、出力ベクトルへの線形マップを実行する自己注意を採用しています。トランスフォーマーモデルは、MLP ブロックまたは多層パーセプトロンブロックの間に積み重ねられたセルフアテンションブロックで構成され、さらに特定の活性化関数を備えた 2 層 MLP で構成されます。

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

状態空間モデル

状態空間モデルは、入力シーケンスの長さに関して線形複雑度を持つシーケンスモデルのグループに属します。状態空間モデルのアーキテクチャは、注意ベースのアーキテクチャではなく、リカレントニューラルネットワークや畳み込みニューラルネットワークと一致しており、暗黙的な潜在空間を通じて 1 次元関数をマッピングする連続動的システムからインスピレーションを得ています。線形動的システムでは、連想スキャンまたは畳み込みスキャンを使用して並列計算を効率化します。実際のシナリオでは、状態空間モデルの反復的な性質が、GPU などの並列性の高い AI ハードウェアに依然として採用されている理由です。しかし、RWKV や RWKV などの SSM の出現により、マンバらは、並列スキャンカーネルを使用して反復操作を効率的に GPU にマッピングし、トランスフォーマーモデルによって達成される効率に匹敵する効率で新しいアーキテクチャのトレーニングを容易にしました。

トランスフォーマー内のシーケンスの長さに関連する固有の 2 次の複雑さは、非常に長いコンテキストにわたる推論と理解を妨げるよく知られた制限です。最近の技術革新では、コンテキストの長さを拡張するというアイデアが導入され、推論中にはるかに長いコンテキストに適用される前に、トランスフォーマーを実行可能なスケールでトレーニングできるようになりました。これらの進歩にもかかわらず、推論プロセスは依然として、特に Key-Value (KV) キャッシュの維持にかなりの量の計算リソースとメモリを必要とし、リソースを大量に消費する作業となっています。最近の研究努力は、アテンションメカニズムに見られるクエリ、キー、値 (QKV) 行列に似た、入力依存のゲートメカニズムを組み込むことによって、状態空間モデルの表現能力を強化することに焦点を当てています。

これらの取り組みは、状態空間再帰の本質的な線形進行を維持し、畳み込みまたは選択的スキャンプロセスによる効率的な実行を可能にすることを目的としています。このアプローチにより、実際のアプリケーションにおける変圧器との性能差が大幅に縮小されます。これらの進歩の中で、Mamba は先行研究の目的を反映する状態空間モデルとして際立っており、最大 2.8 億パラメータのスケールで変圧器に匹敵する優れたパフォーマンスレベルを示します。これは、状態空間モデル (SSM) 再帰の入力に入力依存のゲートを適用することでこれを実現し、同時に特注の選択スキャンカーネルの使用を通じて効率的な計算を保証します。

エキスパートモデルの混合

Mixture of Expert (MoE) モデルは、フォワードパス中にパラメーターを選択的にアクティブにすることで、推論コストと合計パラメーター数の分離を実現します。これらのモデルは、すべてのパラメーターを使用する代わりに、特定のマルチレイヤーパーセプトロン (MLP) エキスパートにトークンを送信します。理想的には、各エキスパートは、本質的にコンパクトなニューラルネットワークであるルーティングメカニズムを使用して、特定の種類の入力を処理するように調整され、各トークンに最適なエキスパートを決定します。このアプローチは、より高密度な構成で同等の数のパラメーターを使用してモデルの包括的な表現力を維持しながら、計算量を大幅に削減することを目的としています。通常、ルーターはトークンからエキスパートインデックスまでの線形層のマッピングであり、各エキスパートは単に標準のトランスフォーマーである多層パーセプトロンです。ただし、エキスパートの割り当て問題は微分不可能であり、エキスパートの混合モデルでは、ハードウェア効率を高めるために、異なるエキスパート間の負荷分散とトレーニングの安定性に苦労することが多いため、開発者はルーターの最適なトレーニング方法をまだ見つけていません。

アーキテクチャ

BlackMamba はその中核として、インターリーブされた MLP ブロックと残差ストリームに沿って順番に追加されるアテンションブロックで構成される標準トランスフォーマーモデルを採用しています。現在、Mixture of Expert モデルの大部分は、多層パーセプトロンブロックを配線されたエキスパート層に単純に置き換えています。一方、BlackMamba フレームワークは、トランスフォーマーの多層パーセプトロンブロックをルーティングされたエキスパート層に置き換えるだけでなく、アテンション層を Mamba 状態空間モデル層に置き換えます。 BlackMamba フレームワークのアーキテクチャを次の図に示します。

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

BlackMamba モデルは、カスタムデータセット上の 300 億を超えるトークンでトレーニングされ、エキスパートの多層パーセプトロンに SwiGLU アクティベーション関数を使用します。このフレームワークは 8 人の専門家でトレーニングします。この数は、開発者がメモリフットプリントとモデルの推論コストの間の適切なバランスとトレードオフであると判断した数値です。 BlackMamba フレームワークのトレーニングに使用されるカスタムデータセットは、Starcoder、SlimPajama、Pile などの既存のオープンソースデータセットの混合で構成されます。次の表は、BlackMamba フレームワークのトレーニングに使用される各データセットの重みを示しています。全体として、データセットには 1.8 兆 XNUMX 億のトークンがあります。

BlackMamba : 結果

Mamba と BlackMamba を公平に比較するために、開発者は同じトレーニングデータで同じトレーニングパラメーターを使用して両方のモデルをトレーニングしました。 BlackMamba フレームワークは、推論時の同一のフォワードパスモデルサイズおよび 1 秒あたりの浮動小数点演算のトレーニングにおいて、Mamba モデルとトランスフォーマーモデルの両方を上回るパフォーマンスを発揮できます。次の図は、最初の 1 トークンプロンプトから所定の長さのシーケンスを自己回帰的に生成するのにかかる時間をシーケンス長の関数として示しています。

さらに、Expert モデルと Mamba モデルの混合のレイテンシーの利点が BlackMamba フレームワークで組み合わされ、トランスフォーマーモデル、純粋な Mamba モデル、および MoE モデルと比較した場合、推論時間が大幅に高速化されます。さらに、BlackMamba フレームワークの推論の利点はシーケンスの長さに直接比例するため、BlackMamba は長いシーケンスの生成において非常に効果的になります。次に、次の図は、それぞれ 340 億 640 万個と XNUMX 億 XNUMX 万個のパラメータを持つ BlackMamba モデルに割り当てられたトークンの数を示しています。見てわかるように、BlackMamba モデルによって実装された改善された Sinkhorn アルゴリズムの結果として、レイヤーの大部分が高度な専門的バランスを示しています。

次の表は、さまざまなオープンソースの事前トレーニング済み言語モデルと比較した BlackMamba フレームワークの評価スコアを示しています。ご覧のとおり、BlackMamba フレームワークは、すべてのベースラインにわたって大部分のフレームワークと競合し、優れたパフォーマンスを発揮できます。さらに、BlackMamba を上回るパフォーマンスを示すモデルはパラメーターの数がかなり多く、パフォーマンスの差が最小限であることは注目に値します。これは、パラメーターが少ない BlackMamba フレームワークの能力を示しています。

最終的な考え

この記事では、Mamba State Space Model と Expert モデルの混合を組み合わせて、これら両方のフレームワークによってもたらされる利点を享受する新しいアーキテクチャである BlackMamba について説明しました。 BlackMamba の実験では、トレーニング FLOP と推論の両方において、既存の Mamba フレームワークとトランスフォーマーベースラインよりも優れたパフォーマンスを発揮することが実証されました。 BlackMamba フレームワークの卓越したパフォーマンスは、MoE からの安価で高速な推論と Mamba からの線形複雑さの生成を組み合わせているため、BlackMamba フレームワークが Mamba と MoE フレームワークの機能を非常にうまく継承して組み合わせることができることを示しています。 BlackMamba フレームワークのアーキテクチャが、トレーニング FLOP と推論コストの点で、強力にトレーニングされた大規模言語モデル、既存の Mamba フレームワーク、およびエキスパートの混合モデルをどのように上回るパフォーマンスを発揮できるかについて説明してきました。さらに、BlackMamba フレームワークは、エキスパートモデルの混合と Mamba フレームワークの両方から生成 FLOP と削減されたトレーニングも同時に継承します。