大規模モデルの推論速度が 3.6 倍に向上しました。「Medusa」の論文はこちらです。Jia Yangqing: 最もエレガントな加速推論ソリューションの 1 つ

ご存知のとおり、大規模言語モデル (LLM) の動作ロジックでは、サイズが大きくなるにつれて言語生成の品質が向上します。しかし、これにより推論の遅延も増加し、実際のアプリケーションに大きな課題が生じます。

システムの観点から見ると、LLM 推論は主にメモリに依存しており、主なレイテンシのボトルネックは算術計算ではなくアクセラレータのメモリ帯域幅に起因します。このボトルネックは、自己回帰デコードの順次的な性質に固有のものであり、各フォワードパスでは、完全なモデル パラメーターを高帯域幅メモリからアクセラレータ キャッシュに転送する必要があります。このプロセスでは単一のトークンしか生成されず、最新のアクセラレータの算術計算能力が十分に活用されないため、効率が低下します。

この問題に対処するために、LLM 推論を高速化する方法が提案されており、これにより、デコード プロセスの演算強度 (FLOP と総データ移動量の比率) を高めるか、デコード手順の数を減らすことができます。このタイプの方法は、投機的デコードによって表されます。投機的デコードでは、小さなドラフト モデルを使用して各ステップでトークン シーケンスを生成し、その後、より大きな元のモデルによってトークン シーケンスを精製して、許容可能な継続を取得します。ただし、適切なスクラッチ モデルを取得することは依然として困難であり、スクラッチ モデルを分散システムに統合することはさらに困難です。

この論文では、プリンストン大学、Together.AI、イリノイ大学アーバナ・シャンペーン校などの研究者らが、別個のスクラッチモデルを使用して候補出力を順次生成する代わりに、バックボーンモデルの上に複数のデコードヘッドを使用して推論を加速するという概念を再検討し、改良しました。この技術を効果的に適用すれば、投機的復号化の課題を克服し、既存の LLM システムにシームレスに統合できることが分かりました。

具体的には、研究者らは、複数のトークンを同時に予測できる追加のデコードヘッドを統合することで LLM 推論を強化する方法である MEDUSA を提案しました。これらのヘッドはパラメータ効率の高い方法で微調整されており、既存のモデルに追加できます。したがって、新しいモデルを必要とせずに、MEDUSA を現在の LLM システム (分散環境を含む) に簡単に統合して、使いやすいユーザー エクスペリエンスを確保できます。

注目すべきは、この論文の著者の 1 人である Tri Dao 氏が、非常に人気のある Transformer 代替アーキテクチャである Mamba の 2 人の著者のうちの 1 人であるということです。彼は Together.AI の主任科学者であり、近々プリンストン大学のコンピューターサイエンスの助教授に就任する予定です。

• 論文アドレス: https://arxiv.org/pdf/2401.10774.pdf
• GitHub アドレス: https://arxiv.org/pdf/2401.10774.pdf

具体的な実装では、研究者は 2 つの重要な洞察を通じて MEDUSA をさらに強化しました。まず、各デコード ステップで単一の候補継続を生成する現在のアプローチでは、許容できる長さが制限され、計算リソースの使用が非効率的になります。この問題に対処するために、彼らは MEDUSA ヘッドを使用して複数の候補継続を生成し、注意マスクの簡単な調整を通じてそれらを検証することを提案しています。 2 番目に、投機的デコードで使用されるものと同様の拒否サンプリング方式を使用して、元のモデルと同じ分布を持つ応答を生成できますが、これは通常、多くの LLM アプリケーションでは必要ありません。

そのため、研究者らは、MEDUSA の出力から妥当な候補を選択するという、典型的な許容可能なスキームを導入できる可能性があると考えました。彼らは、元のモデル予測の偏りを管理するための閾値として温度を使用し、拒否サンプリングに代わる効果的な方法を提供しました。このアプローチは、高温での速度低下などの拒否サンプリングの制限に効果的に対処します。

