# GMM derivation for diagonal covariance matrices

I was trying to understand the derivation of M step in the EM algorithm for GMM. All the resources available consider only "full covariance" matrices. I wanted to implement GMM for "diagonal covarinace" matrices, but I'm unable to proceed with the proof because it's kind of difficult to represent diagonal matrix in a form which can be differentiated. Can I get some hints on how to proceed with the derivation for this case?

• This is a Gaussian Naive Bayes classifier. See for example slides 22-24 (pages 37-48) here. Commented Mar 17 at 20:58
• I agree with @Spätzle's TL;DR conclusion for slightly different reasons. EM is by its nature empirical (data driven) and often exploratory. It seems a very harsh form of regularization to assert the variables are independent prior to EM. While the question appears to be a theoretical analytic exercise, if implemented, you may want to consider the lack of fit versus larger models - maximum entropy / Gaussian graphical models with more off-diagonal terms in the precision matrix to improve goodness-of-fit; or full Bayesian models where data is permitted to eventually overwhelm a diagonal prior. Commented Mar 18 at 12:56
• But wouldn't a naïve Bayes classifier require defining $\begin{pmatrix}k\\2\end{pmatrix}$ boundaries? Commented Mar 18 at 17:02

TL;DR while the diagonal form of the covariance matrix indeed has some effect on the final estimates for the variance terms (which are not easily available in the general case), the formulae are still very close to the general case and thus IMHO it is better to remain with the well-known formulae.

Generally in EM we want to find $$\theta$$ which maximizes the density $$p(x|\theta)$$. We also assume that given the unobserved variable $$z$$, the log density $$\log(p(x,z|\theta))$$ is maximized easily. The EM algorithm has two steps:

• E step: $$Q(\theta|\theta_{t-1})=E_{z|x,\theta_{t-1}}[log(p(x,z|\theta))]$$

• M step: $$\theta_t=\arg\max_\theta((\theta|\theta_{t-1}))$$

Given some $$d$$-dimensional data $$D$$ of size $$N$$ and parameter vector $$\theta$$, the log-likelihood of a GMM with $$K$$ classes (that is, $$y\in\{1,...,K\}$$) is

$$\log(p(D|\theta))=\sum_{i=1}^N\log\left( \sum_{k=1}^K \pi_k\mathcal{N}(x_i|\mu_k,\Sigma_k) \right)$$

Where $$\mathcal{N}(x_i|\mu_k,\Sigma_k)=\frac{1}{\sqrt{(2\pi)^d|\Sigma|}}\exp\left(-0.5(x-\mu)^T\Sigma^{-1}(x-\mu)\right)$$ is the multivariate normal distribution.

Now let's consider the effect of diagonality. In practice, for each class $$k$$ we have the covariance matrix $$\Sigma_k=\begin{pmatrix} \sigma^2_{k1}&&0\\&\ddots&\\0&&\sigma^2_{kd} \end{pmatrix}\Rightarrow \Sigma_k^{-1}=\begin{pmatrix} \sigma^{-2}_{k1}&&0\\&\ddots&\\0&&\sigma^{-2}_{kd} \end{pmatrix}$$

This means that for the $$i^{th}$$ sample, the density under class $$k$$ is

$$p(x_i|\mu_k,\Sigma_k)=\left((2\pi)^d\prod_{j=1}^d\sigma^2_{kj}\right)^{-0.5}\exp\left(-0.5\sum_{j=1}^d \sigma^{-2}_{kj}(x_{ij}-\mu_{kj})^2\right)$$

# E Step

Upon constructing the E step, we use $$y$$ as our hidden variable. Should we know $$y_i$$, we can write $$\log(p(D|\theta))$$ as

$$\log(p(D|\theta))=\log(p(x,y|\theta))=\sum_{i=1}^N\log\left( \sum_{k=1}^K \left(\pi_k\mathcal{N}(x_i|\mu_k,\Sigma_k)\right)^{I\{y_i=k\}} \right)$$

That is,

