複数の大規模モデルを効率的に組み合わせて「互いの長所を補完する」という斬新なアプローチが、トップカンファレンス NeurIPS 2024 で承認されました。

RouterDCと呼ばれるこのルーティング アーキテクチャは、パラメータ効率 (1 億パラメータ未満) と計算効率 (LLM の勾配バックプロパゲーションを実行する必要がない) の利点を持つ、デュアル対照学習に基づくルーティング アーキテクチャです。

言語理解、コード生成、数学的推論などの難しい推論タスクの実験では、RouterDC は分布内 (+2.76%) と分布外 (+1.90%) の両方の設定で既存のルーティング方法を大幅に上回りました。

よく知られているように、LLM は通常、さまざまなデータセットで事前トレーニングおよび微調整されるため、さまざまなタスク間でパフォーマンスが異なります。

LLMルーティングは、複数のLLMを組み合わせる革新的なアプローチです。ルータから学習し、各クエリに最適なLLMを選択します。推論処理中は、選択されたLLMのみを呼び出すため、複数のLLMの相補的な機能を活用しながら、計算効率を維持できます。

新しい RouterDC アプローチには、エンコーダーとしての小さな言語モデルと、候補 LLM に対応する学習可能なLLM 埋め込みのセットが含まれています。

トレーニング データ内の各クエリについて、まず候補 LLM の予測を実際のラベルと比較して、パフォーマンスが最も高い LLM と最も低い LLM を取得し、次に 2 つの対照的な損失を構築します。

• サンプル LLM 対照損失により、クエリ埋め込み (エンコーダーによって抽出) は、パフォーマンスが最も高い LLM 埋め込みと類似しますが、パフォーマンスが最も低い LLM 埋め込みとは類似しなくなります。
• サンプル間対照損失:すべてのトレーニング クエリを複数のグループにクラスタ化し、同じグループ内のクエリ間の類似性を最大化し、異なるグループ内のクエリ間の類似性を最小限に抑えることで、トレーニングの安定性が向上します。

この研究は、南方科技大学と香港科技大学の研究チームによって提案されました。以下は、より詳細な紹介です。

二重対照学習によるルーターのトレーニング

ルーターのアーキテクチャ

図1に示すように、RouterDCは、エンコーダεとして小規模言語モデル（mDeBERTaV3ベース）と、候補LLMに対応する一連の学習可能なLLM埋め込みkTから構成されます。各クエリxiに対して、RouterDCはT個のLLMについて以下の選択確率を生成します。

ここでsim(·,·)はコサイン類似度を表します。

△図1：RouterDC方式の概略図

サンプルLLMコントラスト損失

ルータを学習させるために、研究者らはクエリのサンプル埋め込みをK+個のLLM対応埋め込みのうち、最もパフォーマンスの高いものまで絞り込み、K-個のLLM対応埋め込みのうち、最もパフォーマンスの低いものまで広げました。したがって、サンプルLLM対比損失は次のように表すことができます。

サンプル間比較損失

研究者は実験を通じて、ルーティング問題でサンプル LLM 対照損失のみを使用すると、類似したクエリが異なる埋め込みを持つ可能性があるため、安定しないことを発見しました。

学習の堅牢性を向上させるため、学習サンプルは異なるグループにクラスタリングされます。これにより、学習中に同じグループ内のサンプル同士が接近し、異なるグループからのサンプル同士が離れます。サンプル対LLMの対比損失と同様に、サンプル対サンプルの対比損失は次のように定式化できます。

トレーニングと推論

最終的な最適化の目的は、サンプル-LLM 対照損失とサンプル-サンプル対照損失の組み合わせを最小化することです。

推論中、各テスト クエリでは、トレーニング済みのルータを通じて最も高い確率の LLM を選択し、選択した LLM を使用してクエリに回答するだけで済みます。

RouterDC は、トレーニング中に LLM による勾配バックプロパゲーションを必要とせず、推論中に LLM を 1 回実行するだけで済むため、トレーニングと推論の両方で高い効率を実現します。

実験の結果はどうでしたか?

主な結果

表1は、分散データセットにおけるRouterDCのテスト精度の結果を示しています。以下のことがわかります。

RouterDCは、平均3.98%のパフォーマンス向上を達成し、最優秀単一モデルを大幅に上回りました。個々のタスクレベルでは、RouterDCは3つのタスクにおいて最優秀単一モデルよりも精度が向上しました。GSM8K（0.51%）、ARC-C（0.57%）、HumanEval（1.63%）です。

既存のルーティング手法であるCosineClassifierおよびZOOTERと比較して、RouterDCはすべてのタスクにおいて優れたパフォーマンスを発揮します。LoraRetrieverと比較した場合、RouterDCは平均2.77%の精度向上を実現します。

△表1: 分散タスクのテスト精度（％）

RouterDC の一般化能力を評価するために、表 2 に 3 つの分布外データセット (PreAlgebra、MBPP、C-EVAL) での RouterDC のテスト精度を示します。

ご覧のとおり、RouterDC は再び最高のテスト精度を達成し、最高のパフォーマンスを示した単一の LLM (dolphin-2.9-llama3-8b) を 1.9% 大幅に上回りました。

△表2：分布外課題のテスト精度（％）

サンプル間損失の役割

サンプル間損失の影響を調査するため、図3はサンプル間損失ありとなしのトレーニングおよびテストの精度曲線を示しています。RouterDC（Lサンプル間損失なし）は大きな振動を示しているのに対し、RouterDCははるかに安定していることがわかります。

△図2: GSM8KタスクにおけるRouterDCのトレーニングとテストの精度曲線

図3(a)は、RouterDC（Lsample-sampleなし）を用いて抽出した学習サンプルのTSNE特徴を可視化したものです。異なるタスクに属する学習サンプルが大まかに混在していることがわかります。しかし、Lsample-sampleを組み合わせると、学習サンプルは明確なクラスタリング構造を示します（図3(b)）。

△図3: 学習済みルータによって抽出されたトレーニングサンプル埋め込みのt-SNE可視化

RouterDCはコスト効率に優れています

LLMを評価する上で価格も重要な指標であるため、研究者はRouterBench上で2つのタスクについて実験を行い、コストの影響を具体的に検討しました。図16に示すように、RouterDCはCosineClassifierやZOOTERよりも費用対効果が高いことが分かりました。

△図4: RouterBenchで異なるコストを使用して得られたテスト精度

論文リンク: https://arxiv.org/abs/2409.19886 コードリンク: https://github.com/shuhao02/R...