技術要約: TriRoute

問題提起

大規模言語モデル（LLM）のデプロイには多大な推論コストがかかるが、現在のリソース割り当ては統計的に無駄が多い。高密度なトランスフォーマーは、トークンの意味的な重要性に関わらず、一律の浮動小数点演算（FLOPs）とKey/Value（KV）キャッシュメモリを適用してしまう。文法的に些細な機能語（function word）は、段落の意味を決定づける希少な固有名詞と同じリソースを消費する。

既存の条件付き計算手法は、この一様性を打破しようと試みているが、単一の軸に対して孤立して動作している：

1. Mixture-of-Experts (MoE): トークンをサブセットの専門家（expert）にルーティングすることで、フィードフォワードネットワーク（FFN）を疎化する。
2. Mixture-of-Depths (MoD): 予測可能であると判断されたトークンに対して、トランスフォーマーのブロック全体（アテンション＋FFN）をスキップする。
3. KV-Cache Quantization: キーとバリューをより低いビット幅で保存することで、アテンションメモリを圧縮する。

本論文は、これらの決定は強く結合していると主張している。例えば、希少なエンティティのトークンは、将来のクエリによって忠実に想起されるために、フル解像度のアテンションと高精度なキャッシュ保存を必要とする可能性がある一方で、そのFFN変換は予測可能でスキップ可能なものであるかもしれない。逆に、機能語はこれら3つの次元すべてにおいてスキップ可能かもしれない。これらを個別のハイパーパラメータを持つ独立した最適化問題として扱うことは、これらの相関関係を捉えることができず、サブオプティマルなリソース割り当てや、「テールケース（稀なケース）」への堅牢性（希少なエンティティ、コード、算術に関する性能）の低下を招く。

手法: TriRoute

著者らは、3つの孤立したメカニズムを単一の軽量な学習済みコントローラーに置き換える統一アーキテクチャ、TriRouteを提案する。すべてのレイヤー $\ell$ におけるすべてのトークン $t$ に対して、コントローラーは3つの異種混合の軸にわたる協調されたポリシー $\pi^{(\ell)}_t = (a_t, e_t, b_t)$ を出力する：

1. アテンション解像度 ($a_t$): スキップ（コスト0）、ローカル（ウィンドウ化されたアテンション）、またはフル（高密度の因果的アテンション）の選択。
2. エキスパート選択 ($e_t$): トークンを疎なエキスパート集合、またはヌル・エキスパート ($f_0 \equiv 0$) へルーティングする。ヌル・エキスパートにより、モデルはフル・アテンションを維持しながらFFNをスキップすることが可能となる。これはMoDでは表現できないレジームである。
3. KV-キャッシュ・ビット幅 ($b_t$): 将来のクエリのために、トークン自身のKey/Valueペアがキャッシュに書き込まれる際の精度（2, 4, 8, または16ビット）を決定する。

技術的構成要素

• 共有コントローラー: 2層のMLPトランクが、正規化された残差状態と安価な因果的サイド特徴量（位置、距離、予測エントロピー）を処理し、3つの独立したヘッド（Attention, Experts, Bits）のためのロジットを生成する。この設計により、共通の「トークンの重要性」という概念を用いて全軸に情報を伝えつつ、個別の出力サブスペースを維持できる。
• 異種混合の緩和 (Heterogeneous Relaxation): 学習には離散的な決定の微分が必要である。著者らは以下を使用する：
  • カテゴリカルなアテンションおよびビット幅の決定には、Gumbel-SoftmaxとStraight-Through (ST) 推定器を使用。
  • エキスパート選択には、Top-k Softmax Gatingを使用。
  • 軸間勾配バランシング (Cross-Axis Gradient Balancing): アテンションをスキップすることの影響は、キャッシュのビット数を減らすことよりもはるかに大きいため、著者らは（勾配ノルムのEMAを用いた）軸ごとの勾配再スケーリングメカニズムを導入し、共有トランクにおけるアテンション軸による支配を防いでいる。
• 結合を考慮したバランシング (Coupling-Aware Balancing): 一つの軸での崩壊が他の軸の崩壊を引き起こす「崩壊カスケード（collapse cascade）」（低分散の入力によるもの）を防ぐため、以下の手法を採用している：
  • 軸ごとのホワイトニング (Per-axis Whitening): 各ヘッドの前に、各ヘッドのトランク特徴量をデタッチ（分離）して正規化する。
  • エントロピー・フロア (Entropy Floor): オプションの使用エントロピーが閾値を下回ったときにのみ作動するペナルティ。これにより、予算制約と衝突することなく多様性を確保する。
