UniMoMo: エキスパートマージングに基づくMoE高速化手法 for 大規模推薦モデル

要旨

スパース混合エキスパート（MoE）レイヤーは、条件付き計算によって推薦能力を拡張するが、学習済みチェックポイントは依然として全エキスパートバンクを格納し、それらすべてにルーティングを行う。本研究では、デプロイメント問題、すなわち、圧縮専用のオンラインモジュールを追加することなく、明示的なエキスパート予算の下で、そのチェックポイントをより小さな標準MoEへ変換する問題を扱う。この課題に対処するため、制約付きグラフ粗化問題として定式化された学習後圧縮フレームワークであるUniMoMoを導入する。UniMoMoはパラメータ距離に依存せず、ラベルなしキャリブレーションセットを用いて、共有された推薦状態に対してエキスパートがどの程度類似した応答を示すかを測定し、機能的類似性に基づいてエキスパートをグループ化する。性能低下を防ぐため、ルーティング露出に基づいて高トラフィックエキスパートのマージを制限するレイヤー適応型保護メカニズムを導入する。Amazon Beauty、KuaiRec、TenRecにおいて、2、4、6個のMoEブロックを用いた場合、最終的な4エキスパートのチェックポイントは、ソース相対の5回実行平均NDCG@10比が99.92%～102.30%、A100での実測高速化が1.28倍～1.63倍となった。より積極的な2エキスパート、top-1動作点では、98.36%～104.24%の比率と1.47倍～2.21倍の高速化が得られた。これらのエンドポイント結果は、変換と適応の完全なワークフローを評価するものであり、学習済み推薦MoEが複数のサービング予算でエクスポート可能であることを示している。

English

Sparse mixture-of-experts (MoE) layers expand recommendation capacity through conditional computation, yet a trained checkpoint still stores and routes over its full expert bank. We study a deployment problem: convert that checkpoint to a smaller standard MoE under an explicit expert budget, without adding a compression-specific online module. To address this, we introduce UniMoMo, a post-training compression framework formulated as a constrained graph coarsening problem. Rather than relying on parameter distance, UniMoMo groups experts based on their functional similarity, using an unlabeled calibration set to measure how similarly experts respond to shared recommendation states. To prevent performance degradation, we introduce a layer-adaptive protection mechanism that restricts the merging of high-traffic experts based on their routing exposure. Across Amazon Beauty, KuaiRec, and TenRec with 2, 4, and 6 MoE blocks, the final four-expert checkpoints obtain source-relative five-run mean NDCG@10 ratios of 99.92%--102.30% and measured A100 speedups of 1.28times--1.63times. An aggressive two-expert, top-1 operating point obtains ratios of 98.36%--104.24% and speedups of 1.47times--2.21times. These endpoint results evaluate the complete conversion-and-adaptation workflow and show that a trained recommendation MoE can be exported at multiple serving budgets.