Transformer ではまだ注意が必要ですか?

ここ数週間、AI コミュニティでは、注意を必要としないアーキテクチャを使用して言語モデルを実装するという話題が注目されています。つまり、機械学習コミュニティにおける長年の研究方向がようやく大きな進歩を遂げ、Mamba の 2 つの強力な新しいモデル、Mamba と StripedHyena が誕生したのです。多くの点で、これらは Llama 2 や Mistral 7B などのよく知られた強力なモデルに匹敵します。この研究の方向性は注意力のないアーキテクチャであり、現在ではますます多くの研究者や開発者がこれを真剣に受け止め始めています。

最近、機械学習科学者の Nathan Lambert 氏が「状態空間 LLM: 注意は必要か?」というタイトルの記事を発表しました。 》では、2023 年に注目度フリー モデルの開発について詳しく説明しています。彼はまた、次のようにも述べています。「2024年には、さまざまな言語モデル アーキテクチャから選択できるようになります。」この記事には多くの数学的内容が含まれていますが、それを深く理解することは価値があることに留意してください。この記事はかなり長いので、簡単に索引付けできるように目次を載せておきます。

• はじめに: 注意を使用したくない理由と、リカレント ニューラル ネットワークとは何かについて説明します。
• Mamba モデル: この新しい状態空間モデルは、将来の言語モデルの多くのクラスに機能とハードウェア アクセラレーションを提供します。
• StripedHyena モデル: Together AI のこの新しい 7B モデルは、RNN とトランスフォーマーの最新の研究を組み合わせたもので、非常に優れたパフォーマンスを発揮します。
• Monarch Mixers の研究: この新しい論文では、これがどのように機能するか、また注意や MLP なしでなぜ成功できるのかの例が示されています。
• 動物学研究: これは、効率的な LLM アーキテクチャと、これらの研究に基づくモデルに関する研究のリポジトリです。

他にもたくさんのリンク、読み物、リソースがあります。

これらの内容にご興味がおありの方は、著者の Nathan Lambert 氏とこの分野の第一人者である 2 人の研究者とのインタビューをお読みいただくか、Machine Heart のレポート「トランスフォーマーの優位性を揺るがすのは誰か?」を参照してください。 Mamba の作者が LLM の将来のアーキテクチャについて語ります。

注意 vs 再帰型および状態空間モデル (SSM)

この記事の核心は、さまざまな計算方法によってモデルにどのような異なる機能がもたらされるかを理解することです。この記事は言語に焦点を当てていますが、その考え方は他の多くのモダリティにも適用できます (実際、これらの新しいアーキテクチャが最初に成功を収めたモダリティはオーディオです)。モデルの内部が異なると、一般化の方法、トレーニングの新しいスケーリング法則、推論時間のコスト、表現力の新しいレベル（つまり、モデルが学習できるタスクの複雑さ）などが異なります。アーキテクチャによって、データが同じであっても、モデルの表現方法がすべて変わります。

いつものように、さまざまなアーキテクチャの選択にはそれぞれ長所と短所があります。現在最も人気のある Transformer アーキテクチャのコア コンポーネントである Attention は、さまざまな理由から優れたパフォーマンスと使いやすさを備えています。この記事では、これらの理由をすべて列挙するわけではありません。簡単に言うと、言語タスクを処理するときに自然な帰納的バイアスを持つモデル、GPU および TPU でのトレーニング用に簡単にスケーリングできるモデル、大量の入力バッチを効率的に処理できるモデル (キー値行列の格納など) などには、注意が有益です。

本質的に、注意には過去のすべてのトークンから現在のトークンへのマッピングがあります。この高密度のアーキテクチャにより、モデルはさまざまなコンテンツを表現し、長いコンテキスト サンプルに焦点を当てることができるようになります。