\begin{aligned}Q(\theta|\theta_{t})=E_{y|x,\theta_{t-1}}[\log(p(x,z|\theta))] \\=E_{y|x,\theta_{t}}\left[ \sum_{i=1}^N\sum_{k=1}^K I\{y_i=k\}(\log\pi_k+\log\mathcal{N}(x_i|\mu_k,\Sigma_k))\right] \\=\sum_{i=1}^N\sum_{k=1}^KE_{y|x,\theta_{t}}\left[ I\{y_i=k\}(\log\pi_k+\log\mathcal{N}(x_i|\mu_k,\Sigma_k))\right] \\=\sum_{i=1}^N\sum_{k=1}^KE_{y|x,\theta_{t}}\left[ I\{y_i=k\}\right](\log\pi_k+\log\mathcal{N}(x_i|\mu_k,\Sigma_k)) \end{aligned}

Skipping the full derivation, we end up with the responsibility variable $$r_{ik}^{(t)}$$ which represents the relative probability of sample $$i$$ belonging to category $$k$$ at the $$t^{th}$$ iteration:

$$r^{(t)}_{ik}=E_{y|x,\theta_{t}}\left[ I\{y_i=k\}\right]=\frac{\pi_k^{(t)}\mathcal{N}(x_i|\mu_k^{(t)},\Sigma_k^{(t)})}{\sum_{k'=1}^K\pi_{k'}^{(t)}\mathcal{N}(x_i|\mu_{k'}^{(t)},\Sigma_{k'}^{(t)})}$$

Now I know we talk here about diagonal covariance matrices but let's give it a few more steps and then use this information. The full E step is then:

$$Q(\theta|\theta_{t-1})=\sum_{i=1}^N\sum_{k=1}^Kr_{ik}(\log\pi_k+\log\mathcal{N}(x_i|\mu_k,\Sigma_k))$$

# M Step

For the M step we need to use the constraint $$\sum_k\pi_k=1$$ and so we define the Lagrangian $$\mathcal{L}=Q(\theta|\theta_{t-1})-\lambda\left(\sum_k\pi_k-1\right)$$ and derivate it with respect to $$\pi_k$$:

$$\frac{\partial}{\partial\pi_k}\mathcal{L}=\left(\sum_{i=1}^N r_{ik}\left(\frac{1}{\pi_k}+0\right)\right)-\lambda\cdot 1=0\Rightarrow \pi_k=\frac{1}{\lambda}\sum_{i=1}^N r_{ik}$$

With the condition on $$\sum_k\pi_k$$ we get

$$\sum_k\pi_k=1=\sum_k\left(\frac{1}{\lambda}\sum_{i=1}^N r_{ik}\right)=\frac{N}{\lambda}\Rightarrow \lambda=N$$

And finally $$\hat\pi_k^{ML}=\frac{1}{N}\sum_{i=1}^N r_{ik}$$

The derivation of $$\mu_k$$ is similar, this time without using the Lagrangian:

$$\frac{\partial}{\partial\mu_k}Q(\theta|\theta_{t-1})=\sum_{i=1}^N r_{ik}\frac{\partial}{\partial\mu_k}\left( -0.5(x_i-\mu_k)^T\Sigma_k^{-1}(x_i-\mu_k) \right)=0\Rightarrow \sum_{i=1}^Nr_{ik}\mu_k=\sum_{i=1}^Nr_{ik}x_i$$ and finally $$\hat\mu_k^{ML}=\frac{\sum_{i=1}^Nr_{ik}x_i}{\sum_{i=1}^Nr_{ik}}$$

In the general case, the covariance estimate is

$$\hat{\Sigma}^{ML}_{k}=\frac{\sum_{i=1}^N r_{ik}(x_i-\hat\mu^{ML}_k)(x_i-\hat\mu^{ML}_k)^T}{\sum_{i=1}^N r_{ik}}$$

