Reproduce “DOT-MoE: Differentiable Optimal Transport for MoEfication” on OpenResearch

About this paper

Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta

Amazon researchers developed DOT-MoE, a framework that leverages Differentiable Optimal Transport to convert dense Large Language Model (LLM) layers into sparse Mixture-of-Experts (MoE) architectures. The method significantly reduces active parameters during inference by 50% while retaining approximately 90% of the original dense model's performance, outperforming previous pruning and MoEfication techniques across various LLM benchmarks.

Key insights

• Reformulating neuron decomposition as a Differentiable Optimal Transport problem enables a principled, output-aware approach to constructing functionally coherent and balanced experts.
• Jointly and differentiably optimizing neuron assignment and token routing allows for co-adaptation between expert structure and routing policy, leading to superior initial sparse models.
• Using entropic regularization and Straight-Through Estimators overcomes the combinatorial intractability and non-differentiability of sparse assignment and routing decisions.

Summary by alphaXiv.