ACE、tokenごとに寄与度の低いMoE expertをスキップし、A100のprefillを最大2.25倍高速化

MoEモデルは各tokenについて一部のexpertしか有効化しないものの、一般的な固定top-k routingでは、すべてのtokenが同じ数のexpert slotを実行する必要がある。新たな研究であるACEは、この冗長性をさらに削減しようとするものだ。元のcheckpointとrouterを維持し、別途configuratorを学習することなく、推論時にtokenごとに、選択済みの副次的なexpertのうちどれを省略できるかを判定する。一方、重みが最大のtop-1 expertは必ず維持される。

ACEは2種類のオフライン信号を使用する。Global Spectral Proxy（GSP）は、各SwiGLU expertのgate、up projection、down projection、および前段のRMSNormを統合的に解析し、その全体的な変換能力を推定する。Router-Conditioned Refinement（RCR）は、同一layer内のrouter weightを中心化してexpertが選好する入力方向を構築し、その方向における応答強度を測定する。オンライン段階では、lookup tableを参照し、現在のtokenのgate値と組み合わせて少量のscalar演算を行うだけでよい。GSPとRCRの両方が寄与度を低いと判定した場合にのみ、システムはそのexpertをスキップする。

著者らはQwen3-30B-A3B-Instruct-2507、Qwen3.6-35B-A3B、Gemma-4-26B-A4B-itを対象に、WikiText-2のほか、ARC、PIQA、MATH-500、GPQA-Diamond、HumanEval、LiveCodeBenchなど計7つのタスクで評価した。Qwen3.6-35B-A3Bでexpertスキップ率を50%に設定した例では、ACEは最も強力なベースラインと比べてWikiText-2のperplexityを7.96%低減し、7タスクの平均accuracyを4.15 percentage points向上させた。これは同一のスキップ予算における比較であり、完全なBF16モデルを上回って全般的に性能が向上したという意味ではない。

最適化済みのexpert dispatchを単一のNVIDIA A100で使用した場合、スキップ率60%でfirst-token latencyは1.72～2.25倍、per-token latencyは1.31～1.41倍改善した。入力長1,024、batch size 1では、TTFTが270.9ミリ秒から120.3ミリ秒に短縮された。公開リポジトリにはすでにend-to-end pipelineとテストが含まれているが、一部の補助実験は除外されており、READMEのキャリブレーション手順でも、目標スキップ率をマッピングするためにWikiTextのtokenが依然として使用されている。論文は、結果が単一のdeterministic runとワークロードに適合させたしきい値に基づくことを認めている。今後は、しきい値のクロスドメイン転移、multi-GPU expert parallelにおける通信、さらにvLLMやSGLangなどの本番runtimeへ統合した際の効果を検証する必要がある。