Reviewer wqTg,

We sincerely thank the reviewer for their thoughtful evaluation and insightful questions.

Q.1

As the reviewer correctly noted, the conditions for the theorem rely on the inequality $M \gtrsim C \log (C / 0.001)$ in Lemma 4.9. In connection with this, while we initially described $M$ and $C$ as polylogarithmic in the input dimension $d$, we have since realized that ensuring reliable initialization requires both to be $O(1)$. We will revise the final version accordingly. This refinement is made for theoretical clarity and does not affect the main conclusions or practical relevance of our work. Aside from this adjustment, no other major changes to assumptions or results are necessary. As discussed in Oko et al. (2024a) in the context of feature learning, the Hermite coefficients of the student model can be perturbed by random initialization, a factor not considered in prior work on MoE classification (Chen et al., 2022). Our analysis builds on this observation to study robust specialization under such randomness. However, if the number of clusters $C$ is too large, the number of expert-to-router combinations grows rapidly, making the specialization problem significantly more challenging. We stress that this assumption is mainly to ensure theoretical reliability. In practice, MoE models often incorporate techniques such as auxiliary diversity losses in routing to promote expert diversity and avoid premature collapse (Cai et al., 2024). Our assumption reflects these practical strategies, and the insights into learning dynamics and specialization in MoE remain valid and relevant.

Implications from our theoretical results

While our study is theoretical in nature, it also carries practical implications. These implications are discussed in the "Implications" part of our response to Reviewer YoCg, which we kindly invite the reviewer to refer to.

Due to character limits, we respectfully make use of this space to provide the references.

References

• Allen-Zhu & Li. (2023), Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning, ICLR.
• Arous et al. (2021), Online stochastic gradient descent on non-convex losses from high-dimensional inference, JMLR.
• Bietti et al. (2022), Learning single-index models with shallow neural networks, NIPS.
• Cai et al. (2024). A survey on mixture of experts. arXiv preprint arXiv:2407.06204. Chen et al. (2022), Towards understanding the mixture-of-experts layer in deep learning, NIPS.
• Chowdhury et al. (2023), Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks, ICML.
• Damian et al. (2022), Neural networks can learn representations with gradient descent, COLT.
• Dayi & Chen. (2024), Gradient dynamics for low-rank fine-tuning beyond kernels, arXiv preprint arXiv:2411.15385.
• Dudeja et al. (2018), Learning Single-Index Models in Gaussian Space, PMLR.
• Ge et al. (2018), Learning One-hidden-layer Neural Networks with Landscape Design, ICLR.
• Huang et al. (2024), Harder Tasks Need More Experts: Dynamic Routing in MoE Models, ACL.
• Komatsuzaki et al. (2023), Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints, ICLR.
• Li et al. (2025), Theory on Mixture-of-Experts in Continual Learning, ICLR.
• Liu et al. (2024), GRIN: GRadient-INformed MoE, arXiv preprint arXiv:2409.12136, 2024.
• Mousavi-Hosseini et al. (2023), Gradient-Based Feature Learning under Structured Data, NIPS.
• Oko et al. (2024a), Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations, COLT.
• Oko et al. (2024b), Pretrained transformer efficiently learns low-dimensional target functions in-context, NIPS.
• Pan et al. (2024), Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models, arXiv preprint arXiv:2404.05567.
• Roller et al. (2021), Hash Layers For Large Sparse Models, NIPS.
• Shazeer et al. (2017), Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR.
• Simsek et al. (2025), Learning Gaussian Multi-Index Models with Gradient Flow: Time Complexity and Directional Convergence, CPAL.
• Tang et al. (2025), Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model, ICLR.
• Vural & Erdogdu. (2024), Pruning is Optimal for Learning Sparse Features in High-Dimensions, COLT.
• Wei et al. (2024), Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models, arXiv preprint arXiv:2406.06563.
• Zeng et al. (2024), AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models, ACL.
• Zhang et al. (2024), Diversifying the Expert Knowledge for Task-Agnostic Pruning in Sparse Mixture-of-Experts, arXiv preprint arXiv:2407.09590.
• Zhou et al. (2022), Mixture-of-Experts with Expert Choice Routing, NIPS.

Reviewer bEDR,

