🤖 machine learning

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

This paper introduces MESH, a memory-efficient Sinkhorn optimization method that incorporates hidden-momentum updates and optional block preconditioning to successfully enable low-memory training of Mixture-of-Experts models by addressing the temporal instability of routed expert gradients that causes standard stateless Sinkhorn methods to fail.

Original authors: Masato Fujitake

Published 2026-08-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Masato Fujitake

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, super-smart robot to write stories. To make this robot really good, you need a massive brain made of billions of tiny switches. But here's the catch: the robot's brain is so huge that it barely fits in the computer's memory. It's like trying to run a blockbuster movie on a phone with a tiny battery; the phone gets hot, slows down, and might crash.

To fix this, scientists use special "trainers" called optimizers. The most famous trainer, named AdamW, is like a very careful coach who keeps a detailed notebook for every single switch in the robot's brain, remembering exactly how to nudge it next time. This works great, but those notebooks take up a ton of space. So, researchers invented a lighter trainer called Sinkhorn. It's like a coach who doesn't keep a notebook at all; instead, it just looks at the current mistake and makes a quick, memory-free adjustment. This saves a massive amount of space, but it turns out that for a specific type of robot brain called a "Mixture-of-Experts" (MoE), this memory-free coach gets confused and the robot stops learning well. The question scientists are asking is: Can we keep the memory-saving benefits of the quick coach without losing the robot's ability to learn?

This paper introduces a new method called MESH (Memory-Efficient Sinkhorn for Experts) to solve exactly that problem. The researchers discovered that the quick, memory-free coach fails because the "experts" in the robot's brain don't always get to work. In a Mixture-of-Experts model, the robot only picks a few specific experts to handle each sentence, kind of like a restaurant where only a few chefs are called to the kitchen for a specific order. Because the chefs change from order to order, the quick coach sees a jumbled, noisy signal and gets lost.

The paper suggests that the solution isn't to give the coach a full notebook again, but to give it a "hidden memory." Instead of looking at just the current order, the coach now remembers the average of the last few orders before making a decision. The author calls this "hidden momentum." They tested this on a small 110-million-parameter model (a "nanowhale") and found that this simple trick—smoothing out the noise before adjusting the weights—restores the robot's learning ability.

Here is what they found and what they ruled out:

