Improving MoE Compute Efficiency by Composing Weight and Data Sparsity

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Kilian, Maciej, Mkrtchyan, Oleg, Zettlemoyer, Luke, Shrivastava, Akshat, Aghajanyan, Armen
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866908780559597568
author Kilian, Maciej
Mkrtchyan, Oleg
Zettlemoyer, Luke
Shrivastava, Akshat
Aghajanyan, Armen
author_facet Kilian, Maciej
Mkrtchyan, Oleg
Zettlemoyer, Luke
Shrivastava, Akshat
Aghajanyan, Armen
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_15370
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Improving MoE Compute Efficiency by Composing Weight and Data Sparsity
Kilian, Maciej
Mkrtchyan, Oleg
Zettlemoyer, Luke
Shrivastava, Akshat
Aghajanyan, Armen
Machine Learning
Artificial Intelligence
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.
title Improving MoE Compute Efficiency by Composing Weight and Data Sparsity
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2601.15370