Transformerは時代遅れ？次世代AIアーキテクチャMamba

目次

• 初めに
• 参考論文
• 結論
• Mambaとは
• 前提：Foundation Models/基礎モデルとは
• Attentionの欠点
• 欠点を克服しようと模索された構造化状態空間シーケンスモデル
• 早くて賢いMambaの核心：S6
• Mambaを利用する上での課題
• 感想

初めに

ChatGPTやGeminiの成功を支えているのは、「Attention（アテンション）」という画期的なメカニズムです。

この技術は、文章の中の「どの単語が重要か」をAIが自ら判断し、文脈を深く理解することを可能にしました 。
これにより、AIは人間のように自然な対話ができるようになり、世界中に「AI革命」を巻き起こしたのです。

しかし、最強に見えるAttentionにも「文章が長くなると、計算量が爆発的に増える」という致命的な弱点がありました。

文章の長さが2倍になれば、計算コストは4倍、さらに長くなれば16倍……と「2次関数的」に膨れ上がってしまいます。
このため、超長文を読み込ませようとすると、膨大なメモリと時間が必要になるという「物理的な限界」に直面していました。

この問題を解決する基礎モデルとして登場したのが『Mamba』です。
Mambaは、これまでのモデルの常識を覆す数々の革命を起こしています。

• 圧倒的なスピード： 同サイズのTransformerに比べ、推論スピードは5倍に達します
• 驚異のコストパフォーマンス： Mambaの30億パラメータモデルは、2倍のサイズ（70億）の既存モデルに匹敵する賢さを示します
• 100万トークンの壁を突破： これまでのAIがパンクしていた100万という膨大な情報の塊も、精度を落とさず高速に処理できます

なぜMambaは、これほどまでに速くて賢いのでしょうか？
その秘密は、入力された情報から必要なものだけを瞬時に選別する 「選択メカニズム」と、それを支える革新的なアルゴリズムにあります。

ここからは、AIの未来を書き換える『Mamba』の驚異的なアルゴリズムの正体に迫ります。

参考論文

本記事は下記論文をもとに作成しています。
https://arxiv.org/pdf/2312.00752

結論

• Attentionに代表されるシーケンスモデルやそれらの欠点を克服しようとした構造化状態空間シーケンスモデルというものが研究されていたが、それらを凌駕する選択的状態空間モデルであるMambaを開発した
• Attentionと比較して性能が2倍、計算速度は5倍、コンテキスト長は10倍になっている
• Transformerに最適化されたエコシステムが存在するため、実用化にはまだまだ課題が残っていると考えられる

Mambaとは

ここから論文の主題であるMambaについて深ぼっていきます。

前提となる概念の説明や、比較対象になっているAttentionの説明を交えつつ、Mambaとは何か、何が優れているのか、今後の課題は何かを書いていきます。

前提：Foundation Models/基礎モデルとは

一言でいうと「膨大なデータで学習済みの汎用的なAIの土台」のことを指します。

従来、AIを作るには目的ごとに専用モデルをゼロから作る必要がありました。
しかし現在はFMを巨大なデータで作り、ファインチューニングなどの調整をして用途ごとに使用することが主流となっています。

chatGPTもGeminiもAttentionという技術を使用したTransformerという非常に有名なFMを使用していますが、出力の癖が異なるのは調整が異なるからです。

本記事のMambaは、Transformerのような基礎モデルであり、基礎モデルの性能向上はそのまま世の中のAIの性能向上につながるため、世界を揺るがす技術といえるでしょう。

Attentionの欠点

Attenitonは並列学習による高速学習、文脈同士を結び付ける能力にたけています。
そのおかげでまるで人間と話しているかのような出力が可能になったのですが、その一方文脈を結び付けるからこその欠点が生まれてしまっています。

細かい話は置いておくとして、現在使われているAttentionの欠点は大きく3つです。

1. 文章に比例して計算量が爆発してしまうこと
2. 記憶を圧縮できないこと
3. 参照できる範囲に限界があること

まず計算量について。
これが最大の弱点なのですが、Attentionは分量の長さの2乗で計算コストが増えてしまいます。

文章が2倍になれば計算量は4倍、10倍になれば100倍という形で膨れ上がってしまうため、長い文章や大規模データを処理しようとすると計算が追い付かなくなります。

次に記憶について。
Attentionは過去の情報を要約して覚えることが苦手です。

推論する際には過去のやり取りをすべてそのまま保存している必要があり、メモリも食うし上記の計算量向上にもつながってしまっています。

最後に見える範囲について。
AIが一度に参照する範囲をコンテキストウィンドウと言ったりしますが、そのウィンドウに物理的な限界があります。

