トランスフォーマーに挑むマンバの起源とは？著者の博士論文はSSMの進化の道筋を明らかにしている

大型模型の分野では、トランスフォーマーが全容を一手に引き受けています。しかし、モデルのサイズが拡大し、処理する必要のあるシーケンスが長くなるにつれて、Transformer の限界がますます明らかになります。たとえば、自己注意メカニズムの計算の複雑さは、コンテキストの長さの増加に伴って 2 乗的に増加します。これらの欠点を克服するために、研究者は注意メカニズムの効率的な変種を数多く開発してきましたが、成功は限られています。

最近、「Mamba」と呼ばれる研究がこの状況を打破したようで、言語モデリングにおいてTransformerに匹敵するか、あるいはそれを上回る可能性があります。これはすべて、著者が提案した新しいアーキテクチャ、つまり選択的状態空間モデルのおかげです。これは、Mamba 論文の著者である Albert Gu が以前主導した S4 アーキテクチャ (シーケンスモデリングのための構造化状態空間) を単純に一般化したものです。

Mamba 論文が発表された後、多くの研究者が SSM (状態空間モデル) や S4 などの関連研究に興味を持つようになりました。その中には、飛行機の中ですべての論文を読みたいと言った研究者もいました。この点に関して、アルバート・グーはより良い提案をしました。彼の博士論文は実際にこれらすべての展開を整理しており、読みやすく整理されているかもしれません。

論文の要約では、シーケンスモデルはディープラーニングモデルのバックボーンであり、科学的アプリケーションで幅広い成功を収めていると著者らは述べています。しかし、既存の方法では、さまざまなタスク、モダリティ、機能に対して広範な特殊化が必要であり、計算効率のボトルネックが発生し、より複雑なシーケンスデータをモデル化することが困難です (長い依存関係が関係する場合など)。したがって、一般的なシーケンスをモデル化するための原理的かつ実用的な方法を継続的に開発することが、依然として根本的に重要です。

論文リンク: https://stacks.stanford.edu/file/druid:mb976vf9362/gu_dissertation-augmented.pdf

この論文では、著者らは状態空間モデルを使用したディープシーケンスモデリングへの新しいアプローチを提示しています。これは、理論的根拠を持ち、計算効率が高く、さまざまなデータモダリティとアプリケーションにわたって優れた結果を達成する柔軟なアプローチです。

まず、著者らは、リカレントニューラルネットワークや畳み込みニューラルネットワークなどの標準的なディープシーケンスモデルの長所を一般化する、幅広い表現と特性を備えたモデルのクラスを紹介します。しかし、著者らは、これらのモデルの計算は困難を伴う可能性があることを示し、長いシーケンスにスケーリングする場合と自己回帰推論などの他の設定の両方で、現在のハードウェアで非常に高速に実行される新しい構造化状態空間を開発しています。最後に、彼らは連続信号の増分モデリングのための新しい数学的フレームワークを提案しています。このフレームワークは状態空間モデルと組み合わせることで、原理的な状態表現を付与し、長距離依存性をモデリングする能力を向上させることができます。要約すると、この新しいクラスのメソッドは、特に大規模な共通シーケンスデータを処理するための、機械学習モデルのための効率的で多用途なビルディングブロックを提供します。

以下は論文の各部分の簡単な紹介です。

ディープシーケンスモデル

シーケンスデータのディープラーニングモデルは、再帰、畳み込み、アテンションなどの単純なメカニズムに基づいて構築されたシーケンス間の変換として考えることができます。

これらのプリミティブは、標準的なディープニューラルネットワークアーキテクチャに組み込むことで、ディープシーケンスモデルの主なファミリ (再帰型ニューラルネットワーク (RNN)、畳み込みニューラルネットワーク (CNN)、および勾配降下法バックプロパゲーションなどの標準的なディープラーニング手法を使用して学習できる強力なパラメーター化された変換を表現するトランスフォーマー) を形成できます。図 1.1 と定義 1.1 は、この論文で使用されている抽象シーケンスモデルを示しており、セクション 2.1 では、例を挙げてより正式な定義を示します。

定義1.1（非公式）。著者らは、シーケンスモデルを使用して、シーケンス y = f_θ(x) 上のパラメーター化されたマッピングを参照します。ここで、入力と出力 x、y は、R^D の長さ L の特徴ベクトルのシーケンスであり、θ は勾配降下法によって学習されたパラメーターです。

