Improving MoE Compute Efficiency by Composing Weight and Data Sparsity

TL;DR

This paper introduces a method to improve compute efficiency in Mixture-of-Experts models by combining weight and data sparsity, using null experts to maintain causality and achieve better training and performance outcomes.

Contribution

It proposes a novel approach to incorporate data sparsity within causal MoE models using null experts, enhancing compute efficiency without violating causality.

Findings

Composing weight and data sparsity improves compute efficiency.

Model learns implicit modality-aware routing, favoring vision tokens.

Achieves better training loss and downstream performance at matched FLOPs.

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.