Statistical Perspective of Top-K Sparse Softmax Gating Mixture of Experts

Response to Reviewer dwg3 (Part 1)

Dear Reviewer dwg3, Thanks for providing your valuable reviews, and giving **good (3)** grades to the **soundness, presentation** and **contribution** of our paper. We hope that we address your concerns regarding our paper with the responses below. **Q1: The theoretical results of this paper lack experimental verification.** Thanks for raising your concern. Actually, we already conducted numerical experiments to empirically verify our theoretical results in Appendix D. In that appendix, we show that the empirical convergence rates of parameter estimation totally match the theoretical rates under both the exact-specified and over-specified settings. Therefore, we would like to refer the reviewer to Appendix D for further details.

Response to Reviewer dwg3 (Part 2)

**Q2: The paper only considers Gaussian mixture-of-experts models. What are the difficulties in generalizing the results to MoE methods without Gaussian assumption?** Thanks for your question. We would like to clarify that the proof technique in the paper can be extended to other types of models. The reason that we consider only top-K sparse softmax Gaussian MoE is because it is the most popular one in practice. Per your suggestion, we have generalized the results of the paper to the other settings of MoE when $f$ belongs to the family of location-scale distributions $\mathcal{F}:=\{f(Y|h_1(X,a,b),\sigma):(a,b,\sigma)\in\Theta\}$, where the expert function $h_1(X,a,b):=a^{\top}X+b$ stands for the location, $\sigma$ denotes the scale and $\Theta$ is a compact subset of $\mathbb{R}^d\times\mathbb{R}\times\mathbb{R}_+$. Due to the space limit, we will summarize our findings here, and leave further details, including rigorous proofs, in Appendix E of the current manuscript. Firstly, we characterize the family $\mathcal{F}$ based on the following notion of strong identifiability: **Definition 1 (Strong Identifiability).** We say that the family $\mathcal{F}$ is strongly identifiable if $f(Y|h_1(X,a,b),\sigma)$ is twice differentiable w.r.t its parameters and the following assumption holds true: For any $k\geq 1$ and $k$ pairwise different tuples $(a_1,b_1,\sigma_1),\ldots,(a_k,b_k,\sigma_k)\in\Theta$, if there exist real coefficients $\alpha^{(i)}_{\ell_1,\ell_2}$, for $i\in[k*]$ and $0\leq \ell_1+\ell_2\leq 2$, such that $$ \sum_{i=1}^{k}\sum_{\ell_1+\ell_2=0}^{2} \alpha^{(i)}_{\ell_1,\ell_2}\cdot\dfrac{\partial^{\ell_1+\ell_2}f}{\partial h_1^{\ell_1}\partial \sigma^{\ell_2}}(Y|(a_i)^{\top}X+b_i,\sigma_i)=0, $$ for almost surely $(X,Y)$, then we obtain that $\alpha^{(i)}_{\ell_1,\ell_2}=0$ for any $i\in[k*]$ and $0\leq \ell_1+\ell_2\leq 2$. **Example.** The families of Student's t-distributions and Laplace distributions are strongly identifiable, while the family of location-scale Gaussian distributions is not. Therefore, it suggests that the location-scale Gaussian distribution is weakly identifiable. Since the convergence rates of maximum likelihood estimation when $\mathcal{F}$ is a family of Gaussian distributions, which is not strongly identifiable, have already been studied, we focus on the scenario when $\mathcal{F}$ is strongly identifiable. Under that assumption, the density $f(Y|h_1,\sigma)$ is twice differentiable w.r.t its parameters, thus, it is also Lipschitz continuous. As a consequence, the density estimation rates under both the exact-specified and over-specified in Theorem 1 and Theorem 3 still hold true. Therefore, we consider only the parameter estimation problem in the sequel. **Parameter estimation under the strongly identifiable MoE:** In high level, we need to establish the Total Variation lower bound $$E_X[V(g_{G}(\cdot|X),g_{G*}(\cdot|X))]\gtrsim \mathcal{D}(G,G*),$$ for any $G\in\mathcal{O}_k(\Omega)$. Then, this bound together with the density estimation rate in Theorem 1 (resp. Theorem 3) leads to the parameter estimation rates in Theorem 2 (resp. Theorem 4). Here, the key step is to decompose the difference between the density $g_{\widehat{G}_n}(Y|X)$ and the true density $g_{G*}(Y|X)$ into a combination of linearly independent terms using Taylor expansions. Therefore, we have to involve the above notion of strong identifiability, and separate our convergence analysis based on that notion. **Under the exact-specified settings**, we demonstrate in Theorem 5 that the rates for estimating true parameters $\exp(\beta^{\ast}_{0i})$, $\beta^{\ast}_{1i}$ , $a^{\ast}_i$, $b^{\ast}_i$, $\sigma^{\ast}_i$ are of order $O(n^{-1/2})$, which match those in Theorem 2. **Under the over-specified settings**, we show in Theorem 6 that the rates for estimating true parameters $\beta^{\ast}_{1i},a^{\ast}_i,b^{\ast}_i,\sigma^{\ast}_i$, which are fitted by more than one component, are of order $\widetilde{\mathcal{O}}(n^{-1/4})$. Notably, these rates are no longer determined by the solvability of the system of polynomial equations in Eq. (7). Thus, they are significantly faster than those in Theorem 4. The main reason for this improvement is when $\mathcal{F}$ is strongly identifiable, the interaction among expert parameters via the second partial differential equation in Eq. (6) does not occur. Finally, we would like to emphasize that we consider only Gaussian MoE models in our original manuscript since Gaussian is the most challenging distributions. In particular, as the family of location-scale Gaussian distributions does not satisfy the strong identifiability conditions, we encounter an interaction among parameters via the PDE in Eq. (6), which subsequently leads to the complex system of polynomial equations in Eq. (7). As a result, the convergence analysis for the Gaussian MoE models are much more challenging than that for the strongly identifiable MoE models.