これらの各モデルファミリは、機械学習に大きな成功をもたらしました。たとえば、RNN は機械翻訳にディープラーニングをもたらし、CNN は最初のニューラルオーディオ生成モデルであり、Transformer は NLP の広大な分野に革命をもたらしました。

ただし、これらのモデルには、順次メカニズムから生じるトレードオフも残っています。たとえば、RNN はシーケンシャルデータ用の自然なステートフルモデルであり、時間ステップごとに一定の計算/ストレージのみを必要としますが、トレーニングに時間がかかり、最適化の難しさ (勾配消失問題など) に悩まされるため、長いシーケンスを処理する能力が制限されます。 CNN はローカルコンテキストに焦点を当て、シフト等価性などの機能をエンコードし、高速で並列トレーニングが可能ですが、シーケンス推論コストが高く、コンテキストの長さが本質的に制限されています。トランスフォーマーは長距離の依存関係と並列処理を処理できるため大きな成功を収めていますが、シーケンスの長さの 2 次スケーリングの問題に悩まされています。最近のモデルのもう 1 つのファミリーは、ニューラル微分方程式 (NDE) です。これは、連続時間の問題と長期依存関係を理論的に解決できる、十分に根拠のある数学モデルですが、非常に非効率的です。

これらの問題は、ディープシーケンスモデルが直面している 3 つの大きな課題を示しています。

課題1: 一般的な能力

ディープラーニングの大きな目標は、さまざまな問題に使用できる一般的な構成要素を開発することです。シーケンスモデルは、このような多くの問題を解決するための一般的なフレームワークを提供します。これらは、シーケンスとして投影できるあらゆるコンテキストに適用できます。ただし、現在のモデルでは、特定のタスクやドメインの問題を解決したり、特定の機能をターゲットにしたりするために、依然として大量の特殊な機能が必要になることがよくあります。各モデルの利点は次のように分析されます。

RNN: オンライン処理タスクや強化学習など、隠れ状態の高速更新が必要なステートフル設定。
CNN: オーディオ、画像、ビデオなどの均一にサンプリングされた知覚信号をモデル化します。
トランスフォーマー: 言語などの領域における高密度で複雑な相互作用をモデル化します。
NDE: 欠損データや不規則にサンプリングされたデータなど、非典型的な時系列設定を処理します。

逆に、各モデルラインは、得意としていない機能で苦戦する可能性があります。

課題2: 計算効率

実際にディープシーケンスモデルを適用するには、それらを定義する関数を計算する必要があります (つまり、シーケンス間のマッピングをパラメーター化します)。これにはさまざまな形式があります。トレーニング中、タスクは一般に入力シーケンス全体の損失関数によって記述できます。アルゴリズムの中心的な問題は、フォワードパスを効率的に計算する方法です。推論時（トレーニング後のモデルの展開）には設定が変わることがあります。たとえば、オンライン処理や自己回帰生成設定では、入力は一度に 1 つのタイムステップで提示され、モデルはこれらの入力を効率的に順番に処理できる必要があります。

どちらの状況も、異なるモデルファミリに課題をもたらします。たとえば、RNN は本質的に順次的であり、GPU や TPU などの最新のハードウェアアクセラレータでトレーニングするのは困難ですが、並列処理のメリットを享受できます。一方、CNN とトランスフォーマーはステートフルではないため、効率的な自己回帰推論を行うのが困難です。単一の新しい入力を処理するコストは、モデルのコンテキストサイズ全体に応じて拡大する可能性があります。より特殊なモデルでは追加の機能が追加される場合もありますが、通常は計算がより困難になり、遅くなります (たとえば、高価な微分方程式ソルバーの呼び出しが必要になるなど)。

課題3: 長期依存

現実世界のシーケンスデータでは、数万の時間ステップにわたる推論が必要になる場合があります。長い入力を処理するために必要な計算上の問題に加えて、この問題を解決するには、このような長距離依存関係 (LRD) に存在する複雑な相互作用をモデル化する能力が必要です。具体的には、モデルのコンテキストウィンドウが限られているなど、データ内の相互作用をキャプチャできないことや、再帰型モデルで長い計算グラフを逆伝播するときに発生する勾配消失問題などの最適化の問題によって、困難が生じる可能性があります。

