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

TL;DR

This paper provides a theoretical analysis of the top-K sparse softmax gating mixture of experts, revealing how it affects density and parameter estimation, especially in over-specified models, with convergence rates depending on expert number and input region complexity.

Contribution

It offers the first theoretical insights into the behavior of top-K sparse softmax gating in mixture of experts, including convergence rates and conditions for accurate estimation.

Findings

Density estimation converges at a parametric rate when the true number of experts is known.

Over-specified models require selecting more experts than the true number to ensure density estimation convergence.

Parameter estimation rates slow down significantly in over-specified models due to gating-expert interactions.

Abstract

Top-K sparse softmax gating mixture of experts has been widely used for scaling up massive deep-learning architectures without increasing the computational cost. Despite its popularity in real-world applications, the theoretical understanding of that gating function has remained an open problem. The main challenge comes from the structure of the top-K sparse softmax gating function, which partitions the input space into multiple regions with distinct behaviors. By focusing on a Gaussian mixture of experts, we establish theoretical results on the effects of the top-K sparse softmax gating function on both density and parameter estimations. Our results hinge upon defining novel loss functions among parameters to capture different behaviors of the input regions. When the true number of experts $k_{\ast}$ is known, we demonstrate that the convergence rates of density and parameter estimations are both parametric on the sample size. However, when $k_{\ast}$ becomes unknown and the true model is over-specified by a Gaussian mixture of $k$ experts where $k > k_{\ast}$, our findings suggest that the number of experts selected from the top-K sparse softmax gating function must exceed the total cardinality of a certain number of Voronoi cells associated with the true parameters to guarantee the convergence of the density estimation. Moreover, while the density estimation rate remains parametric under this setting, the parameter estimation rates become substantially slow due to an intrinsic interaction between the softmax gating and expert functions.