Demystifying Softmax Gating Function in Gaussian Mixture of Experts

Summary

This paper addresses the long-standing challenge of parameter estimation in softmax gating Gaussian mixture of experts. The authors identify three theoretical challenges related to softmax gating: (i) the non-uniqueness of parameter estimation, (ii) the interaction between softmax gating and expert functions, and (iii) the complex dependence in the conditional density. They propose Vononoi loss functions to overcome these challenges and establish convergence rates for the maximum likelihood estimator (MLE). Additionally, they discover a connection between the MLE rate and a polynomial equation solvability problem when the number of experts is unknown or over-specified.

Strengths

The paper proposes a novel solution to a long-standing problem in the literature, addressing three fundamental challenges associated with softmax gating. By introducing Vononoi loss functions and establishing convergence rates for the maximum likelihood estimator, the authors offer a fresh perspective on parameter estimation in these models. The paper contributes to the theoretical understanding of softmax gating Gaussian mixture of experts. It tackles the issues of non-uniqueness, interaction between gating and expert functions, and complex dependence in the conditional density. By resolving these challenges, the authors enhance our knowledge of the underlying principles and dynamics of these models. The paper establishes a connection between the MLE rate and a solvability problem of a system of polynomial equations. This connection adds an interesting aspect to the study, highlighting the potential relationships between statistical estimation and mathematical problem-solving.

Weaknesses

Firstly, the absence of a comparison or discussion of popular techniques such as Gumbel-softmax [1] and max-propagate [2] raises concerns. These techniques are widely used in probabilistic MOE methods, and their exclusion from the paper may limit its completeness and applicability. Secondly, the lack of experimental validation is another potential issue. The practical relevance of the proposed methods is claimed by the authors, but without empirical evidence or experiments, it is challenging to assess the true impact and effectiveness of the proposed approach in real-world scenarios. [1] E. Jang, et al. "Categorical Reparameterization with Gumbel-Softmax." International Conference on Learning Representations (2017). [2] J. Ren, et al. "Probabilistic mixture-of-experts for efficient deep reinforcement learning." arXiv preprint arXiv:2104.09122 (2021).

Questions

No

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

No

Summary

The authors explore the theoretical complexity of standard Gaussian Mixture of Experts and propose a Vononoi loss which estimates different components of the MLE using different Voronoi cells. Update: I have read the rebuttal and the other reviewers' responses, and will keep the score at 7 (accept).

Strengths

The authors establish versions of the Voronoi loss in multiple scenarios, both in fixed settings and over-fitted settings. The Voronoi loss function in both settings is computationally efficient up to a linear order of the number of experts, assuming that the true number of experts is fixed.

Weaknesses

Unfortunately, there does not seem to be any practical exploration of the Voronoi loss in the real world setting to verify the theoretical findings, that should be fine due to the posit of the paper as a more theoretical work.

Questions

Minor notes: there seems to be a switching between Vononoi vs Voronoi in the paper. Make sure to standardize them!

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Yes.

Summary

This paper establishes convergence rates for parameters (Gaussian mean / variance and expert weight) for softmax Gaussian mixture of experts. The rates depend on the setting: exact-fitted vs. over-fitted. The solutions are in terms of Voronoi losses and Hellinger distances.

Strengths

**Originality** To the best of my knowledge, this paper is the first to provide convergence rates for softmax Gaussian mixture of experts, making it an original paper. Likewise, the methodology for obtaining this result involves a rather involved setup, which struck me as highly non-trivial, again, supporting claims of originality. Of course, this is not my area of expertise, so it is difficult for me to accurately assess these points. **Quality** I believe the quality of the paper is strong. The authors provide a thorough description of the setup for softmax Gaussian mixture of experts, with rigorous proofs for the cases of exact-fitted and over-fitted number of experts. I have not assessed the quality of the proofs in the appendix. **Clarity** Despite tackling a complex theoretical problem and proposing a fairly involved solution, I found the paper was presented well. Terms are defined clearly, and the authors walk the reader through the main results before diving in to more detail. **Significance** The results seem marginally significant, as the authors claim that they have some practical implications for setting the number of experts and post-processing techniques, like merge-truncate-merge. However, fully demonstrating these practical insights, as well as empirical verification of the convergence rates, are not included in the paper, somewhat limiting their significance.

Weaknesses

I have several concerns regarding the **significance** of the results. First, these results target the convergence of the softmax mixture of expert parameters, however, in practical settings, I would imagine that these parameters are the outputs of deep networks, which may complicate the convergence picture. Second, the convergence results are derived by considering the number of “true” experts, however, it’s unclear to me whether there is a notion of a “true” expert, particularly if the parameterization is used inside of a neural network / model to model a latent variable. Finally, while I understand that this is an entirely theoretical paper, connecting it more closely to the practical settings in which softmax mixture of experts are used would help to more clearly demonstrate the significance of the paper. For instance, even if it’s in a toy setting, is there any way to empirically verify the convergence rates for exact-fitted vs. over-fitted setups? Regarding **presentation**, I found the paper overall quite clear given the complexity of the result. However, some diagrams (there are none in the main paper) could possibly fill in any gaps in the reader’s understanding. For instance, diagrams showing the softmax Gaussian mixture of experts setting, as well as other concepts, like Voronoi cells, Hellinger distance, etc., would be useful additions to the paper.

Questions

My main question is regarding the significance of the result and connecting the result to practical empirical settings — what can be done to make this connection clearer?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

2 fair

Limitations

The authors have adequately addressed the main limitations of their work in the Discussion section of the paper.

Summary

