推理力が2倍にアップ！プリンストン大学と北京大学の卒業生がロング「メデューサ」を提供、33Bモデルは13Bと同等の速さ

LLM アーキテクチャに固有のメモリ制限により、生成は遅く、コストがかかります。

この点に関して、多くの有力者たちはそれを救う方法を模索している。 Karpathy 氏は、大規模モデルの「投機的サンプリング」が推論時間を最適化する優れたソリューションであると繰り返し提案しています。

しかし、投機的デコードは生成プロセスを高速化できますが、複雑すぎるため、多くの開発者に採用されていません。

今日、プリンストン大学、イリノイ大学カリフォルニア大学（UIUC）などの中国のチームが、新しいシンプルなフレームワーク「Medusa」を提案した。

研究者たちは、追加のスクラッチモデルを使用せずに、いくつかの追加のデコードヘッドを導入し、「Medusa ヘッド」を微調整するだけで、単一の A100-8G GPU で 1 日でトレーニングできるようになりました。

結果は、Medusa がモデル推論を約 2 倍直接的に高速化することを示しています。

ビクーニャ-7bとメデューサ

LLM 生成が非効率的なのはなぜですか?

システムの観点から見ると、LLM 生成は「メモリバウンド」コンピューティングモデルに従います。このモデルでは、主なレイテンシのボトルネックは算術計算ではなくメモリの読み取り/書き込みから生じます。

この問題の根本は、自己回帰デコード処理の固有の順次的な性質にあります。

つまり、各フォワードパスでは、モデル全体のパラメータを高帯域幅メモリ (HBM) からアクセラレータのコンピューティングユニットに転送する必要があります。

各サンプルに対して 1 つのトークンのみが生成されますが、この操作では最新のアクセラレータの演算処理能力を十分に活用できず、非効率的なモデルが生成されます。

LLM が登場する前は、この非効率性に対処する一般的な方法は、単に「バッチサイズを増やす」ことで、より多くのトークンを並行して生成できるようにすることでした。

しかし、大型モデルの出現により状況はさらに複雑になっています。

この場合、バッチサイズを増やすと、レイテンシが長くなるだけでなく、Transformer モデルのキー値キャッシュのメモリ要件も大幅に増加します。

さらに、この非効率性は「コスト構造」にも反映されています。

2023 年 9 月現在、プロンプトのみを処理する場合と比較して、GPT-4 の生成コストは約 2 倍、Claude 2 の生成コストは約 3 倍高くなります。

研究者の主な焦点は LLM 生成のレイテンシの改善ですが、Medusa はレイテンシとスループットのバランスを取る必要がある LLM サービスにも適用できます。

トークンあたりの価格

「推測的デコード」は究極の解決策でしょうか?

上記の課題を考慮すると、テキスト生成を高速化する魅力的な戦略は、コンピューティングリソースをより効率的に活用することです。

具体的には、より多くのトークンを並列処理します。

この方法では、簡略化された「ドラフト」モデルを使用して、各ステップでトークン候補のバッチをすばやく生成します。

これらの候補トークンは、元の本格的な言語モデルに対して検証され、テキストの最も妥当な続きが決定されます。

この基本的なロジックは、興味深い仮定に基づいています。つまり、「ドラフト」モデルは、小さくても、元のモデルに受け入れられるシーケンスを生成できるほど十分に優れているはずです。

この仮定が正しい場合、「ドラフト」モデルはトークンシーケンスを迅速に生成でき、元のモデルは複数のトークンを効率的に並列検証できるため、計算スループットが最大化されます。

最近の研究では、慎重に調整されたドラフトモデルを使用すると、投機的デコードによってレイテンシを最大 2.5 倍も削減できることが示されています。

ただし、このアプローチには課題がないわけではありません。

1. 理想的な「ドラフトモデル」を見つける:元のモデルとうまく調和する「小さいながらも強力な」ドラフトモデルを見つけるのは、言うほど簡単ではありません。

2. システムの複雑さ: 1 つのシステムで 2 つの異なるモデルをホストすると、特に分散環境では、計算面と運用面の両方で複数の層の複雑さが生じます。

3. サンプリング効率が低い:サンプリングに投機的デコードを使用する場合は、重要度サンプリング方式を使用する必要があります。これにより、特にサンプリング温度が高い場合に、追加の生成オーバーヘッドが発生します。

