Research map / Optimization / Implicit bias and generalization
Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Theoretically shows that SGD-trained mixture-of-experts with router and adaptive top-k routing can detect and learn a latent cluster structure of single-index models, whereas a vanilla two-layer network fails due to rising information exponent of the shared global task.
Problem
Understanding why MoE experts specialize and how the router learns to dispatch inputs, specifically sample and runtime complexity of SGD for MoE in nonlinear regression with a latent cluster structure of single-index models.
Approach
Analyzes a four-phase gradient-based training (normalized SGD for experts, gating training, reinitialized expert training with adaptive top-k routing, convex second-layer fit) and shows that the information exponent of each cluster is low, so experts weakly recover their local tasks, enabling the router to partition clusters, whereas a single network's effective information exponent increases due
Key results
- Vanilla neural network cannot align with the hidden global feature w_g*; alignment stays O(d^{-1/2}) for all time.
- MoE learns the teacher with total sample complexity tilde O(d^{k*-1}) and polynomial time complexity in d, matching single-index learning by a vanilla network.
- Synthetic experiment with 8 experts, C=2, dimension 200 shows experts specialize and router dispatches clusters as predicted.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a MoE to partition the clustered task into cluster-specific expert subproblems | Gradient-based learning of clustered single-index regression · Synthetic clustered nonlinear regression | expected absolute prediction error: \leq \epsilon | Helped |
| Use a vanilla neural network instead of a MoE | Gradient-based learning of clustered single-index regression · Synthetic clustered nonlinear regression | alignment with hidden global feature: \tilde{O}(d^{-1/2}) | Hurt |
| Train the router after experts weakly recover cluster-specific features | MoE (four-phase SGD) · Synthetic clustered nonlinear regression | cluster routing | Helped |
| Use adaptive top-k routing in Phases III and IV | MoE (four-phase SGD) · Synthetic clustered nonlinear regression | routing to non-corresponding experts | Helped |
| Reinitialize expert weights before Phase III | MoE (four-phase SGD) · Synthetic clustered nonlinear regression | learning the global task | Helped |
Surprising findings
- A vanilla neural network with polynomial width fails to learn the shared global task even though each cluster's task has low information exponent.
Limitations
- Analysis restricted to population gradient flow and idealized assumptions (orthogonality-like conditions, sign invariance of Hermite coefficients).
- Requires k* > 2 (and k* >= 5 for the vanilla-network lower bound) and a specific four-phase training schedule with reinitialization.
- Experiments limited to a small synthetic setting with C=2.
Experimental setup
Datasets: synthetic single-index teacher data
Metrics: alignment of expert neurons with teacher feature vectors sample complexity bound
Baselines: vanilla two-layer neural network kernel ridge regression
Builds on: information exponent Damian et al. feature learning with SGD Dandi et al. additive models sparsely-gated MoE
Closest papers by meaning
- Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient (ICML 2025 · similarity 0.90)
- Learning Mixtures of Experts with EM: A Mirror Descent Perspective (ICML 2025 · similarity 0.90)
- On the Expressive Power of Mixture-of-Experts for Structured Complex Tasks (NeurIPS 2025 · similarity 0.90)
- On Least Square Estimation in Softmax Gating Mixture of Experts (ICML 2024 · similarity 0.89)
- MoDE: A Mixture-of-Experts Model with Mutual Distillation among the Experts (AAAI 2024 · similarity 0.89)
- Tight Clusters Make Specialized Experts (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.