さらに、LLM に予測可能な MEDUSA ヘッドを装備するために、研究者らはさまざまなシナリオに合わせて調整された 2 つの微調整手順を提案しました。計算リソースが限られている場合や、パフォーマンスに影響を与えずに MEDUSA を既存のモデルに組み込むことが目的の場合は、MEDUSA-1 の使用を推奨しています。このアプローチではメモリが最小限で済み、固定されたバックボーン モデルによる生成品質を損なうことなく、QLoRA と同様の量子化技術を使用してさらに最適化できます。

しかし、MEDUSA-1 では、バックボーン モデルの潜在能力を十分に活用できませんでした。したがって、さらに微調整を行うことで、MEDUSA ヘッドの予測精度が向上し、直接的に高速化につながります。そのため、研究者らは、十分なコンピューティング リソースがあるシナリオや、ベース モデルからの直接的な教師付き微調整に適した MEDUSA-2 を提案しました。 MEDUSA-2 の鍵となるのは、モデルの次のトークン予測能力と出力品質を損なうことなく、MEDUSA ヘッド モデルとバックボーン モデルの共同トレーニングを可能にするトレーニング プロトコルです。

実験部分では、バッチ サイズ 1 のシナリオに焦点を当てます。これは、LLM が個人使用のためにローカルにホストされるユース ケースを表します。彼らは、Vicuna-7B および 13B (パブリック データセットを使用してトレーニング)、Vicuna -33B (プライベート データセットを使用してトレーニング)、Zephyr-7B (教師ありの微調整とアライメントを使用してトレーニング) など、さまざまなサイズとトレーニング設定で MEDUSA をテストしました。

結果は、MEDUSA が生成品質を損なうことなく、さまざまなプロンプト タイプで 2.3 ～ 3.6 の推論速度向上を達成できることを示しています。次のアニメーション画像は、Vicuna-7b で Medusa-1 を使用した場合と使用しない場合の推論速度を比較したものです。

論文の共同筆頭著者であるTianle Cai氏は、Medusaプロジェクトの開始以来、TensorRT、TGI、および多くのオープンソースプロジェクトや企業で採用されていると述べた。新しい技術論文では、フルモデルの微調整のための Medusa-2 スキーム、Medusa を任意の微調整 LLM に統合するための自己蒸留、およびその他の加速技術を紹介します。

この研究に関して、Lepton AI の創設者 Jia Yangqing 氏は、Medusa はこれまで見た中で最もエレガントな加速推論ソリューションの 1 つであり、int8/fp8、コンパイルなどを補完し、実際には 2 倍のパフォーマンス向上を達成できる可能性があると述べました。

さらに、Medusa は多くの既存の最適化手法やハイブリッド アクセラレーション ソリューションと統合されています。その結果、アクセラレーションは適度な同時実行性の下でプラスの効果を維持し、A100 や H100 などのカードで特に効果的です。さらに、彼らはラマ モデル用に汎用のメデューサの頭を訓練しました。

方法の概要

MEDUSA は推測的デコードフレームワークに従っており、各デコードステップは主に 3 つのサブステップ (1) 候補の生成、(2) 候補の処理、(3) 候補の受け入れ) で構成されます。 MEDUSAの場合、(1)はMEDUSAヘッドによって実装され、(2)はツリーアテンションによって実装され、MEDUSAヘッドは元のバックボーンモデルの上に配置されるため、(2)で計算されたロジットはサブステップ(1)の次のデコードステップで使用できます。最後のステップ（3）は、不合格サンプル検査または通常合格検査のいずれかによって達成することができる。 MEDUSA の全体的なプロセスを以下の図 1 に示します。

主要コンポーネント

MEDUSA の主要コンポーネントには、主に MEDUSA ヘッドとツリー アテンションが含まれます。

