AIモデルとプラットフォーム
マンバ：シーケンスモデリングを再定義し、トランスフォーマーアーキテクチャを上回る
Unite.AI を Google の優先ソースに追加
この記事では、革新的な状態空間モデル（SSM）であるマンバについて探り、シーケンスモデリングをどのように革命的に変えるのかを見ていきます。Albert GuとTri Daoによって開発されたマンバは、言語処理、ゲノミクス、オーディオ分析などの分野で複雑なシーケンスを効率的に処理する能力で知られています。選択的な状態空間を使用したマンバの線形時間シーケンスモデリングは、多様なモダリティーで優れたパフォーマンスを発揮します。
マンバは、従来のトランスフォーマーに比べて、特に長いシーケンスで計算上の課題を克服する能力があります。マンバの選択的なアプローチは、状態空間モデルで高速な推論とシーケンス長との線形スケーリングを可能にし、スループットを大幅に改善します。
マンバのユニークさは、従来の注意とMLPブロックからの逸脱にあります。この簡素化により、モデルは軽量化され、シーケンス長と線形にスケーリングするようになり、先行するモデルに匹敵することができます。
トランスフォーマー vs マンバ
トランスフォーマー、たとえばGPT-4は、自然言語処理でベンチマークを設定しています。しかし、長いシーケンスではその効率が低下します。ここでマンバが優位性を発揮します。マンバは長いシーケンスをより効率的に処理し、プロセスを簡素化する独自のアーキテクチャを持っています。
トランスフォーマーは、シーケンスデータ（例：言語モデル用のテキスト）を処理することに特化しています。以前のモデルがデータを順次的に処理していたのとは異なり、トランスフォーマーはシーケンス全体を同時に処理し、データ内での複雑な関係を捉えることができます。
トランスフォーマーは、注意メカニズムを使用します。これにより、モデルはシーケンスの異なる部分に焦点を当てることができ、予測を行う際にそれらの重要性を考慮することができます。
この注意は、入力データから導かれた3つの重みセット（クエリ、キー、値）を使用して計算されます。シーケンス内の各要素は、他のすべての要素と比較され、シーケンス内の次の要素を予測する際に各要素に与えるべき「注意」の重みを提供します。
トランスフォーマーには、2つの主要なブロックがあります。入力データを処理するエンコーダーと、出力データを生成するデコーダーです。エンコーダーは、複数の層で構成されており、各層には2つのサブ層があります。マルチヘッドセルフアテンションメカニズムと、シンプルな位置依存の完全接続フィードフォワードネットワークです。各サブ層では、正規化と残差接続が使用され、深いネットワークのトレーニングを支援します。
デコーダーも、エンコーダーと同様の層とサブ層を持ちますが、エンコーダーの出力に対するマルチヘッドアテンションを実行する3番目のサブ層が追加されています。デコーダーの順序性は、ある位置の予測が前にある位置のみを考慮できることを保証し、自己回帰的特性を維持します。
一方、 マンバモデルは異なるアプローチを取ります。トランスフォーマーは長いシーケンスの問題に対処するためにより複雑な注意メカニズムを使用するのに対し、 マンバは選択的な状態空間を使用して、より計算効率の高いアプローチを提供します。
トランスフォーマーの基本的な動作の概要は以下のとおりです：
- 入力処理: トランスフォーマーは、モデルが理解できる形式に入力データをエンコードします。通常、各要素の位置も含む埋め込みを使用します。
- 注意メカニズム: 注意メカニズムの核となるのは、入力シーケンスの他の部分に焦点を当てるスコアを計算することです。
- エンコーダーとデコーダーのアーキテクチャ: トランスフォーマーモデルは、入力を処理するエンコーダーと、出力を生成するデコーダーで構成されています。各コンポーネントには、入力を理解するために複数の層があります。
- マルチヘッドアテンション: エンコーダーとデコーダーの両方に、マルチヘッドアテンションがあり、モデルはシーケンスの異なる部分から同時に情報を取得することができます。
- 位置依存のフィードフォワードネットワーク: 注意の後、シンプルなニューラルネットワークが各位置の出力を個別に処理し、入力と組み合わせて残差接続と層の正規化を行います。
- 出力生成: デコーダーは、エンコーダーのコンテキストと生成した出力に基づいて、出力シーケンスを予測します。
トランスフォーマーのシーケンスを並列に処理し、強力な注意メカニズムを備えることで、翻訳やテキスト生成などのタスクで強力なものとなります。
一方、マンバモデルはシーケンスを選択的な状態空間で処理することで異なります。このアプローチは、トランスフォーマーが長いシーケンスを扱う際の計算的非効率性を解決します。マンバの設計により、推論が速くなり、シーケンス長と線形にスケーリングするため、シーケンスモデリングの新たなパラダイムを実現し、特にシーケンスが長くなるときに効率的になります。
マンバ
マンバの特徴は、従来の注意とMLPブロックからの逸脱にあります。この簡素化により、モデルは軽量化され、シーケンス長と線形にスケーリングするようになり、先行するモデルに匹敵することができます。
マンバの主な特徴は以下のとおりです：
- 選択的なSSM: マンバは無関係な情報をフィルタリングし、関連するデータに焦点を当てることができます。これはシーケンスの処理を効率化するために不可欠です。
- ハードウェアに最適化されたアルゴリズム: マンバは、特にGPUを対象とした並列アルゴリズムを使用し、計算を高速化し、従来のモデルに比べてメモリ要件を削減します。
- 簡素化されたアーキテクチャ: 選択的なSSMと注意およびMLPブロックの統合により、マンバはよりシンプルで均一な構造を提供し、拡張性とパフォーマンスの向上につながります。
マンバは、言語、オーディオ、ゲノミクスを含むさまざまなドメインで優れたパフォーマンスを発揮しています。特に、事前トレーニングとドメイン固有のタスクで優れています。
マンバのコードと事前トレーニング済みモデルは、GitHubでコミュニティの利用に公開されています。
標準のコピー作業は線形モデルにとって簡単です。選択的なコピーと帰納的ヘッドは、LLMにとって動的でコンテンツに依存したメモリを必要とします。
構造化状態空間（S4）モデルは、最近、シーケンスモデルの一種として注目されています。これは、RNN、CNN、古典的な状態空間モデルからの特徴を包含しています。S4モデルは、連続システム、特に一次元の関数またはシーケンスを暗黙の潜在的な状態を通じてマッピングするシステムからインスピレーションを得ています。深層学習の文脈では、これらはシーケンスモデルを設計するための新しい方法論を提供し、効率的で適応性の高いものとなります。
S4モデルの動態
SSM（S4） これは基本的な構造化状態空間モデルです。シーケンス x を入力として、学習されたパラメータ A、B、C、および遅延パラメータ Δ を使用して出力 y を生成します。変換には、パラメータを離散化（連続関数を離散関数に変換）し、時間不変のSSM操作を適用します。
離散化の重要性
離散化は、連続パラメータを固定の式を通じて離散パラメータに変換するプロセスであり、S4モデルが連続時間システムとのつながりを維持することを可能にします。これにより、モデルは解像度不変性などの特性を持ち、適切な正規化を実現し、モデル安定性とパフォーマンスを向上させます。離散化は、RNNで情報の流れを管理するために重要なゲーティングメカニズムにも似ています。
線形時間不変性（LTI）
S4モデルの重要な特徴は、その線形時間不変性です。これは、モデルの動態が時間の経過とともに一貫して変わらないことを意味します。パラメータはすべてのタイムステップで固定されています。LTIは、再帰と畳み込みの基礎であり、シーケンスモデルを構築するためのシンプルながら強力なフレームワークを提供します。
基本的な限界を克服する
S4フレームワークは、従来、LTI性によって制限されていました。これは、適応的な動態を必要とするデータをモデル化する際に課題をもたらします。最近の研究論文では、時間変化パラメータを導入するアプローチを提示し、LTIの制限を解消します。これにより、S4モデルはより多様なシーケンスとタスクを処理できるようになり、適用可能性が大幅に拡大します。
「状態空間モデル」という用語は、広義には潜在的な状態を伴う再帰プロセスを指し、さまざまな分野でさまざまな概念を説明するために使用されてきました。深層学習の文脈では、S4モデルまたは構造化SSMは、計算効率を維持しながら複雑なシーケンスをモデル化するために最適化された、特定のモデルクラスを指します。
S4モデルは、エンドツーエンドのニューラルネットワークアーキテクチャに統合されて、スタンドアロンのシーケンス変換として機能します。CNNの畳み込み層と同様の役割を果たします。さまざまなニューラルネットワークアーキテクチャのシーケンスモデリングの基盤を提供します。
SSM vs SSM + 選択
シーケンスモデリングにおける選択性の動機
構造化SSM
論文では、シーケンスモデリングの基本的な側面は、コンテキストを管理可能な状態に圧縮することであると主張しています。入力またはフィルタリングを選択的に行うモデルは、コンテキストを効率的に維持するためのより効果的な手段を提供します。これは、言語モデリングを含む複雑なタスクで、シーケンスの次元に沿って情報の流れを適応的に制御する能力が不可欠であるため、非常に重要です。
選択的なSSMは、従来のSSMにパラメータの入力依存性を追加することで、以前には達成できなかった適応性を導入します。これにより、時間変化のSSMが可能になり、従来のモデルとは異なり、畳み込みではなく線形帰復メカニズムに依存するようになります。
SSM + 選択（S6） このバリアントには、選択メカニズムが含まれており、パラメータ B と C、および遅延パラメータ Δ に入力依存性が追加されます。これにより、モデルは入力シーケンス x の特定の部分に焦点を当てることができます。パラメータは選択を考慮して離散化され、SSM操作は時間依存的にスキャン操作を使用して適用され、時間の経過とともに焦点を動的に調整します。
マンバのパフォーマンスのハイライト
マンバはすべての評価結果で最高クラス
マンバのパフォーマンスは、推論速度と精度の両面で優れています。マンバの設計により、長いコンテキストの利用が改善され、DNAやオーディオモデリングを含む複雑なタスクで、長距離依存性を必要とするタスクで先行するモデルを上回ります。マンバの多様性は、複数のタスクに対するゼロショット評価でも実証されており、効率性と拡張性の新しい基準を設定しています。
マンバの開始
マンバを利用したい場合は、Linux OS、NVIDIA GPU、PyTorch 1.12+、およびCUDA 11.6+が必要です。インストールには、Mambaリポジトリから必要なパッケージをインストールするためのシンプルなpipコマンドが含まれます。PyTorchのバージョンに互換性の問題が発生した場合は、pipで–no-build-isolationフラグを使用することで解決できます。マンバは、PileやSlimPajamaデータセットなどの大規模なデータセットでトレーニングされたモデルを提供し、さまざまな計算要件とパフォーマンスベンチマークに対応しています。
マンバは、選択的なSSM層からマンバブロック、完全な言語モデル構造まで、さまざまなインターフェイスを提供します。マンバブロックは、因果Conv1d層を使用し、ニューラルネットワーク設計に簡単に統合できます。Pythonの使用例では、マンバモデルをインスタンス化し、データを処理する方法を示しています。これは、システムのシンプルさと柔軟性を強調しています。
事前トレーニング済みのマンバモデルは、Hugging Faceで利用可能です。パラメータ数は130Mから2.8Bまであり、PileデータセットとSlimPajamaデータセットでトレーニングされています。これらのモデルは、GPT-3の次元基準に従って、さまざまな計算要件とパフォーマンスベンチマークに対応するように設計されています。ユーザーは、これらのモデルから高いスループットと精度を期待できます。これにより、マンバは、言語モデリングを含むさまざまなアプリケーションで競争力のある選択肢となります。
マンバの影響
マンバはシーケンスモデリングにおける大きな進歩を表し、情報密度の高いデータを処理するためのトランスフォーマーアーキテクチャの強力な代替手段を提供します。マンバの設計は、現代のハードウェア、特にGPUの要件に合致しており、メモリ使用量と並列処理能力の両方で最適化されています。マンバのコードベースと事前トレーニング済みモデルのオープンソース化により、AIと深層学習の分野の研究者や開発者にとって、利用しやすく堅牢なツールとなっています。