長距離依存性は、効率性、表現力、トレーニング能力の制限により、長い間、シーケンスモデルの課題となってきました。 NDE、RNN、CNN、Transformer などのすべての標準モデルファミリには、これらの問題に対処するために設計された多くの特殊なバリアントが含まれています。例としては、消失勾配に対抗する直交 RNN と Lipschitz RNN、コンテキストサイズを増やすアトラス畳み込み、シーケンスの長さへの二次依存性を減らす効率的な注意バリアントのますます拡大するファミリーなどがあります。ただし、これらのソリューションは長距離依存関係向けに設計されているにもかかわらず、Long Range Arena などの難しいベンチマークではパフォーマンスが依然として低くなります。

状態空間シーケンシャルモデル

この論文では、線形状態空間モデル (SSM) に基づく新しいディープシーケンスモデルファミリを紹介します。著者らは、SSMを、1次元関数またはシーケンスを暗黙の潜在状態x(t)∈R^Nにマッピングする単純なシーケンスモデルとして定義しています。

SSM は、サイバネティクス、計算神経科学、信号処理などの分野で広く使用されている基本的な科学モデルです。広い意味では、SSM という用語は、潜在変数が状態空間でどのように進化するかをモデル化するあらゆるモデルを指します。これらの一般化 SSM にはさまざまな種類があり、x の状態空間 (連続、離散、混合など)、y の観測空間、遷移ダイナミクス、加法ノイズプロセス、システムの線形性などを変更できます。 SSM は歴史的に、階層的ディリクレ過程 (HDP-HMM) やスイッチング線形動的システム (SLDS) など、隠れマルコフモデル (HMM) や線形動的システム (LDS) のバリエーションを指すことが多かった。

式(1.1)の状態空間モデルは状態空間とダイナミクスの両方で連続しており、完全に線形かつ決定論的であるが、定義1.1の意味でのディープシーケンスモデルとしてはまだ使用されていない。この論文では、状態空間シーケンスモデルの多くの利点と、それらを活用して、独自の制限を克服しながら一般的なシーケンスモデリングの課題に対処する方法について説明します。

一般的なシーケンスモデル

SSM は、多くの豊富なプロパティを備えたシンプルで基本的なモデルです。これらは、NDE、RNN、CNN などのモデルファミリと密接に関連しており、実際には複数の形式で記述して、通常は特殊なモデルを必要とするさまざまな機能を実現できます (課題 1)。

SSM は連続的です。 SSM 自体は微分方程式です。したがって、連続プロセスのシミュレーション、欠損データの処理、さまざまなサンプリングレートへの対応など、連続時間モデルの独自のアプリケーションを実行できます。
SSM は周期的です。 SSM は標準的な手法を使用して線形再帰に離散化でき、推論中に状態再帰モデルとしてシミュレートできます。この場合、時間ステップあたりのメモリと計算量は一定のままです。
SSM は畳み込みシステムです。 SSM は、連続畳み込みとして明示的に表現できる線形時間不変システムです。さらに、離散時間バージョンは、離散畳み込みを使用してトレーニングするときに並列化できるため、効率的なトレーニングが可能になります。

したがって、SSM は、並列環境とシーケンシャル環境の両方で、またさまざまなドメイン (オーディオ、ビジョン、時系列など) で効率的に動作する一般的なシーケンスモデルです。論文の第 2 章では、SSM の背景を紹介し、状態空間シーケンスモデルのこれらの特性について説明します。

ただし、SSM の汎用性にはコストも伴います。オリジナルの SSM には、おそらく他のモデルよりも深刻な 2 つの追加の課題がまだ存在しており、これがディープシーケンスモデルとしての使用を妨げています。課題としては、(1) SSM は一般に同じサイズの RNN や CNN よりもはるかに遅い、(2) 長い依存関係を記憶するのが難しい、たとえば RNN の勾配消失問題を継承している、などが挙げられます。

著者らは、SSM の新しいアルゴリズムと理論を通じてこれらの課題に取り組んでいます。

構造化 SSM を使用した効率的なコンピューティング (S4)

残念ながら、一般的な SSM は、状態表現 x(t)∈R^N の計算要件とメモリ要件が高いため、実際にはディープシーケンスモデルとして使用することはできません (課題 2)。

