arxiv Post-Trained MoE Can Skip Half Experts via Self-Distillation

https://huggingface.co/TsinghuaC3I Tsinghua University

Mixture‑of‑Experts (MoE) basics
An MoE layer contains many “expert” feed‑forward sub‑networks, but for each token only a small subset is activated by a router. This sparse activation lets the overall model grow very large while keeping the per‑token compute bounded. Dynamic MoE variants go further by letting the router decide, input‑dependently, how many experts to use, reducing computation for easy tokens.

Problem the paper addresses
Existing dynamic‑MoE techniques usually require training the model from scratch or performing task‑specific fine‑tuning. Consequently, a fully trained static MoE model cannot be easily turned into a dynamic one without risking loss of the routing knowledge that was already learned. This limits practical deployment because inference costs remain high even when many tokens could be handled with fewer experts. <citation src=“1”></citation>

How the new method (ZEDA) works

  1. Zero‑Expert injection – a parameter‑free “zero‑output” expert is added to every MoE layer. These experts produce no contribution unless selected, allowing the model to skip computation for certain tokens.
  2. Two‑stage self‑distillation –
    • Stage 1 (SFT): the augmented model is fine‑tuned while staying close to the original outputs.
    • Stage 2 (OPD): the original static MoE acts as a frozen teacher; the student (augmented model) learns via self‑distillation, guided by a group‑level balancing loss that keeps expert loads even.
      This stabilises the conversion from static to dynamic architecture without needing a new pre‑training run. <citation src=“2”></citation>

Benefits emphasized by the authors

  • Computation reduction: > 50 % of expert FLOPs are eliminated, meaning many tokens bypass expert computation entirely.
  • Small accuracy loss: performance drops only marginally across a suite of 11 benchmarks (math, code, instruction following).
  • Speed‑up: about 1.2 × end‑to‑end inference acceleration is observed.
  • Low adaptation cost: ZEDA requires only a brief self‑distillation fine‑tuning phase, avoiding the expense of full re‑training.
  • Compatibility: because zero‑experts are parameter‑free, the method can be applied to existing large MoE models without architectural redesign. <citation src=“4”></citation>