We sincerely thank the reviewer for the thoughtful and constructive feedback. We will incorporate the suggestions into the final version. Below, we address the reviewer’s comments. For references, please see our response to Reviewer wqTg due to character limits.

Multi-phase training

We agree that the multi-phase training procedure deviates from practical applications. However, multi-stage optimization is often employed to reveal the intrinsic nature of complex optimization problems and is commonly used in studies of feature learning (Arous et al., 2021; Ge et al., 2018; Gudeja et al., 2018; Bietti et al., 2022; Damian et al., 2022; Mousavi-Hosseini et al., 2024;Oko et al., 2024ab). Joint training of the router and experts is promising, but in nonlinear regression, it poses intractable non-convex challenges and is left for future work.

Assumptions and practical relevance

Since obtaining negative results for vanilla neural networks is mathematically challenging, we followed the approach of using gradient flow, as in Mousavi-Hosseini et al. (2024) and Simsek et al. (2025). In addition, while Chen et al. (2022) assume all signals are orthogonal, our analysis only requires orthogonality to the mean vector. Our theoretical results remain rich in implications. Please refer to the "Implications" section of our response to Reviewer YoCg.

Router

We would like to clarify that our theoretical routing mechanism is not an oversimplification, but rather a principled approach to addressing fundamental technical challenges in MoE without practical heuristics. In Phase II, training is performed using top-1 routing with added random noise $r_m$. After this phase, a data $x_c$ is no longer routed to experts $m' \notin \mathcal{M}_c$ (where $\mathcal{M}\_c$ is the set of professional experts) due to the negative correlation between its router weight $\theta_{m'}$ and the cluster center vector $v_c$.(Lemma 4.11 and Lemma C.16). However, it remains possible that competition among experts $m \in \mathcal{M}_c$ leads to load imbalance if we employ top-1 routing. This phenomenon differs from the classification setting in Chen et al. (2022) and Chowdhury et al. (2023), and stems from the need to estimate a continuous function in the regression setting. We analytically determine the cause and employ an adaptive top-$k$ routing in Phases III and IV, where an expert $m \in \mathcal{M}_c$ is activated if it satisfies a sufficient condition $h_m(x_c) \geq 0$, ensuring that it covers the data distribution of cluster $c$. The value of $k$ is adaptively determined based on the data, a strategy that is also adopted in practice (Huang et al., 2024; Zeng et al., 2024). Importantly, while the adaptive top-$k$ arises from a theoretical and technical challenge, practical implementations often incorporate auxiliary losses or dense training to address load imbalance (see lines 126–131 in the right column), ensuring consistency between theory and practice.

Q.1

• Expert choice routing (ECR): Since ECR (Zhou et al., 2022) uses the same gating network as our token choice routing (TCR) (Shazeer et al., 2017), both methods identify clusters through differences in expert recovery reflected in the gradients. The key difference is that TCR selects top-k experts per token, while ECR selects top-k tokens per expert, potentially dropping tokens unselected by any expert. We believe similar theoretical results are attainable if token drop is mitigated, e.g., via expert capacity constraints.
• Dense MoE: Dense training (Pan et al., 2024) activates all experts without applying top-k selection to the softmax outputs. Since all experts are activated, the gradient for expert $m$ depends on a comparison with a weighted average over all experts, potentially causing gradient competition. Specifically, when weak recovery occurs across experts $m \in \mathcal{M}_c$, differences in gradient magnitudes can distort the update direction and complicate theoretical analysis. Therefore, similar guarantees in the dense setting demand extra assumptions or modified algorithms.
• Hash-based Routing: Since hash-based routing (Roller et al., 2021) is non-learning, misaligned partitioning of the cluster structure may induce gradient interference and hinder theoretical guarantees.

Q.2

If all signals ($w^*_i$, $v_j$) have overlaps of order $\tilde{O}(d^{-1/2})$, then by carefully bounding the resulting overlap terms, we believe the MoE can still reliably learn the underlying cluster structure. Specifically, the gating network's gradient includes additional terms of magnitude $\tilde{O}(d^{-1/2})$ relative to the main signal, arising across task–neuron combinations. In contrast, if the overlap is of order $\tilde{\Omega}(1)$, the dynamics become profoundly more complex. Thus, we are unable to provide a definitive conclusion in this case.

