UniMoMo Compresses Recommendation MoE Models by Merging Experts

Introduction

Sparse mixture-of-experts layers are attractive for large recommendation systems because they increase model capacity while activating only a subset of experts for each input. Yet deployment remains expensive: even if computation is sparse, a trained checkpoint still stores the full expert bank and the router still operates over the original set of experts. For production recommenders, that can translate into higher memory use, more routing overhead, and stricter latency constraints.

UniMoMo addresses this deployment-side question directly: given a trained recommendation MoE and a target expert budget, can we convert it into a smaller standard MoE without adding a special compression module at inference time?

Key Ideas

• Functional similarity instead of parameter similarity: Many expert-compression methods compare weights, but similar weights do not always imply similar behavior. UniMoMo uses an unlabeled calibration set and compares expert outputs on shared recommendation states, grouping experts that behave similarly in practice.

• Expert merging as constrained graph coarsening: The method formulates the conversion process as a constrained graph coarsening problem. This gives the expert-budget constraint a clear role and makes merging more systematic than simply pruning or averaging experts.

• Traffic-aware protection: In real MoE recommenders, experts are not used equally. Some experts receive far more routing exposure and may carry important patterns. UniMoMo introduces a layer-adaptive protection mechanism that restricts the merging of high-traffic experts, reducing the risk of damaging major routing paths.

• No extra serving-time component: The compressed checkpoint is still a conventional top-k MoE. This matters for production systems, because avoiding additional online modules keeps the serving stack simpler and reduces integration risk.

• Calibration-aware reconstruction: For merged SwiGLU experts, UniMoMo applies a least-squares correction to better recover intermediate activations from the original model, helping reduce the representational drift caused by merging.

Results and Impact

Across Amazon Beauty, KuaiRec, and TenRec, with models containing 2, 4, and 6 MoE blocks, UniMoMo’s four-expert checkpoints reach five-run mean NDCG@10 ratios of 99.92% to 102.30% relative to the source model. Measured A100 speedups range from 1.28x to 1.63x. Under the more aggressive two-expert, top-1 setting, NDCG@10 ratios are reported at 98.36% to 104.24%, with speedups from 1.47x to 2.21x.

The broader message is that trained recommendation MoE models may contain mergeable functional redundancy. UniMoMo turns that observation into a practical export workflow: teams can prepare multiple serving variants under different expert budgets while retaining a standard MoE architecture.

There are still deployment considerations. The calibration set must reflect the states the model will see in production, and changes in traffic distribution could affect which experts should be protected. Even so, UniMoMo is notable because it frames MoE compression not just as a model-size problem, but as an inference-serving problem where routing behavior, quality preservation, and operational simplicity all matter.

Source: Hugging Face Daily Papers