But we are looking to find what it looks like when having a diagonal covariance. Now, before we move to the covariance estimate itself, let's see what is the effect of diagonality up to here: it indeed changes $$r^{(t)}_{ik}$$ a little bit, as we replace the traditional density with the diagonal one, but... that's it so far. As we use responsibilities 'as is' (meaning we relate to $$r_{ik}$$ as a constant and don't derivate it), we keep on using the simple, well-known GMM estimates.

For the covariance estimates, we're interested in the actual components of $$\Sigma_k$$, so we'll denote $$z_{kl}=\sigma^2_{kl}$$ (where $$l\in\{1,...,d\}$$) and derivate according to $$z_{kl}$$:

\begin{aligned}\frac{\partial}{\partial z_{kl}}Q(\theta|\theta_{t-1})=\\=\sum_{i=1}^N r_{ik}\frac{\partial}{\partial z_{kl}}\left( \log\pi_k+\log\left(\left((2\pi)^d\prod_{j=1}^d\sigma^2_{kj}\right)^{-0.5}\exp\left(-0.5\sum_{j=1}^d \sigma^{-2}_{kj}(x_{ij}-\mu_{kj})^2\right)\right) \right) \\=\sum_{i=1}^N r_{ik}\frac{\partial}{\partial z_{kl}}\left( \log\pi_k -0.5d\log(2\pi) -0.5\sum_{j=1}^d\log z_{kj} -0.5\sum_{j=1}^d z^{-1}_{kj}(x_{ij}-\mu_{kj})^2\right) \\=\sum_{i=1}^N r_{ik}\left( 0 +0 -0.5\frac{1}{z_{kl}} +0.5(x_{il}-\mu_{kl})^2 z^{-2}_{kl}\right) =0.5\sum_{i=1}^N r_{ik}\left( \frac{(x_{il}-\mu_{kl})^2}{z_{kl}^2} -\frac{1}{z_{kl}} \right) =0 \\\Rightarrow \frac{\sum_{i=1}^N r_{ik}}{z_{kl}}=\frac{\sum_{i=1}^N r_{ik}(x_{il}-\mu_{kl})^2}{z_{kl}^2} \end{aligned}

As per the definition of covariance $$\sigma^2_{kl}>0$$, we get that the estimate for the estimate of the $$l^{th}$$ covariance term in the $$k^{th}$$ class is

$$\hat{z}^{ML}_{kl}=\frac{\sum_{i=1}^N r_{ik}(x_{il}-\mu_{kl})^2}{\sum_{i=1}^N r_{ik}}$$

and the covariance matrix estimate for the $$k^{th}$$ class is

$$\hat{\Sigma}^{ML}_{k}=\begin{pmatrix} \hat{z}^{ML}_{k1}&&0\\&\ddots&\\0&&\hat{z}^{ML}_{kd} \end{pmatrix}$$

Note that the eventual form of $$\hat{z}^{ML}_{kl}$$ is not very different from the general case, however it does give us a closed form for the variance terms which we don't have that easily in the general case.

[Motivation: to fit the parameters $$\theta$$ of any model $$p_\theta(x)$$ to a data distribution $$p(x)$$, we typically maximise (w.r.t. $$\theta$$) the log likelihood $$\int_x p(x)\log p_\theta(x)$$. This is equivalent to minimising the KL divergence $$\text{KL}[p(x)\|p_\theta(x)]$$, which is minimised (to 0) iff $$p_\theta(x)\!=\!p(x)$$. For latent variable models, the likelihood can be intractable so a lower bound on the log-likelihood (the "ELBO") can be maximised instead.]

The EM algorithm maximises the Evidence Lower bound (ELBO) for a discrete latent variable model $$p_{\theta,\pi}(x) = \sum_z p_\theta(x|z)p_\pi(z)$$. What you suggest is simply to model $$p_\theta(x|z)$$ by $$\mathcal{N}(\mu_z,diag(\mathbf{\sigma}^2_z))$$ rather than use a full covariance matrix $$\mathcal{N}(\mu_z,{\Sigma}_z)$$, which barely changes things.

The ELBO ($$L$$) can be derived by: $$\int_x p(x)\log p(x) \geq \int_xp(x)\log p_\theta(x) = \int_xp(x)\sum_zq(z|x)\log p_\theta(x)\qquad\qquad\quad\\ \qquad= \int_xp(x)\sum_zq(z|x)\log \tfrac{p_\theta(x|z)p_\pi(z)}{p_{\theta,\pi}(z|x)}\tfrac{q(z|x)}{q(z|x)} \geq \int_xp(x)\sum_zq(z|x)\log \tfrac{p_\theta(x|z)p_\pi(z)}{q(z|x)} \doteq L$$

The ELBO is maximised w.r.t. parameters $$\theta, \pi$$ and the "approximate posterior" $$q$$ (or equiv. its parameters). The ELBO is optimal w.r.t. $$q$$ iff $$q(z|x) = p_{\theta,\pi}(z|x)$$, the true posterior under the model which can be intractable, but if $$z$$ is discrete (as for a GMM where $$z$$ is a cluster label), $$p_{\theta,\pi}(z|x)$$ can be computed exactly by Bayes' rule. This is essentially the E-step: $$q(z|x) \ =\ p_{\theta^\circ,\pi^\circ}(z|x) \ \doteq\ \tfrac{p_{\theta^\circ}(x|z)p_{\pi^\circ}(z)}{\sum_{z'}p_{\theta^\circ}(x|z')p_{\pi^\circ}(z')}$$ using parameters $$\theta^\circ, \pi^\circ$$ from the previous iteration.