まず、MEDUSA ヘッドを元のバックボーン モデルと一緒にトレーニングします。元のバックボーン モデルは、トレーニング中は固定されたままにしておくことも (MEDUSA-1)、一緒にトレーニングすることもできます (MEDUSA-2)。このアプローチにより、強力な基本モデルによって学習された表現を活用して、単一の GPU 上で大規模なモデルでも微調整が可能になります。

さらに、MEDUSA ヘッドの分布は元のモデルと一致するように保証されているため、分布シフトの問題が軽減され、MEDUSA はサービス システム設計の複雑さを増大させず、分散設定に適しています。

候補が増えると計算要件が増加するため、本研究では複数の候補を同時に処理するためのツリー構造の注意メカニズムを採用しています。この注意メカニズムは、従来の因果的注意パラダイムとは異なります。そのフレームワークでは、同じ継続からのトークンのみが履歴データと見なされます。グラフニューラルネットワークの分野でグラフ構造を注意に埋め込むというアイデアに触発され、研究チームは下の図2に示すように、ツリー構造も注意マスクに組み込みました。

トレーニング戦略

バックボーン モデルをフリーズして MEDUSA ヘッドをトレーニングするのは簡単で、計算リソースもほとんど必要ありませんが、バックボーン ネットワークを MEDUSA ヘッドと一緒にトレーニングすると、MEDUSA ヘッドの精度が大幅に向上します。そのため、研究チームは、コンピューティング リソースとユース ケースの具体的な要件に応じて、MEDUSA ヘッドのトレーニング戦略として、MEDUSA-1: バックボーン ネットワークのフリーズと MEDUSA-2: 共同トレーニングという 2 つのレベルを提案しました。

最後に、この研究では、MEDUSA に利用可能なトレーニング データがない場合に対処し、デコード プロセスの効率を向上させるために使用される自己蒸留と典型的な受け入れを含む、MEDUSA の 2 つの拡張機能を提案しました。

実験

異なる設定での MEDUSA の有効性を実証するために、本研究では 2 セットの実験を実施しました。まず、MEDUSA を Vicuna-7B/13B モデルで評価し、MEDUSA-1 と MEDUSA-2 のパフォーマンスを実証しました。次に、Vicuna-33B モデルのトレーニング データセットは公開されておらず、Zephyr-7B モデルは RLHF を使用してトレーニングされているため、自己蒸留の有効性を調べるために MEDUSA を Vicuna-33B および Zephyr-7B モデルで評価しました。

ユースケーススタディ 1: Vicuna-7B/13B モデルでの MEDUSA の評価

Vicuna-7B/13B モデルで MEDUSA-1 と MEDUSA-2 を評価した結果を以下の図 4 に示します。

ユースケーススタディ 2: Vicuna-33B および Zephyr-7B でのトレーニングに自己蒸留を使用する

研究者らは、Vicuna-33B と Zephyr-7B を例に挙げ、自己蒸留が必要な状況に焦点を当てました。彼らはまずシードプロンプトを使用してデータセットを生成し、次に ShareGPT と UltraChat をシードデータセットとして使用し、上記の 2 つの例について約 10 万個のサンプルを含むデータセットを収集しました。

以下の表 1 は、MT-Bench ベンチマークにおけるさまざまな MEDUSA-2 モデルの高速化、オーバーヘッド、および品質を示しています。

下の図 5 は、MEDUSA-2 を使用した場合のさまざまなモデルの高速化を示しています。

アブレーション実験

下の図 6a は、ランダムにサンプリングされた密なツリー設定 (青い点) と最適化された疎なツリー設定 (赤い星) の高速化を比較しています。 6b は密なツリー設定と疎なツリー設定の速度を比較します。

下の図 7 は、異なるサンプリング設定でのモデル パフォーマンスの比較分析を示しています。

2段階微調整の有効性。研究者らは、Vicuna-7B モデルの 2 つの微調整戦略のパフォーマンスの違いを評価しました。