これらの複雑さとトレードオフにより、投機的デコードの広範な採用が制限されます。したがって、投機的デコードは大きな可能性を秘めているものの、広く採用されていません。

Medusa: シンプルさと効率性を兼ね備える

よりユーザーフレンドリーでありながら強力なソリューションを求める声に応えるため、プリンストン大学の研究チームは革新的なフレームワーク Medusa を立ち上げました。

これにより、モデル生成が高速化されるだけでなく、LLM をより多くの人がアクセスして使用できるようになります。

報道によると、最新の方法は、論文「深層自己回帰モデルのブロック単位の並列デコード」で過小評価されていた宝物を再検討するものである。

Transformer モデルの発明に戻ると、後続のトークンを予測するためにまったく新しい「ドラフト」モデルを導入するのではなく、元のモデル自体を単純に拡張してみてはどうでしょうか。

論文アドレス: https://arxiv.org/abs/1811.03115

ここで「メデューサの頭」が登場します。

これらの追加のデコードヘッドは元のモデルとシームレスに統合され、生成された各キーポイントでトークンブロックを生成します。

スクラッチモデルとは異なり、メデューサのヘッドは、トレーニング中は固定されたままの元のモデルと一緒にトレーニングできます。

このアプローチにより、研究者は強力な基本モデルによって学習された表現を活用して、単一の GPU 上で大規模なモデルを微調整できます。

さらに、新しいヘッダーは元の言語モデルヘッダーと同様に 1 つのレイヤーのみで構成されているため、Medusa はサービスシステム設計の複雑さを増大させず、分散環境に適しています。

Medusa ヘッドだけでは、処理速度を 2 倍にする目標を達成することはできません。

しかし、ここにちょっとしたコツがあります。ツリーベースのアテンションメカニズムと組み合わせると、Medusa ヘッドによって生成された複数の候補を並行して検証できます。

このようにして、メデューサの頭の予測能力は2〜3倍に増加します。

さらに、研究者らは従来の重要度サンプリング方式を放棄し、メデューサの頭部生成専用の効率的で高品質な代替方式を作成しました。

この新しいアプローチは、サンプリングのオーバーヘッドを完全に回避し、Medusa をさらに 1 段階高速化します。

つまり、研究者たちは単純なシステムで推測的復号化の課題を解決したのです。

1. 別個のモデルなし:新しいスクラッチモデルを導入する代わりに、複数のデコードヘッドが同じモデルでトレーニングされます。

2. 既存のシステムへの簡単な統合:トレーニングパラメータは効率的で、GPU のパフォーマンスが低い場合でも実行できます。また、追加モデルがないため、分散コンピューティングの設定を調整する必要はありません。

3. サンプリングを緩和として扱う:元のモデル分布に一致するという要件を緩和すると、「非貪欲な生成」が「貪欲なデコード」よりもさらに高速になります。

メデューサの概要

具体的には、Medusa は LLM の最後の隠し状態の上に複数のヘッドを導入し、複数の後続のトークンを並行して予測できるようにします。

Medusa ヘッドを使用してモデルを拡張する場合、元のモデルはトレーニング中に固定され、Medusa ヘッドのみが微調整されます。このアプローチにより、単一の GPU 上で大規模なモデルを微調整することが可能になります。

推論中、各ヘッドは割り当てられた位置に対して複数のトップレベルの予測を生成します。これらの予測は候補に結合され、ツリーベースのアテンションメカニズムを使用して並列処理されます。

最後のステップは、一般的な受け入れスキームを使用して適切な継続を選択することです。受け入れられた最長の候補プレフィックスが、次のデコード段階で使用されます。

このように、Medusa はより多くのトークンを同時に受け入れることでデコードプロセスの効率を向上させ、必要なデコード手順の数を減らします。

次に、Medusa の 3 つのコンポーネント、Medusa ヘッド、ツリーアテンション、および一般的な受け入れスキームについて詳しく見ていきましょう。

メデューサの頭

それで、メデューサの頭とは一体何なのでしょうか?

これらは、元のアーキテクチャの言語モデルヘッド (因果 Transformer モデルの最後のレイヤー) に似ていますが、次のトークンだけではなく、複数の今後のトークンを予測するという点が 1 つ異なります。

研究者たちは、ブロック並列デコード法にヒントを得て、各メデューサヘッドを単層フィードフォワードネットワークとして実装し、残差接続を追加しました。

