Structure Aware Neural Architecture Search for Mixture of Experts
cs.LG
Submitted: 2026-08-30
Updated: 2026-08-30
Comments: 22 pages, 3 figures
Code: https://github.com/intsystems/nas-for-moe
License: http://creativecommons.org/licenses/by/4.0/
The gist: Neural Architecture Search (NAS) has so far rarely been applied to Mixture-of-Experts (MoE) models, and existing MoE designs leave the alignment between experts and the structure of the data to
Terminology
Abstract
Neural Architecture Search (NAS) has so far rarely been applied to Mixture-of-Experts (MoE) models, and existing MoE designs leave the alignment between experts and the structure of the data to emerge on its own. We propose an architecture search framework that makes this alignment an explicit search variable: the assignment of data clusters to experts is optimised jointly with the per-expert architectures. We cast the joint problem as a cluster-aware likelihood maximisation, show that it coincides with the incomplete-data maximum likelihood of a latent-variable mixture, and solve it by a generalised Expectation-Maximisation procedure whose otherwise intractable expert-quality term is supplied by an adaptively refined surrogate. We prove that the iterates converge whenever the surrogate errors are summable, and that at every limit point no candidate the search produces improves the true objective. On a heterogeneous image-classification mixture the method recovers the underlying domain partition on 95% of clusters without ever observing domain labels, and on that benchmark and a four-domain time-series forecasting one alike it outperforms the MoE and NAS baselines that likewise use no label information.
Sources
- A Comprehensive Survey of Mixture-of-Experts: Algorithms, Theory, and Applications
- Mixtral of Experts
- DARTS: Differentiable Architecture Search
- Mixture of experts models for multilevel data: modelling framework and approximation theory
- Multi-Source Domain Adaptation with Mixture of Experts
- A General Theory for Softmax Gating Multinomial Logistic Mixture of Experts
- Divide, Specialize, and Route: A New Approach to Efficient Ensemble Learning
- Learning Factored Representations in a Deep Mixture of Experts
- Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
- Scaling Vision with Sparse Mixture of Experts
- DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
- Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models
- Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
- Scaling Expert Language Models with Unsupervised Domain Discovery
- Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
- OLMoE: Open Mixture-of-Experts Language Models
- Time-MoE: Billion-Scale Time Series Foundation Models with Mixture of Experts
- Moirai-MoE: Empowering Time Series Foundation Models with Sparse Mixture of Experts
- Neural Architecture Search with Reinforcement Learning
- Designing Neural Network Architectures using Reinforcement Learning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks