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

• URL: http://arxiv.org/abs/2608.04407v1
• Date: Wed, 05 Aug 2026 03:23:31 GMT
• Title: MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training
• Abstract summary: Memory-efficient matrixs such as Sinkhorn gradient descent remove most AdamW state for dense Transformer matrices, but direct application to Mixture-of-Experts (MoE) training is unreliable.<n>A SAGE/Sinkhorn hybrid reduces state from 0.883GB to 0.331GB but loses evaluation loss to 3.8265, far above the AdamW baselines observed in the same setup.
• Score: 8.832969171530054
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability [57.03393882933515]
  We describe the architecture and ablations of Qwen3.8-Flash-Next, a sparse mixture-of-experts model with 125B parameters, 6B activated per token, and 51B parameters of n-gram embedding tables held off the accelerator.<n>On fourteen pre-training benchmarks the model leads the 397B-A17B predecessor on eight and trails it on the rest by at most 2.6 points, at 1/3 the activated parameters, 1/3 the training tokens, and roughly 1/9 the training FLOPs.
  arXiv  Detail & Related papers  (2026-08-31T06:35:07Z)
• GROM: Gradient-Free Rapid One-Shot Machine Unlearning [12.335697153382966]
  We propose a novel one-shot unlearning approach, abandoning iterative optimization in favor of a direct, exact analytical solution.<n>GROM applies the weight edit in mere seconds, which makes it orders of magnitude faster than traditional fine-tuning.<n>Because the update removes the targeted content from the weights instead of masking it, GROM also withstands the low-bit quantization attack.
  arXiv  Detail & Related papers  (2026-08-06T09:16:53Z)
• Forgetful Attention: A Trainable Support-Vector Memory with Certified Selection and Exact Unlearning [0.16921396880325776]
  We introduce Support Vector Attention (SV-Attention), a memory whose weights are support coefficients of a one-class Sfp with fixed box parameter C.<n>Its active-set partition gives reserve tokens exactly zero weight, output weights; a incremental incremental solver token to recover the state produced by retraining without it under the same C.<n>We also demonstrate surgical eviction, exact editing, patient-record deletion, and a forgettable retrieval memory over real sentence embeddings.
  arXiv  Detail & Related papers  (2026-07-13T23:13:57Z)
• Closed-Form Spectral Regularization for Multi-Task Model Merging [96.82449201305234]
  Model merging combines several independently fine-tuned experts into a single multi-task model without any training data.<n>State-of-the-art merging methods formulate merging as a layer-wise interference problem.<n>We propose SWUDI, a closed-form method that combines a soft exponential filter, which matches the gradient-flow trajectory of iterative descent.
  arXiv  Detail & Related papers  (2026-06-05T14:00:47Z)
• SAGE: Sign-Adaptive Gradient for Memory-Efficient LLM Optimization [1.675857332621569]
  We propose SAGE (Sign Adaptive GradiEnt) to replace AdamW in hybrid structures.<n>SAGE combines a Lion-style update direction with a new, memory-efficient $O(d)$ adaptive scale.<n>On Llama models up to 1.3B parameters, our SAGE-based hybrid achieves new state-of-the-art perplexity, outperforming all baselines.
  arXiv  Detail & Related papers  (2026-04-09T00:07:38Z)
• Memory-Efficient Structured Backpropagation for On-Device LLM Fine-Tuning [10.913120072779193]
  On-device fine-tuning enables privacy-preserving personalization of large language models.<n>Mobile devices impose severe memory constraints, typically 6--12GB shared across all workloads.<n>We propose Memory-efficient Structured Backpropagation (MeSP)<n>MeSP reduces peak memory from 361MB to 136MB for Qwen2.5-0.5B, enabling fine-tuning scenarios previously infeasible.
  arXiv  Detail & Related papers  (2026-02-13T16:24:33Z)
• Retrieval-Aware Distillation for Transformer-SSM Hybrids [56.85859614817908]
  State-space models (SSMs) offer efficient sequence modeling but lag behind Transformers on benchmarks that require in-context retrieval.<n>We propose *retrieval-aware distillation*, which converts a pretrained Transformer into a hybrid student by preserving only these retrieval-critical heads.<n>We show that preserving **just 2% of attention heads recovers over 95% of teacher performance on retrieval-heavy tasks.
  arXiv  Detail & Related papers  (2026-02-11T21:05:00Z)
• ECO: Quantized Training without Full-Precision Master Weights [58.97082407934466]
  Error-Compensating (ECO) eliminates master weights by applying updates directly to quantized parameters.<n>We show that ECO converges to a constant-radius neighborhood of the optimum, while naive master-weight removal can incur an error that is inversely proportional to the learning rate.
  arXiv  Detail & Related papers  (2026-01-29T18:35:01Z)
• Robust and Efficient Zeroth-Order LLM Fine-Tuning via Adaptive Bayesian Subspace Optimizer [4.6561758107970395]
  Fine-tuning large language models (LLMs) with zeroth-order (ZO) optimization reduces memory by approximating gradients through function evaluations.<n>We introduce BSZO, an adaptive textbfBayesian textbfSubspace textbfZeroth-Order textbfOptimizer.
  arXiv  Detail & Related papers  (2026-01-04T09:35:11Z)
• Kernelized Sparse Fine-Tuning with Bi-level Parameter Competition for Vision Models [80.50996301430108]
  sparse tuning achieves remarkable performance by adjusting only the weights most relevant to downstream tasks.<n>We propose a one-stage method named SNELLA to overcome the above limitations.<n> SNELLA achieves SOTA performance with low memory usage.
  arXiv  Detail & Related papers  (2025-10-28T03:39:18Z)
• AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
  We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models. AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
  arXiv  Detail & Related papers  (2023-10-16T09:04:28Z)
• Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost [53.746169882193456]
  Recent works have proposed various sparse attention modules to overcome the quadratic cost of self-attention. We propose a model that resolves both problems by endowing each attention head with a mixed-membership Block Model. Our model outperforms previous efficient variants as well as the original Transformer with full attention.
  arXiv  Detail & Related papers  (2022-10-27T15:30:52Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.