UniMoMo：推薦 MoE を専門家マージで軽量化する手法

導入

疎な mixture-of-experts（MoE）は、推薦モデルの容量を拡張する有力な方法です。入力ごとに一部の専門家だけを使うため、計算量を抑えながら多様なユーザー行動やアイテム文脈を表現できます。しかし、学習済み checkpoint には依然として全専門家が保存され、ルータも元の専門家集合を対象に動作します。大規模な推薦サービスでは、これはメモリ、遅延、運用コストの問題になります。

UniMoMo が扱うのは、この現実的なデプロイ課題です。すでに学習した推薦 MoE を、指定された専門家数の範囲内で、追加のオンライン圧縮モジュールなしに小さな標準 MoE へ変換します。

核心ポイント

• パラメータではなく機能で比べる：重みの距離が近い専門家が、必ずしも同じように振る舞うとは限りません。UniMoMo は無ラベル校正データを使い、共通の推薦状態に対する各専門家の出力を比較して、機能的に近い専門家を見つけます。

• 制約付きグラフ粗化として定式化：専門家統合を、専門家数の予算を持つグラフ粗化問題として扱います。単純な削除や平均化ではなく、どの専門家を同じグループにまとめるかを体系的に決める設計です。

• 高トラフィック専門家を保護：MoE では専門家の利用頻度が均一ではありません。多くルーティングされる専門家は重要な経路を担っている可能性があります。UniMoMo は層ごとに適応する保護機構を導入し、ルーティング露出の高い専門家の無理な統合を制限します。

• 推論時の構造は標準 MoE のまま：圧縮後のモデルは通常の top-k MoE として動作します。特別な推論モジュールを必要としないため、既存の serving パイプラインに組み込みやすい点が実用上の強みです。

• 校正を使った再構成：SwiGLU 専門家の統合後には、最小二乗による補正で元モデルの中間活性をよりよく再現し、統合による表現のずれを抑えます。

結果と意義

Amazon Beauty、KuaiRec、TenRec の各データセット、および 2、4、6 個の MoE block を持つ設定で、4 専門家 checkpoint は元モデルに対する 5 回平均 NDCG@10 比率が 99.92%〜102.30% となりました。A100 上の実測高速化は 1.28 倍〜1.63 倍です。さらに積極的な 2 専門家・top-1 設定では、NDCG@10 比率が 98.36%〜104.24%、高速化が 1.47 倍〜2.21 倍と報告されています。

この結果は、推薦 MoE の専門家群に機能的な冗長性が存在し得ることを示しています。UniMoMo の重要性は、それを単なる圧縮率の話にせず、複数の serving 予算に合わせて標準 MoE checkpoint をエクスポートできる実務フローに落とし込んだ点にあります。

一方で、校正データが本番分布をどれだけ反映するかは重要です。トラフィックが大きく変化すれば、保護すべき専門家も変わる可能性があります。それでも、品質、ルーティング、推論単純性を同時に扱う UniMoMo は、推薦 MoE の運用負荷を下げる有望な方向と言えます。

出典：Hugging Face Daily Papers