ただし、リカレント ニューラル ネットワーク (RNN) は、時間をモデルに非常に異なる方法で組み込んでおり、この記事で説明する主な競合アプローチです。これらのモデルは、新しい入力データに遭遇するたびに、内部状態変数 (以下で x で表示) を更新します。原理的には、この内部状態は、データへの直接的な計算リンクがなくても、あらゆるシステムの関連する長期的な動作を捉えることができます。これにより、モデルは長いシーケンスの計算に非常に効率的になりますが、最近まで、注意ベースのモデルのパフォーマンスに匹敵することは実証されていませんでした。次の図は、Attention と RNN の計算グラフを比較したものです。

これらのモデルについて議論するときには、多くの専門用語に出会うことになります。研究コミュニティが望んでいるのは、アテンションや畳み込みなどのアーキテクチャの効率的なトレーニング機能を維持しながら、RNN の時間依存機能を備えたモデルを作成することです。このため、状態の連続時間または離散時間の変化に従う状態空間モデル (SSM) に関する最近の研究が数多く行われています: x'(t) = Ax (t) + Bu (t)、y (t) = Cx (t) + Du (t)。巧みな線形代数または微分方程式を使用すると、この状態の進化を支配する行列は、連続時間か離散時間かに応じて、1 次元畳み込みとして表すことができます。畳み込みアーキテクチャは非常に効率的であるため、それは良い兆候ですが、それ以外に、この記事では難しい数学は扱いません。

