Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning

TL;DR

This paper provides a theoretical analysis of Mixture of Experts (MoE), demonstrating its ability to detect and learn latent cluster structures in regression tasks, outperforming vanilla neural networks in such settings.

Contribution

It offers the first theoretical study of MoE's sample and runtime complexity under SGD for nonlinear regression with latent clusters, highlighting its advantages over standard neural networks.

Findings

MoE can effectively identify latent clusters in regression tasks.

Vanilla neural networks fail to detect underlying cluster structures.

MoE leverages expert specialization to solve subproblems more efficiently.

Abstract

Mixture of Experts (MoE), an ensemble of specialized models equipped with a router that dynamically distributes each input to appropriate experts, has achieved successful results in the field of machine learning. However, theoretical understanding of this architecture is falling behind due to its inherent complexity. In this paper, we theoretically study the sample and runtime complexity of MoE following the stochastic gradient descent (SGD) when learning a regression task with an underlying cluster structure of single index models. On the one hand, we prove that a vanilla neural network fails in detecting such a latent organization as it can only process the problem as a whole. This is intrinsically related to the concept of information exponent which is low for each cluster, but increases when we consider the entire task. On the other hand, we show that a MoE succeeds in dividing this problem into easier subproblems by leveraging the ability of each expert to weakly recover the simpler function corresponding to an individual cluster. To the best of our knowledge, this work is among the first to explore the benefits of the MoE framework by examining its SGD dynamics in the context of nonlinear regression.