大規模言語モデルの推論効率と品質のバランスは、現在の研究ホットトピックです。Mixture-of-Experts(MoE)、Mixture-of-Depths(MoD)、KVキャッシュ量子化などの既存の条件付き計算手法は、それぞれFFNの間引き、Transformerブロックのスキップ、アテンションメモリの圧縮を独立して行います。しかし、これらの決定は実際には強く結合しています。まれなトークンが完全なアテンションを必要とする場合、どのエキスパートが処理するかに関わらず高精度キャッシュも必要となる可能性があります。この問題に対処するため、Andrii BalashovとOlena PonomarovaはTriRouteを提案しました。これは、各トークンの各層において、(i) アテンションモード(スキップ/ローカル/フル)、(ii) スパースなFFNエキスパート集合(MoDを回復するヌルエキスパートを含む)、(iii) KVキャッシュビット幅を協調的に出力する軽量な統一コントローラです。
TriRouteの学習は、カテゴリカル決定にはGumbel-Softmaxとストレートスルー推定(straight-through estimator)を用い、エキスパート選択には負荷分散top-kゲーティングを採用した異種緩和手法で行われます。ラグランジュ予算制約の下で平均計算・メモリコストを制御可能なノブに変換し、トレードオフを柔軟に調整できます。研究チームは、素朴な統合学習における軸間ルーティング崩壊カスケード(1つの軸の崩壊が他に伝播する現象)を特定し、軸ごとの正規化と結合認識バランス損失で対処しました。これにより、各軸の出力分布が多様性を保ち、軸間の過度な相関が抑制されます。
実験では、160Mから1.3Bパラメータのデコーダのみモデルを計算最適トークン数で学習し、同一推論FLOPs・メモリにおいてTriRouteが最高の独立MoD+MoE+KV量子化組み合わせをPareto支配することを確認しました。特に、希少エンティティ、コード、算術におけるテールケースのロバスト性が大幅に向上し、純粋なパープレキシティ最適化では失われがちな性能を保持します。事後分析により解釈可能な構造が明らかになりました。コントローラは文頭位置、希少サブワード、固有名詞に完全アテンションと高精度キャッシュを割り当て、機能語には安価なルーティング(スキップアテンションや低精度キャッシュ)を行います。
この研究は、アテンション、エキスパート、キャッシュ決定の統合最適化がリソース制約下での言語モデルの性能を大幅に向上させ、将来のより効率的な適応型推論システムの基盤を築くことを示しています。研究チームは論文とコードの一部を公開し、大規模モデルやマルチモーダルシナリオへの適用を今後の課題としています。