以下に、Mamba 論文 (https://arxiv.org/abs/2312.00752) からの式を示します。これについて専門家になりたいのでなければ、これが連続時間で構築され (1a と 1b)、多くの場合離散化され (2a と 2b)、カーネル K になる (3a と 3b) ことを知っておくだけで十分です。技術的には、これは 1D 畳み込みです。

マンバのSSM方程式

著者は、2024年までにすべてが変わるとは予想されないものの、2～4年以内に劇的な変化をもたらす可能性が高いと述べています。タスクによって使用される LLM アーキテクチャは異なります。著者らはまた、ChatGPT のようなシステムが日常的なタスクを実行するために複数のタイプの言語モデルを使用するだろうと予想しています。本稿で説明するように、この RNN 構造に基づいて構築されたモデル (技術的な理由により多くの変更が加えられています) には、長いコンテキストのタスクにおける潜在的な精度と推論コストの点で明らかなスケーリングの利点があります。

言語モデルと機械学習の難しい数学に興味があるなら、12 月は間違いなく始めるのに良い月です。多くの合理的な人は、注目度が置き換えられる可能性が高いことを知っていますが、それがどのように、いつ起こるのか疑問に思うだけです。注意に特化したインフラストラクチャへの大規模な投資を考えると、これがすぐに GPT-N や Gemini と同じレベルに達するとは予想されません。もしそれが成功して注目が薄れれば、Google は大きな問題に直面することになるだろう。なぜなら、TPU はこれらの新しいテクノロジーに使用できなくなる可能性があるからだ (TPU がすでに MoE でうまく動作していないのと同じように)。しかしながら、SSM と関連技術は依然として多くの課題に直面しており、概念的にはまだ実証されていないものも多くあります。たとえば、

• 効果的なスケーリングに必要な GPU のパワーを効率的に活用します。
• モデルを簡単に微調整し、パフォーマンスの大部分を維持する機能。
• コンテキスト学習やシステムプロンプトなどの機能を実行する機能。
• 実際、大規模な Transformer モデルのパラメーターと計算のほとんどは依然としてフォワード フィード ネットワーク (FFN) であり、これは SSM が使用するか変更しない部分でもあります。
• RNN の隠れ状態に必要な容量のボトルネック。
• 特に長い文書の場合、検索メモリなどの機能を統合する機能。これは、既存の長いテキストを拡張するのではなく、複雑な情報ソースを統合することに重点を置いています。

リストは無限です。この記事を読む際には、この新しい注意力を必要としない技術には大きな可能性があるものの、非常に印象的な結果がいくつかあったとしても、それが現時点で最先端の技術と競合できるという意味ではないことに留意してください。それは、まだ非常に初期段階にあるということを意味します。

最近発売されたモデル

MambaとRNNの効率的な計算

12月4日、Albert Gu氏とTri Dao氏は、注意力のないLLMのパフォーマンスをTransformerと同等にし、同時に長いコンテキストのシナリオにおける計算上の制限の問題を解決することを目指した新しいモデルであるMambaを発表しました。 Mamba には 3 つの重要な機能があります。

1. (データ) 選択メカニズム: 「入力に基づいて SSM のパラメータをパラメータ化できるシンプルな選択メカニズムを設計します。」

2. ハードウェア対応アルゴリズム: 畳み込みを特徴スキャンに切り替えるスイッチにより、モデルを既存のハードウェア上でより効率的に実行できます。

3. アーキテクチャ: 以前の SSM ループとトランスフォーマー フォワード モジュール スタイルを組み合わせます。

これら 3 つの側面には、多くの数学的知識が関係します。結局のところ、これらは計算効率を低下させることなく SSM の表現力を向上させるために設計された手法です。

データ選択メカニズムは、循環空間 B と C の処理行列を入力テキスト x の関数として表すことができます (これは、除去行列の線形時間不変性 (LTI) とも呼ばれます)。これにより、入力シーケンスの動作がドメインによって大きく異なる可能性があるため、一般性が犠牲になって表現力が向上します。これらのマトリックスは、どの入力トークンが最も重要であるかを学習できるため、「選択」という名前が付けられています。

Mamba論文で示されたアルゴリズム

ハードウェア対応コンポーネントは、隠し状態 h をメモリ内に最も効率的に保存する方法に重点を置いています。 SSM 更新で使用されるコア パラメーター (線形化された A、B、C 行列) は SRAM と呼ばれるキャッシュに保存されるため、重みを移動しても大きな計算ボトルネックが発生することはありません。次の図は、使用されるメモリの種類を示しています。

最後に、Mamba の論文には、SSM モデルと Transformer モデルにヒントを得たデザインの新しいモデル モジュールも含まれています。この記事の著者は、この設計の理由を明確に述べていませんが、現在の機械学習におけるアクティベーションと MLP の重要性、および現在の最良の LLM が基本的にこれらを組み込んでいるという事実を考慮すると、このアプローチは合理的であるはずです。

このプロジェクトは、SSM コミュニティによる多くの研究の集大成と見ることができます (StripedHyena の場合は確かにそうではありません)。カスタマイズされた CUDA コアを使用することで、GPU は完全にパワーアップしたと言えます。

CUDA コアの特別に設計されたアーキテクチャにより、以下に示すように推論速度が大幅に向上します。

最後に、Pythia スイートと比較すると、モデルのサイズを考慮すると平均評価パフォーマンスは低くなります。 Pythia はパラメータ効率の点ではもはや最良のモデルではないことに留意する必要がありますが、Mistral (y 軸のさらに上) などのモデルと比較されるため、これは実際には賛辞です。さらに、モデルは必ずしも堅牢かつ柔軟というわけではありませんが、下の図に示すように、知っておく価値はあります。

以前のインタビューで述べたように、Tri Dao 氏は、アーキテクチャは拡張法則曲線をより適合性を持って上方に動かすだけであり、データ駆動型 LLM は依然として最良のモデルを作成するための最大の要素であると述べました。著者らは、これにより、利用可能なコンピューティングと関連タスクでより優れたパフォーマンスを達成するために、モデル アーキテクチャがさらに制約されると主張しています。これは素晴らしいものです。 Mamba のモデルとコードのアドレス: https://github.com/state-spaces/mamba

さらに、GitHub には Mamba の最小限の実装があります: https://github.com/johnma2006/mamba-minimal

実世界パフォーマンス: StripedHyena-7B

前の 2 つのプロジェクトは LLM のアーキテクチャを進化させることを目的としていましたが、新しい LLM アーキテクチャという点では StripedHyena (SH) が最も注目を集めていますが、その目標は多くの高度な方法とアーキテクチャ (注意を含む) を組み合わせて LLM のパフォーマンスを進化させることです。

12月8日、Together AIは最初のモデル「StripedHyena-7B」をリリースしました。 StripedHyena (SH) は衝撃的です。この新しい言語モデルは、一般的に使用されている多くの言語モデルと同等のパフォーマンスを発揮します。ブログの説明 (https://www.together.ai/blog/stripedhyena-7b) によると、このモデルはグラフティングと呼ばれる手法を使用しています。本質的には、Together AI はさまざまな事前トレーニング済みモデルからモジュールを取得し、それらを接続して、モデルのトレーニングを継続し、パフォーマンスを安定させます。このブログ投稿には次のように書かれています:

「Transformer と Hyena のアーキテクチャ コンポーネントを統合し、より長いコンテキスト データで拡張された RedPajama データセットのハイブリッド データセットでトレーニングしました。

Hyena という名前は、論文「Hyena Hierarchy: Towards Larger Convolutional Language Models」に由来しています。

OpenLLM リーダーボード タスクでは、StripedHyena が Llama 2 と Yi 7B を簡単に破りました。

多くの注意不要のアーキテクチャと同様に、このモデルの主なセールス ポイントの 1 つは、ロング コンテキスト パフォーマンスです。この論文では、ZeroScrolls ベンチマークを使用して、StripedHyena がタスクで Mistral 7b v0.1 よりも平均 f1 スコアが 3 ポイント高いことを示しています (ただし、すべてのサブクラスで勝っているわけではありません)。 Mistral を上回った StripedHyena のスコアは 27.5 にとどまりましたが、GPT-4 の平均スコアが 41.7 に過ぎないことを考えると、この結果はまだ良好です。

パフォーマンス以外にも、この研究の大部分は LLM のさまざまな概念の計算効率に関するものです。このブログ投稿では、さまざまなアーキテクチャのモデルに対する Chinchilla スタイルのスケーリング規則について詳しく説明します。下の図の左側は Llama-2 と StripedHyena を比較したもので、右側は予算に基づいた最適な注意比率を示しています。

Mamba と同様に、StripedHyena リリースには推論の改善に関する詳細な情報が含まれています。まず最初に、エンドツーエンドの完全な速度について説明します。

2 番目に、コンテキストの長さが長くなるにつれて、その合計メモリ使用量は 2 次的に増加する傾向を示します。

このモデルは現実世界のパフォーマンスに重点を置いているため、著者らはこれを Mistral 7b と比較しました。作者は、Mistral の回答の方が好みだが、StripedHyena の回答は正しく、もっと早くリリースされていれば、しばらくの間はベストモデルになっていたかもしれない、と述べました。これは、これらの新しいアーキテクチャがそれほど遅れていないことを示しています。

ただし、いくつかの制限もあります。研究チームは、基本モデルで使用されたデータを共有せず、「より長いコンテキストデータで強化された、RedPajamaデータセットのハイブリッドデータセット」であるとだけ述べた。

ブログ投稿の最後で、彼らは、Together のこの方向への次の動きに注目する必要があることを明確にしています。

• より長いコンテキストを持つより大きなモデル
• マルチモーダルサポート
• さらなるパフォーマンスの最適化
• より長いコンテキストを活用するために、StripedHyena を検索プロセスに統合します。

最近の研究

このセクションでは、Mamba や StripedHyena と同じ影響を与える可能性がある新しい論文を紹介しますが、この段階に到達するにはさらなる研究が必要になる可能性があります。

モナークミキサー: 注意や多層パーセプトロンのないモデル

• 論文タイトル: Monarch Mixer: シンプルなサブ二次 GEMM ベースのアーキテクチャ
• 論文: https://arxiv.org/abs/2310.12109
• コード: https://github.com/HazyResearch/m2
• ブログ: https://hazyresearch.stanford.edu/blog/2023-07-25-m2-bert

この論文では、Transformer から注意を取り除くことだけでなく、パラメータの大部分を占める MLP を取り除くことも研究しています。このような研究は、今後 6 ～ 12 か月以内に Mamba のようなモデルで発表される予定です。

GEMM は、FFN、RNN、LSTM、GRU など、多くの一般的なアーキテクチャの基礎となる一般化された行列乗算アルゴリズムです。 GEMM は GPU 上で非常に効率的に実行できます。行列乗算は現代の機械学習の中心であるため、これは当然のことですが、行列乗算を新しい方法で設定しても、必ずしもうまく拡張できるとは限りません。

これが何を意味するのかを理解するために、抽象的な説明を見てみましょう。

「機械学習モデルは、より長いコンテキストをサポートし、より優れたパフォーマンスを実現するために、シーケンスの長さとモデルの次元の点でスケーリングし続けています。しかし、Transformer などの既存のアーキテクチャは、これらの軸に沿って 2 乗的にスケーリングします。私たちは、シーケンスの長さとモデルの次元に対して 2 乗未満でスケーリングする高性能アーキテクチャは存在するのか、と問いました。」

現在の長いコンテキストの長さはすでに非常に長いですが、その推論効率は高くありません。著者らは、モデルのサイズはここでは中心的な要素ではないと主張しています。言い換えれば、スケーリング則が指数べき乗則よりも小さい法則に従うアーキテクチャを探してください。以下に続きます:

「私たちは、シーケンスの長さとモデル次元の両方に沿って二次的にスケーリングする新しいアーキテクチャである Monarch Mixer (M2) を提案します。」 Monarch マトリックスは、多くの線形変換をキャプチャし、GPU で高い計算効率を実現し、二次的にスケーリングできる、シンプルで表現力豊かな構造化マトリックスのクラスです。 「

下の画像は Monarch マトリックスの画像を示しています。これは、入力を最初にシーケンスの長さで混合し、次にモデル次元で混合します (両方を同時に混合するわけではありません)。

「概念実証のために、非因果的 BERT スタイルの言語モデリング、ViT スタイルの画像分類、因果的 GPT スタイルの言語モデリングという 3 つの領域で M2 のパフォーマンスを調査しました。」

この投稿では、GPT スタイルの言語モデリングに焦点を当てています。この手法は勢いを増しているからです。ただし、注意を必要としないアーキテクチャは、多くの分野でメリットをもたらす可能性があります (拡散モデルと同様に)。最近機械学習に着手した人にとって、Bidirectional Encoder Representations from Transformers (BERT) について知っておくことは価値があります。BERT は、微調整された多数のバージョンを生成した最初の Transformer モデルであり、多くの利点をもたらします。 BERT スタイルのモデルと ViT (Visual Transformer) のパフォーマンスを続けます。

「非因果的な BERT スタイルのモデリングでは、M2 は下流の GLUE 品質において BERT-base および BERT-large に匹敵しますが、使用するパラメータは 27% 少なく、4k シーケンス長で 9.1 倍のスループットを達成しています。ImageNet では、M2 はパラメータ数が半分であるにもかかわらず、ViT-b の精度を 1% 上回っています。」

さて、GPT スタイルのモデルに戻りましょう。

「因果的 GPT スタイルのモデルは技術的な課題を提起します。マスキングによる因果関係の強制は、二次の計算ボトルネックをもたらします。このボトルネックを軽減するために、多変量多項式の評価と補間に基づく Monarch 行列の新しい理論的分析を実行します。これにより、サブ二次特性を維持しながら、M2 を因果モデルとしてパラメーター化できます。このパラメーター化を使用すると、M2 は事前トレーニング済みの難しさの点で、The PILE 上の 360M パラメーターの GPT スタイルの Transformer に匹敵します。これは、注意や MLP を使用せずに Transformer のような品質を実現できる可能性があることを初めて示しています。」

この段落には多くの情報が含まれています。基本的に、トランスフォーマーを使用して推論を実行するときは、生成された各トークンが過去のトークンのみを参照するように、アテンション マトリックスを上三角マトリックスにマスクする必要があります。これはモデルのデコーダー部分にあり、BERT のようなエンコーダーを見ると、完全にアクティブ化されたアテンション マトリックスが表示されます。 M2 論文の計算は、この二次のボトルネックを軽減します (大まかに言えば、N 個の生成されたトークンに対して M 個のコンテキスト トークンに直接焦点を当てます)。これはすごいですが、数学的には理解しにくいです。しかし、そのインタビューを見れば、理解している人たちがそれをどのように説明しているかが分かります。

その論文のモデルには 3 億 6000 万のパラメータがあるため、これは多くの GPT2 モデルの範囲内ですが、まだ遠い道のりです。

モデル動物学とモデルに基づく

この一連のモデルリリースには、合成タスクで言語モデル アーキテクチャを理解し、テストするためのソフトウェア ライブラリである Zoology も含まれています。アドレス: https://github.com/HazyResearch/zoology

Hazy Research の Christopher Re 氏が、関連する 2 つのブログ投稿を共有しました。

最初の記事では、言語モデリングにおける連想想起の問題をさまざまなアーキテクチャがどのように管理するかを研究します。連想想起とは、さまざまなデータ ソースまたは概念から信号を取得して組み合わせるモデルの能力です。アドレス: https://hazyresearch.stanford.edu/blog/2023-12-11-zoology1-analysis

つまり、これは注意ベースのモデルが非常に優れたパフォーマンスを発揮するタスクです。ここに印象的な例があります:

「7,000 万パラメータの注意モデルが、14 億パラメータのゲート畳み込みモデルよりも優れていることがわかりました。コンテキスト リコールは専門家のタスクのように聞こえるかもしれませんが、機械学習では長い歴史があり、以前の研究では、これらのタスクを解決するにはコンテキスト学習などの魅力的な機能に関連していることが示されています。」

2番目の論文は、彼らの研究結果に基づいた新しいモデルアーキテクチャを公開したため、さらに注目を集めました。ベースアドレス: https://hazyresearch.stanford.edu/blog/2023-12-11-zoology2-based

記事にはこう書かれている。

これらの特性を 3 つの次元にわたって実証し、Based が以下を提供できることを示します。

シンプルで直感的な理解: 単純な畳み込みとアテンションは、さまざまな種類のシーケンスをモデル化するのに適しているという考えに基づいています。それぞれの欠点を克服するために新たな複雑さを導入するのではなく、一般的に使用されているそれぞれのバージョン (短い 1D 畳み込み、"スパイキング" 線形アテンション) を直感的な方法で組み合わせて、両方の長所を最大限に活用します。

高品質なモデリング: シンプルであるにもかかわらず、Based は言語モデリングの難しさの点で、複数のスケールで完全な Llama-2 スタイルのトランスフォーマー (回転埋め込み、SwiGLU MLP など) や最新の状態空間モデル (Mamba、Hyena) よりも優れていると評価されています。

効率的な高スループット推論: 純粋な PyTorch で実装されたモデルの場合、Based は競合する Transformer (スライディング ウィンドウ アテンションと FlashAttention 2 を使用する同等のパラメーター サイズの Mistral) よりも 4.5 倍高い推論スループットを実現します。 LLM がバッチ処理タスクを実行するには、高いスループットが不可欠です。

つまり、異なるアーキテクチャにはそれぞれ長所と短所があります。