The Main Fix: The paper shows that the failure happens because the "routed experts" (the chefs who only work sometimes) need their signals smoothed out over time before the memory-free adjustment happens. By averaging the recent history of these experts (using a hidden buffer that doesn't count as permanent memory), the new MESH method works much better.
What Didn't Work: The researchers tried many other ideas to fix the problem, but they ruled them out. They found that simply making the adjustments smaller, changing how the experts are grouped, or using only "sign" information (just knowing if the error went up or down, not how much) did not fix the issue. The key was specifically the time-smoothing, not just any kind of extra data.
How Sure Are They? The results are based on experiments with a specific small model. In these tests, the new method reduced the memory needed for the optimizer's state by 62.5% (dropping from 0.883GB to 0.331GB) and lowered the peak computer memory usage by about 12.6%. However, the robot's final test score (evaluation loss) was still slightly worse than the heavy, notebook-keeping coach (AdamW), with scores around 3.64 for the new method compared to 3.59 for the old one. The author suggests this small gap might be because other parts of the robot's brain (like the vocabulary) still need the heavy coach, but they are confident that the "hidden momentum" is the missing piece for the experts.

In short, the paper argues that you don't need a full notebook to train these special robot brains; you just need a way to remember the recent past without writing it down permanently. This "hidden momentum" approach, called MESH, gets you most of the way there, saving a lot of memory while keeping the robot smart enough to learn.

Technical Summary: MESH – Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

1. Problem Statement

The paper addresses the challenge of optimizing Mixture-of-Experts (MoE) language models with memory efficiency. While AdamW is the standard optimizer for pretraining due to its combination of temporal first-moment smoothing and coordinate-wise adaptive scaling, it incurs a high memory cost by storing two full state tensors per parameter.

Memory-efficient alternatives, such as Sinkhorn-style matrix optimizers and SAGE/Sinkhorn hybrids, reduce this cost by removing moment states for dense matrices. However, the author demonstrates that directly applying these stateless methods to routed MoE experts fails. In a controlled 110M-parameter "nanowhale" DeepSeek-style MoE setting, a standard SAGE/Sinkhorn hybrid reduced optimizer state significantly (from 0.883GB to 0.331GB) but resulted in a severe degradation in evaluation loss (3.8265 vs. the AdamW baseline of 3.58–3.64).

The core failure is identified as an ordering problem specific to routed experts:

Routed Gradients are Conditional: Unlike dense MLP gradients, routed expert gradients are conditional samples based on router decisions and gate weights, which vary temporally.
Nonlinearity of Sinkhorn: Sinkhorn normalization is a nonlinear operation. Applying it to an instantaneous, high-variance routed gradient ( $Sinkhorn(G_t)$ ) behaves differently than applying it to a temporally averaged signal ( $Sinkhorn(E[G_t])$ ).
Noise Amplification: Stateless Sinkhorn normalizes routing-induced noise at every step, whereas the failure of stateless methods stems from the lack of temporal smoothing before the matrix normalization occurs.

2. Methodology: MESH and MESH-B

The author proposes MESH (Memory-Efficient Sinkhorn for Experts), a hybrid optimizer designed specifically for MoE architectures.

Core Mechanism: Hidden Momentum

MESH replaces the stateless Sinkhorn update for MoE experts with a hidden-momentum approach. Instead of storing an explicit first-moment tensor in the optimizer state (which would defeat the memory-saving goal), MESH maintains the first moment within the gradient-buffer lifecycle:

Accumulation: An unnormalized expert momentum accumulator is maintained: $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$ .
Buffer Lifecycle: The accumulator $H_{e,t}$ is stored in the gradient buffer rather than the optimizer state dictionary.
Update Cycle: Before the backward pass, the buffer is reattached as the gradient ( $p.grad \leftarrow \beta_1 H_{e,t}$ ). The backward pass adds the current gradient, effectively updating the accumulator for the next step.
Normalization: The Sinkhorn update is computed on the smoothed signal: $U_{e,t} = Sinkhorn(H_{e,t})$ .

This approach restores the temporal first-moment signal required for routed experts without storing an explicit optimizer-state tensor for them.

Optional Extension: MESH-B

The paper also evaluates MESH-B, which adds a coarse block/neuron inverse-RMS multiplier to the Sinkhorn direction. This acts as a block-level preconditioner. The block statistics are computed from the smoothed source ( $H_{e,t}$ ) rather than a separate raw-gradient stream to maintain implementation simplicity.

Role Assignment Strategy

MESH maintains a hybrid strategy for different parameter classes:

Embeddings / LM Head: Handled by SAGE (light-state updates).
Non-Expert Dense Matrices: Handled by stateless Sinkhorn.
Routed/Shared Experts: Handled by MESH (hidden-momentum Sinkhorn).
Norms / Biases: Handled by SAGE.

3. Key Contributions

Identification of the Failure Point: The paper identifies routed MoE expert matrices as the specific failure point of direct SAGE/Sinkhorn optimization in pretraining, distinguishing them from dense matrices and embeddings.
MESH Algorithm: Introduction of a hidden-momentum Sinkhorn update that restores a first-moment signal for experts without adding explicit optimizer state, utilizing the gradient buffer lifecycle.
Ablation Analysis: Systematic negative ablations showing that second-moment-only, sign/momentum-only, confidence-only, and scale-only interventions do not explain the performance gap. The primary causal ingredient is identified as temporal smoothing before matrix normalization.
MESH-B Evaluation: Evaluation of a block-preconditioned variant, showing that while it can improve the memory-quality frontier in specific seeds, it is not universally necessary compared to the temporal smoothing provided by MESH.

4. Experimental Results

Experiments were conducted on a 110M-parameter DeepSeek-style MoE model trained on FineWeb-Edu for 5,000 steps.

Memory Efficiency:
- AdamW: 0.883GB optimizer state.
- MESH / MESH-B: 0.331GB optimizer state (a 62.5% reduction relative to AdamW).
- Peak CUDA Allocation: Reduced by approximately 12.6% (from ~4.37GB to ~3.82GB).
Performance (Evaluation Loss):
- AdamW Baseline: ~3.5877 (average across seeds 100/512).
- Stateless SAGE/Sinkhorn: 3.8265 (Significant degradation).
- MESH: 3.6430 (Close to AdamW, but slightly higher).
- MESH-B: 3.6365 (Slightly better than MESH, but still trails AdamW).
Diagnostic Findings (Seed 42):
- Full-state variants (MESH-FULL, MESH-B-FULL) that explicitly store the first moment recover AdamW-like quality (e.g., MESH-B-FULL with V(H) reached 3.6102) with roughly half the optimizer state of AdamW.
- This confirms that the temporal smoothing is the critical missing ingredient, while the hidden-momentum implementation successfully approximates this without the full state cost.

5. Significance and Claims

The paper modestly claims that MESH provides a substantially lower-state point on the memory-quality frontier for MoE training.

Primary Conclusion: MoE experts require temporal smoothing before matrix normalization to function effectively with Sinkhorn-style optimizers. This smoothing is the dominant causal factor for closing the performance gap with AdamW.
Secondary Conclusion: Block/neuron preconditioning (MESH-B) can improve results in specific settings but is not established as a universally necessary component.
Limitations: The author acknowledges a residual performance gap between MESH/MESH-B and AdamW. They hypothesize this is due to non-expert role assignments (specifically vocabulary matrices and non-expert dense matrices), which remain on SAGE or stateless Sinkhorn in their setup, whereas AdamW adapts them with full moments.
Scope: The study is intentionally narrow, focusing on a specific 110M architecture and 5,000-step runs. The author does not claim MESH outperforms AdamW in all metrics but rather that it recovers most of the quality lost by stateless methods while drastically reducing memory overhead.

In summary, MESH demonstrates that the failure of memory-efficient optimizers in MoE training is not inherent to the MoE structure itself, but rather to the lack of temporal smoothing for routed gradients. By implementing "hidden momentum," MESH resolves this ordering problem, offering a viable path toward memory-efficient MoE pretraining.