Response to Reviewer X5io

Dear Reviewer X5io We would like to thank you for the wealth of your comments, and giving **good (3)** grades to the **soundness, presentation** and **contribution** of our paper. Below are our responses to your questions. We hope that we address your concerns regarding our paper. **Q1: Chen et al. [1] also analyze MoE, but in the deep learning setting. How does the analysis presented here differ from this work?** Thanks for your question. We would like to highlight the differences between Chen et al. [1] and our work in terms of goals, settings and results as follows: **1. Goals.** - Chen et. al. [1] studied the mechanism of the mixture of experts (MoE) layers for deep learning from the perspective of a classification problem. Firstly, they would like to know whether using the mixture of experts yields a higher test accuracy than using a single expert. Secondly, they considered the top-K sparse softmax gating MoE model, where $K=1$. They aimed to explore how the router learns to dispatch the data to the right expert. - In our work, we aim to understand the effects of the top-K sparse softmax gating function, for arbitrary $K\geq 1$, on the MoE models via the density and parameter estimation problems. Then, we compare the performance of the top-K sparse softmax gating MoE to that of the MoE with dense softmax gating function. Finally, we provide insight into finding the best choices of $K$ based on the theoretical results. **2. Settings.** - Chen et. al. [1] sampled the data from a distribution with multiple clusters, and generated their label from the set {$1,-1$} uniformly. Then, they trained a top-1 sparse softmax gating MoE layer based on the data using the gradient descent method to minimize an empirical logistic loss. Here, they used a linear gating network and formulated each expert as a two-layer convolutional neural network. - In our work, we assume that the data are generated from the ground-truth top-K sparse softmax gating Gaussian MoE model, where the mean of each Gaussian distribution is a linear expert network (an extension to general expert network is possible but will require more complex calculations and we leave that for the future work). Then, we estimate the true density function and true parameters of that model by using the maximum likelihood method. Note that, parameter estimation is important to understand as they directly lead to the estimation of experts and softmax weights. **3. Results.** - Chen et. al. [1] demonstrated that the top-K sparse softmax gating mixture of non-linear experts can achieve nearly 100\% test accuracy, while a single expert can only reach a test accuracy of no more than 87.5\% on the data distribution. Moreover, they proved that each expert will be specialized to a specific portion of the data, which is determined by the initialization of the weights. Finally, they figured out that the router can learn the cluster-center features and route the input data to the right experts. - In our work, we establish the convergence rates of density estimation and parameter estimation under the top-K sparse softmax gating Gaussian MoE model when the true number of experts is known and unknown. From our theoretical results, we realize that using the top-K sparse softmax gating function not only helps scale up the model capacity but also keeps the model performance comparable to that when using the dense softmax gating function. Additionally, we also find out that using the top-K sparse softmax gating function, for $K=1$, leads to faster convergence rates of parameter estimation than other values of $K$. For further details, please refer to our response to your Question 2. **Q2: What are the practical implications of the work, if any, for the use of Top-K sparse MoE as a means of conditional computation?** Thanks for your question. We would like to refer the reviewer to our response to the Common Question 1 in the General Response (part 2) for further details. **Q3: It would have been helpful to outline high level ideas before delving into notation-heavy math.** Thanks for your suggestion. We would like to refer the reviewer to the Common Question 2 in the General Response (part 2) for further details. **References** [1] Z. Chen, Y. Deng, Y. Wu, Q. Gu, Y. Li. Towards understanding mixture of experts in deep learning. In NeurIPS, 2022.

