🤖 AI Summary
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.
📝 Abstract
Memory-efficient matrix optimizers such as Sinkhorn gradient descent remove most AdamW optimizer state for dense Transformer matrices, but direct application to Mixture-of-Experts (MoE) training is unreliable. We study this failure in a controlled 110M-parameter nanowhale DeepSeek-style MoE pretraining setting. A SAGE/Sinkhorn hybrid reduces optimizer state from 0.883GB to 0.331GB but degrades evaluation loss to 3.8265, far above the AdamW baselines observed in the same setup (3.58--3.64 across the seeds we study). We show that routed MoE expert matrices are the dominant failure point: their gradients are conditional, temporally varying, and poorly served by stateless Sinkhorn normalization. We propose MESH, a hidden-momentum Sinkhorn update for MoE experts. MESH restores a temporal first-moment signal through the gradient-buffer lifecycle, without storing the expert first moment as optimizer state. MESH is an optional block-preconditioned variant that adds a coarse neuron/block inverse-RMS multiplier. Across ablations, temporal smoothing before matrix normalization is the primary causal ingredient; block/neuron preconditioning can improve the memory-quality frontier, but is not established as universally necessary. In two additional seeds, MESH and MESH-B reduce optimizer-state memory by 62.5\% and peak PyTorch CUDA allocation by about 12.6\% relative to AdamW, with a modest evaluation-loss gap. Full-state diagnostic variants recover AdamW-like performance in ablations, supporting the conclusion that MoE experts need temporal smoothing, but not necessarily full coordinate-wise AdamW state.