# Breaking the gridlock in Mixture-of-Experts: Consistent and Efficient Algorithms Ashok Vardhan Makkuva Sewoong Oh Sreeram Kannan Pramod Viswanath ###### Abstract Mixture-of-Experts (MoE) is a widely popular model for ensemble learning and is a basic building block of highly successful modern neural networks as well as a component in Gated Recurrent Units (GRU) and Attention networks. However, present algorithms for learning MoE, including the EM algorithm and gradient descent, are known to get stuck in local optima. From a theoretical viewpoint, finding an efficient and provably consistent algorithm to learn the parameters remains a long standing open problem for more than two decades. In this paper, we introduce the first algorithm that learns the true parameters of a MoE model for a wide class of non-linearities with global consistency guarantees. While existing algorithms jointly or iteratively estimate the expert parameters and the gating parameters in the MoE, we propose a novel algorithm that breaks the deadlock and can directly estimate the expert parameters by sensing its echo in a carefully designed cross-moment tensor between the inputs and the output. Once the experts are known, the recovery of gating parameters still requires an EM algorithm; however, we show that the EM algorithm for this simplified problem, unlike the joint EM algorithm, converges to the true parameters. We empirically validate our algorithm on both the synthetic and real data sets in a variety of settings, and show superior performance to standard baselines. Machine Learning, ICML ## 1 Introduction In this paper, we study a popular gated neural network architecture known as Mixture-of-Experts (MoE). MoE is a basic building block of highly successful modern neural networks like Gated Recurrent Units (GRU) and Attention networks. A key interesting feature of MoE is the presence of a gating mechanism that allows for specialization of experts in their respective domains. MoE allows for the underlying expert models to be simple while allowing to capture complex non-linear relations between the data. Ever since their inception more than two decades ago (Jacobs et al., [1991](#bib.bib10)), they have been a subject of great research interest (Tresp, [2001](#bib.bib31); Collobert et al., [2002](#bib.bib6); Ng & Deisenroth, [2014](#bib.bib22); Theis & Bethge, [2015](#bib.bib30); Le et al., [2016](#bib.bib17); Gross et al., [2017](#bib.bib9); Sun et al., [2017](#bib.bib28); Wang et al., [2018](#bib.bib33)) across multiple domains such as computer vision, natural language processing, speech recognition, finance, and forecasting. The basic MoE model is the following: let $\boldsymbol{x}\in\mathbb{R}^{d}$ be the input feature vector and $y\in\mathbb{R}$ be the corresponding label. Then the discriminative model $P_{y|\boldsymbol{x}}$ for the $k$-mixture of experts ($k$-MoE) in the regression setting is: [TABLE] Figure [1](#S1.F1) details the architecture for $k$-MoE. [FIGURE:S1.F1] The interpretation behind ([1](#S1.E1)) is that for each input $\boldsymbol{x}$, the gating network chooses an expert based on the outcome of a multinomial random variable $z\in[k]$, whose probability depends on $\boldsymbol{x}$ in a parametric way, i.e. $z|\boldsymbol{x}\sim\mathrm{softmax}(\langle{\boldsymbol{w}_{1}^{\ast}},{\boldsymbol{x}}\rangle,\ldots,\langle{\boldsymbol{w}_{k}^{\ast}},{\boldsymbol{x}}\rangle)$. The chosen expert then generates the output $y$ from a Gaussian distribution centred at a non-linear activation of $\boldsymbol{x}$, i.e. $g(\langle{\boldsymbol{a}_{z}^{\ast}},{\boldsymbol{x}}\rangle)$, with variance $\sigma^{2}$. We want to learn the expert parameters $\boldsymbol{a}_{i}^{\ast}\in\mathbb{R}^{d}$ (also referred to as the regressors) and the gating parameters $\boldsymbol{w}_{i}^{\ast}\in\mathbb{R}^{d}$, assuming we know the non-linear activation $g:\mathbb{R}\to\mathbb{R}$. This problem of learning MoE has been a long standing open problem for more than two decades, even though it is a fundamental building block of several state-of-the-art gated neural network architectures. Gated neural networks such as GRUs and Sparsely-gated-MoEs have been widely successful in challenging tasks like machine translation (Chung et al., [2014](#bib.bib5); Shazeer et al., [2017](#bib.bib26); Vaswani et al., [2017](#bib.bib32)). Parameters are typically learnt through (stochastic) gradient descent on a non-convex loss function. However, these methods do not possess any theoretical guarantees, even for the simplest gated neural network, which is the MoE. On the other hand, existing guarantees for simpler models without gating units do not extend to MoEs. Consider the mixture of generalized linear models (M-GLMs) (Sedghi et al., [2014](#bib.bib24); Sun et al., [2014](#bib.bib29); Yi et al., [2016](#bib.bib36); Zhong et al., [2016](#bib.bib38)), which is a strict simplification of the $k$-MoE model in ([1](#S1.E1)), where $\boldsymbol{w}_{i}^{\ast}=0$ for all $i\in\{1,\ldots,k\}$. The learning in M-GLMs is usually done through a combination of spectral methods and greedy methods such as EM. A major limitation of these methods is that they rely critically on the fact that the mixing probability is a constant and hence they do not generalize to MoEs (see Section [2](#S2)). In addition, the EM algorithm, which is the workhorse for learning in parametric mixture models, is prone to bad local minima (Sedghi et al., [2014](#bib.bib24); Balakrishnan et al., [2017](#bib.bib2); Zhong et al., [2016](#bib.bib38)) (we independently verify this for MoEs in Section [4](#S4)). These theoretical shortcomings and practical relevance of the MoE models lead to the following fundamental question: Can we find an efficient and a consistent algorithm (with global initializations) that recovers the true parameters of the model with theoretical guarantees? In this paper, we address this question precisely and make the following contributions: **1) First theoretical guarantees:** We provide the first (poly-time) efficient algorithm that recovers the true parameters of a MoE model with global initializations (Theorem [1](#Thmtheorem1) and Theorem [2](#Thmtheorem2)). We allow for a wide class of non-linearities which includes the popular choices of identity, sigmoid, and ReLU. To the best of our knowledge, ours is the first work to give global convergence guarantees for MoE. **2) Algorithmic innovations:** Existing algorithms jointly or iteratively estimate the expert parameters and the gating paramters in the MoE and can get stuck in local minima. In this paper, we propose a novel algorithm that breaks the gridlock and can directly estimate the expert parameters by sensing its echo in a cross-moment tensor between the inputs and the output (Algorithm [1](#alg1) and Algorithm [2](#alg2)). Once the experts are known, the recovery of gating parameters still requires an EM algorithm; however, we show that the EM algorithm for this simplified problem, unlike the joint EM algorithm, converges to the true parameters. The proofs of global convergence of EM as well as the design of the cross-moment tensor are of independent mathematical interest. **3) Novel transformations:** In this paper, we introduce the novel notion of “Cubic and Quadratic Transform (CQT)". These are polynomial transformations on the output labels tailored to specific non-linear activation functions and the noise variance. The key utility of these transforms is to equip MoEs with a supersymmetric tensor structure in a principled way (Theorem [1](#Thmtheorem1)). [FIGURE:S1.F2] **Related work.** While there is a huge literature on MoEs ((Yuksel et al., [2012](#bib.bib37); Masoudnia & Ebrahimpour, [2014](#bib.bib20)) are detailed surveys), there are relatively few works on its learning guarantees. (Jordan & Xu, [1995](#bib.bib15)) is the first work to analyze the local convergence of joint-EM for both the gating and the expert parameters. As noted earlier, however, EM is prone to bad local minima. In contrast, our algorithms have *global convergence* guarantees. It is important to note that even for the simpler problem of mixtures of Gaussians, it is known that EM gets stuck in local minima, whenever number of mixtures, $k$, is at least $3$ (Jin et al., [2016](#bib.bib13)), whereas we can handle $2k-12$. **Notation.** In this paper, we denote Euclidean vectors by bold face lowercase letters $\boldsymbol{a},\boldsymbol{b}$, etc., and scalars by plain lowercase letters $y,z$, etc. We use $\mathcal{N}(y|\mu,\sigma^{2})$ either to denote the density or the distribution of a Gaussian random variable $y$ with mean $\mu$ and variance $\sigma^{2}$, depending on the context. $[d]\triangleq\{1,\ldots,d\}$. $\mathrm{Perm}[d]$ denotes the set of all permutations on $[d]$. We use $\otimes$ to denote the tensor outer product of vectors in $\mathbb{R}^{d}$. $\boldsymbol{x}^{\otimes 3}$ denotes $\boldsymbol{x}\otimes\boldsymbol{x}\otimes\boldsymbol{x}$, where $(\boldsymbol{x}\otimes\boldsymbol{x}\otimes\boldsymbol{x})_{ijk}=x_{i}x_{j}x_{k}$. $\mathrm{sym}(\boldsymbol{x}\otimes\boldsymbol{y}\otimes\boldsymbol{z})$ denotes the symmetrized version of $\boldsymbol{x}\otimes\boldsymbol{y}\otimes\boldsymbol{z}$, i.e. $\mathrm{sym}(\boldsymbol{x}\otimes\boldsymbol{y}\otimes\boldsymbol{z})_{ijk}=\sum_{\sigma\in\mathrm{Perm}[d]}x_{\sigma(i)}y_{\sigma(j)}z_{\sigma(k)}$. $\boldsymbol{e}_{i},i\in[d]$ denotes the standard basis vectors for $\mathbb{R}^{d}$. Through out the paper, we assume that $\boldsymbol{w}_{k}^{\ast}=0$, without loss of generality. ## 2 Algorithms In this section, we present our algorithms to learn the regression and gating parameters *separately*. Figure [2](#S1.F2) summarizes our algorithm. First we take a moment to highlight the issues of the existing approaches. For illustration purposes, we suppose that $k=2$ in ([1](#S1.E1)). We assume without loss of generality that $\boldsymbol{w}^{*}_{k}=\boldsymbol{w}^{*}_{2}=0$ and denote $\boldsymbol{w}_{1}^{\ast}=\boldsymbol{w}^{\ast}$. Thus the $2$-MoE model is given by $P_{y|\boldsymbol{x}}$: [TABLE] **Issues with traditional tensor methods.** In the far simplified setting of the absence of the gating parameter, i.e. $\boldsymbol{w}^{\ast}=0\in\mathbb{R}^{d}$, we see that $2$-MoE reduces to $2$-uniform mixture of GLMs. In this case, for $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$, the standard approach is to construct a $3^{\mathrm{rd}}$-order tensor $\mathcal{T}$ by regressing the output $y$ on the score transformation $\mathcal{S}_{3}(\boldsymbol{x})\triangleq\boldsymbol{x}\otimes\boldsymbol{x}\otimes\boldsymbol{x}-\sum_{i\in[d]}\mathrm{sym}\left(\boldsymbol{x}\otimes\boldsymbol{e}_{i}\otimes\boldsymbol{e}_{i}\right)$, i.e. [TABLE] Here the second equality follows from the generalized Stein’s lemma that $\mathbb{E}[f(\boldsymbol{x})\cdot\mathcal{S}_{3}(\boldsymbol{x})]=\mathbb{E}[\nabla_{\boldsymbol{x}}^{(3)}f(\boldsymbol{x})]$ under some regularity conditions on $f:\mathbb{R}^{d}\mapsto\mathbb{R}$ (see Lemma [2](#Thmlemma2) in Appendix [A](#A1)). Then the regressors can be learned through spectral decomposition on $\mathcal{T}$, where the uniqueness of decomposition follows from (Kruskal, [1977](#bib.bib16)). If we apply a similar technique for $2$-MoE in ([2](#S2.E2)), we obtain that [TABLE] where $\alpha_{i},\beta_{i},\gamma_{i},\delta$ are some scalar constants depending on the parameters $\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast},\boldsymbol{w}^{\ast}$ and $g$ (see Appendix [D.1](#A4.SS1) for the proof). Thus ([4](#S2.E4)) reveals that traditional spectral methods do not yield a supersymmetric tensor of the desired parameters for MoEs. In fact, ([4](#S2.E4)) contains all the $3^{\mathrm{rd}}$-order rank-$1$ terms formed by $\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast}$ and $\boldsymbol{w}^{\ast}$. Hence we cannot recover these parameters uniquely. Note that the inherent coupling between the regressors $\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast}$ and the gating parameter $\boldsymbol{w}^{\ast}$ in ([2](#S2.E2)) manifests as a cross tensor in ([4](#S2.E4)). This coupling serves as a key limitation for the traditional methods which critically rely on the fact that the mixing probability $p=\frac{1}{2}$ in ([4](#S2.E4)) is a constant. In fact, we recover ([3](#S2.E3)) by letting $\boldsymbol{w}^{\ast}=0$ in ([4](#S2.E4)). **Issues with EM algorithm.** EM algorithm is the workhorse for parameter learning in both the $k$-MoE and HME models (Jordan & Jacobs, [1994](#bib.bib14)). However, it is well known that EM is prone to spurious minima and existing theoretical results only establish local convergence for the regressors and the gating parameters. Indeed, our numerical experiments in Section [4.3](#S4.SS3) verify this fact. Figure [3(b)](#S4.F3.sf2) and Figure [3(c)](#S4.F3.sf3) highlight that joint-EM often gets stuck in bad local minima. ### 2.1 The proposed algorithm for learning MoE In order to tackle these challenges, we take a different route and propose to estimate the regressors and gating parameters *separately*. To gain intuition about our approach, let us consider 2-MoE model in ([2](#S2.E2)) with $\sigma=0$ and linear $g$. Then we have that $y$ either equals $\langle{\boldsymbol{a}_{1}^{\ast}},{\boldsymbol{x}}\rangle$ with probability $\sigma(\langle{\boldsymbol{w}^{\ast}},{\boldsymbol{x}}\rangle)$ or equals $\langle{\boldsymbol{a}_{2}^{\ast}},{\boldsymbol{x}}\rangle$ with probability $1-\sigma(\langle{\boldsymbol{w}^{\ast}},{\boldsymbol{x}}\rangle)$, where $\sigma(\cdot)$ is the sigmoid function. If we exactly know $\boldsymbol{w}^{\ast}$, we can recover $\boldsymbol{a}_{1}^{\ast}$ and $\boldsymbol{a}_{2}^{\ast}$ by solving a simple linear regression problem since we can recover the true latent variable $\boldsymbol{z}\in\{1,2\}$ with high probability. Similarly, if we know $\boldsymbol{a}_{1}^{\ast}$ and $\boldsymbol{a}_{2}^{\ast}$, it is easy to see that we can recover $\boldsymbol{w}^{\ast}$ by solving a binary linear classification problem. Thus knowing either the regressors or the gating parameters makes the estimation of other parameters easier. However, how do we first obtain one set of parameters without any knowledge about the other? Our approach precisely addresses this question and breaks the *grid lock*. We show that we can extract the regressors $\boldsymbol{a}_{1}^{\ast}$ and $\boldsymbol{a}_{2}^{\ast}$ without knowing $\boldsymbol{w}^{\ast}$ at all, just using the samples. Although we explain our approach with two mixtures, all claims are made precise for general $k$ in Theorems [1](#Thmtheorem1) and [2](#Thmtheorem2), and the algorithms are written for general $k$ as well in Algorithms [1](#alg1) and [2](#alg2). #### Step 1: Estimation of regressors To learn the regressors, we first pre-process $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$ using the score transformations $\mathcal{S}_{3}$ and $\mathcal{S}_{2}$, i.e. [TABLE] These score functions can be viewed as higher-order feature extractors from the inputs. As we have seen in ([3](#S2.E3)), these transformations suffice to learn the parameters in M-GLMs. However this approach fails in the context of MoE, as highlighted in ([4](#S2.E4)). Can we still construct a supersymmetric tensor for MoE? To answer this question in a principled way, we introduce the notion of “Cubic and Quadratic Transform (CQT)" for the labels, i.e. [TABLE] The coefficients $(\alpha,\beta,\gamma)$ in these polynomial transforms are obtained by solving a linear system of equations (see Appendix [C](#A3)). For the special case of $g=$linear, we obtain $\mathcal{P}_{3}(y)=y^{3}-3(1+\sigma^{2})\,y$ and $\mathcal{P}_{2}(y)=y^{2}$. These special transformations are specific to the choice of non-linearity $g$ and the noise variance $\sigma$. The key intuition behind the design of these transforms is that we can nullify the cross moments and obtain supersymmetric tensor in ([3](#S2.E3)) if we regress $\mathcal{P}_{3}(y)$ instead of $y$, for properly chosen constants $\alpha$ and $\beta$. This is made mathematically precise in Theorem [1](#Thmtheorem1). A similar argument holds for $\mathcal{P}_{2}(y)$ too. In addition, the choice of these polynomials is unique in the sense that any other polynomial transformations fail to yield the desired tensor structure. Using these transforms, we construct two special tensors $\hat{\mathcal{T}_{3}}\in(\mathbb{R}^{d})^{\otimes 3}$ and $\hat{\mathcal{T}_{2}}\in(\mathbb{R}^{d})^{\otimes 2}$. Later we use the robust tensor power method (Anandkumar et al., [2014](#bib.bib1)) on these tensors to learn the regressors. Algorithm [1](#alg1) details our learning procedure. Theorem [1](#Thmtheorem1) establishes the theoretical justification for our algorithm. [FIGURE:alg1] #### Step 2: Estimation of gating parameters To gain intuition for estimating the gating parameters, let $g=\mathrm{linear}$ in ([2](#S2.E2)) for simplicity. Moreover, assume that we know both $\boldsymbol{a}_{1}^{\ast}$ and $\boldsymbol{a}_{2}^{\ast}$. Then taking conditional expectation on $y$, we obtain from ([2](#S2.E2)) that [TABLE] where $f$ is the sigmoid function. Thus, [TABLE] Note that since $\boldsymbol{x}$ is Gaussian, $\langle{\boldsymbol{a}_{1}^{\ast}-\boldsymbol{a}_{2}^{\ast}},{\boldsymbol{x}}\rangle$ is non-zero with probability $1$. Hence, to recover $\boldsymbol{w}^{\ast}$, in view of Stein’s lemma, we may write [TABLE] However, it turns out that the above chain of equalities does not hold. Surprisingly, the first equality, which essentially is the law of iterated expectations, is not valid in this case as $\frac{y-\langle{\boldsymbol{a}_{2}^{\ast}},{\boldsymbol{x}}\rangle}{\langle{\boldsymbol{a}_{1}^{\ast}-\boldsymbol{a}_{2}^{\ast}},{\boldsymbol{x}}\rangle}$ is not integrable since it is a mixture of two Cauchy distributions, as proved in Appendix [D.4](#A4.SS4). Thus the above analysis highlights the difficulty of learning the gating parameters even in the simplest setting of two linear mixtures. Can we still learn $\boldsymbol{w}^{\ast}$ using method of moments (MoM)? In Theorem [3](#Thmtheorem3), we precisely address this question and show that we can still provably recover the gating parameters using MoM, by designing clever transformations on the data to infer the parameters of a Cauchy mixture distribution. While Theorem [3](#Thmtheorem3) highlights that gating parameters can be learnt using the method of moments for $2$-MoE, we still need a principled approach to learn these parameters for a more generic setting of $k$-MoE. Recall that the traditional joint-EM algorithm randomly initializes both the regressors and the gating parameters and updates them iteratively. Figure [3(b)](#S4.F3.sf2) and Figure [3(c)](#S4.F3.sf3) highlight that this procedure is prone to spurious minima. Can we still learn the gating parameters with *global initializations*? To address this question, we utilize the regressors learnt from Algorithm [1](#alg1). In particular, we use EM algorithm to update *only* the gating parameters, while fixing the regressors $\hat{\boldsymbol{a}}_{1},\ldots,\hat{\boldsymbol{a}}_{k}$. We show in Theorem [2](#Thmtheorem2) that, with *global/random* initializations, this variant of EM algorithm learns the true parameters. To the best of our knowledge, this is the first global convergence result for EM for $k>2$ mixtures. This motivates the following algorithm ($\varepsilon>0$ is some error tolerance): [FIGURE:alg2] ## 3 Theoretical analysis In this section, we provide the theoretical guarantees for our algorithms in the population setting. We first formally state our assumptions and justify the rationale behind them: 1. 1. $\boldsymbol{x}$ follows standard Gaussian distribution, i.e. $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$. 2. 2. $\left\|\boldsymbol{a}_{i}^{\ast}\right\|_{2}=1$ for $i\in[k]$ and $\left\|\boldsymbol{w}_{i}^{\ast}\right\|_{2}\leq R$ for $i\in[k-1]$, with some $R>0$. 3. 3. $\boldsymbol{a}_{i}^{\ast},i\in[k]$ are linearly independent and $\boldsymbol{w}_{i}^{\ast}$ is orthogonal to $\mathrm{span}\{\boldsymbol{a}_{1}^{\ast},\ldots,\boldsymbol{a}_{k}^{\ast}\}$ for $i\in[k-1]$. 4. 4. The non-linearity $g:\mathbb{R}\to\mathbb{R}$ is $(\alpha,\beta,\gamma)$-valid, which we define in Appendix [C](#A3). For example, this class includes $g=$linear, sigmoid and ReLU. **Remark.** We note that the Gaussianity of the input distribution and norm constraints on the parameters are standard assumptions in the learning of neural networks literature (Janzamin et al., [2015](#bib.bib12); Li & Yuan, [2017](#bib.bib18); Ge et al., [2017](#bib.bib8); Zhong et al., [2017](#bib.bib39); Du et al., [2017](#bib.bib7); Safran & Shamir, [2017](#bib.bib23)) and also that of M-GLMs (Sedghi et al., [2014](#bib.bib24); Yi et al., [2016](#bib.bib36); Zhong et al., [2016](#bib.bib38); Balakrishnan et al., [2017](#bib.bib2)). An interpretation behind Assumption $3$ is that if we think of $\boldsymbol{x}$ as a high-dimensional feature vector, distinct sub-features of $\boldsymbol{x}$ are used to perform the two distinct tasks of classification (using $\boldsymbol{w}_{i}^{\ast}$’s) and regression (using $\boldsymbol{a}_{i}^{\ast}$’s). We note that we need the above assumptions only for the technical analysis. In Section [4.1](#S4.SS1) and Section [4.2](#S4.SS2), we empirically verify that our algorithms work well in practice even under the relaxation of these assumptions. Thus we believe that the assumptions are merely technical artifacts. We are now ready to state our results. ###### **Theorem 1**** **(Recovery of regression parameters)**.** *Let $(\boldsymbol{x},y)$ be generated according the true model ([1](#S1.E1)). Under the above assumptions, we have that* [TABLE] *where $c^{\prime}_{g}$ and $c_{g,\sigma}$ are two non-zero constants depending on $g$ and $\sigma$. Hence the regressors $\boldsymbol{a}_{i}^{\ast}$’s can be learnt through tensor decomposition on $\mathcal{T}_{2}$ and $\mathcal{T}_{3}$.* ###### Proof. (Sketch) To highlight the central ideas behind the proof, first let $g=$linear. From ([1](#S1.E1)) we get that [TABLE] where $p_{i}^{\ast}(\boldsymbol{x})\triangleq P_{i|\boldsymbol{x}}$ for $i\in[k]$. Taking the cross moment of $y$ with $\mathcal{S}_{3}(\boldsymbol{x})$ and using Lemma [2](#Thmlemma2) we obtain that [TABLE] Notice that had $p_{i}^{\ast}(\boldsymbol{x})$ been a constant in the above equation, we would obtain a supersymmetric tensor easily as is the case with M-GLMs. However, $\mathbb{E}[\nabla_{\boldsymbol{x}}^{(3)}(p_{i}^{\ast}(\boldsymbol{x})\langle{\boldsymbol{a}_{i}^{\ast}},{\boldsymbol{x}}\rangle)]$ now contains all the third-order rank-$1$ terms involving the tensor product of $\boldsymbol{w}_{1}^{\ast},\ldots,\boldsymbol{w}_{k-1}^{\ast}$ and $\boldsymbol{a}_{i}^{\ast}$ for any fixed $i$. Our key insight is that this issue can be avoided if we cleverly transform $y$. In particular, we consider a cubic transformation $\mathcal{P}_{3}(y)=y^{3}-3y(1+\sigma^{2})$ and obtain that [TABLE] Now it turns out that after using the orthogonality of $\boldsymbol{w}_{i}^{\ast}$ and $\boldsymbol{a}_{i}^{\ast}$, and the fact $\mathbb{E}[p(Z)]=\mathbb{E}[p^{\prime}(Z)]=\mathbb{E}[p^{\prime\prime}(Z)]=0$ for $3^{\mathrm{rd}}$-Hermite polynomial $p(z)=z^{3}-3z$ and $Z\sim\mathcal{N}(0,1)$, we can nullify the cross-moments between $\boldsymbol{w}_{i}^{\ast}$’s and $\boldsymbol{a}_{i}^{\ast}$’s to obtain that [TABLE] Similarly, we can show that $\mathbb{E}[\mathcal{P}_{2}(y)\cdot\mathcal{S}_{2}(\boldsymbol{x})]=2\sum_{i\in[k]}\mathbb{E}[p_{i}^{\ast}(\boldsymbol{x})].(\boldsymbol{a}_{i}^{\ast})^{\otimes 2}$. For a general non-linearity $g:\mathbb{R}\to\mathbb{R}$, we can similarly design cubic and quadratic polynomials $\mathcal{P}_{3}=y^{3}+\alpha y^{2}+\beta y$ and $\mathcal{P}_{2}=y^{2}+\gamma y$ such that we can still construct supersymmetric tensors involving the regressors. In order to obtain the unique set of coefficients $(\alpha,\beta,\gamma)$, we need to solve a linear system of equations, which we describe in Appendix [C](#A3). ∎ Once we obtain $\mathcal{T}_{2}$ and $\mathcal{T}_{3}$, the recovery gurantees for the regressors $\boldsymbol{a}_{i}^{\ast}$ follow from the standard tensor decomposition guarantees, for example, Theorem 4.3 and Theorem 5 of (Anandkumar et al., [2014](#bib.bib1)). We assume that the learnt regressors $\boldsymbol{a}_{i}$ are such that $\max_{i\in[k]}\left\|\boldsymbol{a}_{i}-\boldsymbol{a}_{i}^{\ast}\right\|_{2}=\sigma^{2}\varepsilon$ for some $\varepsilon>0$. Now we present our theoretical results for global convergence of EM. First we briefly recall the algorithm. Let $\Omega$ denote the domain of our gating parameters, defined as [TABLE] Then the population EM for the mixture of experts consists of the following two steps: - • **E-step:** Using the current estimate $\boldsymbol{w}_{t}$ to compute the function $Q(\cdot|\boldsymbol{w}_{t})$, - • **M-step:** $\boldsymbol{w}_{t+1}=\mathrm{argmax}_{\boldsymbol{w}\in\Omega}Q(\boldsymbol{w}|\boldsymbol{w}_{t})$, where the function $Q(.|\boldsymbol{w}_{t})$ is the expected log-likelihood of the complete data distribution with respect to current posterior distribution. Mathematically, [TABLE] where $\mathrm{const}$ refers to terms not depending on $\boldsymbol{w}$, $P_{\boldsymbol{w}}(z=i|x)=\exp(\boldsymbol{w}_{i}^{\top}\boldsymbol{x})/\sum_{j}\exp(\boldsymbol{w}_{j}^{\top}\boldsymbol{x})$ and $p_{\boldsymbol{w}_{t}}^{(i)}\triangleq\mathbb{P}\left[z=i|\boldsymbol{x},y,\boldsymbol{w}_{t}\right]$ corresponds to the posterior probability for the $i^{\text{th}}$ expert, given by [TABLE] In ([8](#S3.E8)), the expectation is with respect to the true distribution of $(\boldsymbol{x},y)$, given by ([1](#S1.E1)). Thus the EM can be viewed as a deterministic procedure which maps $\boldsymbol{w}_{t}\mapsto M(\boldsymbol{w}_{t})$ where [TABLE] When the estimated regressors $\boldsymbol{a}_{i}$ equal the true parameters $\boldsymbol{a}_{i}^{\ast}$, it follows from the self-consistency property of the EM that the true parameter $\boldsymbol{w}^{\ast}$ is a fixed-point for the EM operator $M$, i.e. $M(\boldsymbol{w}^{\ast})=\boldsymbol{w}^{\ast}$ (McLachlan & Krishnan, [2007](#bib.bib21)). However, this does not guarantee that EM converges to $\boldsymbol{w}^{\ast}$. In the following theorem, we show that even when the regressors are known approximately, EM algorithm converges to the true gating parameters at a geometric rate upto an additive error, under *global* initializations. For the error metric, we define $\left\|\boldsymbol{w}-\boldsymbol{w}^{\prime}\right\|\triangleq\max_{i\in[k-1]}\left\|\boldsymbol{w}_{i}-\boldsymbol{w}^{\prime}_{i}\right\|_{2}$ for any $\boldsymbol{w},\boldsymbol{w}^{\prime}\in\Omega$. We assume that $R=1$ for simplicity. (Our results extend straightforwardly to general $R$). ###### **Theorem 2****.** *Let $\varepsilon>0$ be such that $\max_{i}\|\boldsymbol{a}_{i}-\boldsymbol{a}_{i}^{\ast}\|_{2}=\sigma^{2}\varepsilon$. There exists a constant $\sigma_{0}>0$ such that whenever $0<\sigma<\sigma_{0}$, for any random initialization $\boldsymbol{w}_{0}\in\Omega$, the population-level EM updates on the gating parameter $\{\boldsymbol{w}\}_{t\geq 0}$ converge almost geometrically to the true parameter $\boldsymbol{w}^{\ast}$ upto an additive error, i.e.* [TABLE] *where $\kappa_{\sigma},\kappa$ are dimension-independent constant depending on $g$ and $\sigma$ such that $\kappa_{\sigma}\xrightarrow{\sigma\rightarrow 0}0$ and $\kappa\leq(k-1)\frac{\sqrt{6(2+\sigma^{2})}}{2}$ for $g=$linear, sigmoid and ReLU.* ###### Proof. (Sketch) One can show that the $Q(\cdot|\boldsymbol{w}_{t})$ defined in ([8](#S3.E8)) is a strongly concave function. Moreover, if we let $\varepsilon=0$ and $\boldsymbol{w}_{t}=\boldsymbol{w}^{\ast}$, we have from the self-consistency of EM that $\mathrm{argmax}Q(\cdot|\boldsymbol{w}^{\ast})=\boldsymbol{w}^{\ast}$. Thus if we can show that the functions are $Q(\cdot|\boldsymbol{w}_{t})$ and $Q(\cdot|\boldsymbol{w}^{\ast})$ “sufficiently close" whenever $\boldsymbol{w}_{t}$ and $\boldsymbol{w}^{\ast}$ are close, we can use the EM convergence analysis tools from (Balakrishnan et al., [2017](#bib.bib2)) to show that their corresponding maximizers also stay close upto a scaling factor determined by $\kappa_{\sigma}$ above. Then it follows that the EM updates converge geometrically. ∎ **Remark.** In the M-step of the EM algorithm, the next iterate is chosen so that the function $Q(\cdot|\boldsymbol{w}_{t})$ is maximized. Instead we can perform an ascent step in the direction of the gradient of $Q(\cdot|\boldsymbol{w}_{t})$ to produce the next iterate, i.e. $\boldsymbol{w}_{t+1}=\Pi_{\Omega}(\boldsymbol{w}_{t}+\alpha\nabla Q(\boldsymbol{w}_{t}|\boldsymbol{w}_{t}))$, where $\Pi_{\Omega}(\cdot)$ is the projection operator. This variant of EM algorithm is known as *Gradient EM*. In Appendix [G](#A7), we show that Gradient EM also enjoys similar convergence guarantees. **MoM to learn gating parameters.** In Theorem [2](#Thmtheorem2), we proved that EM algorithm provably recovers the true gating parameters for any $k\geq 2$ mixtures. In this section, we show that for the special case of $k=2$, we can learn $\boldsymbol{w}^{\ast}$ (upto the unit direction) using an alternative procedure involving MoM. First we define [TABLE] The following theorem establishes that the the CDF of the random variable $\mathrm{Ratio}(\boldsymbol{x},y)$, when regressed on input $\boldsymbol{x}$, is proportional to $\boldsymbol{w}^{\ast}$. ###### **Theorem 3****.** *Suppose that $(\boldsymbol{a}_{1},\boldsymbol{a}_{2})=(\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast})$. Then we have that* [TABLE] *where $\alpha\in\mathbb{R}$ is a scalar given by $\alpha=\mathbb{E}[f^{\prime}(\langle{\boldsymbol{w}^{\ast}},{\boldsymbol{x}}\rangle)\cdot\left(1-2\Phi\left(\frac{|\langle{\boldsymbol{a}_{1}-\boldsymbol{a}_{2}},{\boldsymbol{x}}\rangle|}{2\sigma}\right)\right)]$.* ###### Proof. (Sketch) We first show that $\mathrm{Ratio}(\boldsymbol{x},y)$ is a mixture of Cauchy distributions. Then we show that $\mathbb{E}[\mathds{1}\{\mathrm{Ratio}(\boldsymbol{x},y)\leq z\}|\boldsymbol{x}]=\mathbb{P}\left[\mathrm{Ratio}\leq z|\boldsymbol{x}\right]=f(\boldsymbol{w}^{\top}\boldsymbol{x})\Phi\left((z-1)\frac{|\Delta_{x}|}{\sigma}\right)+(1-f(\boldsymbol{w}^{\top}\boldsymbol{x}))\Phi\left(z\frac{|\Delta_{x}|}{\sigma}\right)$ where $\Delta_{x}=(\boldsymbol{a}_{1}-\boldsymbol{a}_{2})^{\top}\boldsymbol{x}$. Then our result follows from taking the first moment of the indicator random variable with $\boldsymbol{x}$ and Stein’s lemma. ∎ ## 4 Experiments In this section, we empirically validate our algorithm in various settings and compare its performance to that of EM on both synthetic and real world datasets 111Codes are available at this repository [MoE codes](https://github.com/Ashokvardhan/Breaking-the-gridlock-in-MoE-Consistent-and-Efficient-Algorithms).. In both the scenarios, we found that our algorithm consistently outperforms the existing approaches. For the tensor decomposition in our Algorithm [1](#alg1), we use the Orth-ALS package by (Sharan & Valiant, [2017](#bib.bib25)). In all the synthetic experiments, we first draw the regressors $\{\boldsymbol{a}_{i}^{\ast}\}_{i=1}^{k}$ i.i.d uniformly from the unit sphere $\mathbb{S}^{d-1}$. The input distribution $P_{\boldsymbol{x}}$ and the generation of $\boldsymbol{w}_{i}^{\ast}$’s are detailed for each experiment. Then the labels $y_{i}$ are generated according to the true $k$-MoE model in ([1](#S1.E1)) for linear activation. Additional experiments in this setting with non-linear activations are detailed in Appendix [H.1](#A8.SS1). Experiments with real world data are provided in Section [4.4](#S4.SS4). ### 4.1 Non-gaussian inputs In this section we let the input distribution to be mixtures of Gaussians (GMM). We let $k=2,d=10$ and $\sigma=0.1$. The gating parameter $\boldsymbol{w}^{\ast}\in\mathbb{R}^{10}$ is uniformly chosen from the unit sphere $\mathbb{S}^{9}$. To generate the input features, we first randomly draw $\mu_{1},\mu_{2}\in\mathbb{S}^{9}$, and generate $n$ i.i.d. samples $\boldsymbol{x}_{i}\sim p\mathcal{N}(\mu_{1},I_{d})+(1-p)\mathcal{N}(\mu_{2},I_{d})$, where $p\in\{0.1,0.3,0.5,0.7,0.9\}$. Here $n=2000$. Since $\boldsymbol{x}$ is a $2$-GMM, its score functions $\mathcal{S}_{3}(\boldsymbol{x}),\mathcal{S}_{2}(\boldsymbol{x})$ are computed using the densities of Gaussian mixtures (Janzamin et al., [2014](#bib.bib11)). To gauge the performance of our algorithm, we measure the correlation of our learned parameters $\boldsymbol{a}_{1},\boldsymbol{a}_{2}$ and $\boldsymbol{w}$ with the ground truth, i.e. [TABLE] where $\pi:\{1,2\}\to\{1,2\}$ is a permutation. Similarly, for the gating parameter, we define [TABLE] Here we assume that all the parameters are unit-normalized. The closer the values of fit are to $1$, the closer the learnt parameters are to the ground truth. As shown in Table [2](#S4.T2), our algorithms are able to learn the ground truth very accurately in a variety of settings, as indicated by the measured fit. This highlights the fact that our algorithms are robust to the input distributions. ### 4.2 Non-orthogonal parameters In this section we verify that our algorithms still work well in practice even under the relaxation of Assumption $3$. For the experiments, we consider the similar setting as before with $k=2,d=10,\sigma=0.1$ and the gating parameter $\boldsymbol{w}^{\ast}$ is drawn uniformly from $\mathbb{S}^{9}$ without the orthogonality restriction. We let $\boldsymbol{x}_{i}\stackrel{{\scriptstyle i.i.d.}}{{\sim}}\mathcal{N}(0,I_{d})$. We choose $n=2000$. We use $\mathsf{RegressorFit}$ and $\mathsf{GatingFit}$ defined in ([11](#S4.E11)) and ([12](#S4.E12)) respectively, as our performance metrics. From Table [2](#S4.T2), we can see that the performance of our algorithms is almost the same across both the settings. In both the scenarios, our fit is consistently greater than $0.9$. In Figure [3(a)](#S4.F3.sf1), we plotted $\mathsf{GatingFit}(\boldsymbol{w}_{t})$ vs. the number of iterations $t$, as $\boldsymbol{w}_{t}$ is updated according to Algorithm [2](#alg2), over $10$ independent trials. We observe that the learned parameters converge to the true parameters in less than $5$ iterations. [FIGURE:S4.F3] [FIGURE:S4.T2] ### 4.3 Comparison to joint-EM Here we compare the performance of our algorithm with that of the joint-EM. We let the number of mixture components be $k=3$ and $k=4$. We let $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$ and the gating parameters are drawn uniformly from $\mathbb{S}^{9}$. If $\boldsymbol{A}=[\boldsymbol{a}_{1}\ldots\boldsymbol{a}_{k}]$ and $\boldsymbol{W}=[\boldsymbol{w}_{1}\ldots\boldsymbol{w}_{k-1}\ 0]$ denote the estimated expert and gating parameters respectively, our evaluation metric is $\mathcal{E}$, the Frobenious norm of the parameter error accounting for the best possible permutation $\pi:[k]\to[k]$, i.e. $\mathcal{E}(\boldsymbol{A},\boldsymbol{W})=\inf_{\pi}\|\boldsymbol{A}-\boldsymbol{A}_{\pi}^{\ast}\|_{F}+\|\boldsymbol{W}-\boldsymbol{W}_{\pi}^{\ast}\|_{F}$, where $\boldsymbol{A}_{\pi}^{\ast}=[\boldsymbol{a}^{\ast}_{\pi(1)}\ldots\boldsymbol{a}^{\ast}_{\pi(k)}]$ denotes the permuted regression parameter matrix and similarly for $\boldsymbol{W}_{\pi}^{\ast}$. In Figure [3(b)](#S4.F3.sf2) and Figure [3(c)](#S4.F3.sf3), we compare the performance of our algorithm with the joint-EM algorithm for $n=8000,d=10,\sigma=0.5$. The plotted estimation error $\mathcal{E}(\boldsymbol{A},\boldsymbol{W})$ is averaged for $10$ trials. It is clear that our algorithm is able to recover the true parameters thus resulting in much smaller parameter error than the joint-EM which often gets stuck in local optima. In addition, our algorithm is able to learn these parameters in very few iterations, often less than $10$ iterations. We also find that our algorithm consistently outperforms the joint-EM for different choices of non-linearities, number of samples, number of mixtures, etc. (details provided in Appendix [H](#A8)). Note that the above error metric $\mathcal{E}(\boldsymbol{A},\boldsymbol{W})$ is close to zero if and only if $\mathsf{Regressor\ Fit}$ and $\mathsf{Gating\ Fit}$ is close to one. ### 4.4 Real data To highlight the generalizability of our algorithm, in Appendix [H.2](#A8.SS2) of the supplement, we compare the performance of our algorithm to that of the standard approaches on a variety of real world datasets. Results from these experiments highlight the fact that in the real world scenario, where the underlying data is not generated according to a MoE model, our approach still learns a superior set of parameters as opposed to the existing algorithms. This fact is reflected in the lowest prediction errors obtained by our algorithm. ## 5 Discussion In this paper we provided the first provable and globally consistent algorithm that can learn the true parameters of a MoE model. We believe that ideas from (Sedghi et al., [2014](#bib.bib24)) can be naturally extended for the finite sample complexity analysis of the tensor decomposition to learn the regressors and similarly, techniques from (Balakrishnan et al., [2017](#bib.bib2)) can be extended to the finite sample EM convergence analysis for the gating parameters. While we have focused here on parameter recovery, however, there are no statistical bounds on output prediction error when the data is not generated from the model. MoE models are known to be capable of fitting general functions, and getting statistical guarantees on learning in such regimes is an interesting direction for future work. ## Acknowledgements This work is partly supported by NSF grants 1927712 and 1815535, NSF awards CNS-1718270, 1651236, 1703403, and the Army Research Office under grant W911NF1810332. ## References - Anandkumar et al. (2014) Anandkumar, A., Ge, R., Hsu, D., Kakade, S. M., and Telgarsky, M. Tensor decompositions for learning latent variable models. *J. Mach. Learn. Res.*, 15(1):2773–2832, January 2014. ISSN 1532-4435. - Balakrishnan et al. (2017) Balakrishnan, S., Wainwright, M. J., and Yu, B. Statistical guarantees for the EM algorithm: From population to sample-based analysis. *The Annals of Statistics*, 45(1):77–120, 2017. - Brooks et al. (1989) Brooks, T., Pope, D., and Marcolini., A. Airfoil self-noise and prediction. Technical report, NASA, 1989. URL [https://archive.ics.uci.edu/ml/datasets/Airfoil+Self-Noise](https://archive.ics.uci.edu/ml/datasets/Airfoil+Self-Noise). - Chen et al. (2014) Chen, Y., Yi, X., and Caramanis, C. A convex formulation for mixed regression with two components: Minimax optimal rates. In *Conference on Learning Theory*, pp. 560–604, 2014. - Chung et al. (2014) Chung, J., Gülçehre, Ç., Cho, K., and Bengio, Y. Empirical evaluation of gated recurrent neural networks on sequence modeling. abs/1412.3555, 2014. - Collobert et al. (2002) Collobert, R., Bengio, S., and Bengio, Y. A parallel mixture of SVMs for very large scale problems. *Neural Computing*, 2002. - Du et al. (2017) Du, S. S., Lee, J. D., Tian, Y., Poczos, B., and Singh, A. Gradient descent learns one-hidden-layer cnn: Don’t be afraid of spurious local minima. *arXiv preprint arXiv:1712.00779*, 2017. - Ge et al. (2017) Ge, R., Lee, J. D., and Ma, T. Learning one-hidden-layer neural networks with landscape design. *arXiv preprint arXiv:1711.00501*, 2017. - Gross et al. (2017) Gross, S., Szlam, A., et al. Hard mixtures of experts for large scale weakly supervised vision. In *Computer Vision and Pattern Recognition (CVPR), 2017 IEEE Conference on*, pp. 5085–5093. IEEE, 2017. - Jacobs et al. (1991) Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. *Neural Computation*, 1991. - Janzamin et al. (2014) Janzamin, M., Sedghi, H., and Anandkumar, A. Score function features for discriminative learning: Matrix and tensor framework. abs/1412.2863, 2014. URL [http://arxiv.org/abs/1412.2863](http://arxiv.org/abs/1412.2863). - Janzamin et al. (2015) Janzamin, M., Sedghi, H., and Anandkumar, A. Beating the perils of non-convexity: Guaranteed training of neural networks using tensor methods. *arXiv preprint arXiv:1506.08473*, 2015. - Jin et al. (2016) Jin, C., Zhang, Y., Balakrishnan, S., Wainwright, M. J., and Jordan, M. Local maxima in the likelihood of gaussian mixture models: Structural results and algorithmic consequences. *arXiv preprint arXiv:1609.00978*, 2016. - Jordan & Jacobs (1994) Jordan, M. I. and Jacobs, R. A. Hierarchical mixtures of experts and the EM algorithm. *Neural computation*, 6(2):181–214, 1994. - Jordan & Xu (1995) Jordan, M. I. and Xu, L. Convergence results for the EM approach to mixtures of experts architectures. *Neural Networks*, 8(9):1409–1431, 1995. - Kruskal (1977) Kruskal, J. B. Three-way arrays: rank and uniqueness of trilinear decompositions, with application to arithmetic complexity and statistics. *Linear algebra and its applications*, 18(2):95–138, 1977. - Le et al. (2016) Le, P., Dymetman, M., and Renders, J.-M. Lstm-based mixture-of-experts for knowledge-aware dialogues. *arXiv preprint arXiv:1605.01652*, 2016. - Li & Yuan (2017) Li, Y. and Yuan, Y. Convergence analysis of two-layer neural networks with relu activation. In *Advances in Neural Information Processing Systems*, pp. 597–607, 2017. - Liu & Yeh (2017) Liu, Y.-C. and Yeh, I.-C. Using mixture design and neural networks to build stock selection decision support systems. *Neural Computing and Applications*, 28(3):521–535, 2017. doi: 10.1007/s00521-015-2090-x. URL [https://archive.ics.uci.edu/ml/datasets/Stock+portfolio+performance](https://archive.ics.uci.edu/ml/datasets/Stock+portfolio+performance). - Masoudnia & Ebrahimpour (2014) Masoudnia, S. and Ebrahimpour, R. Mixture of experts: a literature survey. *Artificial Intelligence Review*, 42(2):275, 2014. - McLachlan & Krishnan (2007) McLachlan, G. and Krishnan, T. *The EM algorithm and extensions*, volume 382. John Wiley & Sons, 2007. - Ng & Deisenroth (2014) Ng, J. W. and Deisenroth, M. P. Hierarchical mixture-of-experts model for large-scale gaussian process regression. *arXiv preprint arXiv:1412.3078*, 2014. - Safran & Shamir (2017) Safran, I. and Shamir, O. Spurious local minima are common in two-layer relu neural networks. *arXiv preprint arXiv:1712.08968*, 2017. - Sedghi et al. (2014) Sedghi, H., Janzamin, M., and Anandkumar, A. Provable tensor methods for learning mixtures of classifiers. *arXiv preprint arXiv:1412.3046*, 2014. - Sharan & Valiant (2017) Sharan, V. and Valiant, G. Orthogonalized ALS: A theoretically principled tensor decomposition algorithm for practical use. In *Proceedings of the 34th International Conference on Machine Learning*, volume 70 of *Proceedings of Machine Learning Research*, pp. 3095–3104, 06–11 Aug 2017. URL [http://proceedings.mlr.press/v70/sharan17a.html](http://proceedings.mlr.press/v70/sharan17a.html). - Shazeer et al. (2017) Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. *arXiv preprint arXiv:1701.06538*, 2017. - Stein (1972) Stein, C. A bound for the error in the normal approximation to the distribution of a sum of dependent random variables. In *Proceedings of the Sixth Berkeley Symposium on Mathematical Statistics and Probability*, volume 2, pp. 583–602. University of California Press, 1972. - Sun et al. (2017) Sun, X., Peng, X., Ren, F., and Xue, Y. Human-machine conversation based on hybrid neural network. In *Computational Science and Engineering (CSE) and Embedded and Ubiquitous Computing (EUC), 2017 IEEE International Conference on*, volume 1, pp. 260–266. IEEE, 2017. - Sun et al. (2014) Sun, Y., Ioannidis, S., and Montanari, A. Learning mixtures of linear classifiers. In *Proceedings of the 31st International Conference on Machine Learning*, volume 32, pp. 721–729, 2014. - Theis & Bethge (2015) Theis, L. and Bethge, M. Generative image modeling using spatial lstms. In *Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2*, NIPS’15, pp. 1927–1935, Cambridge, MA, USA, 2015. MIT Press. - Tresp (2001) Tresp, V. Mixtures of gaussian processes. NIPS, 2001. - Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In *Advances in Neural Information Processing Systems*, pp. 5998–6008, 2017. - Wang et al. (2018) Wang, X., Yu, F., Wang, R., Ma, Y.-A., Mirhoseini, A., Darrell, T., and Gonzalez, J. E. Deep mixture of experts via shallow embedding. *arXiv preprint arXiv:1806.01531*, 2018. - Yeh (1998) Yeh, I.-C. Modeling of strength of high performance concrete using artificial neural networks. *Cement and Concrete Research*, 28(12):1797–1808, 1998. URL [https://archive.ics.uci.edu/ml/datasets/Concrete+Compressive+Strength](https://archive.ics.uci.edu/ml/datasets/Concrete+Compressive+Strength). - Yi et al. (2014) Yi, X., Caramanis, C., and Sanghavi, S. Alternating minimization for mixed linear regression. In *International Conference on Machine Learning*, pp. 613–621, 2014. - Yi et al. (2016) Yi, X., Caramanis, C., and Sanghavi, S. Solving a mixture of many random linear equations by tensor decomposition and alternating minimization. *arXiv preprint arXiv:1608.05749*, 2016. - Yuksel et al. (2012) Yuksel, S. E., Wilson, J. N., and Gader, P. D. Twenty years of mixture of experts. *IEEE Transactions on Neural Networks and Learning Systems*, 23(8):1177–1193, 2012. - Zhong et al. (2016) Zhong, K., Jain, P., and Dhillon, I. S. Mixed linear regression with multiple components. pp. 2190–2198. 2016. - Zhong et al. (2017) Zhong, K., Song, Z., Jain, P., Bartlett, P. L., and Dhillon, I. S. Recovery guarantees for one-hidden-layer neural networks. *arXiv preprint arXiv:1706.03175*, 2017. **Organization.** The appendix is organized as follows: - • Appendix [A](#A1) and Appendix [B](#A2) contain the requisite material for method of moments and the convergence analysis of EM respectively. - • Appendix [C](#A3) details the class of non-linearities for which our results hold. - • Appendix [D](#A4) contains all the proofs of Section [3](#S3). Two technical lemmas needed to prove Theorem [2](#Thmtheorem2) are relegated to Appendix [E](#A5) and Appendix [F](#A6). - • Appendix [G](#A7) provides convergence guarantees for Gradient EM. - • Appendix [H](#A8) contains additional experiments for the comparison of joint-EM and our algorithm for the synthetic data. ## Appendix A Toolbox for method of moments In this section, we introduce the key techniques that are useful in parameter estimation of mixture models via the method of moments. Stein’s identity (Stein’s lemma) is a well-known result in probability and statistics and is widely used in estimation and inference taks. A refined version of the Stein’s lemma (Stein, [1972](#bib.bib27)) for higher-order moments is the key to parameter estimation in mixture of generalized linear models. We utilize this machinery in proving Theorem [1](#Thmtheorem1). We first recall the Stein’s lemma. ###### **Lemma 1**** **(Stein’s lemma (Stein, [1972](#bib.bib27)) )**.** *Let $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$ and $g:\mathbb{R}^{d}\to\mathbb{R}$ be a function such that both $\mathbb{E}[\nabla_{\boldsymbol{x}}g(\boldsymbol{x})]$ and $\mathbb{E}[g(\boldsymbol{x})\cdot\boldsymbol{x}]$ exist and are finite. Then* [TABLE] The following lemma, which can be viewed as an extension of Stein’s lemma for higher-order moments, is the central technique behind parameter estimation in M-GLMs. ###### **Lemma 2**** **((Sedghi et al., [2014](#bib.bib24)))**.** *Let $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$ and $\mathcal{S}_{3}(\boldsymbol{x})$ be as defined in ([6](#S2.E6)) and let $\mathcal{S}_{2}(\boldsymbol{x})\triangleq\boldsymbol{x}\otimes\boldsymbol{x}-I_{d}$. Then for any $g:\mathbb{R}^{d}\to\mathbb{R}$ satisfying some regularity conditions, we have* [TABLE] ## Appendix B Toolbox for EM convergence analysis Recall that the domain of our gating parameters is $\Omega=\{\boldsymbol{w}:\left\|\boldsymbol{w}\right\|\leq 1\}$. Then the population EM for the mixture of experts consists of the following two steps: - • **E-step:** Using the current estimate $\boldsymbol{w}_{t}$ to compute the function $Q(\cdot|\boldsymbol{w}_{t})$. - • **M-step:** $\boldsymbol{w}_{t+1}=\mathrm{argmax}_{\left\|\boldsymbol{w}\right\|\leq 1}Q(\boldsymbol{w}|\boldsymbol{w}_{t})$. Thus the EM can be viewed as a deterministic procedure which maps $\boldsymbol{w}_{t}\mapsto M(\boldsymbol{w}_{t})$ where [TABLE] Our convergence analysis relies on tools from (Balakrishnan et al., [2017](#bib.bib2)) where they provided local convergence results on both the EM and gradient EM algorithms. In particular, they showed that if we initialize EM in a sufficiently small neighborhood around the true parameters, the EM iterates converge geometrically to the true parameters under some strong-concavity and gradient stability conditions. We now formally state the assumptions in (Balakrishnan et al., [2017](#bib.bib2)) under which the convergence guarantees hold. We will show in the next section that these conditions hold *globally* in our setting. ###### **Assumption 1**** **(Convexity of the domain)**.** $\Omega$ is convex. ###### **Assumption 2**** **(Strong-concavity)**.** $Q(\cdot|\boldsymbol{w}^{\ast})$ is a $\lambda$-strongly concave function over a $r$-neighborhood of $\boldsymbol{w}^{\ast}$, i.e. $\mathcal{B}(\boldsymbol{w}^{\ast},r)\triangleq\{\boldsymbol{w}\in\Omega:\left\|\boldsymbol{w}-\boldsymbol{w}^{\ast}\right\|\leq r\}$. ###### **Remark 1****.** An important point to note is that the true parameter $\boldsymbol{w}^{\ast}$ is a fixed point for the EM algorithm, i.e. $M(\boldsymbol{w}^{\ast})=\boldsymbol{w}^{\ast}$. This is also known as *self-consistency* of the EM algorithm. Hence it is reasonable to expect that in a sufficiently small neighborhood around $\boldsymbol{w}^{\ast}$ there exists a unique maximizer for $Q(\cdot|\boldsymbol{w}^{\ast})$. ###### **Assumption 3**** **(First-order stability condition)**.** Assume that [TABLE] ###### **Remark 2****.** Intuitively, the gradient stability condition enforces the gradient maps $\nabla Q(\cdot|\boldsymbol{w})$ and $\nabla Q(\cdot|\boldsymbol{w}^{\ast})$ to be close whenever $\boldsymbol{w}$ lies in a neighborhood of $\boldsymbol{w}^{\ast}$. This will ensure that the mapped output $M(\boldsymbol{w})$ stays closer to $\boldsymbol{w}^{\ast}$. ###### **Theorem 4**** **(Theorem 1, (Balakrishnan et al., [2017](#bib.bib2)))**.** *If the above assumptions are met for some radius $r>0$ and $0\leq\gamma<\lambda$, then the map $\boldsymbol{w}\mapsto M(\boldsymbol{w})$ is contractive over $\mathcal{B}(\boldsymbol{w}^{\ast},r)$, i.e.* [TABLE] *and consequently, the EM iterates $\{\boldsymbol{w}_{t}\}_{t\geq 0}$ converge geometrically to $\boldsymbol{w}^{\ast}$, i.e.* [TABLE] *whenever the initialization $\boldsymbol{w}_{0}\in\mathcal{B}(\boldsymbol{w}^{\ast},r)$.* ## Appendix C Class of non-linearities In this section, we characterize the class of non-linearities for which our theoretical results for the recovery of regressors hold. Let $Z\sim\mathcal{N}(0,1)$ and $Y|Z\sim\mathcal{N}(g(Z),\sigma^{2})$, where $g:\mathbb{R}\to\mathbb{R}$. For $(\alpha,\beta,\gamma)\in\mathbb{R}^{3}$, define [TABLE] and [TABLE] ###### **Condition 1****.** $\mathbb{E}[\mathcal{S}_{3}^{\prime}(Z)]=\mathbb{E}[\mathcal{S}_{3}^{\prime\prime}(Z)]=0$ and $\mathbb{E}[\mathcal{S}_{3}^{\prime\prime\prime}(Z)]\neq 0$. ###### **Condition 2****.** $\mathbb{E}[\mathcal{S}_{2}^{\prime}(Z)]=0$ and $\mathbb{E}[\mathcal{S}_{2}^{\prime\prime}(Z)]\neq 0$. We are now ready to define the $(\alpha,\beta,\gamma)$-valid class of non-linearities. ###### **Definition 1****.** *We say that the non-linearity $g$ is $(\alpha,\beta,\gamma)$-valid if there exists $(\alpha,\beta,\gamma)\in\mathbb{R}^{3}$ such that both Condition [1](#Thmcondition1) and Condition [2](#Thmcondition2) are satisfied.* We have that [TABLE] Thus $\mathbb{E}[\mathcal{S}_{3}^{\prime}(Z)]=\mathbb{E}[\mathcal{S}_{3}^{\prime\prime}(Z)]=0$ implies that [TABLE] To ensure Condition [1](#Thmcondition1), we need the pair $(\alpha,\beta)$ obtained by solving the above linear equation to satisfy $\mathbb{E}[\mathcal{S}_{3}^{\prime\prime\prime}(Z)]\neq 0$. Similarly, $\mathbb{E}[\mathcal{S}_{2}^{\prime}(Z)]=0$ implies that [TABLE] Thus Condition [2](#Thmcondition2) stipulates that $\mathbb{E}[\mathcal{S}_{2}^{\prime\prime}(Z)]\neq 0$ with this choice of $\gamma$. It turns out that these conditions hold for a wide class of non-linearities and in particular, when $g$ is either the identity function, or the sigmoid function, or the ReLU. For these three choices of popular non-linearities, the values of the tuple $(\alpha,\beta,\gamma)$ are provided below (which are obtained by solving the linear equations mentioned above). ###### **Example 1****.** If $g$ is the identity mapping, then $\mathcal{P}_{3}(y)=y^{3}-3y(1+\sigma^{2})$ and $\mathcal{S}_{2}(y)=y^{2}$. ###### **Example 2****.** If $g$ is the sigmoid function, i.e. $g(z)=\frac{1}{1+e^{-z}}$, then $\alpha$ and $\beta$ can be obtained by solving the following linear equation: [TABLE] The second-order transformation is given by $\mathcal{S}_{2}(y)=y^{2}-y$ (since $\gamma=-1$ when $g$ is sigmoid). ###### **Example 3****.** If $g$ is the ReLU function, i.e. $g(z)=\mathrm{max}\{0,z\}$, then $\alpha=-3\sqrt{\frac{2}{\pi}},\beta=3\left(\frac{4}{\pi}-\sigma^{2}-1\right)$ and $\gamma=-2\sqrt{\frac{2}{\pi}}$. ## Appendix D Proofs of Section [3](#S3) In this section, for the simplicity of the notation we denote the true parameters as $\boldsymbol{w}_{i}$’s and $\boldsymbol{a}_{i}$’s dropping the $\ast$ sign. ### D.1 Proof of Theorem [1](#Thmtheorem1) for $k=2$ ###### Proof. Suppose that $g$ is the linear activation function. For $k=2$, ([1](#S1.E1)) implies that [TABLE] where $f(\cdot)$ is the sigmoid function. Using the fact $\mathbb{E}[Z^{3}]=\mu^{3}+3\mu\sigma^{2}$ for any Gaussian random variable $Z\sim\mathcal{N}(\mu,\sigma^{2})$, we get [TABLE] Moreover, [TABLE] Thus, [TABLE] If we define $\mathcal{P}_{3}(y)\triangleq y^{3}-3y(1+\sigma^{2})$, in view of Lemma [2](#Thmlemma2) we get that [TABLE] Using the chain rule for multi-derivatives, the first term simplifies to [TABLE] Since $f(z)=\frac{1}{1+e^{-z}}$, $f^{\prime}(\cdot),f^{\prime\prime\prime}(\cdot)$ are even functions whereas $f^{\prime\prime}(\cdot)$ is an odd function. Furthermore, both $\boldsymbol{x}\mapsto(\boldsymbol{a}_{1}^{\top}\boldsymbol{x})^{3}-3(\boldsymbol{a}_{1}^{\top}\boldsymbol{x})$ and $\boldsymbol{x}\mapsto\boldsymbol{a}_{1}^{\top}\boldsymbol{x}$ are odd functions whereas $\boldsymbol{x}\mapsto 3(\boldsymbol{a}_{1}^{\top}\boldsymbol{x})^{2}-3$ is an even function. Since $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$, $-\boldsymbol{x}\stackrel{{\scriptstyle(d)}}{{=}}\boldsymbol{x}$. Thus all the expectation terms in ([15](#A4.E15)) equal zero except for the last term since $\mathbb{E}[f(\boldsymbol{w}^{\top}\boldsymbol{x})]=\frac{1}{2}>0$. We have, [TABLE] Similarly, [TABLE] Together, we have that [TABLE] Now consider an arbitrary link function $g$ belonging to the class of non-linearities described in Appendix [C](#A3). Then [TABLE] implies that [TABLE] and [TABLE] If we define $\mathcal{P}_{3}(y)\triangleq y^{3}+\alpha y^{2}+\beta y$, we have that [TABLE] where $(a)$ follows from the choice of $\alpha$ and $\beta$ and the fact that $\boldsymbol{w}\perp\{\boldsymbol{a}_{1},\boldsymbol{a}_{2}\}$, and $c_{g,\sigma}\triangleq\mathbb{E}\left[\left(g(Z)^{3}+\alpha g(Z)^{2}+g(Z)(\beta+3\sigma^{2})\right)^{\prime\prime\prime}\right]$ where $Z\sim\mathcal{N}(0,1)$ . The proof for $\mathcal{T}_{2}$ is similar. ∎ ### D.2 Proof of Theorem [1](#Thmtheorem1) for general $k$ ###### Proof. The proof for general $k$ closely follows that of $k=2$, described in Appendix [D.1](#A4.SS1). For the general $k$, we first prove the theorem when $g$ is the identity function, i.e. [TABLE] Denoting $P_{i|\boldsymbol{x}}$ by $p_{i}(\boldsymbol{x})$, we have that [TABLE] Hence [TABLE] If we let $\mathcal{P}_{3}(y)\triangleq y^{3}-3y(1+\sigma^{2})$, we get [TABLE] Since $\boldsymbol{x}\sim\mathcal{N}(0,I_{d})$ and $\boldsymbol{a}_{i}\perp\mathrm{span}\{\boldsymbol{w}_{1},\ldots,\boldsymbol{w}_{k-1}\}$, we have that $\boldsymbol{a}_{i}^{\top}\boldsymbol{x}\perp(\boldsymbol{w}_{1}^{\top}\boldsymbol{x},\ldots,\boldsymbol{w}_{k-1}^{\top}\boldsymbol{x})$. Moreover, $\mathbb{E}[(\boldsymbol{a}_{i}^{\top}\boldsymbol{x})^{3}-3(\boldsymbol{a}_{i}^{\top}\boldsymbol{x})]=\mathbb{E}[(\boldsymbol{a}_{i}^{\top}\boldsymbol{x})^{2}-1]=\mathbb{E}[\boldsymbol{a}_{i}^{\top}\boldsymbol{x}]=0$ for each $i$. Using the chain-rule for multi-derivatives, the above equation thus simplifies to [TABLE] For a generic $g:\mathbb{R}\to\mathbb{R}$ which is $(\alpha,\beta,\gamma)-$valid, let $\mathcal{P}_{3}(y)=y^{3}+\alpha y^{2}+\beta y$. Then it is easy to see that the same proof goes through except for a change in the coefficients of rank-$1$ terms, i.e. [TABLE] where $\alpha_{i}\triangleq\mathbb{E}\left[\left(g(Z)^{3}+\alpha g(Z)^{2}+g(Z)(\beta+3\sigma^{2})\right)^{\prime\prime\prime}\right]$ where $Z\sim\mathcal{N}(0,1)$ and *′′′* denotes the third-derivative with respect to $Z$. Note that Condition [2](#Thmcondition2) together with the fact that $\mathbb{E}[p_{i}(\boldsymbol{x})]>0$ ensures that $\alpha_{i}\neq 0$ and thus the coefficients of the rank-$1$ terms are non-zero. The proof for $\mathcal{T}_{2}$ is similar. ∎ ### D.3 Proof of Theorem [2](#Thmtheorem2) The following two lemmas are central to the proof of Theorem [2](#Thmtheorem2). Let $\boldsymbol{A}^{\top}=[\boldsymbol{a}_{1}|\ldots|\boldsymbol{a}_{k}]\in\mathbb{R}^{d\times k}$ denote the matrix of regressor parameters whereas $\boldsymbol{W}^{\top}=[\boldsymbol{w}_{1}|\ldots|\boldsymbol{w}_{k-1}]\in\mathbb{R}^{d\times(k-1)}$ denote the matrix of gating parameters. With a slight change of notation, when $\boldsymbol{A}=\boldsymbol{A}^{\ast}$, we denote the EM operator $M(\boldsymbol{W})$ as either $M(\boldsymbol{W},\boldsymbol{A}^{\ast})$ or $M(\boldsymbol{w})$, introduced in Section [3](#S3). For the general case, we simply denote it by $M(\boldsymbol{W},\boldsymbol{A})$. In the following lemmas, we use the norm $\|\boldsymbol{A}\|=\max_{i\in[k]}\|\boldsymbol{A}_{i}^{\top}\|_{2}$ where $\boldsymbol{A}\in\mathbb{R}^{k\times d}$ is a matrix of regressors, similarly for any matrix of classifiers $\boldsymbol{W}\in\mathbb{R}^{(k-1)\times d}$. ###### **Lemma 3**** **(Contraction of the EM operator)**.** *Under the assumptions of Theorem [2](#Thmtheorem2), we have that* [TABLE] *Moreover, $\boldsymbol{W}=\boldsymbol{W}^{\ast}$ is a fixed point for $M(\boldsymbol{W},\boldsymbol{A}^{\ast})$.* ###### **Lemma 4**** **(Robustness of the EM operator)**.** *Let the matrix of regressors $\boldsymbol{A}$ be such that $\max_{i\in[k]}\|\boldsymbol{A}_{i}^{\top}-(\boldsymbol{A}^{\ast}_{i})^{\top}\|_{2}=\sigma^{2}\varepsilon_{1}$. Then for any $\boldsymbol{W}\in\Omega$, we have that* [TABLE] *where $\kappa$ is a constant depending on $g,k$ and $\sigma$. In particular, $\kappa\leq(k-1)\frac{\sqrt{6(2+\sigma^{2})}}{2}$ for $g=$linear, sigmoid and ReLU.* We are now ready to prove Theorem [2](#Thmtheorem2). ###### Proof. We first note that the EM iterates $\{\boldsymbol{W}_{t}\}_{t\geq 1}$ evolve according to [TABLE] Thus [TABLE] where the last inequality follows from Lemma [3](#Thmlemma3) and Lemma [4](#Thmlemma4). Recursively using the above inequality, we obtain that [TABLE] ∎ ### D.4 Proof of Theorem [3](#Thmtheorem3) ###### Proof. We are given that $(\boldsymbol{a}_{1},\boldsymbol{a}_{2})=(\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast})$. Denoting $\boldsymbol{w}^{\ast}$ with $\boldsymbol{w}$, from ([13](#A4.E13)), we have that [TABLE] Thus, [TABLE] Notice that in the above equation we have $(\boldsymbol{a}_{1}-\boldsymbol{a}_{2})^{\top}\boldsymbol{x}$ in the denominator. But this equals zero with zero probability whenever $\boldsymbol{x}$ is generated from a continuous distribution; in our case $\boldsymbol{x}$ is Gaussian. Thus we may write [TABLE] However, it turns out that the above chain of equalities does not hold. Surprisingly, the first equality, which essentially is the law of iterated expectations, is not valid in this case as $\frac{y-\boldsymbol{a}_{2}^{\top}\boldsymbol{x}}{(\boldsymbol{a}_{1}-\boldsymbol{a}_{2})^{\top}\boldsymbol{x}}$ is not integrable. To see this, notice that the model in ([13](#A4.E13)) can also be written as [TABLE] Thus, [TABLE] Since $Z$ is independent of $N$ and $\frac{N}{(\boldsymbol{a}_{1}-\boldsymbol{a}_{2})^{\top}\boldsymbol{x}}$ is a Cauchy random variable, it follows that the random variable $\mathrm{Ratio}$ is not integrable. To deal with the non-integrability of $\mathrm{Ratio}$, we look at its conditional cdf, given by [TABLE] where $\Phi(\cdot)$ is the standard Gaussian cdf. Substituting $z=0.5$ and using the fact that $\Phi(z)+\Phi(-z)=1$, we obtain [TABLE] Since $\Phi\left(\frac{|(\boldsymbol{a}_{1}-\boldsymbol{a}_{2})^{\top}\boldsymbol{x}|}{2\sigma}\right)$ is a symmetric function in $\boldsymbol{x}$ its first moment with $\boldsymbol{x}$ equals zero. Furthermore, if we assume that $\boldsymbol{w}$ is orthogonal to $\boldsymbol{a}_{1}$ and $\boldsymbol{a}_{2}$, we have [TABLE] Thus, if $\left\|\boldsymbol{w}\right\|=1$, we have that [TABLE] In the finite sample regime, $\mathbb{E}\left[\mathds{1}\left\{\mathrm{Ratio}\leq 0.5\right\}\cdot\boldsymbol{x}\right]$ can be estimated from samples using the empirical moments and its normalized version will be an estimate of $\boldsymbol{w}$. ∎ ## Appendix E Proof of Lemma [4](#Thmlemma4) We need the following lemma which establishes the stability of the minimizers for strongly convex functions under Lipschitz perturbations. ###### **Lemma 5****.** *Suppose $\Omega\subseteq\mathbb{R}^{d}$ is a closed convex subset, $f:\Omega\to\mathbb{R}$ is a $\lambda$-strongly convex function for some $\lambda>0$ and $B$ is an $L$-Lipschitz continuous function on $\Omega$. Let $\boldsymbol{w}_{f}=\operatorname*{\arg\!\min}_{\boldsymbol{w}\in\Omega}f(\boldsymbol{w})$ and $\boldsymbol{w}_{f+B}=\operatorname*{\arg\!\min}_{\boldsymbol{w}\in\Omega}f(\boldsymbol{w})+B(\boldsymbol{w})$. Then* [TABLE] ###### Proof. Let $\boldsymbol{w}^{\prime}\in\Omega$ be such that $\left\|\boldsymbol{w}^{\prime}-\boldsymbol{w}_{f}\right\|>\frac{L}{\lambda}$. Let $\boldsymbol{w}_{\alpha}=\alpha\boldsymbol{w}_{f}+(1-\alpha)\boldsymbol{w}^{\prime}$ for $0<\alpha<1$. From the fact that $\boldsymbol{w}_{f}$ is the minimizer of $f$ on $\Omega$ and that $f$ is strongly convex, we have that [TABLE] Furthermore, the strong-convexity of $f$ implies that [TABLE] Since $B$ is $L$-Lipschitz, we have [TABLE] Adding ([18](#A5.E18)) and ([19](#A5.E19)), we get [TABLE] By the assumption that $\left\|\boldsymbol{w}^{\prime}-\boldsymbol{w}_{f}\right\|>\frac{L}{\lambda}$, the term $\frac{L}{\lambda}-\left(1-\frac{\alpha}{2}\right)\left\|\boldsymbol{w}^{\prime}-\boldsymbol{w}_{f}\right\|$ will be negative for sufficiently small $\alpha$. This in turn implies that $f(\boldsymbol{w}_{\alpha})+B(\boldsymbol{w}_{\alpha})\frac{L}{\lambda}$. The conclusion follows. ∎ We are now ready to prove Lemma [4](#Thmlemma4). Fix any $\boldsymbol{W}\in\Omega$ and let $\boldsymbol{A}=\begin{bmatrix}\boldsymbol{a}_{1}^{\top}\\ \ldots\\ \boldsymbol{a}_{k}^{\top}\end{bmatrix}\in\mathbb{R}^{k\times d}$ be such that $\max_{i\in[k]}\left\|\boldsymbol{a}_{i}-\boldsymbol{a}_{i}^{\ast}\right\|_{2}=\sigma^{2}\varepsilon_{1}$ for some $\varepsilon_{1}>0$. Let [TABLE] where, [TABLE] and, [TABLE] Here $p^{(i)}(\boldsymbol{A},\boldsymbol{W})\triangleq\frac{p_{i}(\boldsymbol{x})N_{i}}{\sum_{i\in[k]}p_{i}(\boldsymbol{x})N_{i}}$ denotes the posterior probability of choosing the $i^{\mathrm{th}}$ expert, where [TABLE] Since both $Q(\cdot|\boldsymbol{W},\boldsymbol{A})$ and $Q(\cdot|\boldsymbol{W},\boldsymbol{A}^{\ast})$ are strongly concave functions over $\Omega$ with some strong-concavity parameter $\lambda$, Lemma [5](#Thmlemma5) implies that [TABLE] where $L$ is the Lipschitz-constant for the function $l(\cdot)\triangleq Q(\cdot|\boldsymbol{W},\boldsymbol{A})-Q(\cdot|\boldsymbol{W},\boldsymbol{A}^{\ast})$. We have that [TABLE] Without loss of generality let $i=1$. Since $l(\cdot)$ is linear in $\boldsymbol{W}^{\prime}$, it suffices to show for each $i$ that [TABLE] We show that $L=\kappa\varepsilon_{1}$, or equivalently, [TABLE] Let [TABLE] By hypothesis, we have that $\left\|\Delta_{i}\right\|_{2}\leq\sigma^{2}\varepsilon_{1}$ for all $i\in[k]$. Thus in order to show that [TABLE] it suffices to show that [TABLE] Or equivalently, [TABLE] We can rewrite the difference of the posteriors as [TABLE] Since $N_{i}=\mathcal{N}(y|g(\boldsymbol{a}_{i}^{\top}\boldsymbol{x}),\sigma^{2})=\frac{1}{\sqrt{2\pi\sigma^{2}}}e^{-(y-g(\boldsymbol{a}_{1}^{\top}\boldsymbol{x}))^{2}/2\sigma^{2}}$, we have that [TABLE] Thus, [TABLE] Hence, [TABLE] where we denoted $(\boldsymbol{a}_{i})_{t}$ by $\boldsymbol{a}_{i}$ in the integrals above(with a slight abuse of notation) for the sake of notational simplicity. For any $i\neq 1$, we have that [TABLE] For $g=$linear, sigmoid and ReLU, we have that $|g^{\prime}(\cdot)|\leq 1$. Moreover, $\frac{p_{i}(\boldsymbol{x})p_{1}(\boldsymbol{x})N_{i}N_{1}}{(p_{1}(\boldsymbol{x})N_{1}+p_{i}(\boldsymbol{x})N_{i})^{2}}\leq 1/4$. Thus we have [TABLE] We thus get [TABLE] Now it remains to bound $\sqrt{\mathbb{E}[(y-g(\boldsymbol{a}_{i}^{\top}\boldsymbol{x}))^{2}]}$. Since $\left\|\boldsymbol{a}_{i}\right\|_{2}\leq 1$, one can show that $\mathbb{E}[g(\boldsymbol{a}_{i}^{\top}\boldsymbol{x})^{2}]\leq 1$ for the given choice of non-linearities for $g$. Also, we have that [TABLE] where we used the following facts: (i) $\langle{\boldsymbol{a}_{i}^{\ast}},{\boldsymbol{x}}\rangle$ is independent of the random variable $p_{i}^{\ast}(\boldsymbol{x})$ for each $i\in[k]$, (ii) $\langle{\boldsymbol{a}_{i}^{\ast}},{\boldsymbol{x}}\rangle\stackrel{{\scriptstyle(d)}}{{=}}\langle{\boldsymbol{a}_{1}^{\ast}},{\boldsymbol{x}}\rangle$ and (iii) $\mathbb{E}[g(\langle{\boldsymbol{a}_{1}^{\ast}},{\boldsymbol{x}}\rangle)^{2}]\leq 1$. Since $\mathbb{E}[(y-g(\boldsymbol{a}_{i}^{\top}\boldsymbol{x}))^{2}]\leq 2\mathbb{E}[y^{2}]+\mathbb{E}[g(\boldsymbol{a}_{i}^{\top}\boldsymbol{x})^{2}]$, after substituting these bounds in ([26](#A5.E26)), we get [TABLE] Similarly, [TABLE] Substituting the above two inequalities in ([23](#A5.E23)), we obtain that [TABLE] Defining $\kappa\triangleq(k-1)\frac{\sqrt{6(2+\sigma^{2})}}{2}$ and using the fact that $\left\|\Delta/\sigma^{2}\right\|_{2}\leq\varepsilon_{1}$, we thus obtain [TABLE] ## Appendix F Proof of Lemma [3](#Thmlemma3) ### F.1 Proof for $k=2$ ###### Proof. We first prove the lemma for $k=2$. We show that the assumptions in Appendix [B](#A2) hold *globally* in our setting yielding a geometric convergence. Here we simply denote $M(\boldsymbol{W},\boldsymbol{A}^{\ast})$ as $M(\boldsymbol{w})$ dropping the explicit dependence on $\boldsymbol{A}^{\ast}$. Recall that [TABLE] where [TABLE] For simplicity we drop the subscript in the above expectation with respect to the distribution ${p_{\boldsymbol{w}^{\ast}}(\boldsymbol{x},y)}$. Now we verify each of the assumptions. - • Convexity of $\Omega$ easily follows from its definition. - • We have that [TABLE] Note that the strong-concavity of $Q(\cdot|\boldsymbol{w}^{\ast})$ is equivalent to the strong-convexity of $-Q(\cdot|\boldsymbol{w}^{\ast})$. Denoting the sigmoid function by $f$, we have that for all $\boldsymbol{w}\in\Omega$, [TABLE] where $(a)$ follows from finding the two possible eigenvalues of the positive-definite matrix in the previous step and considering the minimum among them to ensure strong-convexity. Here the value of $\lambda$ is found numerically to be approximately around $0.1442$. - • For any $\boldsymbol{w},\boldsymbol{w}_{t}\in\Omega$, [TABLE] Thus, [TABLE] where we want to prove in $(a)$ that $\gamma_{\sigma}$ is smaller than $0.14$ for all $\boldsymbol{w}\in\Omega$. Intuitively, this means that the posterior probability in ([27](#A6.E27)) is smooth with respect to the parameter $\boldsymbol{w}$. We will now show that this can be achieved in the high-SNR regime when $\sigma$ is sufficiently small. This will ensure that $\kappa_{\sigma}\triangleq\frac{\gamma_{\sigma}}{\lambda}<1$. In particular, the value of $\gamma_{\sigma}$ is dimension-independent and depends only on the choice of the non-linearity $g$. To prove that [TABLE] it suffices to show [TABLE] Or equivalently, [TABLE] Let $\Delta\triangleq\boldsymbol{w}-\boldsymbol{w}^{\ast}$ and $f(u)\triangleq p_{1}(\boldsymbol{x},y,\boldsymbol{w}_{u})$ where $\boldsymbol{w}_{u}=\boldsymbol{w}^{\ast}+u\Delta,u\in[0,1]$. Thus $f(1)=p_{1}(\boldsymbol{x},y,\boldsymbol{w})$ and $f(0)=p_{1}(\boldsymbol{x},y,\boldsymbol{w}^{\ast})$. So we get [TABLE] where the gradient is evaluated with respect to $\boldsymbol{w}_{u}$. Differentiating ([27](#A6.E27)) with respect to $\boldsymbol{w}$, we get that [TABLE] Thus, [TABLE] where the last inequality follows from Lemma $5$ of (Balakrishnan et al., [2017](#bib.bib2)). Our goal is to now prove that $\gamma_{\sigma}\rightarrow 0$ as $\sigma\rightarrow 0$. First observe that [TABLE] where the key observation is that irrespective of the sign of $(y-g(\boldsymbol{a}_{1}^{\top}\boldsymbol{x}))^{2}-(y-g(\boldsymbol{a}_{2}^{\top}\boldsymbol{x}))^{2}$, the ratio still goes to zero and hence by dominated convergence theorem $\mathbb{E}[R(\boldsymbol{x},y,\boldsymbol{w}_{u},\sigma)^{2}]\rightarrow 0$ for each $u\in[0,1]$. Now we show that this convergence is uniform in $u$ and thus $\gamma_{\sigma}\rightarrow 0$. For simplicity, define [TABLE] Thus, [TABLE] Similarly, [TABLE] Thus, we get [TABLE] Hence [TABLE] where $(a)$ follows from the fact $\left\|\boldsymbol{w}_{u}\right\|\leq 1$ and $\mathbb{E}[e^{4\boldsymbol{w}_{u}^{\top}\boldsymbol{x}}]=e^{8\left\|\boldsymbol{w}_{u}\right\|^{2}}\leq e^{8}$, for each $u\in[0,1]$. Now we analyze the convergence rate of the last term $\mathbb{E}[e^{-2n^{2}|\Delta_{1}-\Delta_{2}|}]$ for the case of linear regression, i.e. $g(z)=z$. Notice that for the two-mixtures, we have [TABLE] Thus, [TABLE] Since $Z$ can equal either [math] or $1$, we have [TABLE] ∎ ### F.2 Proof for general $k$ ###### Proof. The proof strategy for general $k$ is similar. First let $\varepsilon_{1}=0$. Our task is to show that the assumptions of Appendix [B](#A2) hold globally in our setting. The domain $\Omega$ is clearly convex since [TABLE] Now we verify Assumption [2](#Thmassumption2). The function $Q(.|\boldsymbol{w}_{t})$ is given by [TABLE] where $p_{\boldsymbol{w}_{t}}^{(i)}\triangleq\mathbb{P}\left[z=i|\boldsymbol{x},y,\boldsymbol{w}_{t}\right]$ corresponds to the posterior probability for the $i^{\text{th}}$ expert, given by [TABLE] Throughout we follow the convention that $\boldsymbol{w}_{k}=0$. Thus the gradient of $Q$ with respect to the $i^{\text{th}}$ gating parameter $\boldsymbol{w}_{i}$ is given by [TABLE] Thus the $(i,j)^{\text{th}}$ block of the negative Hessian $-\nabla^{(2)}_{\boldsymbol{w}}Q(\boldsymbol{w}|\boldsymbol{w}^{\ast})\in\mathbb{R}^{d(k-1)\times d(k-1)}$ is given by [TABLE] where $p_{i}(\boldsymbol{x})=\frac{e^{\boldsymbol{w}_{i}^{\top}\boldsymbol{x}}}{1+\sum_{j\in[k-1]e^{\boldsymbol{w}_{j}^{\top}\boldsymbol{x}}}}$. It is clear from ([52](#A6.E52)) that $-\nabla^{(2)}_{\boldsymbol{w}}Q(\boldsymbol{w}|\boldsymbol{w}^{\ast})$ is positive semi-definite. Since we are interested in the strong convexity of $-Q(\cdot|\boldsymbol{w}^{\ast})$ which is equivalent to positive definiteness of the negative Hessian, it suffices to show that [TABLE] Since the Hessian is continuous with respect to $\boldsymbol{w}$ and consequently the minimum eigenvalue of it, there exists a $\boldsymbol{w}^{\prime}\in\Omega$ such that [TABLE] where $\boldsymbol{a}=(\boldsymbol{a}_{1}^{\top},\ldots,\boldsymbol{a}_{k-1}^{\top})^{\top}\in\mathbb{R}^{d(k-1)}$. In view of ([52](#A6.E52)), the above equation can be further simplified to [TABLE] where $\boldsymbol{a}_{\boldsymbol{x}}=(\boldsymbol{a}_{1}^{\top}\boldsymbol{x},\ldots,\boldsymbol{a}_{k-1}^{\top}\boldsymbol{x})^{\top}\in\mathbb{R}^{k-1}$ and $M_{\boldsymbol{x}}$ is given by [TABLE] Let the infimum in ([53](#A6.E53)) is attained by $\boldsymbol{a}^{\ast}$, i.e. $\lambda=\mathbb{E}[(\boldsymbol{a}_{\boldsymbol{x}}^{\ast})^{\top}M_{\boldsymbol{x}}\boldsymbol{a}_{\boldsymbol{x}}^{\ast}]$. For each $\boldsymbol{x}$, $M_{\boldsymbol{x}}$ is strictly diagonally dominant since $|M_{\boldsymbol{x}}(i,i)|=p_{i}(\boldsymbol{x})(1-p_{i}(\boldsymbol{x}))=p_{i}(\boldsymbol{x})\left(\sum_{j\neq i,j\in[k]}p_{j}(\boldsymbol{x})\right)>p_{i}(\boldsymbol{x})\left(\sum_{j\neq i,j\in[k-1]}p_{j}(\boldsymbol{x})\right)=\sum_{j\neq i}M(i,j)$. Thus $M_{\boldsymbol{x}}$ is positive-definite and $(\boldsymbol{a}_{\boldsymbol{x}}^{\ast})^{\top}M_{\boldsymbol{x}}\boldsymbol{a}_{\boldsymbol{x}}^{\ast}>0$ whenever $\boldsymbol{a}^{\ast}_{\boldsymbol{x}}\neq 0$. Since $x$ follows a continuous distribution it follows that $\boldsymbol{a}^{\ast}_{\boldsymbol{x}}\neq 0$ with probability $1$ and thus $\lambda=\mathbb{E}[(\boldsymbol{a}_{\boldsymbol{x}}^{\ast})^{\top}M_{\boldsymbol{x}}\boldsymbol{a}_{\boldsymbol{x}}^{\ast}]>0$. Now it remains to show that Assumption [3](#Thmassumption3) too holds, i.e. [TABLE] Note that $\boldsymbol{w}=(\boldsymbol{w}_{1}^{\top},\ldots,\boldsymbol{w}_{k-1}^{\top})^{\top}\in\mathbb{R}^{d(k-1)}$. We will show that [TABLE] where $(\nabla Q(M(\boldsymbol{w})|\boldsymbol{w}))_{i}\in\mathbb{R}^{d}$ refers to the $i^{\text{th}}$ block of the gradient and $\gamma_{\sigma}\rightarrow 0$. Observe that [TABLE] Let $\Delta=\boldsymbol{w}-\boldsymbol{w}^{\ast}$ and correspondingly $\Delta=(\Delta_{1}^{\top},\ldots,\Delta_{k-1}^{\top})^{\top}$ where $\Delta_{i}=\boldsymbol{w}_{i}-\boldsymbol{w}_{i}^{\ast}$. Thus it suffices to show that [TABLE] Or equivalently, [TABLE] We consider the case $i=1$. The proof for the other cases is similar. Recall that [TABLE] For simplicity we define $N_{i}=\mathcal{N}(y|g(\boldsymbol{a}_{1}^{\top}\boldsymbol{x}),\sigma^{2})$. It is straightforward to verify that [TABLE] Thus [TABLE] Similarly, [TABLE] Let $\boldsymbol{w}_{u}\triangleq\boldsymbol{w}^{\ast}+u\Delta,u\in[0,1]$ and $f(u)\triangleq p_{\boldsymbol{w}_{u}}^{(1)}$. Thus [TABLE] So we get [TABLE] Now our goal is to show that $\mathbb{E}[R_{i}(\boldsymbol{x},y,\boldsymbol{w}_{u},\sigma)^{2}]\rightarrow 0$ as $\sigma\rightarrow 0$. For $i=1$, we have [TABLE] Similarly, [TABLE] For $\boldsymbol{w}=\boldsymbol{w}_{u}$ and $i\neq 1$, we have that [TABLE] Thus, by Dominated Convergence Theorem, $\mathbb{E}[R_{i}(\boldsymbol{x},y,\boldsymbol{w}_{u},\sigma)^{2}]\rightarrow 0$ for each $u\in[0,1]$. To show that $\int_{0}^{1}\mathbb{E}[R_{i}(\boldsymbol{x},y,\boldsymbol{w}_{u},\sigma)^{2}]du\rightarrow 0$, we can now follow the same analysis as in the proof of Theorem [2](#Thmtheorem2) from ([29](#A6.E29)) on-wards (replacing $\boldsymbol{w}$ there with $\boldsymbol{w}_{1}-\boldsymbol{w}_{i}$) which ensures that $\gamma^{(1)}_{\sigma}$ in our case converges to zero. Similarly for other $i\in[k-1]$, we get that $\gamma^{(i)}\rightarrow 0$. Taking $\gamma_{\sigma}=\gamma^{(1)}_{\sigma}+\ldots+\gamma^{(k-1)}_{\sigma}$ and $\kappa_{\sigma}=\frac{\gamma_{\sigma}}{\lambda}$ completes the proof. ∎ ## Appendix G Gradient EM algorithm In this section, we provide the convergence guarantees for the gradient EM algorithm. For simplicity, we prove the results for $k=2$ and $(\boldsymbol{a}_{1},\boldsymbol{a}_{2})=(\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast})$. Thus we want to learn the gating parameter $\boldsymbol{w}^{\ast}$ in this setting. The results for the general case follow essentially the same proof as that of Theorem [2](#Thmtheorem2). In particular, our Theorem [5](#Thmtheorem5) can be viewed as a generalization of Lemma [3](#Thmlemma3). Together with Lemma [4](#Thmlemma4), extension to general $k$ is straightforward. Note that in the M-step of the EM algorithm, instead of maximizing $Q(\cdot|\boldsymbol{w}_{t})$, we can chose an iterate so that it increases the $Q$ value instead of fully maximizing it, i.e. $Q(\boldsymbol{w}_{t+1}|\boldsymbol{w}_{t})\geq Q(\boldsymbol{w}_{t}|\boldsymbol{w}_{t})$. Such a procedure is termed as *generalized EM*. *Gradient EM* is an example of generalized EM in which we take an ascent step in the direction of the gradient of $Q(\cdot|\boldsymbol{w}_{t})$ to produce the next iterate, i.e. [TABLE] where $\alpha>0$ is a suitably chosen step size and the gradient is with respect to the first argument. To account for the constrained optimization, we can include a projection step. Mathematically, [TABLE] where $\Pi_{\Omega}$ refers to the projection operator. Our next result establishes that the iterates of the gradient EM algorithm too converge geometrically for an appropriately chosen step size $\alpha$. ###### **Theorem 5****.** *Suppose that the domain $\Omega=\{\boldsymbol{w}\in\mathbb{R}^{d}:\|\boldsymbol{w}\|_{2}\leq 1\}$ and $(\boldsymbol{a}_{1},\boldsymbol{a}_{2})=(\boldsymbol{a}_{1}^{\ast},\boldsymbol{a}_{2}^{\ast})$. Then there exist constants $\alpha_{0}>0$ and $\sigma_{0}>0$ such that for any step size $0<\alpha\leq\alpha_{0}$ and noise variance $\sigma<\sigma_{0}$, the gradient EM updates on the gating parameter $\{\boldsymbol{w}\}_{t\geq 0}$ converge geometrically to the true parameter $\boldsymbol{w}^{\ast}$, i.e.* [TABLE] *where $\rho_{\sigma}$ is a dimension-independent constant depending on $g$ and $\sigma$.* ###### **Remark 3****.** The condition $\sigma<\sigma_{0}$ ensures that the Lipschitz constant $\rho_{\sigma}$ for the map $G$ is strictly less than $1$. The constant $\alpha_{0}$ depends only on two universal constants which are nothing but the strong-concavity and the smoothness parameters for the function $Q(\cdot|\boldsymbol{w}^{\ast})$. ###### Proof. In addition to the assumptions of Appendix [B](#A2), if we can ensure that the map $-Q(\cdot|\boldsymbol{w}^{\ast})$ is $\mu$-smooth, then the proof follows from Theorem $3$ of (Balakrishnan et al., [2017](#bib.bib2)) if we choose $\alpha_{0}=\frac{2}{\mu+\lambda}$ where $\lambda$ is the strong-convexity parameter of $-Q(\cdot|\boldsymbol{w}^{\ast})$. The strong-convexity is already established in Appendix [D.3](#A4.SS3). To find the smoothness parameter, note that [TABLE] The contraction parameter is then given by [TABLE] Since $\gamma_{\sigma}\xrightarrow{\sigma\rightarrow 0}0$, $\rho_{\sigma}<1$ whenever $\sigma<\sigma_{0}$ for a constant $\sigma_{0}$. ∎ ## Appendix H Additional experiments ### H.1 Synthetic data [FIGURE:A8.F4] In Figure [4](#A8.F4), we varied the number of samples our data set and fixed the other set of parameters to $k=3,d=5,\sigma=0.5$. [FIGURE:A8.F5] In Figure [5](#A8.F5) we repeated our experiments for the choice of $n=10000,d=5,k=3$ for two different popular choices of non-linearities: sigmoid and ReLU. The same conclusion as in the linear setting holds in this case too with our algorithm outperforming the EM consistently. ### H.2 Real data [FIGURE:A8.F6] For real data experiments, we choose the $3$ standard regression data sets from the UCI Machine Learning Repository: Concrete Compressive Strength Data Set, Stock portfolio performance Data Set, and Airfoil Self-Noise Data Set (Yeh, [1998](#bib.bib34); Liu & Yeh, [2017](#bib.bib19); Brooks et al., [1989](#bib.bib3)). In all the three tasks, the goal is to predict the outcome or the response $y$ for each input $\boldsymbol{x}$, which typically contains some task specific attributes. For example, in the concrete compressive strength, the task is to predict the compressive strength of the concrete given its various attributes such as the component of cement, water, age, etc. For this data, the input $\boldsymbol{x}\in\mathbb{R}^{8}$ corresponds to $8$ different attributes of the concrete and the output $y\in\mathbb{R}$ corresponds to its concrete strength. Similarly, for the stock portfolio data set the input $\boldsymbol{x}\in\mathbb{R}^{6}$ contains the weights of several stock-picking concepts such as weight of the Large S/P concept, weight of the Small systematic Risk concept, etc,. and the output $y$ is the corresponding excess return. The airfoil data set is obtained from a series of aerodynamic and acoustic tests of two and three-dimensional airfoil blade sections and the goal is predict the scaled sound pressure level (in dB) given the frequency, angle of attack, etc,. For all the tasks, we pre-processed the data by whitening the input and scaling the output to lie in $(-1,1)$. We randomly allotted $75\%$ of the data samples for training and the rest for testing. Our evaluation metric is the prediction error on the test set $(\boldsymbol{x}_{i},y_{i})_{i=1}^{n}$ defined as [TABLE] where $\hat{y}_{i}$ corresponds to the predicted output response using the learned parameters. In other words, [TABLE] We ran the joint-EM algorithm (with $10$ different trails) on these tasks with various choices for $k\in\{2,\ldots,10\},\sigma\in\{0.1,0.4,0.8,1\},g\in\{\text{linear},\text{sigmoid},\text{ReLU}\}$ and found the best hyper-parameters to be $(k=3,\sigma=0.1$ and $g=\text{linear})$, $(k=3,\sigma=0.4,g=\text{sigmoid})$ and $(k=3,\sigma=0.1,g=\text{linear})$ for the three datasets respectively. For this choice of best hyper-parameters found for joint-EM, we ran our algorithm. Figure [6](#A8.F6) highlights the predictive performance of our algorithm as compared to that of the EM. We also plotted the variance of the test data for reference and to gauge the performance of our algorithm. In all the settings our algorithm is able to obtain a better set of parameters resulting in smaller prediction error.