状態次元 N およびシーケンス長 L の SSM の場合、完全な潜在状態 x を計算するだけで O(N^2L) の演算と O(NL) の空間が必要になります。これは、全体の出力を計算する場合の下限値 Ω(L + N) と比較すると小さくなります。したがって、妥当なサイズのモデル (例: N ≈ 100) の場合、SSM は同等のサイズの RNN または CNN よりも数桁多くのメモリを使用するため、汎用シーケンスモデリングソリューションとしては計算上非実用的になります。

この計算上のボトルネックを克服するには、効率的なアルゴリズムに適した方法で状態行列 A に構造を課す必要があります。著者らは、構造行列 A のさまざまな形式を持つ構造化状態空間シーケンスモデル (S4) (または略して構造化状態空間) のファミリと、任意の表現 (再帰型や畳み込み型など) で S4 モデルを効率的に計算できる新しいアルゴリズムを紹介します。

論文の第 3 章では、これらの効率的な S4 モデルのさまざまなタイプを紹介しています。最初の構造は、状態行列の対角パラメータ化を使用します。これは、ほぼすべての SSM を表すのに十分なほど単純かつ汎用的です。次に、後で導入される SSM の特殊なクラスを捕捉するために必要な低ランク補正項を許可することで、これを一般化します。著者らは、生成関数、線形代数変換、構造化行列乗算の結果などの多数の技術的アイデアを組み合わせることで、時間計算量O(N + L) と空間計算量 O(N + L) の両方の構造に対するアルゴリズムを開発しており、これはシーケンスモデルに対して本質的に厳密です。

HIPPOを使用して長距離依存関係を解決する

計算上の問題を無視しても、基本的な SSM は実験でのパフォーマンスが依然として低く、長距離の依存関係をモデル化することができません (課題 3)。直感的には、1 つの説明として、線形 1 次 ODE は指数関数として解かれるため、勾配がシーケンスの長さに応じて指数関数的に増加するという問題が発生する可能性があります。これは線形回帰としての解釈にも見られ、これは回帰行列を繰り返し指数関数的に増加させることを含み、RNN におけるよく知られた勾配消失/爆発問題の原因です。

第 4 章では、著者らは SSM から一歩後退し、代わりに循環モデルを使用して第一原理から LRD をモデル化する方法を調査します。彼らは、オンライン関数近似（またはメモ化）と呼ばれる問題を形式化して解決する HIPPO と呼ばれる数学的フレームワークを開発しました。このアプローチは、連続関数の履歴を圧縮したままにして、連続関数を徐々に記憶することを目的としています。これらのアプローチは完全に独立した動機を持っていますが、すべて SSM の特定の形式です。これらの最終的な方法は、その動機は完全に独立しているものの、SSM の特定の形式であることが示されています。

論文の第 5 章では、このフレームワークを改良し、SSM 抽象化にさらに厳密に結び付けています。直交 SSM の概念を導入し、HIPPO を広く一般化し、すべての SSM パラメータを原則的に初期化する方法など、より多くの例と理論的結果を導き出します。

HIPPOの概要

入力関数u(t)、固定確率測度ω(t)、およびN個の直交基底関数（多項式など）のシーケンスを考えます。各時刻 t において、時刻 t より前の u の履歴をこの基底に投影することができ、その結果、提供されたメトリック ω に対する u の履歴の最良近似を表す係数ベクトル x(t)∈R^N が得られます。関数u(t)∈Rを係数x(t)∈R^Nに写像することを、計量ωに関する高次多項式射影演算子（HIPPO）と呼びます。多くの場合、(A, B) には x ′ (t) = Ax (t) + Bu (t) という形式の閉じた式が存在します。

HIPPOとS4の組み合わせ

HIPPO は重要な特性を持つ SSM を構築するための数学的ツールを提供しますが、S4 は計算表現に関するものです。第 6 章では、この 2 つを正式に関連付け、組み合わせることで両方の長所を活かすことができることを示します。この論文では、長距離依存性を処理するために HIPPO によって生成される特殊な行列が、実際には第 3 章で開発された特定の構造形式で記述できることを示しています。これは、S4 と HIPPO を組み合わせた具体的な例であり、機能が豊富で非常に効率的で、長距離推論に優れた一般的なシーケンスモデルを実現します。

