Provably Learning Mixtures of Linear Experts

Introduction

Mixture-of-experts (MoE) models decompose a prediction problem across a collection of specialized experts, with a gate (or router) that decides which experts are responsible for each input (Jacobs et al., [10]; Jordan & Jacobs, [12]). The idea is now standard in large models: sparsely gated MoE layers let the parameter count of a network grow by orders of magnitude while the compute per input stays fixed (Shazeer et al., [23]; Fedus et al., [6]). Often, many more experts than are necessary are initialized, the gate is initialized at random, and expert load-balancing penalties are implemented to incentivize broad expert utilization. Each of these choices is motivated by experience rather than theory. Why gradient-based training of a gated mixture should succeed at all is only partly understood.

An MoE can be seen as a latent-variable model whose likelihood is nonconvex. The nonconvexity comes from the coupling of gate and experts: the gate decides which expert sees which data, and the experts’ errors feed back into the gate. The theory of expectation–maximization (EM) has shed considerable light on the specific latent-variable model of a Gaussian mixture model (GMM). When the GMM has exactly as many components as the ground-truth, EM and its gradient variants enjoy global guarantees only for two components (Xu et al., 2016; Daskalakis et al., [3]) or from a good initialization (Balakrishnan et al., [1]; Yan et al., [29]), and for three or more Gaussian components the population likelihood provably has spurious local maxima that trap EM from random initialization with high probability (Jin et al., [11]). A recent line of work has shown that this obstruction disappears when the model is over-parameterized: with $n = \Omega(m \log m)$ fitted components, randomly initialized gradient EM recovers an $m$-component Gaussian mixture in polynomial time, first for $m = 1$ (Xu et al., [28]) and then for general well-separated mixtures (Zhou et al., 2020). Recent work on MoEs over-provision of experts so that every latent cluster is covered at initialization is the mechanism behind recent gradient-based analyses (Chen et al., [2]; Kawata et al., [13]; Liao & Kyrillidis, 2026). Whether it carries the joint likelihood dynamics of gate and experts has remained open.

In this paper we give a global convergence guarantee for gradient-based training of a Gaussian-gated mixture of linear experts. Data are generated by drawing a latent expert index $\ell$ with probability $\pi_\ell$, then an input $x \sim\mathcal{N}(\mu_\ell, I_d)$, and finally a response $y = x^{\mathsf{T}}\theta_\ell^*+\epsilon$ with $\epsilon\sim\mathcal{N}(0,1)$, independent of $(\ell,x)$. Although this is a generative model, its predictive function is exactly an MoE in the sense used in practice. A quick calculation shows that

that is, a softmax router with weights $\mu_i$ and biases $b_i$ with linear experts $\theta_i$. The mixing weights $\pi_i$, which are usually fixed or ignored in MoE analysis, play the role of learnable router biases. We train this model by gradient EM on the population log-likelihood, which we show coincides with gradient descent on the Kullback–Leibler (KL) divergence between the truth and the model.

Our main result (Theorem 4.1) states that when the fitted model has $n \gtrsim m\log m$ experts, the gates are initialized at random from the data, the experts are initialized by a windowed least-squares fit, and the true gate centers are well separated, gradient EM drives the population KL loss below any target $\epsilon\ge e^{-cdn}$ within a polynomial number of steps. Global guarantees for gradient-based MoE training exist for staged algorithms, purpose-built losses and specific teacher models (§ 2); to our knowledge this is the first for joint training of gate and experts on the likelihood itself, and the first in which the mixing weights are free. The proof establishes a gradient-dominance inequality. A central difficulty is controlling the expert parameters: the expert errors (residuals) $y-x^{\mathsf{T}}\theta_i$ have to remain bounded within a shrinking budget throughout training. The least-squares warm start supplies the initial bound on the residuals, and an expert-drift argument ensures the budget is not violated.

Our contributions are as follows.

• A global guarantee for joint gate–expert training (Theorem 4.1): with $n \gtrsim m\log m$ experts (for balanced mixing weights), random gates and least-squares experts, population gradient EM reaches any accuracy $\epsilon\ge e^{-cdn}$ in polynomially many steps.

• A proof technique that decouples the loss into per-cluster problems, controls each by a Gaussian log-Sobolev inequality, and keeps the expert residuals inside a shrinking budget along the whole trajectory.

• Experiments showing (i) a sharp success transition along $n \sim m\log m$ (Figure 1), (ii) that randomly initialized experts succeed almost as often as least-squares experts once the gates cover the clusters (Figure 2), and (iii) that free mixing weights concentrate routing on about $m$ experts and make the model prunable, while load-balancing penalties destroy both properties (Figure 3 and Figure 4).

Related Work

EM for exactly parameterized mixtures. The convergence of EM for Gaussian mixtures was studied early by Xu & Jordan [27]. Modern convergence theory begins with Balakrishnan et al. [1], who showed that population EM converges locally at a geometric rate, and verified this for two-component Gaussian mixtures and two-component mixed linear regression. Subsequent work extended these guarantees to mixtures of many linear regressions, to two-component mixed linear regression from random initialization, and to well-separated Gaussian mixtures ([15, 16, 14]). For three or more components the picture changes: Jin et al. [11] showed that the population log-likelihood has spurious local maxima to which EM from random initialization converges with high probability.

