技術要約: TRACE-Router

問題提起

現在の大規模言語モデル（LLM）のルーティングシステムは、主にリクエストレベルで動作しており、個々のプロンプトに対して独立したルーティング決定を行っています。この粒度はシングルターンの推論には適していますが、タスクがツール利用、検索、環境遷移を伴う複数のLLM呼び出しを含む、長期間のステートフルな実行トレース（execution trace）で構成されるエージェント的ワークロードとは乖TL（ミスマッチ）が生じています。

核心的な不一致は、**フィードバックの帰属（feedback attribution）**にあります：

1. 遅延したフィードバック: エージェント型アプリケーションにおいて、成功または失敗は、個々の中間ターンの品質ではなく、最終的なタスクレベルの結果（例：テストスイートの通過や目標の達成）によってのみ決定されます。
2. クレジット割り当ての失敗: リクエストレベルのルーターは、この遅延したタスクレベルの報酬を、特定の中間ルーティング決定に正しく帰属させることができません。一見すると最適なローカルの決定（例：中間ステップをより小さく高速なモデルにルーティングすること）であっても、それがモデル特有の状態の蓄積喪失や複雑なタスクの完了失敗を引き起こす場合、グローバルな劣最適化を招く可能性があります。
3. 静的なポリシー: 既存のルーターは、デプロイ前にルーティングポリシーを固定するオフライン学習済みの分類器に依存することが多く、ワークロードの変化への適応や、タスク開始後に誤った割り当てを修正することを妨げます。

手法: TRACE-Router

著者らは、ルーティングおよび学習の単位を個別のリクエストからタスク全体のトレースへとシフトさせるフレームワークであるTRACE-Routerを提案しています。このシステムは、「ルーティングの単位はフィードバックの単位と一致しなければならない」という原則に基づいて動作します。

コアメカニズム

1. タスク一貫性のあるルーティング（Task-Consistent Routing）:

   • タスク受理時、ルーターは実行トレース全体に対して単一のバックエンドモデルを選択します。
   • トレース内のすべてのリクエストには、永続的なタスク識別子が付与されます。後続のリクエストはこの保存されたバックエンド割り当てを再利用し、実行中を通じてモデル固有の状態と一貫性を維持します。
   • これにより、単一のクレジット割り当て対象が作成されます。すなわち、タスク全体の報酬は、トレース全体にサービスを提供した単一のバックエンドに帰属されます。

2. コンテキスト条件付きオンライン学習（Context-Conditioned Online Learning）:

   • 多様なワークロードレジーム（例：レイテンシが極めて重要なタスク vs 正解率が極めて重要なタスク）を扱うため、システムはコンテキスト・バンディットのアプローチを採用しています。
   • タスクは、受理時に利用可能な特徴量（例：正規表現ベースの複雑性分類器）に基づき、追加のLLM呼び出しを必要とせずに、粗いルーティング・コンテキスト（例：Easy, Medium, Hard）に割り当てられます。
   • システムは各コンテキストに対して独立したバンディット・ポリシー（$B_c$）を保持し、異なるタスクタイプに特化した精度とレイテンシのトレードオフを学習します。

3. コールドスタートUCBポリシー（Cold-Start UCB Policy）:

   • ルーティング・ポリシーは、コールドスタート初期化を備えた**上限信頼限界（UCB）**アルゴリズムを利用します。
   • 各コンテキストにおける疎なフィードバックの統計的な困難に対処するため、システムはUCB選択ルールに移行する前に、強制的なラウンドロビン探索（各アームにつきk回のプル）を実行します。
   • 著者らは、「ウォームスタート」による事前分布の導入は精度を向上させず、むしろシステムを低速なモデルへと偏らせる傾向があることを発見したため、コールドスタートをデフォルトとしています。