We are happy to clarify any remaining questions during the discussion.

Reviewer VrRr,

We sincerely thank the reviewer's insightful and constructive comments. In the final version, we will separate the presentation of theoretical results and experiments, and further elaborate on the writing in the supplementary material. For the theoretical results, we will include a high-level summary, detail implications, and highlight key differences from prior works. We clarify each point and question below. Owing to the character limits, we respectfully refer the reviewer to our author response to Reviewer wqTg for the references.

Implication

Please refer to the "Implications" part in our response to Reviewer YoCg.

Key differences from existing works

Prior MoE optimization studies (Chen et al., 2022; Chowdhury et al., 2023) focus on binary classification with noisy features. In contrast, we analyze regression and highlight how gradient interference (see lines 125-127 in the left column) across clusters weakens the signal and effectively increases the information exponent, revealing a clear separation between vanilla neural networks and MoE in Theorem 4.3. Although Li et al. (2025) also tackle regression, their linear setting does not account for the non-convex exploration phase we uncover. Regarding information exponents in feature learning, Oko et al. (2024a) assume additive task structure within single data points, whereas we study additive structure across clusters. Theorem 4.7 is the first to apply nonlinear sample complexity analysis to MoE, showing how the router’s gradients reflect expert specialization via weak recovery differences.

Real-data experiments

We certainly appreciate the importance of experiments on real-world data, as the reviewer has rightly pointed out. However, our study is theoretical in nature. Real-world datasets inevitably contain noise, which makes it challenging to isolate specific phenomena in a controlled and interpretable manner. Therefore, we conducted controlled experiments using synthetic data. Within the line of work on the complexity of learning low-dimensional nonlinear functions in specific architectures, Dayi & Chen (2024) and Oko et al. (2024b) also rely on synthetic experiments based on Gaussian data, while Vural & Erdogdu (2024) do not provide empirical results. In the final revision, we will include the more detailed learning process of MoE based on synthetic experiments.

Q.

Structure of MoE

• Gradient-aware routing: While our analysis shows that the MoE mitigates gradient interference by explicitly partitioning the network, the number of experts is typically selected heuristically in practice. This raises the possibility that integrating gradient-aware routing mechanisms could lead to more efficient and adaptive MoE architectures. Recent works also explore the role of gradient interference in MoE (Liu et al., 2024; Tang et al., 2025).

• Adaptive routing: To prevent competition among professional experts ($\in \mathcal{M}_c$), we adopted top-$k$ routing to avoid potential load imbalance. Thus, dynamically adjusting the number of activated experts $k$ is expected to stabilize training. This insight stems from our theoretical studies in nonlinear regression and aligns with recent NLP research (Huang et al., 2024; Zeng et al., 2024), where $k$ is adaptively varied based on tokens.

• Freezing (pruning) redundant experts: To mitigate competition among professional experts, reducing redundant experts is also worth exploring, potentially lowering deployment cost. Recent work has investigated expert merging in this context (Zhang et al., 2024).

Training process of MoE

• Upcycling from dense checkpoints: We argued that meaningful router learning requires differences in weak recovery among experts, which in turn necessitates a long exploration stage due to the non-convex optimization. In this light, upcycling dense MLP checkpoints trained on distinct domains may offer a practical approach to accelerate convergence. This idea has seen several empirical successes in recent LLMs (Komatsuzaki et al., 2023; Wei et al., 2024).

• Stage-specific routing: The noise $r_m$ introduced during Phase II ensures uniform gradient flow and sufficient signal for all experts. In contrast, the adaptive top-$k$ routing in Phase III and IV addresses competition among professional experts. This distinction—different challenges occur at each learning stage—suggests that stage-specific routing strategies may be helpful for effective MoE training.

We welcome any further questions or concerns that may arise and would be pleased to provide clarification.

Reviewer YoCg,

We sincerely appreciate the reviewer's thoughtful and constructive feedback. Below, we carefully address the raised concerns and questions. Due to character limits, please see our response to Reviewer wqTg for references.

Model parameters

Following the reviewer's precise feedback, we will revise the manuscript to provide clear explanation of $w^*_c$, $w^*_g$, $s_c$, and $\nu$.