• 微分可能なコストモデルとラグランジュ制御: 単一のラグランジュ制約が、FLOPsとメモリの両方をカバーする目標平均予算 $C^\star$ を強制する。双対変数（ラグランジュ乗数）はオンラインで更新され、実務者が3つの個別のスパース性ハイパーパラメータを調整することなく、単一のスカラーノブで全パレート・フロントをスイープすることを可能にする。

主な貢献

1. 統一された定式化: TriRouteは、アテンション解像度、エキスパート選択、およびKVキャッシュ精度を、単一のコントローラーと予算の下で共同学習する初めてのアーキテクチャである。
2. 安定した共同学習レシピ: 著者らは、軸ごとのホワイトニング、エントロピー・フロア、および勾配バランシングを通じて、「軸間ルーティング崩壊カスケード」を特定し、緩和することで、異種混合の離散決定のEnd-to-End学習を可能にした。
3. 設計空間の探索: 共有トランクと個別のヘッドを持つ構成が最適な粒度であり、完全に分離されたルーターや完全に共有された単一のヘッドよりも優れていることを示している。また、ヘッダーごとのアテンション・ルーティングと、トークンレベルのFFN/ビット・ルーティングの組み合わせが「スイートスポット」であることを示している。
4. テールケースの保存: 平均パープレキシティを最適化する適応型モデルがしばしば希少な入力に対して性能を低下させるのに対し、TriRouteはリソースを必要な場所に正確に割り当てることで、希少なエンティティ、コード、および算術に関する堅牢性を維持する。

結果

計算量最適化されたトークン数で訓練された、160Mから1.3Bパラメータのデコーダーのみのモデルを用いて評価を行った：

• パレート優位性: TriRouteは、独立してチューニングされた最良の組み合わせ（MoD + MoE + KV量子化）を厳密に上回るコスト・品質のフロントを実現する。一致した予算（約55%のFLOPsと40%のKVメモリ）において、TriRouteは高密度モデルのダウンストリーム精度を96〜99%回復する。
• 性能向上: 独立した組み合わせと比較して、TriRouteはパープレキシティを0.3〜0.4改善し、すべてのスケールにおいてダウンストリーム精度を0.7〜1.0ポイント向上させた。
• テールの堅牢性: 希少エンティティ、コード、および数学のバケットにおいて、独立した組み合わせは顕著な劣化（例：希少エンティティで+2.1のパープレキシティ）を示したが、TriRouteは高密度モデルの0.4パープレキシティ以内に留まり、共同的な調整が「サイレントな劣化」を防ぐことを示した。
• 解釈可能性: コントローラーは解釈可能なパターンを学習していることが分析により明らかになった：
  • 読み書きの忠実度: 希少なサブワードや固有名詞は、フル・アテンションと高精度なキャッシュを受け取る。
  • デカップリング: 希少なエンティティは、フル・アテンションと高精度なキャッシュを維持しながら、FFNをスキップ（ヌル・エキスパート）することが多い。これはブロック単位のMoDでは不可能な挙動である。
  • 深さの特化: 初期のレイヤーは全般的に高いアテンションを維持するが、後のレイヤーは選択的になり、非クリティカルなトークンに対してアテンションをスキップする。

意義と主張

本論文は、共有された学習済みコントローラーが、固定された推論予算を最も重要な場所に費やすための原理的な方法であることを主張している。3つの軸を統合することで、TriRouteは手動で調整された孤立したメカニズムの非効率性を回避する。

著者らは、結果から、「正しい」計算量は、トークンあたりの単一のスカラーではなく、異種のリソースにわたる結合された選択のベクトルであるという示唆を得たと強調している。統一されたアプローチは、パレート・フロントを改善するだけでなく、どの入力が節約に貢献するかを根本的に変え、希少で重要なトークンが過小評価されないようにしている。

最後に、著者らは実用的なデプロイメントに関する謙虚な注記を述べている。理論的なFLOP削減は顕著であるが、実際のウォールクロック・スピードアップは、現在、システムレベルのギャップ（不規則なアテンション、混合精度キャッシュ管理、ルーターのオーバーヘッド）によって制限されている。このFLOP-レイテンシ間のギャップを埋めることが、実用的な採用に向けた主要な次のステップであると考えている。