Convergence Rates for Softmax Gating Mixture of Experts

TL;DR

The paper provides a rigorous convergence analysis for softmax gating MoEs and its dense-to-sparse and hierarchical variants, identifying when parameter and expert estimation can be achieved at polynomial (parametric) rates and when they deteriorate to subpolynomial or even exponential data requirements. Central ideas include the strong identifiability condition for non-linear experts, Voronoi-based loss functions to capture parameter discrepancies, and an algebraic independence criterion for routers to avoid gating-expert interactions that slow convergence. Key findings show that strongly identifiable experts (e.g., two-layer FFNs with GELU/sigmoid/tanh) enable fast, input-dependent estimation, while linear experts generally incur severely slower rates due to PDE-type parameter interactions, especially under dense-to-sparse gating unless a router-expert algebraic independence condition holds. The results offer concrete guidance for designing MoE architectures with regard to expert choice and gating/router structure to achieve better sample efficiency in practice, while also noting limitations such as the single-layer MoE focus and the regression-model assumption.

Abstract

Mixture of experts (MoE) has recently emerged as an effective framework to advance the efficiency and scalability of machine learning models by softly dividing complex tasks among multiple specialized sub-models termed experts. Central to the success of MoE is an adaptive softmax gating mechanism which takes responsibility for determining the relevance of each expert to a given input and then dynamically assigning experts their respective weights. Despite its widespread use in practice, a comprehensive study on the effects of the softmax gating on the MoE has been lacking in the literature. To bridge this gap in this paper, we perform a convergence analysis of parameter estimation and expert estimation under the MoE equipped with the standard softmax gating or its variants, including a dense-to-sparse gating and a hierarchical softmax gating, respectively. Furthermore, our theories also provide useful insights into the design of sample-efficient expert structures. In particular, we demonstrate that it requires polynomially many data points to estimate experts satisfying our proposed \emph{strong identifiability} condition, namely a commonly used two-layer feed-forward network. In stark contrast, estimating linear experts, which violate the strong identifiability condition, necessitates exponentially many data points as a result of intrinsic parameter interactions expressed in the language of partial differential equations. All the theoretical results are substantiated with a rigorous guarantee.

Figures (4)

• Figure 1: Illustration of mixture of experts.
• Figure 2: Illustration of two-level hierarchical mixture of experts.
• Figure 3: Illustration of Voronoi cells generated by $k_*=6$ atoms of the ground-truth $G_*$ (red triangles) and $k=10$ fitted atoms of the estimator $\widehat{G}_n$ (blue rounds). By definition, each Voronoi cell is generated by one ground-truth atom, and its cardinality equals the number of corresponding fitted atoms. For instance, the red triangle in cell 4 is fitted by three blue rounds, implying that the cardinality of Voronoi cell 4 is three.
• Figure 4: Illustration of Voronoi cells given in equations \ref{['eq:Voronoi_cells_level_1']} and \ref{['eq:Voronoi_cells_level_2']}. Above, Voronoi cells $\mathcal{A}_{1},\mathcal{A}_{2},\ldots,\mathcal{A}_{k^*_1}$ in the first level are generated by ground-truth parameters $\omega^*_{1},\omega^*_{2},\ldots,\omega^*_{k^*_1}$ (red triangles), respectively. As the number of expert groups $k_1^*$ is known, each Voronoi cell $\mathcal{A}_{j_1}$ contains one fitted parameter $\omega_{i_1}$ (blue round). In the second level, each rectangle represent a set of $k_2^*=3$ Voronoi cells $\{\mathcal{A}_{j_2|j_1}:j_2\in[k_2^*]\}$ generated by ground-truth parameters $\theta^*_{j_2|j_1}:=(\kappa^*_{j_2|j_1},\eta^*_{j_1j_2})$(red squares), for $j_1\in[k_1^*]$, each of which contains a total of $k_2=5$ fitted parameters $\theta_{i_2|j_1}:=(\kappa_{i_2|j_1},\eta_{i_2j_1})$ (blue stars).

Theorems & Definitions (14)

• Proposition 1
• Definition 1: Strong Identifiability
• Theorem 1
• Theorem 2
• Proposition 2
• Theorem 3
• Definition 2: Algebraic Independence
• Theorem 4
• Proposition 3
• Theorem 5