SlimWise、MoEのデコード時だけ専門家を削減して推論を高速化

概要

Mixture-of-Experts（MoE）モデルは、トークンごとに一部の専門家だけを呼び出すことで、単一の巨大な密モデルより計算量を抑えます。しかし、実際のバッチ生成では、この利点がそのままシステム性能になるとは限りません。複数のリクエストが異なる専門家へルーティングされるため、バッチ全体では専門家プールのほぼ全体にアクセスすることがあります。その結果、専門家の重みを読み出して移動するトラフィックが大きなボトルネックになります。

SlimWiseは、推論全体に同じ剪定を適用するのではなく、プリフィルとデコードの性質を分けて扱います。入力コンテキストは完全モデルで処理し、トークンを逐次生成するデコード段階だけを軽量化する設計です。

主な仕組み

• プリフィルでは完全な専門家を維持：プリフィルはコンテキスト全体をまとめて処理する段階です。ここで専門家を削ってもスループット上の効果が限定的な場合があるため、SlimWiseはモデル品質を優先します。
• デコードで専門家プールを縮小：デコードでは各トークンの生成ごとに重みアクセスが繰り返されます。専門家数を減らすことで、こうした重みトラフィックの削減を狙います。
• KVキャッシュを直接引き継ぐ：完全モデルがプリフィルで作ったKVキャッシュを、剪定済みのデコーダーが変換なしで利用します。この学習不要の引き継ぎにより、剪定だけを行う場合よりも品質低下を抑えます。
• 軽量な蒸留を追加可能：剪定後に残る精度差や生成長の変化に対して、完全モデルのKVキャッシュから続きを生成できるようデコーダーを蒸留します。更新対象は少数のパラメーターに限定されます。

論文は、タスクの正解率だけでは剪定の影響を十分に評価できない場合があると指摘します。生成長が変われば、ベンチマーク上のスコアが近くても、実際の計算量やサービスコストは異なり得るためです。SlimWiseはvLLMに実装され、プリフィルとデコードを分離する構成と、同一環境に配置する構成の両方に対応します。2種類のMoEバックボーンと3種類の剪定基準で評価され、Qwen3.6-35B-A3Bでは専門家50%削減時にデコードスループットが最大1.81倍向上しました。

意義と課題

SlimWiseが示すのは、モデル圧縮も推論フェーズごとのボトルネックに合わせるべきだという考え方です。コンテキスト処理では完全モデルを使い、重みアクセスが繰り返される生成ループだけを圧縮することで、品質と効率の中間点を作れます。

一方、効果はバッチサイズ、入力と出力の長さ、ルーティング分布、デプロイ方式によって変わります。蒸留を使う場合は追加の学習と運用も必要です。そのため、SlimWiseは無条件に最大限の剪定を行う手法というより、完全モデルと積極的な圧縮の間を柔軟に調整するサービング設計として捉えるのが適切です。

出典：Hugging Face Daily Papers