Blackwell Tensorコア: メタトークンを使用したHYMBAモデル (小規模言語モデル)

11,134 文字

みなさん、こんにちは。NVIDIAから発表された新しい言語モデルHBAについて、最新の研究成果をご紹介させていただきます。これはMとTransformerアーキテクチャを組み合わせた並列化モデルで、新しいNVIDIAチップのハードウェアアーキテクチャ向けに、美しい新しいメタトークンを活用してます。詳しく見ていきましょうか。
通常、言語モデルには事前学習があって、ファインチューニングがあって、そしてアライメントがあって、最後に推論があります。つまり、プロンプトを送信して、答えが生成されるのを待つわけですな。
新しいトレンドや改善点は、OpenAI-1のようにタスク特化型のクエリに対するテスト時のアルゴリズムをスケールさせることです。前回の動画でお見せしたように、推論実行時にLoRAアダプター調整プロセスを使って推論能力を最適化し、OpenAI-1に近づけることができます。また、この動画では、強化学習のポリシーモデルと報酬モデルを一緒に学習させて、モンテカルロ探索アルゴリズムの探索空間を最適化・拡張するテスト時学習アルゴリズムのソリューションも見てきました。
推論アクターではたくさんのことが起こってますが、もちろんポリシーモデルと報酬モデルも互換性があり、一貫した構造を持っていなければならず、推論を始める前に学習される必要があります。そうなると、システム全体の複雑さが増してきますね。
でも、NVIDIAの新しいアイデアを待ってください。NVIDIAは「今までは単純やったな。self-attentionのマルチヘッドを持つTransformerアーキテクチャだけやったけど、新しいチップを使って、こうしましょう」と言うてます。Mアーキテクチャも導入して、state space modelsとTransformerアーキテクチャを並列で動かそうというわけです。
これはかなりチャレンジングですが、残念ながら新しいBlackwellチップへのアクセス権は持ってません。新しいNVIDIAのハードウェア最適化へのアクセスがあったらええのになぁ思います。モデルの最適化とアイデアの最適化をNVIDIAのGPU上で並列に実行できたらすごいですよね。
ただ、彼らは小規模なモデルに焦点を当ててることに気づきました。Small Language Modelと呼ばれるもので、学習可能なパラメータが20億、あるいは1.x億程度です。ここで新たな課題に直面します。小規模で、高速で、効率的である必要があるからです。そうなると、今では20億以下の学習可能パラメータを持つLLMやSmall LRMについて話すことになります。NVIDIAのアイデアが素晴らしいのは、このような小規模なLRMは特定のタスクに対して追加情報を提供する必要があるからです。特定のタスクに対して、メタトークンでLLMを準備する必要があります。
メタトークンは、プロンプト構造の前に配置され、システムのアーキテクチャの不可欠な部分となります。メタトークンはTransformerアーキテクチャとMambaアーキテクチャの両方に対応します。タスク特化型やということを申し上げましたが、まずはNVIDIAの論文を詳しく分析して、技術的な詳細を見ていきましょう。そして、実装が簡単で美しいアイデアをご紹介させていただきます。
例えば、ベンチマークスタディを行いたい場合、3-4つのベンチマークテストを自分で実行したいとします。このメタトークン構造に、ドメイン特有の知識やタスク特有の知識を組み込みます。数学的推論のベンチマークを実行したい場合は、タスク特有の推論知識についての最適な方法論を少し入れます。医療や化学、物理学、数学のドメインでテストを実行する場合は、私たちの小さなLLMが十分な学習容量を持っていない可能性のある基本的なドメイン知識を入れるわけです。
元のアイデアに戻りましょう。大規模モデルの学習では、事前学習、ファインチューニング、DPOアライメントがありますね。例えばAppleがしているように、事前学習済みの一般的にファインチューニングされたモデルを取り、写真アプリケーションや要約作成、ビデオ修正などのために、モデル全体をファインチューニングするのではなく、アダプターを使用します。LoRAアダプターチューニングや量子化LoRAアダプターチューニングですね。これが通常のやり方です。
NVIDIAはこのアイデアを取り入れて、「素晴らしい。特定のタスク、特定のベンチマーク、特定のドメインでテストを実行して、これが最高のモデルであることを示したい場合、このメタトークンセクターに知識を提供します」と言いました。そして、これを1つのアーキテクチャに統合するのが素晴らしいところです。新しいNVIDIAチップ、Blackwellでもなんでも、完全に新しい構造で実行する必要があるため、事前学習、ファインチューニング、アライメントすべてを行う必要があります。そこで彼らは「事前学習を最適化する最良の方法は何か」と考えました。
そして、HBAまたはHと呼ばれるモデルを構築しました。これは単にTransformerとMambaシステムのハイブリッドシステムのようなものですが、それ以上のものがあります。特に素晴らしいのがメタトークンで、これについて焦点を当てたいと思います。メタトークンは埋め込まれ、Transformerが作用するベクトル空間に与えられるデータ構造として埋め込まれます。また、state space modelのメタトークンも使用します。
理論的には、Mambaのstate space modelの初期化のために、「なぜMambaが必要なのか」と問うこともできます。Mambaはそれほど強力ではありませんし、Transformerアーキテクチャの二次計算量を削減する方法も知っています。この素晴らしいチャンネルの購読者の方々はご存知かと思いますが、1年前にGoogleのRing Attentionや、100万トークンのコンテキスト長を持つInfinityについて、二次計算量に圧倒されないようにする方法をお見せしました。
解決策はありますが、これは小さなLLM向けであることを忘れないでください。携帯電話や腕時計で動作させる可能性もあるので、インフラストラクチャもパワーも持っていません。Mambaについては、以前の動画でS6 modが本当にTransformerより優れているのかを詳しく見て、一緒にハイブリッドMamba Transformerをコーディングしましたが、このハイブリッドにはいくつかの問題がありました。
従来のTransformerは二次計算コストと複雑さ、高メモリ要件に直面する一方で、state spaceは非二次計算効率にもかかわらず、メモリ再現タスクで苦戦しています。このギャップを埋めるためにハイブリッドアーキテクチャを試みましたが、レイヤー統合の逐次的な性質によりボトルネックに直面しました。そこでNVIDIAが登場して「新しいチップで並列に実行しよう」と言っています。ハイブリッド、つまり並列について理解する必要があります。監督付きファインチューニングとDPOアライメントのテストを見たい場合は、この動画をご覧ください。また、Mamba S6 modを超えて、数学と理論物理学に少し触れたい場合は、ベクトル場に進みます。これがMを超えた進化の次のステップです。
でも、シンプルに保ちましょう。NVIDIAはここでシンプルにしたいと考えています。すべてのユーザーが複数のエージェントを持つべきです。シンプルなアイデアとして、複雑なRing Attentionメカニズムは構築せず、長期メモリ用の単純なMシステムを使用します。
これを見て「理論的にはin-context learningの可能性があるのに、なぜNVIDIAは新チップのハードウェアアクセラレーションでin-context learningを採用しないのか」と思われるかもしれません。in-context learningがファインチューニングやrackシステムを上回るパフォーマンスを発揮し、in-context learning plusと呼ばれるものがあるため、もはやファインチューニングは必要ないということを、詳しく説明した動画もあります。
長いコンテキスト用に特別に設計された再帰的システムについて話すのではなく、シンプルに、地に足をつけて考える必要があります。このメタトークンのアイデアは本当に素晴らしいものです。
NVIDIAによるとメタトークンは、学習可能な埋め込み構造として初期化され、事前学習から始まるトレーニング中にモデルパラメータと一緒に最適化されます。新しくすべてを一から構築する必要があるためです。このメタトークンは、小規模なモデルがシーケンスをより効果的に処理するのを助ける、追加のメモリまたはコンテキストのレイヤーとして機能します。NVIDIAチップ上のAIシステムに複雑な質問をする場合、トレーニング中にメタトークンは私の個人的なクエリの入力シーケンスと相互作用し、これがTransformerアーキテクチャの注意演算子とMamba構造のstate space遷移に影響を与えます。
メタトークンの値と埋め込み構造は、トレーニングタスク全体で損失関数を最小化するように調整され、ドメインに関する重要な文脈的知識や推論の方法、タスク特有のパターンを埋め込みます。これは素晴らしいアイデアです。LLMが小規模なため、クエリに対して、システムがメタトークンを追加のシェル、追加の情報本体としてLLMに提供し、すべてが一緒に学習されます。
これらのメタトークンは、私たちのLLMのパラメトリック知識に本来備わっていない一般的な知識、例えば数学的文脈や法的推論をエンコードします。特定のテンプレートに従う必要がありますよね。これらすべてがメタトークン構造に含まれます。LLMが非常に小さいため、LoRAキャッシュやMambaのstate初期化手順として見ることができ、これによりMのメモリ効率とシステム全体のリコール機能が向上します。
代替案もあります。NVIDIAは新しいチップアーキテクチャで多くの素晴らしいテストを行ったと確信しています。in-context learningを採用するか、小規模LLM向けに特にプロンプトエンジニアリングを最適化することもできましたが、彼らはメタトークンを選択しました。これが非常に興味深いポイントです。
ただし、タスクに対してこれを行う場合、システム全体を学習させる必要があります。事前学習、ファインチューニング、すべてを行う必要があります。1.2億のサブ2億というのは小さな数字で、1時間でできると思うかもしれません。はい、わかります。でも、まずは簡単なところから始めましょう。このモデルをどう使いたいかがわかっていて、特定のタスクに向けてTransformerとMambaの並列構造を持つ完全に新しいモデルを学習させるのは素晴らしいことです。でも想像してみてください。前回の2つの動画で紹介した特定のタスクに対するテスト時アルゴリズムについて考えると、理論的には、推論時に特定のタスクに対して、タスク最適化のための一貫したメタトークンを選択する可能性があります。
これを使って、知識グラフ特有の埋め込みを使用してクエリやドメインをマッピングし、タスク特有のクエリに関連するすべてのメタトークンを選択することができます。もっと突飛なアイデアを出すと、他のテスト時アルゴリズムでLoRAアダプターを学習させたように、未知のタスクに対して動的に合成メタトークンを作成する軽量な条件付き生成器を学習させることもできます。
このように、美しい新しいアイデアから多くのアイデアが生まれています。メタトークンはLLMのサイズ制限（2B未満）を補うために、重要な文脈的なガイダンスやタスク特有の初期化をMambaシステムに提供します。これらは学習された埋め込みで、完全なプロセスを実行する必要があり、法的、科学的、数学的など、どんな分野で作業しているかに関わらず、ドメイン特有の情報をエンコードする不可欠な部分です。
また、タスク特有のものにもなり得ます。量子力学のような基本的な物理学の原理といったドメインコンテキストを持つことができ、数学的証明の書き方など構造的なものも可能です。研究論文を要約するための最適な方法論など、関連する知識を持つこともできます。これらはLLMのパラメトリック知識には含まれません。2億未満の学習可能なパラメータしかないからです。しかし、これらをメタトークンに入れるのです。Transformerの注意メカニズムやMamba構造のstate space更新のアンカーとして機能することで、両方を並列で実行します。どのように組み合わさるかについてはお話しします。これらのメタトークンにより、小規模言語モデルが特定のタスクに対する我々のガイダンスを持って、入力の関連する側面に焦点を当てることができます。
これは素晴らしい新しいアイデアです。テスト時の最適化計算の適応に関する3つ目の動画になります。2024年11月22日のHBA、ハイブリッドヘッドアーキテクチャの研究は一見単純に見えますが、論文を読むと素晴らしい内容です。もちろん、結果もお伝えします。
NVIDIAのスーパーチップ（私はそう呼んでいます。これらのチップへのアクセス権は持っていません）上でのこのモデルは、Llama 3.23Bさえも上回るパフォーマンスを示しています。つまり、1.5Bのモデルが、ベースモデルの2倍の大きさのシステムを上回っているのです。印象的ですね。Hugging Faceで見つけることができますが、ハードウェア構成については分かりません。これらのチップすべてへのアクセス権は持っていませんし、もしかしたら古いチップとの下位互換性があるかもしれませんが、本当のパワーは次世代のものから出てくると思います。
もしかしたら、皆さんはすでにこれらのモデルにアクセスできるかもしれません。簡単にまとめると、このNVIDIAの新しいハイブリッドモデルの主なポイントは、並列モデルであり、事前学習段階でモデルの重みとともに新しいメタトークンを共同で最適化することです。メタトークンは、NVIDIAのスライディングビデオアテンションヘッドやMamba構造のstate space modelsのメカニズムと互換性があります。なぜなら、並列化の世界でTransformerとMambaを組み合わせる必要があるため、メタトークンを推論実行中にキー値キャッシュの初期化に変換しますが、アーキテクチャを変更することはありません。アーキテクチャは十分に複雑だからです。
素晴らしいですね。これらの新しいチップと新しいモデルが好きですが、技術論文には興味深いことがあります。論文を深く掘り下げてみると、トークンとアテンションスコアの分布についてよく見てみる必要があります。エントロピー計算を見ることで、いくつかの情報が得られます。
トークンレベルで見ると、メタトークンの行列があり、通常は人間の文章の開始を示す特定のトークン、つまり文頭トークンがあります。これは従来のTransformerモデルや類似のアーキテクチャで使用されている単なる指標ですが、残念ながら何か問題が起きています。
マルチヘッドの自己注意機構で、Llama 3.23Bの学習可能パラメータモデルを調査すると、文頭トークン（内容がなく、単に「ここから人間の最初の文章が始まります」と伝えるだけのもの）にシステム全体がどれだけの注意を払っているかを見ると、完全な注意スパンの56%が単一のトークンに向けられています。これは驚くべきことです。
より正確に言うと、Llama 3.23Bは注意の50%以上が、何かが保存されることを示す単一のトークン、つまりシーケンス開始トークンに集中しています。「文」は人間の言葉ですが、シーケンス開始トークンというのが正しいですね。
興味深いことに、このメタトークンには追加の利点があります。これが素晴らしいのは、特別なトークンや特定のトークンへのこの過度な注目、この注意の負担を軽減するのに役立つからです。シーケンス開始からメタトークンへと注意の焦点をリダイレクトすることで、モデルは完全なシーケンス全体でより意味のある注意計算を分散させることができます。NVIDIAが観察したのは（私ではありません）、プロンプトが数学やコード、その他の異なるドメインからのものである場合、推論時に異なるメタトークンが活性化されるということです。
これは、異なるメタトークンが異なる世界知識をその埋め込みにカプセル化していることを示唆しており、これは理にかなっています。そして、この異なる世界知識は、注意メカニズムを導いて関連情報により焦点を当てるために活用できます。
これはまさに私たちが望んでいることです。推論実行時に単純なクエリを出して「この数学の公式を説明して」と言った場合、注意メカニズムがこれが数学に関するものであり、数学には形式性があることを理解し、数学的基準に基づいた説明が欲しいと思います。あるいはコードの場合、10種類の異なるプログラミング言語間を行き来するのではなく、一貫した形式で持っていたいわけです。
メタトークンの知識が何であれ、メタトークンは新しい魅力的なものです。もう1つ興味深いことがあります。注意マップエントロピーを計算すると（理論物理学に少し詳しい方なら、理論物理学からの洞察をコンピュータサイエンスでも使用できることをご存知でしょう）、NVIDIAが各レイヤーとヘッドの注意マップエントロピーを計算したのですが、これは本当に驚くべきことです。
並列構造があり、Transformerアーキテクチャからの注意ヘッドとMamba構造からのSSMヘッドがあることを覚えておいてください。メタトークンありとなしの場合を見ると、明るいオレンジ色でメタトークンありの場合が見えます。これは私たちが望んでいたことをまさに反映しています。Transformerではこれは完璧に機能します。Mambaでは通常2つの問題領域がありますが、気にしないでください。通常、メタトークンを使用すると（濃い青色）、分析されたすべての異なるレイヤーにおけるエントロピーがかなり小さくなります。
定義として、注意マップエントロピーは単に、すべての異なるトークンにわたる注意スコアの分布を反映します。エントロピーが低いということは、注意スコアがより小さなトークンのサブセットに集中しているため、より強い検索効果を示します。これは知識のより高い局所化を意味し、これはまさに私たちが望んでいることです。これはこのシステムが機能していることの美しい指標であり、メカニズムについて自分で考えると理にかなっています。
さて、ここまでTransformerとMambaの並列構造について話してきましたが、これらを組み合わせる必要があります。TransformerとMambaシステムのキー値キャッシュ、そしてそれらを並列で実行してキャッシュ最適化を行う美しい問題があります。
Transformerのキー値キャッシュから始めましょう。各注意ヘッドは入力シーケンス内のトークンのキーと値を計算し、推論時にこれらのキーと値は再計算を避けるためにキー値キャッシュに保存されます。これは特に長いコンテキストウィンドウを含むタスクで重要です。
キー値キャッシュの共有が行われ、これによりAIシステムのメモリ使用量と計算オーバーヘッドを削減できます。これは素晴らしく、まさに私たちが望んでいることです。2つの可能性があります。同じキー値キャッシュをモデルの複数のレイヤーで共有するか、同じレイヤー内の複数の注意ヘッド間で同じキー値キャッシュを共有するかです。これは、独自のTransformerアーキテクチャを構築する場合の最適化構造の一種です。これら2つのオプションがあり、キー値キャッシュを共有することで、パフォーマンスの大幅な低下なしに冗長な保存を排除できます。
NVIDIAはここでデータを提供し、このHBAでメモリフットプリントを大幅に削減できると言っています。最大で3.8倍のメモリ削減が可能とのことです。ただし、単純化するたびに精度を失うという暗部があることにも注意してください。
レイヤーは、異なる特定のサブタスクや特定の入力シーケンスに対して独立してキー値マッピングを調整する能力を失う可能性があります。つまり、一般的なTransformerアーキテクチャのアイデアを思い出すと、異なる注意ヘッドはシーケンスや文の異なる特性を見ています。レイヤー内の注意ヘッドは、データ内の異なる関係に焦点を当てることが多く、構文的、意味的なパターンなどを見ています。
ヘッド間で共有すると、モデルがこれらの微妙な違い、多様な特徴を捉える能力が低下するリスクがありますが、もちろんメモリ削減と高速な処理が得られます。これは現在NVIDIAの最適化プロセスの一部でもあります。
彼らは、ほとんどのレイヤーで主にスライディングビデオアテンションを使用していると述べています。これにより、グローバルコンテンツではなく、このスライディングウィンドウ内のローカルコンテンツに焦点を当てることで、キー値キャッシュのサイズをさらに削減できます。スライディングウィンドウアテンションメカニズムについては、特別な動画があります。
コーヒーが必要ですね。さて、このNVIDIAの新しいHBAモデルについて、少し驚いたことがあります。もちろんこれはプロトタイプで、今日発表されたばかりですが、長いコンテキストについて話す際に、8,000トークンまで処理できると言っています。今日の基準では、共有値を持つ並列化されたTransformer-Mamba構造にとって8Kトークンは素晴らしいことですが、一般的な世界では8Kは本当に大したことではありません。まあ、どこかから始めなければならないので、8Kから始めるのは良いでしょう。
私はここでMambaの役割と、このHYMBAについて調査したいと思いました。「なぜこのシンプルなシステムの方がずっと簡単なのか、Mambaは本当にハードウェアアクセラレーションのためのアーキテクチャの単純化に貢献しているのか」と疑問に思いましたが、NVIDIAはこのステップを進めており、より複雑な数学的アルゴリズムを使用していないようです。
このHBAシステムについてまとめると、Transformerアーキテクチャで高価な完全シーケンス注意を処理する必要なく、より効率的にグローバルコンテキストを処理します。したがって、このタスクは現在Mambaに与えられており、Mambaはフィーディングメモリのように機能し、重要なグローバル特徴を保持することで、Transformerのメモリ需要も削減します。
NVIDIAは（私には検証できませんが、これらのチップへのアクセス権がないので）、チップ上の線形代数演算をより良い方法で活用するように設計されていると述べています。これはNVIDIAのGPUチップやGoogleのテンソル処理ユニットとうまく整合しますが、GoogleのTPUはTransformerアーキテクチャにより焦点を当てているかもしれません。
いずれにせよ、NVIDIAは、この美しい新しいチップ（何であれ）とのハードウェア互換性からHYMBAが恩恵を受け、8Kシーケンスを処理する際でも、より高速な処理速度とより高いスループットを実現すると述べています。そして今、動画の最後に私のクレイジーなアイデアを紹介しましょう。これが何かを意味するなら（つまり、チップとハードウェアと互換性のあるソフトウェアが必要で、その後で最適化を実行する必要がありますが）、将来的にこの新しいメタトークンを、私たちがよく知っているin-context learningと統合できないでしょうか。
なぜなら、メタトークンは数学的推論手順や法的分析テンプレートなど、ドメインに関する知識を表現でき、これにより新しいLLMが入力プロンプトで提供するin-context learningの例をより良く解釈できるようになるからです。
これは興味深いですね。NVIDIAがここで2億未満の学習可能なパラメータを持つ本当に小さな言語モデル用に開発したこのアイデアを、通常の大規模言語モデルの世界で使用できないでしょうか。これは興味深い実験だと思いますが、残念ながら新しいチップアーキテクチャへのアクセス権がありません。
メタトークンは、モデル内部の固定された埋め込み構造で、事前学習にありますが、ここには適応する方法があるかもしれません。なぜなら、数学的ベクトル空間で作業しているからです。すでに2つの新しいアイデアがあります。一方、in-context learningの例は外部のユーザー提供の入力です。
これらは異なりますが、美しい有益な関係を提供できると思います。なぜなら、メタトークンで提供するドメイン知識について、テスト時トレーニングでモデルに指示できるからです。そして、データベースにあるもの、プロンプト構造の最適化のためのプロンプトで知っているものについて、追加のin-context learningの例を持つことができます。
このように、OpenAIやPerplexityによるこの自動プロンプト最適化は、この新しい統合で限界に達すると思います。互いに補完し合える可能性がありますが、もちろん、コストに関する問題が出てきます。というのも、OpenAIのこの01では、時に答えを待つのに90秒や120秒かかることがありますが、今度は完全に新しいハードウェアアーキテクチャをチップ上で動かし、Mambaアーキテクチャも実行し、キー値キャッシュを共有して最適化できますが、これは本当にコストのかかる作業になると想像できます。
ですが、未来に向けた新しいアイデア、新しいモデル、新しいハードウェア、新しいアーキテクチャについて考えているので、これはチップが私たちにとって利用可能になったときに調査すべきことだと思います。
以上で私の動画は終わりです。新しいアイデアをいくつかお見せできたことを願っています。特にNVIDIAの新しいインフラストラクチャで試せる新しい方法論や新しいモデルについてです。
毎日、私たちには新しいアイデアがあり、それらのアイデアの魅力的な実装があります。これらの企業は並列TransformerとMambaアーキテクチャを持つ完全に新しいモデルを構築しており、私たちが知っているものすべてを上回るパフォーマンスを示しているようです。3日前に公開した古いAIシステムさえも上回っています。
ぜひチャンネル登録をお願いします。次の動画でまたお会いしましょう。