これらのヘッドのトレーニングは非常に簡単です。元のモデルをトレーニングしたのと同じコーパスを使用することも、モデル自体を使用して新しいコーパスを生成することもできます。

このトレーニングフェーズでは、元のモデルは変更されず、メデューサの頭部のみが微調整されます。

このターゲットを絞ったトレーニングにより、パラメータ効率が非常に高いプロセスが実現され、すぐに収束します。

特に、投機的デコード手法で別のスクラッチモデルをトレーニングする際の計算の複雑さと比較すると、利点はより顕著です。

研究者らがテストしたビクーニャモデルでは、メデューサの頭が次のトークンを予測する精度が60％で第1位となった。同時に、まだ改善の余地があります。

木への注目

テスト中、チームはいくつかの驚くべき指標を発見しました。次のトークンを予測する第 1 ランクの精度は約 60% でしたが、第 5 ランクの精度は 80% を超えていました。

この大幅な改善は、メデューサのヘッドを効果的に利用して複数のトップランクの予測を生成できれば、デコード手順ごとに生成されるトークンの数を増やすことができることを示唆しています。

この目標を達成するために、研究者たちはまず、各メデューサの頭から得た上位の予測の直積を取って候補のセットを作成しました。

次に、グラフニューラルネットワークの考え方に従って、依存関係グラフをアテンションメカニズムにエンコードし、複数の候補を並列に処理できるようにします。

たとえば、以下に示すように、最初のメデューサの頭に対して最初の 2 つの予測を使用し、2 番目のメデューサの頭に対して最初の 3 つの予測を使用します。

注意ツリー

上の図の視覚化は、ツリー構造の注意を使用して複数の候補を同時に処理するプロセスを示しています。

この場合、最初のヘッドからの予測は 2 番目のヘッドからの予測とペアになり、複数レベルのツリー構造が形成されます。

このツリーの各層は、メデューサの頭の予測に対応しています。したがって、2 * 3 = 6 個の候補が生成されます。

これらの候補はそれぞれ、ツリー構造内の異なるブランチに対応します。

各トークンが前のトークンにのみアクセスすることを保証するために、研究者は、現在のトークンから前のトークンに注意が戻ることを明確に許可する注意マスクを設計しました。

これを実行し、位置エンコーディングに応じて位置インデックスを設定することで、バッチサイズを増やすことなく、さまざまな候補を同時に処理できます。

研究者らはまた、いくつかの研究でも非常によく似たツリー状の注意の考え方が採用されていると指摘した。

それらと比較して、最新の方法では、ツリーパターンが推論中に規則的かつ固定される、より単純な形式のツリー状注意が好まれ、ツリー状注意マスクの前処理が行われ、効率がさらに向上します。

典型的な受け入れ

投機的デコードに関する初期の研究では、「重要度サンプリング」を使用して、元のモデル予測に密接に関連するさまざまな出力が生成されました。

しかし、その後の研究では、このアプローチはサンプリング温度が上昇するにつれて効率が低下する傾向があることが示されています。

簡単に言えば、ドラフトモデルが元のモデルと同じくらい優れている場合は、理想的にはその出力をすべて受け入れて、プロセスを非常に効率的にする必要があります。

ただし、重要度サンプリングでは途中でこのソリューションが拒否される可能性があります。

現実の世界では、必ずしも元のモデルの分布と一致させるためではなく、モデルの創造性を制御するためだけにサンプリング温度を調整することがよくあります。

では、なぜ受け入れられる信頼できる候補者に焦点を当てないのでしょうか?

この目的のために、研究者らは「典型的承認」制度を導入した。

プリンストン大学の研究者たちは、既存の切り捨てサンプリング法からヒントを得て、元のモデルに基づいて十分に可能性のある候補を選択することを目指しました。

研究者は、元のモデルに従って予測された確率に基づいて閾値を設定し、候補者がこの閾値を超えた場合、その候補者は受け入れられます。

技術的には、研究者は、切り捨てサンプリングの場合と同様に、ハードしきい値とエントロピー関連のしきい値の最小値を使用して、候補者を受け入れるかどうかを決定します。

これにより、デコード中に意味のあるトークンと適切な継続が選択されるようになります。

最初のトークンを受け入れるために常に貪欲なデコードを使用し、各ステップで少なくとも 1 つのトークンが生成されるようにしています。最終出力は、テストに合格した最長のシーケンスです。

