Improving MoE Compute Efficiency by Composing Weight and Data Sparsity

TL;DR

The paper addresses compute efficiency in MoE transformers by combining weight sparsity with data sparsity through a minimal, causality-preserving extension: null experts. By duplicating null slots and applying a standard load-balancing loss plus a stabilizing Z-Loss, the model learns to route low-information tokens to null experts while preserving full computation for informative tokens, enabling data sparsity without train–inference mismatch. Experiments on vision-language model training show that, at matched expected FLOPs, data-sparse MoE configurations outperform dense baselines and exhibit emergent modality-aware compute allocation, routing vision tokens more to nulls and adjusting based on task prompts. This approach yields a more compute-efficient frontier and practical gains in training loss and downstream performance, suggesting broader applicability to multimodal and attention components beyond MoE layers.

Abstract

Mixture-of-Experts layers achieve compute efficiency through weight sparsity: each token activates only a subset of experts. Data sparsity, where each expert processes only a subset of tokens, offers a complementary axis. Expert-choice routing implements data sparsity directly but violates causality in autoregressive models, creating train-inference mismatch. We recover data sparsity within causal token-choice MoE by leveraging zero-compute (null) experts within the routing pool. When a token routes to null experts, those slots consume no compute. The standard load balancing objective trains the model to uniformly use all experts (real and null) therefore creating data sparsity in expectation without the causality violations. We evaluate on vision-language model training, where data heterogeneity is pronounced: vision encoders produce many low-information tokens while text tokens are denser. At matched expected FLOPs, composing weight and data sparsity yields a more compute-efficient frontier than weight sparsity alone, with gains in training loss and downstream performance. The model learns implicit modality-aware allocation, routing vision tokens to null experts more aggressively than text, without explicit modality routing.

Figures (10)

• Figure 1: Top: Weight and data sparsity as dual budget constraints. Weight sparsity bounds experts per token (row budget $\leq \rho_w E$); data sparsity bounds tokens per expert (column budget $\leq \rho_d T$); composing both yields a budget over the full $T \times E$ matrix. Bottom: Implementations. Token-choice achieves weight sparsity causally; expert-choice achieves data sparsity but requires seeing future tokens; null experts compose both while preserving causality.
• Figure 2: Effect of data sparsity ($\rho$) on training loss (Left) and evaluation score (Right) for MoE models upcycled from dense 0.6B and 1.7B models. Lower $\rho$ indicates more null experts. Training loss decreases monotonically with increased sparsity, while evaluation scores peak at $\rho \approx 0.5$.
• Figure 3: Training loss (Left) and average evaluation score (Right) as a function of training FLOPs for data-dense and data-sparse MoE models. Blue circles represent data-dense baselines with fixed top-K routing; red circles represent data-sparse models using null experts, with opacity indicating the sparsity ratio $\rho$. The gray line connects data-dense models to form a Pareto frontier, with the shaded region indicating suboptimal performance. Data-sparse models consistently outperform data-dense baselines at equivalent compute budgets, achieving lower training loss and higher evaluation scores. Labels indicate model scale and expected active experts $\mathbb{E}[K]$; $\rho$ denotes the ratio of real experts to total experts (e.g., $\rho=0.50$ means half of selected experts are real on average).
• Figure 4: MoE compute distribution by modality across sparsity configurations. Left: Token distribution (constant at 74% vision, 26% text). Middle: Compute distribution. Dense configurations ($\mathbf{K}^{2}_{1.0}$) allocate compute proportionally to token count; data-sparse configurations route vision tokens to null experts, reducing vision's share from 74% to 36%. Right: Compute intensity (fraction of top-K slots filled by real experts). Vision intensity drops to 0.04 at $\mathbf{K}^{12}_{0.17}$ while text remains at 0.19.
• Figure 5: Per-token compute utilization for a sample sequence. Left: Original input (text and image patches). Right: Compute overlay where brightness indicates fraction of top-K slots filled by real experts—bright regions receive full computation, dark regions route to null experts. Inset values show mean compute per segment (image or text block) and overall sample compute (top-right). Vision tokens receive less compute on average due to redundant patches.