Router

Please refer to the "Router" part in our response to Reviewer bEDR.

Implications

Chen et al. (2022) present a failure of a single expert in binary classification where noisy features confound the true label. In contrast, our work focuses on a regression setting and reveals a different failure mode where gradient interference across clusters attenuates the signal, increasing the information exponent and hindering optimization. Our findings highlight how MoE mitigates interference by explicitly partitioning the network. Such interference has been recognized as a major challenge in multi-task learning (see lines 125-127 in the left column).

Furthermore, we demonstrate that, for the router to distinguish between experts, a weak recovery of the feature index is necessary. Our analysis reveals that this requires a sample complexity of $\tilde{\Theta}{(d^{k^*-1})}$ during the exploration stage, implying that a sufficiently long exploration phase is essential before the router can begin meaningful learning. This contrasts with the linear expert setting in Li et al. (2025), as our findings arise from non-convex optimization.

Finally, unlike in classification (Chen et al., 2022; Chowdhury et al., 2023), in the regression setting, competition among experts within the set of professional experts $\mathcal{M}_c$ may lead to load imbalance after the router has learned to assign inputs. This observation suggests a practical motivation for several heuristic approaches (see lines 122–131 in the right column).

References

According to the reviewer’s suggestion, we will incorporate references Chen et al. (2022) and Allen-Zhu & Li (2023) as inspiration for our data structure, and Liu et al. (2024) and Tang et al. (2025) as practical works on gradient interference in MoE in Assumption 3.2. We will also include Arous et al. (2021), Ge et al. (2018), Gudeja et al. (2018), Bietti et al. (2022), Damian et al. (2022), and Oko et al. (2024a) in lines 096–098 in the right column.

$v_c$ and $g_j$

We appreciate the reviewer pointing that out. The notation $v_c$ is correct, but $g_j$ is a typo. We will correct the typo in the final version.

Q1.

In our setting, feature vectors are sampled from $\mathbb{S}^{d-1}$ rather than from $\mathbb{R}^d$, in order to normalize signal strength for the analysis of simultaneous learning of multiple signals. This setting is frequently employed in prior works (e.g., Arous et al., 2021; Dudeja et al., 2018; Bietti et al., 2022; Damian et al., 2022; Oko et al., 2024ab). We would also like to note, in high-dimension, a uniform vector on $\mathbb{S}^{d-1}$ converges in distribution to $\mathcal{N}(0,d^{-1}I_d)$, in each coordinate, by concentration of measure.

Q2.

As the reviewer correctly pointed out, in $\hat{F}_m$, the number of activated experts is not fixed and $k$ is a stochastic variable. There have indeed been prior attempts to vary the number of activated experts depending on each token (Huang et al., 2024; Zeng et al., 2024).

Q3.

When analyzing the gradient, the correlation between the nonlinear target and the nonlinear activation arises. To decompose this correlation, we employ the Hermite expansion, which forms an orthogonal basis in the $L^2$ space under the Gaussian measure, as first used in Ge et al. (2018) and Dudeja et al. (2018). Specifically, we have

$$\mathbb{E}\_{z \sim \mathcal{N}(0, I_d)} [ f^*\_c({w^*\_c }^\top z) a_{m,j} \sigma_m({w_{m,j}}^\top z + b_{m,j}) ] = \sum_{i = k^*}^{p^*} \alpha_{m,j,i} \beta_{c,i} {\langle w_{m,j}, w^*_c \rangle}^i.$$

This expresses the interaction as an expansion in powers of alignment. Crucially, the first non-zero order $k^*$ in the expansion (i.e., information exponent) governs the signal strength.

Q4.

In our study, the top-1 routing used in Phase II constitutes a sparse MoE. Although the adaptive top-$k$ routing employed in Phases III and IV allows $k$ to vary, it remains a sparse MoE in the sense that not all experts are activated. While adaptive top-$k$ is motivated by technical challenges in the theoretical analysis of nonlinear regression, we believe that end-to-end training with top-1 routing is feasible with the incorporation of auxiliary losses to control load balancing, which presents an interesting direction for future work.

We will incorporate the revisions identified through this review into the final version. We would be happy to address any further questions during the discussion.