PyTorch 1ファイルでのMamba実装
(github.com/johnma2006)- mamba-minimalは、MambaをPyTorchの1ファイルでシンプルかつ最小限に実装したプロジェクト
- 公式実装とforward/backward passで同じ数値出力を出すことを目標としている
- コードは簡略化されており、読みやすいようにコメント付きで構成されている
- 公式実装の中核的な最適化は含まれていないため速度は提供せず、適切なパラメータ初期化も含まれていない
- デモでは
state-spaces/mamba-370mとEleutherAI/gpt-neox-20bトークナイザを使い、プロンプト補完の例を実行する
プロジェクト概要
- mamba-minimalは、MambaをPyTorchの1ファイルで実装したシンプルな最小実装
- 目的は、公式実装と同じ動作をより読みやすいコードで示すこと
- 主な特徴:
- 公式実装とforward passおよびbackward passで同等の数値出力
- 簡略化されたコード
- 読みやすくコメント付きの実装
含まれていないもの
- 速度は目標ではない
- 公式実装は強力に最適化されている
- その最適化はMamba論文の中核的な貢献に含まれる
- この実装は可読性のため、実装の大半をシンプルに保っている
- 適切なパラメータ初期化は含まれていない
- 可読性を犠牲にせず追加できる項目として示されている
デモの使用例
demo.ipynbでプロンプト補完の例を確認できる- 例では
model.MambaとHugging FacetransformersのAutoTokenizerを使用する - 使用モデルとトークナイザ:
state-spaces/mamba-370mEleutherAI/gpt-neox-20b
- 例のプロンプトは
Mamba is theで、生成結果にはMambaを毒ヘビとして説明する文が含まれる
参考資料
- Mambaアーキテクチャは Mamba: Linear-Time Sequence Modeling with Selective State Spaces で紹介されている
- 論文の著者は Albert Gu と Tri Dao
- 公式実装は
state-spaces/mambaリポジトリにある
1件のコメント
Hacker Newsの意見
以前、同僚と一緒に、共有されるモデルコードの大半を分離したライブラリを作った。これを使うと、多くのモデルを Python の import とコメントを除いて約100行で実装できる
BERT: https://github.com/explosion/curated-transformers/blob/main/...
Llama 1/2: https://github.com/explosion/curated-transformers/blob/main/...
MPT: https://github.com/explosion/curated-transformers/blob/main/...
TorchScript JIT や PyTorch flash attention のような機能にも対応している
xformers も似た問題を扱っているが、Triton を使って高性能な Transformer モジュールを提供することにより重点を置いている。ただ、ライブラリの特定のコンポーネントだけを取り出して使うのは簡単ではなく、ランタイムエラーが続いたので、ひとまず後回しにした。BERT アーキテクチャをベースに何かを作っているので、参考にしてみる
次のプロジェクトで使ってみるつもり
元の Mamba コードは速度最適化やその他の要素が多く、すぐには理解しにくいが、この実装は学習に役立ちそう
トークンを1つずつ推論するときは、すべてがずっと単純になる。自作の Mamba 推論実装もある: https://github.com/rbitr/llm.f90/tree/master/ssm
長年検証されてきた科学計算コードの基盤であり、PyTorch や Numpy のようなライブラリでラップして使われることが多いのは知っているが、今どき人気のある言語ではない。選んだ理由が気になる
Mambaについて、機械学習研究者ではない人にも理解できるように説明してほしい部分がある
Mambaの核心的な洞察は、状態空間モデルの古くからの問題を解くところにある。状態空間モデルは入力コンテキストを圧縮するのには向いているが、入力を隠れ状態に圧縮する過程で、Transformerのようにコンテキストを効果的に活用するために必要な情報が消えてしまう
解決策は、論文で選択メカニズムと呼ばれているものを作ることだ。このメカニズムは入力依存なので、入力が変わるたびにモデルが各ステップの出力を調整できる。そのために、いくつかの状態空間変数を入力不変ではなく入力依存にし、各時点の入力を状態空間変数へ投影する線形層などを付けている
ただし、状態空間変数を入力依存にすると計算オーバーヘッドが生じる。これを、現代のGPUメモリ構造を最大限活用するハードウェア認識アルゴリズムで解決し、HBMの内外へデータを移動することをできるだけ避けている
Tri DaoはFlash Attentionを作った人物で、これもTransformerでハードウェアをより効率的に使う方法だった。こういう分野こそ、まさに彼の専門領域だ
NeurIPS 2023の直前に出たMambaとBasedは、多重クエリ連想想起（MQAR）と、マルチヘッドAttentionに着想を得たゲーティング／選択のデータ依存性を含んでいた。この2つがHyenaやそれ以前の状態空間アーキテクチャに欠けていた重要な要素であり、新しいモデル群は連想想起タスクでAttentionと同じくらい良くなり、検索以外のタスクではおそらくAttentionを少し上回る可能性も見えている
もちろんMambaの大きなディテールは、効率的なCUDA実装だ。それがなければ、Transformerがすでに適している作業では、このアーキテクチャの意義は薄れるかもしれない
コンテキスト長をそれほど心配しなくても、新しい領域は多く開ける。DNA配列解析は長い依存関係を持つ線形のタスクだし、画像・動画・高次元情報をトークンストリームとして見る方法も考えられる。昔のCRTモニターのようにピクセルを走査するやり方だ
AIの初期の夢の1つは、環境と継続的に相互作用するエージェントの単一の学習軌跡が持続的に進化することだった。こうした無限コンテキスト長モデルは、その夢をより実現しやすくするかもしれない
ただし現時点では、この種のモデルの重要な実タスクに対するダウンストリーム応用は、Attentionベースの成熟した応用に比べると、概して検証やチューニングがまだ不十分だ。古い再帰型ニューラルネットワークとの類推はある程度役に立つが、この5年間、人々はAttentionとTransformerに過度に特化してきたため、Transformer側の慣性は大きい
元論文では、パラメータが変換された後、モデルを線形漸化式またはグローバル畳み込みの2つの方法で計算できると説明している。一般に、入力シーケンス全体を事前に見られる学習では並列化しやすい畳み込みモードを使い、入力を1時点ずつ見る自己回帰推論では効率的な再帰モードに切り替える
そのため学習はRetNetの並列順伝播モードのように並列化可能だ。通常の推論は、できるだけ長いコンテキストを得るために再帰モードで行われ、チャンキングがないため、推論中にRAMとVRAMをどれだけ消費するかは判断しにくい
論文を説明しつつ、大きな流れの中でどこに位置づけられるのかについても多くの文脈を与えている。展開を聞くのはかなり興味深い
https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
各トークンが他のすべてのトークンにどれだけ注意を向けるかを計算する二次のAttentionの代わりに、入力と同じ長さの長い畳み込みカーネルを何らかの形で計算してからconv1dを適用する
限定的な理解では、FFTを適用し、行列積を行い、再びIFFTで戻すことに少し関係している。動くことは分かっているが遅い。FFTを計算する方法はいくつもあり、その1つがバタフライ行列だ。おそらく近似にすぎないが十分に良く、現在のハードウェアでは非常に速く効率的なようだ
二次複雑度は悪く聞こえるが、実際にはハードウェア上の制約のために、二乗未満のアルゴリズムのほうが遅いことも多い。だから状態空間モデルへの期待が大きいとしても、Llamaは終わったと言うのは簡単ではない。Mambaがスケールを大きくしてもうまくいくのかもまだ分からず、それを知るには実際に学習に数百万ドルを使ってみる必要がある。それでも楽観的ではある
二乗未満系のもう1つの興味深いモデルはRWKVだ。見てみる価値はあるが、おそらくすでにポッドキャストで取り上げていたと思う
独学で、論文も以前ざっと眺めた程度なので、かなり間違っているかもしれない。またAttentionには通常KVキャッシュがあり、性能に大きく役立つが、Mambaではそれはできないと見ている
「Mamba は推定全長が150mを超える世界最長の毒蛇」という文で笑ってしまった。
それでも本当に素晴らしく、arXiv 論文を参照してくれているので、論文を直接解釈するよりも、こうした記事を読む自分のような人間でも内部を少し覗けるのがよかった。
後でその文に訂正を出すことになったら面白そう。
アルゴリズムの核心は並列プリフィックススキャンだろうと予想していた。それこそが Mamba の要点ではないかと思う。
for i in range(l):x = deltaA[:, :, i] * x + deltaB_u[:, :, i]y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')ys.append(y)愚かな質問かもしれないが、Hugging Face に上がっている Mamba モデルを学習させる難易度が気になる。
最大のモデルは 2.8B に見えるが、The Pile のようなデータセットで学習するには GPU が何枚必要で、どれくらい時間がかかるのだろうか？
推論も RAM を半分しか使わずに3〜5倍速く実行されるように見える。
公式の CUDA 版を読み解こうとして最初の試みが失敗した後、結局手を付けられなかったが、この実装はずっと良さそうに見える。
さらにもう一つの単一ファイル PyTorch 実装とは、本当に素晴らしい。以前の hlb-CIFAR10 や関連プロジェクト、そして minGPT や DawnBench のような先行する影響が、シンプルな単一ファイル形式を少しでも後押しする助けになっていたらよいと思う。
こうした取り組みは効率的な機械学習研究に重要であり、いまこの分野のためにできる最も重要なことの一つかもしれない。
研究はイノベーションの速度で前進し、イノベーションは実験実行時間の逆数に応じて速くなり、これは研究や簡単なハック目的のコードのコルモゴロフ複雑性と明らかに関係している。
こうしたツールが研究にどれほど重要か、個人的に知識発見のプロセスをどれほど速めてくれたかは、いくら強調してもし足りない。アイデアを数分で素早くスケッチし、すぐに信号対雑音比の高い結果を受け取る能力は、研究を進めるうえで不可欠な要素になった。
知識蒸留と MDL(https://en.wikipedia.org/wiki/Minimum_description_length) は、現在の論文投稿・査読プロセスが助長しているように見える不要な飾り、雑多なもの、過度に密集した「取り残されまい」とする低価値なトピック競争を巻き戻すうえで非常に重要だと思う。
最近はこの問題を避け、もう少し良いスケーリング解へ進みたいので、コードを1ファイルの自己完結した短い gist である「コードスケッチ」として配布し始めた。開発時間を減らし、概念を含んだ粗く未整理の動くコードを人々にすぐ渡せる。今のところかなりうまく機能しているようで、続けたい。
こういうコードをもっと見たい。大規模にデータを学習する研究者なら、情報の伝播方法においてもデータ効率的であるべきだ。
人類がかなりの複雑さを持つ何かを、これほど速く発展させたことはおそらくないのではないか。
似た速度が見られるのは SpaceX くらいで、今年も最先端のロケットを2機打ち上げた。2024年には何が出てくるのか楽しみだ。
x_projにバイアスがないので、x_proj と dt_proj の重みをまとめられそうに見える。重み調整の要件があるなら実行時に単純にできるかもしれないし、単一カーネルとバイアスのほうが最終的には速そうだ。確信はない。
元論文の議論があったのか気になる。見逃したようだが、かなり興味深い。
「効率的な実装の不足によりメモリ不足や非現実的な計算要求が生じ、SSM としても解釈できる以前の強力な再帰モデルである RWKV と RetNet ベースラインのコンテキスト長 8k の全結果が欠けている」という部分がよく理解できなかった。
RetNet はメモリを大量に使わず、チャンク単位の順伝播実装を使えば VRAM 使用量はチャンクサイズに制限される。この点こそがコンテキスト長をテストする核心だ。
オリジナルの Mamba モデルをテストした人がいるのか気になる。並列順伝播モードの RetNet と比べて学習速度はどの程度なのだろうか？
https://openreview.net/forum?id=AL1fq05o7H
複雑なものを核心だけ残して整理した実装はいつでもよいものだ。