このアプローチの素晴らしさは、その適応性にあります。

サンプリング温度を 0 に設定すると、最も効率的な形式である貪欲なデコードに戻ります。温度が上昇すると、最新の研究方法はより効率的になり、より長いシーケンスを受け入れることができます。

さらに、研究者たちは厳密なテストを通じてこの主張を確認しました。

つまり、本質的には、最新の標準的な入学許可制度は、LLM がより創造的な成果を生み出すためのより効率的な方法を提供します。

ラマはどれくらいの速さで唾を吐くことができますか？

研究者たちは、チャットアプリケーション向けに特別に調整された Vicuna モデルを使用して Medusa をテストしました。

モデルのサイズはさまざまで、パラメータの数はそれぞれ 7B、13B、33B です。

この研究の目的は、実際のチャットボット環境で Medusa がこれらのモデルをどれだけ高速化できるかを測定することでした。

メデューサの頭部を訓練する際、研究者はシンプルなアプローチを採用し、公開されている ShareGPT データセットを使用しました。これは、元々 Vicuna モデルに使用されたトレーニングデータのサブセットであり、1 エポックのみトレーニングされています。

さらに重要なのは、モデルパラメータのサイズに応じて、トレーニングプロセス全体を数時間から 1 日以内に完了でき、必要なのは A100-80G GPU 1 台だけであることです。

特に、Medusa は量子化された基本モデルと簡単に組み合わせることができ、メモリ要件を削減できます。

研究者たちはこれを利用し、33B モデルのトレーニング時に 8 ビットの量子化を使用しました。実際の環境をシミュレートするために、MTベンチを使用して評価しました。

Medusa はシンプルな設計により、幅広いユースケースで一貫して約 2 倍の絶対時間高速化を実現します。

Medusa の最適化により、33B パラメータの Vicuna モデルが 13B モデルと同じ速度で実行できることは注目に値します。

アブレーション実験

メデューサの頭の構成

研究者たちは、メデューサの頭の予測力を活用して、各頭について考慮すべき最有力候補の数を柔軟に選択しました。

たとえば、最初のヘッドの上位 3 つの予測と、2 番目のヘッドの上位 2 つの予測を選択できます。これらの上位候補の直積をとると、モデルが評価するための 6 つの継続オプションが生成されます。

この構成可能性にはトレードオフが伴います。

一方、より多くの上位予測を選択すると、モデルが生成されたトークンを受け入れる可能性が高まります。その一方で、各デコード手順の計算オーバーヘッドも増加します。

最適なバランスを見つけるために、研究者たちはさまざまな構成を試し、添付の図に示すように最も効率的な設定に落ち着きました。

典型的な受け入れ閾値

一般的な承認シナリオでは、「しきい値」と呼ばれる重要なハイパーパラメータが、モデル自体の予測に基づいて、生成されたトークンが妥当かどうかを研究者が判断するのに役立ちます。

しきい値が高くなるほど、受け入れ基準が厳しくなり、このアプローチによって得られる全体的な高速化に影響します。

研究者たちは、MT ベンチで創造性重視の 2 つのタスクに関する実験を実施することで、品質とスピードアップの間のトレードオフを調査しました。

図に示す結果は、貪欲なデコード方法と比較して、典型的な受け入れが 10% 高速化できることを示しています。この高速化は、貪欲なデコードに比べて実際にはプロセスを遅くするランダムサンプリングを使用した推測的なデコード方法よりも大幅に優れています。

著者について

ティエンレ・カイ

共同筆頭著者の Tianle Cai 氏はプリンストン大学の博士課程の学生であり、指導教員は Kai Li 氏と Jason D. Lee 氏です。

博士課程在学中、彼は Xuezhi Wang 氏と Denny Zhou 氏の指導の下、Google DeepMind でインターンシップを行い、また Sébastien Bubeck 氏と Debadeepta Dey 氏の指導の下、Microsoft Research でもインターンシップを行いました。

彼は北京大学で応用数学を専攻し、コンピュータサイエンスも専攻して学士号を取得しました。Liwei Wang教授の指導の下、機械学習の分野で研究を始めました。

Tianle Cai の学術的関心は、最適化、表現学習、アーキテクチャ設計 (Transformer、グラフニューラルネットワークなどに重点を置く)、最近ではシステムアーキテクチャの共同設計など、機械学習の幅広い分野にわたります。