Learning Mixtures of Experts with EM: A Mirror Descent Perspective

TL;DR

This paper provides a theoretical analysis of the EM algorithm for Mixtures of Experts models, showing its equivalence to mirror descent and establishing convergence guarantees, with empirical results demonstrating its advantages over gradient descent.

Contribution

It offers the first rigorous convergence analysis of EM for MoE models, connecting EM to mirror descent and providing conditions for linear convergence.

Findings

EM converges faster than gradient descent in MoE training.

Theoretical guarantees for local linear convergence of EM in MoE models.

Empirical results confirm EM's superior convergence and accuracy on synthetic and real data.

Abstract

Classical Mixtures of Experts (MoE) are Machine Learning models that involve partitioning the input space, with a separate "expert" model trained on each partition. Recently, MoE-based model architectures have become popular as a means to reduce training and inference costs. There, the partitioning function and the experts are both learnt jointly via gradient descent-type methods on the log-likelihood. In this paper we study theoretical guarantees of the Expectation Maximization (EM) algorithm for the training of MoE models. We first rigorously analyze EM for MoE where the conditional distribution of the target and latent variable conditioned on the feature variable belongs to an exponential family of distributions and show its equivalence to projected Mirror Descent with unit step size and a Kullback-Leibler Divergence regularizer. This perspective allows us to derive new convergence results and identify conditions for local linear convergence; In the special case of mixture of $2$ linear or logistic experts, we additionally provide guarantees for linear convergence based on the signal-to-noise ratio. Experiments on synthetic and (small-scale) real-world data supports that EM outperforms the gradient descent algorithm both in terms of convergence rate and the achieved accuracy.