そのため、圧縮されてない情報の一部のみしか参照できなくなるため、コンテキストが長期化すればするほど精度が落ちてしまいます。

欠点を克服しようと模索された構造化状態空間シーケンスモデル

上記のようなAttentionの欠点を克服するために登場したアーキテクチャに構造化状態空間シーケンスモデル = SSMというものもあります。

SSMの特徴として下記のようなものが挙げられます。

• シーケンス長に対して線形にスケールする効率性
• RNNとCNNのいいとこどりをした再帰と畳み込みの二面性
• 連続的なシステムを離散化する仕組み

このように、SSMはAttentionの最大の欠点であった2次関数的な計算量増加を線形スケーリングに落とし、長いコンテキストにおいても精度を落としにくくしたモデルでした。

ただ、計算効率ではAttentionに圧勝したSSMですが、下記のような欠点もありました。

• コンテンツ認識が欠如しており内容を見て判断できない
• 処理のルールが固定されている
• 音声等の連続データは得意だがテキストのような離散データとの相性が悪い

SSMは処理が早く長いコンテキストに有利だったのですが、テキストへの対応が難しいことからAttentionの代わりにはなりえませんでした。

早くて賢いMambaの核心：S6

そんなSSMの欠点すら克服してAttentionを超えようとしているのが本論文で発表さえたMambaというわけです。

MambaではSSMの進化系であるS6という技術を使用されており、従来のSSMに加えて、

• Selection：入力に応じて情報を選ぶ選択アルゴリズム
• Scan：並列処理で超高速に計算するスキャン手法

が追加されています。
SSM=S4とS6の違いは下記。

	SSM(S4)	S6
パラメータの性質	時間不変：どんな言葉が来ても処理ルールが一定	時間変化：入力される単語に応じてルールが変化
情報の選別	不得意：すべてを同じ重みで流してしまう	得意：重要な情報は残し、不要な情報はフィルタリングする
計算手法	畳み込み：文章全体を一気に処理する	並列スキャン：順番に処理しつつGPUで並列化する
主な弱点	テキストのような情報密度が高いデータに弱かった	(論文に記載なし)

上記のように、SSMの高速処理を引き継ぎつつ、データの内容に応じて処理を変更することでテキストなどにも対応できるようになりました。

結果として、

• Attentionの2倍の性能を発揮し、トークン数が半分でよくなった
• 推論速度が5倍に上がった
• 100倍近いコンテキスト長にも対応できるようになった

などなど、Attentionを凌駕する基礎モデルであることが証明されました。

Mambaを利用する上での課題

Attentionを超える性能を発揮したMambaですが、今すぐTransformerから切り替わるかと言われるとそうとは言えません。

Mambaを採用するために乗り越えなければいけない壁は下記のようになります。

• スケーリングの未知の領域
  • 本論文で性能が証明されているのは30億パラメータ程度のサイズまで
  • chatGPTなどは数千億~兆単位のパラメータを持つといわれているため、そのレベルにおいてもTransformerを圧倒できるかは検証段階です。
• 巨大すぎるTransformerエコシステム
  • 現在のAI開発は、単なるアルゴリズムだけでなく、その周りの道具もすべてTransformer専用に最適化されています。
  • Mambaのほうが高速処理できる可能性はありますが、それらを実現するには今まで作り上げたエコシステムを刷新する必要があるかもしれず莫大なコストがかかります。
• 学習コストとリスク
  • エコシステムを刷新する決断をしたとしても、Mambaを学習し商用利用できる状態にするまでにも電気代と計算資源が必要となり、数十億単位のお金がかかります。
  • スケーリング性能が未知な状態でTransformerを捨ててまでMambaに移行するのはギャンブルになります。

ただ、一定スケールにおいてAttentionを凌駕する性能が出ていることは確かなため、今後は小モデルから実験的に開発が進み、段階的に移行するのではと考えています。

個人的にはリソースに余裕があり将来に投資できるGoogleなどのビッグテックがMambaなどのTransformerと異なるモデルの開発を進め、実績ができたところで爆発的に移行が進むかなと予想しています。

感想

心のどこかでTransformerがAIの最終進化だと思っていた気がします。

AIはここ2，3年で爆発的な進化を遂げてきましたが、Transformer自体は2017年に発表された技術だったと記憶しています。

今のAIの進化速度を見て、Mambaも1~3年後に実用化されて有名になるのではと思っています。

一番に実用化した会社の株価がまたぞろ伸びる気もするので、いろんな意味で実用化が楽しみな技術ですね。