This paper studies convergence rate of Maximum Likelihood Estimate (MLE) of parameters of Gaussian mixture of experts with softmax gating. Unlike some of the previous papers, the authors allow for the gating probabilities to be input dependent. Convergence rate (upper bound) is provided for both the setting where the number of components is known, and when only a lower bound is known. For the latter case, the authors show connection between the MLE convergence rate and solvability of a central system of polynomial equations.

Strengths

* Makes progress on a difficult problem complicated by lack of identifiability and complex interactions between parameters of the gating mechanism and of the individual components. * Proposes a new "Voronoi loss", which lower bound the Hellinger distance, and is crafted to capture properties of the mixture of experts model (the involved Voronoi cells roughly correspond to the subsets of the input space dominated by a particular expert under the ground truth). * AFAICT provides the first MLE rates for models where the gating mixture is input dependent.

Weaknesses

* As someone with only cursory understanding of the area (emergency reviewer), the paper was hard to follow and full of technical detail that to me felt like obfuscating rather than facilitating understanding. * Unclear what can a practitioner take away, except that they should set the number of components as low as possible. (This is a very minor weakness, as I don't think every paper needs to have direct application.)

Questions

* With the problem of identifiability, what happens if we introduce some penalty ($\ell^1$, $\ell^2$)? I was surprised that this (perhaps naive) solution was not discussed in the background section or anywhere else. * Could you make more explicit how the level of separation between individual components of the ground truth impacts the convergence rate? Intuitively, the more separated, the easier the estimation problem should be. * Potentially related question: What do you mean when you say that a particular Voronoi cell contains only one component of the MLE? Isn't the MLE a random variable? (Wouldn't that imply that the identity of cells which have only one component changes? Also what about the cells which have zero MLE components in them?)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Summary

This work considers the problem of parameter estimation for the **softmax-gated mixture of Gaussian experts**. For simplicity, first consider the case in which the mixture weights do not depend on the covariates X. Then, this model can be seen as a mixture of linear regressions where each “regression component” is corrupted by Gaussian noise. The paper considers the much more challenging case where the mixture weights depend on the covariates. The mixing weights are modeled as linear functions which are normalized by applying softmax to the vector of mixing weights. The paper proposes a novel "Voronoi distance" for the parameter space, which takes into account the invariances of the parametric model, and is used to determine convergence rates for the maximum likelihood estimator (MLE). In the well-specified case where $k$, the number of components of the model used, is equal to $k^*$, the true number of components, the MLE parameters converge at the parametric rate $n^{-1/2}$. In the overparametrized case where $k > k^*$, the rates of the component parameters are slower and depend on how many components of the overparametrized model fall inside the Voronoi cell of the true model’s components.

Strengths

- The paper overcomes significant technical challenges that come from using softmax gates for the mixing weights. The main contribution of the paper is showing that the expected total variation distance of the MLE’s conditional distribution $\hat{p}(Y|X)$ from the true conditional distribution $p(Y|X)$ is *lower bounded* by the proposed Voronoi distance. The Voronoi distance takes care of the inherent symmetries of the parameter space. For example, the mixing weight parameters can only be identified up to global translations due to the softmax operation. - Since one can show that the expected TV distance of the MLE converges at the parametric rate $n^{-1/2}$ by computing the bracketing number of the parametric family [van de Geer, 2000; Theorem 7.2], this leads to an upper bound on the Voronoi distance and thus the $\ell_2$ distance of the MLE parameters. - To be more precise, the paper proposes a family of Voronoi distances, for both the well-specified case and overparameterized case. The convergence rates for overparametrized cases are determined implicitly by a system of polynomial equations which in turn come from analyzing the Taylor expansion of the model’s conditional density w.r.t. the parameters.

Weaknesses

**Exposition.** Though the technical contributions of this paper are solid, its exposition could be improved to make it more accessible to non-expert readers. - **System of polynomial equations (Section 1).** The full details of the system of polynomial equations in the introduction is difficult to follow and seems unnecessary since it is reproduced in Section 3. It would be easier and less intimidating for the readers if the authors omitted the full details and simply noted that $r(m)$ is related to whether some system of polynomial equations admits no non-trivial solutions. Then, one could end with “we conjecture r(m) = 2m, and establish this for the first few cases m=1,2,3, in Lemma 1.” It would also be great if the authors could provide more explanation for how one arrives at these polynomial equations. - **Computation of the Voronoi loss function (Section 3.2).** This would fit better in the experiment section. In typical estimation settings, we do not have access to the true parameters anyway. This quantity is only computable in synthetic settings where the experimenter knows the true parameters. - **Practical implication paragraphs.** These paragraphs in Section 1 and Section 3.2 are redundant.

Questions

- The authors mention in Section 1 that understanding the parameter estimation problem for mixture of experts models has been a long-standing open problem. Could the authors provide explanation for why this was the case? What were the technical barriers and what recent technical advances or breakthroughs have paved the way for the results established here? - Which aspects of the Voronoi distances are novel? Assigning the model’s components according to the Voronoi cells of true components seems like a natural thing to do for any mixture model. Is it the \inf over translations of the mixing weights? Or the rate functions r(m) related to the system of polynomial equations? **Editorial comment** - In the introduction, it would motivate the paper’s result more if the authors provided simple and practical examples showing the importance of parameter estimation in this particular mixture of experts model rather than just estimating the density. Perhaps there are examples in which the parameters in the mixture of experts model are interpretable and have practical value? - Typo. First paragraph of Section 3.3. Total variation distance should be V not h. - It would be helpful to explain what mixing measures when they are first introduced in Section 1. Perhaps after Eq.(1) in the **Problem setting** paragraph.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Response to Authors

Thank you for your rebuttal. After reading the other reviews and the authors' responses, I have decided to raise my score.