Over-parameterized mixtures. Fitting more components than the data contain changes the picture. Dwivedi et al. [5] showed that over-specification creates singularities in the Fisher information and slows the local rate of EM from geometric to polynomial, and Xu et al. [28] then proved the first global result in this regime: gradient EM fitting an $n$-component mixture to a single Gaussian converges from random initialization at rate $O(1/\sqrt{t})$. Zhou et al. [30] extended this to a general $m$-component, well-separated Gaussian mixture, showing that $n=\Omega(m\log m)$ components suffice for randomly initialized gradient EM to converge globally in polynomial time and with polynomially many samples. Our analysis builds on their framework; the experts add a second parameter set whose errors enter the gate through the regression residuals, and controlling that coupling is the main new difficulty.

Theory of mixtures of experts. Mixtures of experts were introduced by Jacobs et al. [10] and Jordan & Jacobs [12], who already proposed EM for training them. A recent body of work analyzes the statistical properties of MoE estimators: convergence rates of the maximum-likelihood estimator for Gaussian MoEs (Ho et al., [9]) along with the effect of the softmax gate on identifiability and rates (Nguyen et al., [20, 21]). The Gaussian-gated model of Nguyen et al. [22] is closest to ours; however, these results concern the estimator rather than the algorithm that computes it. On the algorithmic side, Fruytier et al. [7] cast EM for mixtures of experts as mirror descent and obtain local guarantees. Makkuva et al. [18] recover the experts by a moment method and then the gate by EM, and Makkuva et al. [19] obtain global guarantees for gradient descent by designing separate losses for the gate and the experts. Closest to us are analyses of gradient-based training under a latent cluster structure: Chen et al. [2] (top-1 routing and classification, where the nonlinearity of the experts is essential), Kawata et al. [13] (Gaussian clusters with single-index targets, trained by a multi-phase SGD schedule that updates the router and the experts separately), and Liao & Kyrillidis [17] (joint gradient flow on the population squared error for a normalized sigmoid router with Hermite experts, an orthonormal teacher and $d \ge\operatorname{poly}(m)$, followed by pruning). All three provision $\Omega(m \log m)$ experts so that every cluster is covered at initialization, as we do. We differ in analyzing the likelihood objective directly, with gate, experts and mixing weights trained jointly by one gradient method, for arbitrary well-separated centers and mixing weights and without a high-dimensional regime.

Routing and load balancing in practice. Sparsely gated MoE layers (Shazeer et al., [23]) and their simplifications (Fedus et al., [6]) route each input to a few experts and add an auxiliary loss that penalizes imbalanced expert utilization, since imbalanced routing wastes capacity and complicates parallel execution. In the idealized case where the fraction of inputs dispatched to expert $i$ equals its average routing probability $\bar{\rho}_i$, the Switch auxiliary loss (Fedus et al., [6]) equals $n \sum_{i=1}^{n} \bar{\rho}_i^2 = n/n_{\mathrm{eff}}$, where $n_{\mathrm{eff}} := \left(\sum_{i=1}^{n} \bar{\rho}_i^2\right)^{-1}$ is the effective number of experts the router uses; minimizing the penalty therefore pushes utilization toward the full fitted width $n$. More recently, balancing has also been pursued through per-expert bias terms that are added to the routing scores used for expert selection and adjusted by a rule during training, instead of through an auxiliary loss (Wang et al., [25]; DeepSeek-AI [4]). These biases still aim at balanced load, but they leave the gradient objective untouched.

Learning a Mixture of Experts

Problem Statement

Our data model is an $m$-component mixed linear regression in which each regression vector has its own Gaussian covariate cluster. A sample $z=(x,y)\in\mathbb{R}^{d}\times\mathbb{R}$ is generated by

We can now define the mixture distribution $p_\star$ over $z$. Let $\varphi(x;\alpha):=(2\pi)^{-d/2}e^{-\|x-\alpha\|^2/2}$ denote the p.d.f. of a Gaussian with mean $\alpha\in\mathbb{R}^d$ and unit covariance and write the one dimensional case as $\varphi_1(y;\alpha)$. Also, let $\mathcal{P}^k\subset\mathbb{R}^{k+1}$ denote the probability simplex. For $\mu\in\mathbb{R}^d$ and $\theta\in\mathbb{R}^d$ the joint density of one component is $\varphi(z;\mu,\theta):=\varphi(x;\mu)\varphi_1(y;x^\mathsf{T}\theta)$. The ground truth is then the $m$-component joint density

and the over-parameterized model has $n\ge m$ experts,

We write $\varphi_i(z) := \varphi(z; \boldsymbol{\mu}_i, \boldsymbol{\theta}_i)$ and $\varphi_\ell^*(z) := \varphi(z; \boldsymbol{\mu}_\ell^*, \boldsymbol{\theta}_\ell^*)$, and drop parameter arguments when they are clear from context. Define $\pi_{\min} := \min_\ell\pi_\ell^*$. We write $\mathbb{E}_\ell:= \mathbb{E}_{z \sim\varphi_\ell^*}$ for expectation under the $\ell$th true component, i.e., $x \mid\ell\sim\mathcal{N}(\boldsymbol{\mu}_\ell^*, I_d)$ and then $y \mid x,\ell\sim\mathcal{N}(x^\top\boldsymbol{\theta}_\ell^*, 1)$, so that $\mathbb{E}_{p_*} = \sum_\ell\pi_\ell^* \mathbb{E}_\ell$.

Let $S_\ell:= \{i \in[n] : \|\boldsymbol{\mu}_i - \boldsymbol{\mu}_\ell^*\| \le\|\boldsymbol{\mu}_i - \boldsymbol{\mu}_j^*\| \ \forall j\}$ (ties broken by a fixed rule) be the Voronoi group of $\boldsymbol{\mu}_\ell^*$, $\ell(i)$ the group of $i$, $\hat{\pi}_\ell:= \sum_{i \in S_\ell} \pi_i$ and

measure how far the fitted components sit from the truth of their group.

EM Algorithm