arXiv (ML)AI
TITLE_JA: TriRoute:注意機構、専門家選択、KVキャッシュ割り当ての統合学習ルーティング
TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation
この記事についてAIに質問する →
日本語要約青い用語にマウスを合わせると解説が表示されます
言語モデルの推論効率を向上させるための複数の技術が存在しますが、従来はそれぞれが独立に機能していました。Mixture-of-Experts(MoE)はフィードフォワードネットワークを疎にし、Mixture-of-Depths(MoD)はトランスフォーマーブロック全体をスキップし、KVキャッシュ量子化は注意メモリを圧縮します。しかし新しい研究では、これら3つの決定(注意の解像度、専門家選択、キャッシュのビット幅)が強く相互に関連しており、共同で最適化されるべきだと主張しています。稀なトークンは完全な注意を必要とするだけでなく、どの専門家で処理されるかに関わらず高精度のキャッシュも必要になる可能性があるからです。
研究チームが提案したTriRouteは、3つの軸すべてに共有される軽量コントローラで、各層の全トークンに対して統一されたポリシーを生成します。具体的には、スキップ・ローカル・完全のいずれかの注意モード、疎な専門家セット、そしてKVキャッシュのビット幅を決定します。このコントローラーはGumbel-Softmaxと直進推定、そしてロードバランスの取れたtop-kゲーティングを用いて複数の軸にわたってエンドツーエンドで訓練されます。
訓練過程では、単純な共同学習において「ルーティング崩壊カスケード」という問題が生じることが判明しました。これは1つの軸での崩壊が他の軸に伝播する現象ですが、軸ごとの正規化とカップリング対応の均衡損失関数により対処されています。160M~1.3Bパラメータのモデルでの実験では、TriRouteは個別のMoD、MoE、KV量子化の組み合わせを上回り、稀なエンティティ、コード、算術演算などの周辺事例の堅牢性をより良く保持します。詳細分析から、コントローラが文頭の位置、稀な部分語、固有表現に完全な注意と高精度キャッシュを割り当て、機能語には安価なルーティングを適用する解釈可能な構造が明らかになりました。