4. 遅延報酬とスカラー化（Delayed Reward and Scalarization）:

   • フィードバックはタスク終了後にのみ観測されます。システムは、勾配付きの精度（$a_{t,m}$）とエンドツーエンドのレイテンシ（$\ell_{t,m}$）を、パラメータ $\alpha$ を用いて以下のスカラー報酬に結合します：
     $$r^{(\alpha)}_{t,m} = (1 - \alpha) a_{t,m} - \alpha \tilde{\ell}_{t,m}$$
   • $\alpha$ をスイープ（変化）させることで、システムは一連のデプロイメント・ポリシーを追跡でき、運用者は精度とレイテンシのトレードオフを調整することが可能になります。

主な貢献

1. タスクレベルのルーティング定式化: 本論文は、モデルルーティングをタスク一貫性のあるオンラインモデル選択として定式化し、ルーティング決定をフィードバックの単位（タスクトレース）と一致させることで、遅延報酬からの原理的なクレジット割り当てを可能にしました。
2. コンテキスト認識型のオンライン適応: TRACE-Routerは、オフライン学習や明示的なタスク複雑度の推定を必要とせず、コンテキストに応じたポリシーをオンラインで学習するコンテキスト・バンディット型ルーターを導入しました。
3. 改善されたパレート・フロント: 本フレームワークは、経験的な精度・レイテンシのパレート・フロントを一貫して学習し、ヒューリスティックなベースライン（例：複雑性ルーター）や学習済みのリクエストレベルのルーターを凌駕します。

実験結果

著者らは、3つのエージェント・ベンチマーク（$\tau^2$-Bench（小売および通信ドメイン）、LiveCodeBench、Terminal-Bench）を用いてTRACE-Routerを評価しました。

• パレート・フロントの優位性: すべてのベンチマークにおいて、TRACE-Routerは非劣なパレート・フロントの点を達成しました。これは、モデル間のレイテンシ一致型ランダム補間を一貫して上回っています。
  • $\tau^2$-Benchにおいて、レイテンシ一致型補間よりも7〜8ポイント精度が高かった。
  • Terminal-Benchにおいて、最強の単一モデル・ベースラインよりも7.1ポイント高い精度を達成しつつ、レイテンシを36%削減した。
• コンテキスト条件付け: コンテキスト条件付けによる総体的な利得は緩やかでしたが、到達可能な動作点を決定する上で極めて重要でした。コンテキスト条件付きのバリアントは、特にモデルの性能特性が分岐するドメイン（例：通信）において、コンテキストフリーのバリアントが劣位となった特定の $\alpha$ 値において、フロントを維持しました。
• コールドスタート vs ウォームスタート: 実験により、コールドスタート初期化がウォームスタート・バリアントよりも優れていることが確認されました。ウォームスタートは事前分布のバイアスを強化し、低速なモデルを選択する傾向がありました。
• スケーラビリティ: 4つのモデルを用いたアブレーション研究において、TRACE-Routerは、ヒューリスティックなルーターが選択し続ける信頼性の低いモデル（例：タイムアウトしやすいモデル）を特定し、回避することに成功しました。これは、候補プールが増大しても堅牢であることを示しています。

意義と主張

本論文は、タスク一貫性のあるコミットメントが、デプロイ時のモデル選択における実用的かつ優れた単位であることを主張しています。ワークフローを独立したリクエストの集合としてではなく、第一級のスケジューリング・オブジェクトとして扱うことで、TRACE-Routerは現在のリクエストレベルのルーターに内在するクレジット割り当てのミスマッチを解決します。

著者らは、自身のアプローチが明示的なタスク複雑度の推定やオフライン学習を必要としないことを強調しています。代わりに、遅延したタスクレベルのフィードバックを活用することで、精度とレイテンシのバランスを取るルーティング・ポリシーを適応的に学習します。これらの結果は、適応的なタスクレベルのルーティングが、同じレイテンシ予算内で確率的なモデル混合を行うよりも大幅に高い精度をもたらすことを示唆しており、より効率的で信頼性の高いエージェントAIデプロイメントへの道筋を示しています。