マンバの解説

トランスフォーマーに挑む状態空間モデル 現在、AIが世界を席巻している。 ここで言うAIとは、トランスフォーマーのことだ。ここ数年におけるAIの大きなブレークスルーのほぼすべては、トランスフォーマーによるものだ。 しかし、Mambaは状態空間モデル（SSM）と呼ばれる別のモデル群の一つである。重要なことに、Mambaは初めて、長いシーケンス長（例えば100万トークン）においても実行可能でありながら、トランスフォーマーと同等の性能（そして決定的に重要なスケーリング則）を約束している。この長いコンテキストを実現するために、Mambaの開発者たちはアテンション・メカニズムにおける「二次のボトルネック」を取り除いた。Mambaはまた高速に動作する - 「トランスフォーマーの最大5倍の速さ」で動作する。 Mambaの開発者であるGuとDaoはこう書いている: Mambaは高速な推論とシーケンス長に対する線形スケーリングを享受し、その性能は実データにおいて最大100万長のシーケンスまで向上する。一般的なシーケンスモデルのバックボーンとして、Mambaは言語、オーディオ、ゲノミクスなど、いくつかのモダリティにおいて最先端の性能を達成する。言語モデリングにおいて、我々のMamba-3Bモデルは、同じサイズのトランスフォーマーを上回り、その2倍のサイズのトランスフォーマーと同等の性能を、事前学習と下流評価の両方で達成する。 ここでは以下について議論する: Mamba (🐍) 対 トランスフォーマー (🤖) の利点（と欠点） Mambaを理解するための類推と直感 解釈可能性、AI安全性、応用にとってMambaが意味するもの トランスフォーマーの問題点 - もしかするとアテンションだけでは不十分かもしれない 我々はまさに歴史のトランスフォーマー時代にいる。機械学習はかつて猫や犬を検出することだった。今や、トランスフォーマーによって、人間のような詩を生成し、平均的な競技プログラマーよりも優れたコーディングを行い、タンパク質折り畳み問題を解決している。 しかし、トランスフォーマーには一つの核心的な問題がある。トランスフォーマーでは、各トークンは予測を行う際に、それ以前のすべてのトークンを振り返って見ることができる。この「振り返り」のために、我々は各トークンに関する詳細な情報をいわゆるKVキャッシュに保存する。 このペアワイズ（対ごとの）通信は、訓練時の順方向パスがO(n²)の時間計算量（恐るべき二次のボトルネック）であることを意味し、自己回帰的に生成される各新しいトークンはO(n)の時間を要する。言い換えれば、コンテキストサイズが増加するにつれて、モデルは遅くなる。 さらに悪いことに、このキー・バリュー（KV）キャッシュを保存するにはO(n)の空間が必要である。その結果、メモリ使用量が拡大するにつれて、恐るべきCUDAのメモリ不足（OOM）エラーが重大な脅威となる。もし空間だけが問題であれば、より多くのGPUを追加することを検討するかもしれない。しかし、遅延が二次関数的に増加するため、単に計算資源を追加するだけでは実行可能な解決策とはならないかもしれない。 限界的には、スライディング・ウィンドウ・アテンションのような技術や、FlashAttentionのような巧妙なCUDA最適化によって、二次のボトルネックを緩和することができる。しかし最終的に、超長いコンテキストウィンドウ（例えば、あなたが共有したすべての会話を記憶するチャットボット）には、別のアプローチが必要だ。 基盤モデルのバックボーン 根本的に、すべての優れた機械学習アーキテクチャのバックボーンは、二つの重要な操作のためのコンポーネントを持っている: トークン間の通信 トークン内の計算 トランスフォーマーでは、これはアテンション（通信）とMLP（計算）である。我々はこれら二つの操作を最適化することでトランスフォーマーを改良する。 我々はアテンション・コンポーネントを、トークン間通信を促進する代替メカニズムに置き換えたい。具体的には、Mambaは通信目的で制御理論に着想を得た状態空間モデル（SSM）を採用し、計算には多層パーセプトロン（MLP）型の射影を保持する。 積み重ねられたトランスフォーマー・ブロックで構成されるトランスフォーマーのように、Mambaは上記のような積み重ねられたMamba・ブロックで構成される。 我々はシーケンス変換におけるSSMの選択を理解し、動機付けたい。 Mambaの動機付け - Temple Runへの回帰 Temple Runエージェントを構築していると想像してみよう。それはランナーがいつ左右に移動すべきかを選択する。 正しい方向を選択するためには、周囲の情報が必要だ。関連する情報の集合を状態と呼ぼう。ここでの状態には、おそらく現在の位置と速度、最も近い障害物の位置、天候条件などが含まれるだろう。 主張1: もし現在の世界の状態と、世界がどのように進化しているかを知っていれば、これを用いて移動する方向を決定できる。 ここで、常に画面全体を見る必要はないことに注意してほしい。走っている間、障害物が画面下方に移動することに気づけば、画面の大部分に何が起こるかを把握できる。新しい情報を理解するには画面の上部だけを見ればよく、残りはシミュレートすればよい。 これは自然な定式化につながる。hを隠れ状態、世界に関する関連知識とする。また、xを入力、つまり毎回得られる観測とする。h'は隠れ状態の微分、すなわち状態がどのように進化しているかを表す。我々はy、最適な次の動き（右か左）を予測しようとしている。 さて、主張1は、隠れ状態h、h'、そして新しい観測xから、yを把握できると述べている。 より具体的には、状態hは微分方程式（式1a）として表現できる: $h’(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$ hを知ることで、次の動きyを決定できる（式1b）: $y(t) = \mathbf{C}h(t) + \mathbf{D}x(t)$ システムの進化は、その現在の状態と新たに獲得された観測によって決定される。状態の大部分は、既知の状態ダイナミクスを以前の状態に適用することで推論できるため、小さな新しい観測で十分である。つまり、画面の大部分は新しいものではなく、単に以前の状態の自然な下方軌道の継続である。状態を完全に理解することで、次の行動yとして最適な選択が可能になる。 画面の上部を観察することで、システムダイナミクスについて多くを学ぶことができる。例えば、この上部の速度が増加していることは、画面の残りの部分も加速していることを示唆するため、ゲームがスピードアップしていると推論できる。このように、ゲームについて何も知らず、限られた観測しか持っていない状態から始めても、画面のダイナミクスを全体的に理解することがかなり迅速に可能になる。 状態とは何か？ ここで状態とは、入力変数と組み合わせることで、将来のシステムの振る舞いを完全に決定する変数を指す。理論的には、状態を一度把握すれば、過去について未来を予測するために他に知る必要はない。この状態の選択により、システムはマルコフ決定過程に変換される。理想的には、状態はシステムの本質的な特性を捉えた、かなり少量の情報である。つまり、状態は過去の圧縮である。 離散化 - 量子化された世界で生きるためにどう対処するか よし、素晴らしい！つまり、ある状態と入力観測が与えられれば、次の行動を決定する自己回帰型のシステムが得られる。すごい！ しかし実際には、ここに少し問題がある。我々は時間を連続的にモデル化している。しかし実生活では、新しい入力は離散的な時間ステップで得られ、新しい行動も離散的な時間ステップで行われる。 この連続時間の微分方程式を、離散時間の差分方程式に変換したい。この変換プロセスは離散化として知られている。離散化は文献においてよく研究された問題である。Mambaは零次ホールド（ZOH）離散化を使用する。道徳的に何が起こっているかの感覚を与えるために、単純な一次近似を考えてみよう。 式1aから、 $h’(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$ そして小さな∆に対して、 $h’(t) \approx \frac{h(t+\Delta) - h(t)}{\Delta}$ これは微分の定義による。 $h_{t+1} = h(t + \Delta)$ と置き、式1aに代入すると: $h_{t+1} - h_t \approx \Delta (\mathbf{A}h_t + \mathbf{B}x_t)$ $\Rightarrow h_{t+1} \approx (I + \Delta \mathbf{A})h_t + (\Delta \mathbf{B})x_t$ したがって、係数の名前を変更し、インデックスを再ラベル付けした後、離散表現が得られる: もし以前にRNNを見たことがあり、これが親しみを感じるなら - あなたの直感を信じてほしい: いくつかの入力xがあり、それは前の...

マンバの解説

背景や根拠まで確認しますか？

関連記事

調べる

選ぶ

サイト