適用、アブレーション、拡張

一般的なシーケンスモデリング機能

第 7 章では、さまざまなドメインとタスクにおける S4 アプローチの包括的な実証的検証を示します。 S4 メソッドを汎用のシンプルなディープニューラルネットワークに組み込むと、多くのベンチマークで SOTA が向上しました。

特別なハイライトと機能は次のとおりです:

一般的なシーケンスモデリング。アーキテクチャを変更することなく、S4 は音声分類ではオーディオ CNN よりも優れ、時系列予測問題では特殊な Informer モデルよりも優れ、シーケンス CIFAR では 90% を超える精度で 2-D ResNet に匹敵します。
長期にわたる依存関係。効率的なシーケンスモデルの LRA ベンチマークでは、S4 はすべてのベースラインと同等の速度でありながら、すべての Transformer バリアントよりも 25% 以上高い平均精度を達成しています。 S4 は、困難な LRA Path-X タスク (長さ 16384) を 96% の精度で解決した最初のモデルです。一方、これまでのすべての研究のランダム推測精度はわずか 50% です。
サンプリング解像度が変更されます。特殊な NDE 手法と同様に、S4 は再トレーニングを必要とせずに時系列のサンプリング頻度の変化に適応できます。
高速な自己回帰生成による大規模な生成モデリング。 CIFAR-10 密度推定の場合、S4 は最良の自己回帰モデル (次元あたり 2.85 ビット) と同等です。 WikiText-103 言語モデリングに関しては、S4 は Transformers とのギャップを大幅に縮小し (0.5 の困惑度以内)、注意のないモデルで SOTA を達成しました。 RNN と同様に、CIFAR-10/WikiText-103 では、S4 は標準の自己回帰モデルよりも 60 倍高速に潜在状態からピクセル/トークンを生成します。

理論的なアブレーション

著者の S4 の扱いでは、各パラメータを慎重に初期化する方法やそれを HIPPO フレームワークに組み込む方法など、SSM のトレーニングに関する多くの理論的な詳細について説明しています。彼らはこれらの詳細について包括的な経験的分析とアブレーション研究を実施し、SSM 理論のさまざまな側面を検証しました。たとえば、HIPPO によって SSM のモデリング機能が大幅に向上し、標準シーケンスモデルベンチマークで元の SSM インスタンスよりも 15% のパフォーマンス向上が達成されることが検証されました。アルゴリズム的には、S4 アルゴリズムは従来の SSM アルゴリズムを数桁改善します (たとえば、速度は 30 倍になり、メモリ使用量は 1/400 に削減されます)。

アプリケーション: オーディオ波形生成

多様な特性を持つシーケンスモデリングプリミティブとして、S4 はさまざまなニューラルネットワークアーキテクチャに組み込んで、さまざまな方法で使用できます。第 8 章では、オーディオ波形のサンプリングレートが高いために難しい問題となる、生のオーディオ波形生成への S4 の適用について説明します。この章では、S4 を中心に構築された SaShiMi マルチスケールアーキテクチャを紹介します。このアーキテクチャは、自己回帰や拡散を含む複数の生成モデリングパラダイムにわたる、制限のないオーディオおよび音声生成の最先端技術を進歩させます。このアプリケーションは、効率的なトレーニング、高速な自己回帰生成、連続信号をモデル化するための強力な誘導バイアスなど、S4 の柔軟な機能を強調しています。

拡張: コンピュータビジョンのための多次元信号

著者らは主に 1 次元のシーケンスに焦点を当てていますが、画像 (2D) やビデオ (3D) など、一部の形式のデータは本質的に高次元です。シーケンスモデルの柔軟性は、これらの環境にも適用されます。第 9 章では、S4 を 1 次元から多次元 (ND) 信号に拡張した S4ND について説明します。 S4ND は、基礎となる連続信号を直接モデル化するなどの S4 の機能を継承し、入力解像度の変更をより適切に処理するなどの関連する利点を備えています。これは、ImageNet などの大規模な視覚タスクで競争力のあるパフォーマンスを発揮する最初の連続モデルです。

詳細については原文論文を参照してください。

最後に、この機会に、皆様の理解と学習のために、いくつかの SSM 研究をレビューし、紹介したいと思います。

論文 1: 注意を怠った事前トレーニング