Liama 3とMambaの強力な組み合わせ！線形RNNに蒸留され、推論速度が1.6倍向上します。

Llama 3 を Mamba に統合すると、推論速度が最大 1.6 倍向上します。

さらに、そのパフォーマンスは低下せず、元のモデルよりも優れています。

Together AI のこの新しい研究では、蒸留を通じて Transformer モデルと Mamba モデルを組み合わせ、ハイブリッドモデル用の推論加速アルゴリズムも組み込んでいます。

Mamba アーキテクチャの設計者であり、FlashAttention の作者でもある Tri Dao もこのプロジェクトに参加しました。

Together AI の創設者兼 CEO は、Transformer と Mamba のハイブリッドは大規模モデルの将来の主要な開発方向を表していると述べました。

トランスフォーマーをマンバに蒸留する

蒸留プロセスが正式に開始される前に、Transformer から線形 RNN への初期化が必要です。

著者らは、Transformer の注目メカニズムと RNN の計算の間には一定の類似点があることを観察しました。

したがって、Transformer の注意を線形化することができ、それによって 2 つの間の接続を確立できます。

この対応関係を利用することで、事前トレーニング済みの Transformer モデルのパラメータを Mamba モデルにコピーできます。

パラメータの初期化を完了した後、著者らは 3 段階の蒸留プロセスを採用して Mamba モデルのパフォーマンスをさらに向上させ、Transformer の知識をより深く学習できるようにしました。

最初の段階は疑似ラベルの蒸留に基づいています。つまり、事前トレーニング済みの Transformer 教師モデルを使用してラベルなしデータに疑似ラベルを生成し、次に Mamba 学生モデルをこれらの疑似ラベルでトレーニングさせます。

このプロセスにおける損失関数は、KL ダイバージェンス損失とクロスエントロピー損失を組み合わせたもので、それぞれ教師モデルの出力分布を模倣し、疑似ラベルを適合させるために使用されます。

第 2 段階では、トレーニングにラベル付き指示データセット (OpenHermes 2.5 など) を使用して、指示データセットの教師あり微調整が行われます。

最終段階では、人間のフィードバックデータを使用して、報酬モデルに基づいてシステムを最適化します。

著者らは、モデルの出力に関する人間からのフィードバックデータを収集し、このデータに基づいて報酬モデルを構築し、RL アルゴリズム (PPO など) を使用して、この報酬モデルの下でモデルのパフォーマンスを最適化しました。

8 個の 80G A100 GPU では、各ハイブリッドモデルの蒸留プロセス全体にかかる時間は 5 日未満です。

上記の蒸留プロセスを通じて、著者らはTransformer-Mambaハイブリッドモデルを取得し、推論プロセスを加速するためのSpeculative Decodingアルゴリズムを提案しました。

投機的デコードアルゴリズムの基本的な考え方は、軽量のドラフトモデルを使用して複数のトークンを予測し、検証モデル（Verifier）を使用してこれらの予測を検証することです。

これにより、デコードの並列性が大幅に向上し、生成プロセスが加速されます。

ドラフトモデルは通常、現在のコンテキストに基づいて次の K 個のトークンを予測する小さなトランスフォーマーです。

予測された K 個のトークンについては、Transformer レイヤーはこれらの K 個のトークンを直接並列に処理し、それらの隠し状態を計算できます。

Mamba レイヤーは、まず現在のトークンの隠し状態を計算し、それを以前の隠し状態と比較しながら、各トークンを順番に処理する必要があります。

シーケンス内の K 個のトークンがすべて受け入れられた場合、それらは出力シーケンスに追加され、次のトークンセットの予測が続行されます。

トークンが拒否された場合、予測シーケンスは最初に拒否されたトークンから切り捨てられ、最初のステップが返されて、その位置から予測が再度開始されます。

テスト結果では、ハイブリッドモデルが、1 ターン (AlpacaEval) と複数ターン (MT-Bench) のチャットダイアログタスクの両方で Llama-3 と同等かそれ以上のパフォーマンスを発揮することが示されています。

さらに、異なる混合比のモデルのパフォーマンスをテストしたところ、混合比が 1:1 のモデルが最も優れたパフォーマンスを発揮することがわかりました。

ゼロサンプルの一般的な NLP タスクの評価では、ハイブリッドモデルは平均して同じサイズの RNN モデルよりも優れています。

サンプルサイズが制限された OpenLLM リーダーボードでは、ハイブリッドモデルは最高のオープンソース RNN モデルと同等のパフォーマンスを示し、GSM8K および CRUX タスクでは対応する Instruct モデルよりも優れたパフォーマンスを示します。

モデルのパフォーマンスに加えて、著者らは投機的デコードアルゴリズムによってもたらされる高速化効果もテストしました。

最初のテストは純粋なMambaモデルで行われました。その結果、2.8Bおよび7Bモデルでは、オリジナルのデコード手法と比較して推論速度が1.7～2.6倍向上しました。

さらに、著者らは、抽出した Zephyr ハイブリッドモデルと Llama ハイブリッドモデルをテストし、Zephyr ハイブリッドモデルでは推論速度が 1.8 倍以上向上し、Llama ハイブリッドモデルでも約 1.6 倍の加速が見られることを発見しました。

論文リンク: https://www.together.ai/blog/...