Your question is about the M-step, which assumes $$q$$ is fixed (as above), and optimises the parameters $$\theta, \pi$$. (Since here $$p_\theta(x|z) = \mathcal{N}(x;\mu_z,diag(\mathbf{\sigma}^2_z))$$, $$\theta$$ refers to the $$\mu_z$$s and $$\mathbf{\sigma}_z$$s.)

Gradients are straightforward to compute (and set to zero) as in the standard case, e.g.: $$\nabla_{\pi_z}L = \int_x p(x) q(z|x) /p_\pi(z) - \lambda \ \implies\ \lambda p_\pi(z) = \int_x p(x) q(z|x) \\ \ \implies\ \lambda = \sum_{z'} \int_x p(x) q(z'|x) \ \implies\ p_\pi(z) = \tfrac{\int_x p(x) q(z|x)}{\sum_{z'} \int_x p(x) q(z'|x)}$$ [we assume $$p_\pi(z)\doteq\pi_z$$ and $$\lambda$$ is a Lagrange multiplier for the constraint $$\sum_zp_\pi(z) = 1$$]

$$\nabla_{\mu_z}L = \nabla_{\mu_z}\!\int_xp(x) q(z|x)\log p_\theta(x|z)\\ = \nabla_{\mu_z}\!\int_xp(x) q(z|x) \big\{\! -\tfrac{D}{2}\log(2\pi) - \tfrac{1}{2}\sum_i\log\sigma_i^2 - \tfrac{1}{2}(x-\mu_{z})^\top diag(\mathbf{\sigma}^2_z)^{-1}(x-\mu_{z}) \big\}\\ = \int_xp(x) q(z|x)(\mu_{z}-x)diag(\mathbf{\sigma}^2_z)^{-1} \ \implies\ \mu_{z}= \tfrac{\int_xp(x) q(z|x)x}{\int_xp(x) q(z|x)}$$

$$\nabla_{\sigma_{z,i}}L = \int_xp(x) q(z|x) \big\{\! - \tfrac{1}{2\sigma_i^2} + \tfrac{1}{2}\tfrac{(x_i-\mu_{z,i})^2}{(\sigma_i^2)^2} \big\}\ \implies \sigma_i^2 = \tfrac{\int_xp(x) q(z|x) (x_i-\mu_{z,i})^2}{\int_xp(x) q(z|x)}$$

Throughout, expectations w.r.t. $$p(x)$$ can be approximated by sampling and so replaced with sums over data points $$\{x_n\}_n$$, giving the standard formulations. The $$q(z|x)$$ are often considered "weights", but are effectively posterior class distributions for each data point.

The only place this differs to standard EM for GMM is in the calculation of cluster variances (as expected). This will affect all final outputs (including class means and the class distribution) as the algorithm aims to fit a mixture of "axis-aligned" Gaussians to the data.