Response to Reviewer qHcA

Dear Reviewer qHcA, We would like to thank you for providing your constructive feedback and giving **good (3)** grade to the **contribution** of our paper. We hope that we can address your concerns with the responses below, and eventually convince you to raise your final score. **Q1: The writing in the universal assumptions section is unclear, and should be revised for readability.** Thanks for your suggestions. We have revised the universal assumptions in the current manuscript (in blue color) as follows: **(U.1) Convergence of MLE:** To ensure the convergence of parameter estimation, we assume that the parameter space $\Omega$ is compact subset of $\mathbb{R}\times\mathbb{R}^d\times\mathbb{R}^d\times\mathbb{R}\times\mathbb{R}_+$, while the input space $\mathcal{X}$ is bounded. **(U.2) Identifiability:** Next, we assume that $\beta^*_{1k_*}=0_d$ and $\beta^*_{0k_*}=0$ so that the top-K sparse softmax gating Gaussian mixture of experts is identifiable. Under that assumption, we show in Appendix C that if $G$ and $G'$ are two mixing measures such that $g_{G}(Y|X)=g_{G'}(Y|X)$ for almost surely $(X,Y)$, then it follows that that $G\equiv G'$. Without that assumption, the result that $G\equiv G'$ does not hold, which leads to unncessarily complicated loss functions (see Proposition 1 in [1]). **(U.3) Distinct Experts:** To guarantee that experts in the mixture [1] are different from each other, we assume that parameters $(a^{\ast}_i,b^{\ast}_i,\sigma^{\ast}_i)$, for $i\in[k*]$, are pairwise distinct. **(U.4) Input-dependent Gating Functions:** To make sure that the gating functions depend on the input $X$, we assume that at least one among parameters $\beta^*_{11},\ldots,\beta^*_{1k_*}$ is different from zero. Otherwise, the gating functions would be independent of the input $X$, which simplifies the problem significantly. In particular, the model (1) would reduce to an input-free gating Gaussian mixture of experts, which was already studied in [2]. **Q2: The implications of the theory for practitioners could be expanded on more in the discussion.** Thanks for your suggestion. We would like to refer the reviewer to our response to the Common Question 1 in the General Response (Part 2) for further details. **Q3: How do these convergence results compare to prior analysis on softmax vs sparse gating?** Thanks for your question. It is worth emphasizing that our paper is the very first work establishing the theoretical guarantee for the effects of top-K sparse gating function on the mixture of experts via the parameter estimation problem. Additionally, the most related work to ours is [1], which established the convergence rates of parameter estimation under the dense softmax gating Gaussian mixture of experts. For the comparison of parameter estimation rates under these two models, please refer to the paragraph "No trade-off between model capacity and model performance" in your Question 2. **Q4: Guidance on model design is lacking. Does this theory suggest one gating approach over the other?** Thanks for your question. We would like to refer the reviewer to our response to the Common Question 1 in the General Response (part 2) for further details. **Q5: Due to the instability of the EM algorithm, there is no error bar provided in Figure 3b. Can we still consider this experiment to be reliable to to justify the theoretical results?** Thanks for your comment. Indeed, as we mentioned in the Appendix, the EM iterates are unstable, which leads to high error bars. That is why we did not include that in the previous version of our manuscript. In the new version of our manuscript, per your suggestion, we have included error bars for the EM iterates in Figure 3b. **Q6: More intuition and examples earlier on could make it more accessible.** Thanks for your suggestion. We would like to refer the reviewer to the Common Question 2 in part 2 of the General Response section for further details. **References** [1] H. Nguyen, TT. Nguyen, and N. Ho. Demystifying softmax gating in Gaussian mixture of experts. Advances in Neural Information Processing Systems, 2023a. [2] N. Ho, C. Y. Yang, and M. I. Jordan. Convergence rates for Gaussian mixtures of experts. Journal of Machine Learning Research, 23(323):1–81, 2022.