Search papers, labs, and topics across Lattice.
This paper investigates the limitations of applying memory-efficient Sinkhorn gradient descent to Mixture-of-Experts (MoE) training, revealing that the method fails due to the conditional and temporally varying nature of routed MoE expert gradients. The authors introduce MESH, a hidden-momentum Sinkhorn update that effectively restores a temporal first-moment signal without the need for extensive optimizer state, achieving a significant reduction in memory usage. Experimental results show that MESH can decrease optimizer-state memory by 62.5% while maintaining competitive evaluation loss compared to traditional AdamW optimizers.
MESH achieves a 62.5% reduction in optimizer-state memory for Mixture-of-Experts training while preserving performance, challenging assumptions about memory efficiency in deep learning.
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.
