MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training
This work addresses the failure of Sinkhorn-based optimization in Mixture-of-Experts (MoE) training, which stems from the expert routing matrix exhibiting a high and highly time-varying gradient condition number. The authors propose MESH, a novel method that identifies the expert matrix as the primary cause of performance degradation and introduces an implicit momentum mechanism. This mechanism provides first-moment signals along the temporal dimension within the gradient buffer’s lifetime, thereby eliminating the need to explicitly store expert-level AdamW optimizer states. MESH optionally incorporates block- or neuron-level inverse RMS preconditioning to further enhance stability. Compared to AdamW, MESH reduces optimizer state memory by 62.5% and peak CUDA memory by 12.6%, with only a minor increase in evaluation loss. Ablation studies confirm that temporal smoothing is the key factor underlying its effectiveness.