# Music Machine Learning - Expectation-Maximization

### Author: Philippe Esling (esling@ircam.fr)

In this course we will cover
1. Unsupervised learning and [clustering](#clustering).
2. Motivating the need for [latent variable](#latent) models.
2. Approaching the problem naïvely with the [k-Means](#kmeans) algorithm.

## Probabilistic clustering

- Usually clustering is done in _hard_ version
    * Each point belongs to one cluster
    * For instance the basic k-Means algorithm
$$
cluster_{id} = f(x)
$$
- We might want to do probabilistic (soft) clustering
    * Each point belongs more or less to all clusters
    * Defines a probability of belonging
$$
p(cluster_{id} \vert x)
$$

Using probabilistic models has several interests
- Hyperparameter tuning 
- Generative modeling

### Maximum likelihood

If we want to optimize our probabilistic model, we have to solve a typical maximum likelihood problem (considering that points $x_{i}$ are independent)
$$
\underset{\theta}{\text{max. }} p(x \vert \theta) = \prod_{i=1}^{N} p(x_{i} \vert \theta)
$$
So we could try to replace our full model inside this problem
$$
\underset{\theta}{\text{max. }} \prod_{i=1}^{N} \sum_{c} \pi_{c} \mathcal{N}\left(x_{i}\vert \mu_{c}, \Sigma_{c}\right)
$$

However, we have to include the following constraints
$$
\sum_{c} \pi_{c} = 1; \pi_{c} \geq 0, \forall c \\
\Sigma_{c} \succcurlyeq 0
$$

We could rely on optimization
- Hard to follow constraints
- Slow and unstable to train
Reasoning is way easier if we introduce a latent variable

### Introducing a latent variable

We will introduce a latent variable $t$, which _explains_ a given datapoint. In our case, this variable tells us _from which group_ the point was generated.

In our case, we can simply say that our latent variable corresponds to the mixing probabilities
$$
\begin{align}
p(z=c\vert\theta) &= \pi_{c} \\
p(x\vert t=c,\theta) &= \mathcal{N}\left(x\vert \mu_{c}, \Sigma_{c}\right)
\end{align}
$$

Therefore, our original model can be retrieved as a marginalization of the cluster latent variable such that
$$
p(x \vert \theta) = \sum_{c} p(x \vert t=c,\theta) p(t=c\vert\theta)
$$

#### Estimating parameters

Now in order to optimize the parameters, we should need to optimize both the $\theta$ and $t$ variables. If we know the cluster belonging (source) variable $t$, then the optimization is easy, since we know that
$$
p(x\vert z=c, \theta) = \mathcal{N}\left(x\vert \mu_{c}, \Sigma_{c}\right)
$$
And computing the parameters of this Gaussian is pretty easy, as we simply need to update the values depending on the corresponding points (weighted mean)
$$
\mu_{c} = \frac{\sum_{i} p(z=c\vert x_{i}, \theta) x_{i}}{\sum_{i} p(z=c\vert x_{i}, \theta)}
$$

#### Estimating sources

Now in order to optimize the source variable $t$, then the optimization is easy if we know the parameters as we have
$$
p(t = c \vert x, \theta) = \frac{p(x \vert t = c, \theta) p(t = c \vert \theta)}{Z}
$$
Note that here the normalization parameter

## Expectation-Maximization

As we just have seen, a natural alternating solution emerges in the form of the Expectation-Maximization algorithm
1. Start with a random set of Gaussians parameters $\theta$
2. Until convergence repeat
    - For each point compute $p(t=c \vert x_{i}, \theta)$
    - Update Gaussian parameters $\theta$

### Jensen inequality

We say that $f(x)$ is concave if, for any $a,b,\alpha$ and $0\leq\alpha\leq 1$ we have
$$
f(\alpha a + (1 - \alpha)b) \geq \alpha f(a) + (1 - a)f(b)
$$

We can generalize this definition to multiple points
$$
f(\alpha_{1}a_{1} + \alpha_{2}a_{2} + \alpha_{3}a_{3}) \geq \alpha_{1}f(a_{1})+\alpha_{2}f(a_{2}) + \alpha_{3}f(a_{3})
$$
with the constraint $\alpha_{1}+\alpha_{2}+\alpha_{3} = 1$. We can see that this formula can generalize if we consider $\alpha$ as probabilities. We obtain Jensen's inequality
$$
f\left(\mathbb{E}_{p(t)}[t]\right) \geq \mathbb{E}_{p(t)}\left[f(t)\right]
$$

### Kullback-Leibler divergence

The Kullback-Leibler divergence allows to compare two probability distributions
$$
\mathcal{KL}(q\Vert p) = \int q(x) \log \frac{q(x)}{p(x)} dx
$$

The KL divergence has several properties
1. $\mathcal{KL}(q\Vert p) \neq \mathcal{KL}(p\Vert q)$
2. $\mathcal{KL}(q\Vert q) = 0$
3. $\mathcal{KL}(q\Vert p) \geq 0$

### General form of EM

We are looking to maximize our model given our parameters, hence if we consider data points as independent we obtain
$$
\begin{align}
\underset{\theta}{\text{max. }} \log p(x \vert \theta) &= \log \prod_{i=1}^{N} p(x_{i} \vert \theta) \\
&= \sum_{i=1}^{N} \log p(x_{i} \vert \theta)
\end{align}
$$

Given that we want to maximize for our mixture model, we can write the final optimization problem as maximizing
$$
\log p(x \vert \theta) = \sum_{i=1}^{N} \log \sum_{c} p(x_{i}, t_{i}=c \vert \theta)
$$

#### Lower bounding the problem

Our goal here is to find a lower bound to this complicated expression so that
$$
\log p(x \vert \theta) \geq \mathcal{L}(\theta)
$$
Hence, if we can maximize the lower bound, it amounts to maximizing the original function. However, the problem is that the quality of our optimization is bound by the quality of our lower bound. Instead, we would like to introduce a lower bound that we could optimize. For instance, if we introduce a density
$$
\log p(x \vert \theta) = \sum_{i=1}^{N} \log \sum_{c} \frac{q(t_{i}=c)}{q(t_{i}=c)} p(x_{i}, t_{i}=c \vert \theta)
$$
Here, we can use Jensen's inequality such that
$$
\log\left(\sum_{c} \alpha_{c} v_{c}\right) \geq \sum_{c} \alpha_c \log(v_{c})
$$
Hence we obtain
$$
\log p(x \vert \theta) = \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log\frac{p(x_{i}, t_{i}=c \vert \theta)}{q(t_{i}=c)}
$$

To summarize, we now have a function $q$ that we can maximize and that should lower-bound so that
$$
\log p(x \vert \theta) \geq \mathcal{L}(\theta,q), \forall q
$$
Hence we can obtain the best optimization at a step as
$$
q^{k+1} = \underset{q}{\text{argmax. }} \mathcal{L}(\theta^{k}, q)
$$

### Summarizing the steps so far
E-step
$$
q^{k+1} = \underset{q}{\text{argmax. }} \mathcal{L}(\theta^{k}, q)
$$
M-step
$$
\theta^{k+1} = \underset{\theta}{\text{argmax. }} \mathcal{L}(\theta, q^{k+1})
$$

### Deriving the E-Step
Note that we want to minimize the gap between our lower bound and the true distribution, which can be written as
$$
\mathcal{D} = \log p(x \vert \theta) - \mathcal{L}(\theta\vert q)
$$
We can develop this expression
$$
\begin{align}
\mathcal{D} &= \sum_{i=1}^{N} \log p(x_{i} \vert \theta) - \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log\frac{p(x_{i}, t_{i}=c \vert \theta)}{q(t_{i}=c)} \\
&= \sum_{i=1}^{N} \left(\log p(x_{i} \vert \theta) - \sum_{c} q(t_{i}=c) \log\frac{p(x_{i}, t_{i}=c \vert \theta)}{q(t_{i}=c)} \right) \\
&= \sum_{i=1}^{N} \left(\log p(x_{i} \vert \theta)\sum_{c} q(t_{i}=c) - \sum_{c} q(t_{i}=c) \log\frac{p(x_{i}, t_{i}=c \vert \theta)}{q(t_{i}=c)} \right) \\
&= \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \left(\log p(x_{i} \vert \theta) - \log\frac{p(x_{i}, t_{i}=c \vert \theta)}{q(t_{i}=c)} \right) \\
&= \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log \frac{p(x_{i} \vert \theta)q(t_{i}=c)}{p(x_{i}, t_{i}=c \vert \theta)} \\
&= \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log \frac{q(t_{i}=c)}{p(t_{i} = c \vert x_{i}, \theta)} \\
&= \sum_{i=1}^{N} \mathcal{KL}(q(t_{i}) \Vert p(t_{i} \Vert x_{i}, \theta)
\end{align}
$$
Hence, we need to minimize the KL divergence, which amounts to setting the approximate function as the posterior distribution
$$
q(t_{i}) = p(t_{i} \vert x_{i}, \theta)
$$

### Deriving the M-step

Based on the approximation that we want to maximize
$$
\mathcal{L}(\theta, q) = \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log\frac{p(x_{i}, t_{i}=c \vert \theta)}{q(t_{i}=c)}
$$
We can expand this expression as
$$
\mathcal{L}(\theta, q) = \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log p(x_{i}, t_{i}=c \vert \theta) - \sum_{i=1}^{N} \sum_{c} q(t_{i}=c) \log q(t_{i}=c)
$$
As we can see the second term is not dependent on $\theta$, so the maximization is 
$$
\mathcal{L}(\theta, q) = \mathbb{E}_{q} \log p(x,t\vert \theta) + const
$$
Note that the fonction is usually concave, which makes it easy to optimize

#### Convergence guaranties

Note that by properties of the lower bound we have
$$
\log p (x \vert \theta^{k+1}) \geq \mathcal{L}(\theta^{k+1}, q^{k+1}) \geq \mathcal{L}(\theta^{k},q^{k+1})
$$
Therefore we have that
$$
\log p (x \vert \theta^{k+1}) \geq \log p (x \vert \theta^{k})
$$
Hence the EM algorithm is guaranteed to converge

### K-Means from GMM perspective

We can show that we can re-obtain the K-Means algorithm from our GMM formulation
- Fix covariances to be identical $\Sigma_{c}=I$
- Fix weights to be uniform $\pi_{c} = \frac{1}{C}$
Set the probability as
$$
p(x_{i}\vert t_{i}=c, \theta) = \frac{1}{Z} \exp (-0.5 \Vert x_{i} - \mu_{c} \Vert^{2})
$$
We can define $q\in Q$, where $Q$ is the set of delta-functions

# Expectation-maximization algorithm

In this assignment, we will derive and implement formulas for Gaussian Mixture Model — one of the most commonly used methods for performing soft clustering of the data. 

### Setup
Loading auxiliary files and importing the necessary libraries.


```python
try:
    import google.colab
    IN_COLAB = True
except:
    IN_COLAB = False
if IN_COLAB:
    print("Downloading Colab files")
    ! shred -u setup_google_colab.py
    ! wget https://raw.githubusercontent.com/hse-aml/bayesian-methods-for-ml/master/setup_google_colab.py -O setup_google_colab.py
    import setup_google_colab
    setup_google_colab.load_data_week2()
```


```python
import numpy as np
from numpy.linalg import slogdet, det, solve
import matplotlib.pyplot as plt
import time
from sklearn.datasets import load_digits
from w2_grader import EMGrader
%matplotlib inline
```

### Grading
We will create a grader instance below and use it to collect your answers. Note that these outputs will be stored locally inside grader and will be uploaded to the platform only after running submitting function in the last part of this assignment. If you want to make a partial submission, you can run that cell anytime you want.


```python
grader = EMGrader()
```

## Implementing EM for GMM

For debugging, we will use samples from a Gaussian mixture model with unknown mean, variance, and priors. We also added initial values of parameters for grading purposes.


```python
samples = np.load('samples.npz')
X = samples['data']
pi0 = samples['pi0']
mu0 = samples['mu0']
sigma0 = samples['sigma0']
plt.scatter(X[:, 0], X[:, 1], c='grey', s=30)
plt.axis('equal')
plt.show()
```

### Reminder

Remember, that EM algorithm is a coordinate descent optimization of variational lower bound $\mathcal{L}(\theta, q) = \int q(T) \log\frac{p(X, T|\theta)}{q(T)}dT\to \max$.

<b>E-step</b>:<br>
$\mathcal{L}(\theta, q) \to \max\limits_{q} \Leftrightarrow \mathcal{KL} [q(T) \,\|\, p(T|X, \theta)] \to \min \limits_{q\in Q} \Rightarrow q(T) = p(T|X, \theta)$<br>
<b>M-step</b>:<br> 
$\mathcal{L}(\theta, q) \to \max\limits_{\theta} \Leftrightarrow \mathbb{E}_{q(T)}\log p(X,T | \theta) \to \max\limits_{\theta}$

For GMM, $\theta$ is a set of parameters that consists of mean vectors $\mu_c$, covariance matrices $\Sigma_c$ and priors $\pi_c$ for each component.

Latent variables $T$ are indices of components to which each data point is assigned, i.e. $t_i$  is the cluster index for object $x_i$.

The joint distribution can be written as follows: $\log p(T, X \mid \theta) =  \sum\limits_{i=1}^N \log p(t_i, x_i \mid \theta) = \sum\limits_{i=1}^N \sum\limits_{c=1}^C q(t_i = c) \log \left (\pi_c \, f_{\!\mathcal{N}}(x_i \mid \mu_c, \Sigma_c)\right)$,
where $f_{\!\mathcal{N}}(x \mid \mu_c, \Sigma_c) = \frac{1}{\sqrt{(2\pi)^n|\boldsymbol\Sigma_c|}}
\exp\left(-\frac{1}{2}({x}-{\mu_c})^T{\boldsymbol\Sigma_c}^{-1}({x}-{\mu_c})
\right)$ is the probability density function (pdf) of the normal distribution $\mathcal{N}(x_i \mid \mu_c, \Sigma_c)$.

### E-step
In this step we need to estimate the posterior distribution over the latent variables with fixed values of parameters: $q_i(t_i) = p(t_i \mid x_i, \theta)$. We assume that $t_i$ equals to the cluster index of the true component of the $x_i$ object. To do so we need to compute $\gamma_{ic} = p(t_i = c \mid x_i, \theta)$. Note that $\sum\limits_{c=1}^C\gamma_{ic}=1$.


<b>Important trick 1:</b> It is important to avoid numerical errors. At some point you will have to compute the formula of the following form: $\frac{e^{y_i}}{\sum_j e^{y_j}}$, which is called _softmax_. When you compute exponents of large numbers, some numbers may become infinity. You can avoid this by dividing numerator and denominator by $e^{\max(y)}$: $\frac{e^{y_i-\max(y)}}{\sum_j e^{y_j - \max(y)}}$. After this transformation maximum value in the denominator will be equal to one. All other terms will contribute smaller values. So, to compute desired formula you first subtract maximum value from each component in vector $\mathbf{y}$ and then compute everything else as before.

<b>Important trick 2:</b> You will probably need to compute formula of the form $A^{-1}x$ at some point. You would normally inverse $A$ and then multiply it by $x$. A bit faster and more numerically accurate way to do this is to directly solve equation $Ay = x$ by using a special function. Its solution is $y=A^{-1}x$, but the equation $Ay = x$ can be solved by methods which do not explicitely invert the matrix. You can use ```np.linalg.solve``` for this.

<b>Other usefull functions: </b> <a href="https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.slogdet.html">```slogdet```</a> and <a href="https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.det.html#numpy.linalg.det">```det```</a>

<b>Task 1:</b> Implement E-step for GMM using template below.


```python
def E_step(X, pi, mu, sigma):
    """
    Performs E-step on GMM model
    Each input is numpy array:
    X: (N x d), data points
    pi: (C), mixture component weights 
    mu: (C x d), mixture component means
    sigma: (C x d x d), mixture component covariance matrices
    
    Returns:
    gamma: (N x C), probabilities of clusters for objects
    """
    N = X.shape[0] # number of objects
    C = pi.shape[0] # number of clusters
    d = mu.shape[1] # dimension of each object
    gamma = np.zeros((N, C)) # distribution q(T)

    ### YOUR CODE HERE
    
    return gamma
```


```python
gamma = E_step(X, pi0, mu0, sigma0)
grader.submit_e_step(gamma)
```

### M-step

In M-step we need to maximize $\mathbb{E}_{q(T)}\log p(X,T | \theta)$ with respect to $\theta$. In our model this means that we need to find optimal values of $\pi$, $\mu$, $\Sigma$. To do so, you need to compute the derivatives and 
set them to zero. You should start by deriving formulas for $\mu$ as it is the easiest part. Then move on to $\Sigma$. Here it is crucial to optimize function w.r.t. to $\Lambda = \Sigma^{-1}$ and then inverse obtained result. Finaly, to compute $\pi$, you will need <a href="https://www3.nd.edu/~jstiver/FIN360/Constrained%20Optimization.pdf">Lagrange Multipliers technique</a> to satisfy constraint $\sum\limits_{i=1}^{n}\pi_i = 1$.

<br>
<b>Important note:</b> You will need to compute derivatives of scalars with respect to matrices. To refresh this technique from previous courses, see <a href="https://en.wikipedia.org/wiki/Matrix_calculus"> wiki article</a> about it . Main formulas of matrix derivatives can be found in <a href="http://www2.imm.dtu.dk/pubdb/views/edoc_download.php/3274/pdf/imm3274.pdf">Chapter 2 of The Matrix Cookbook</a>. For example, there you may find that $\frac{\partial}{\partial A}\log |A| = A^{-T}$.


<b>Task 2:</b> Implement M-step for GMM using template below.


```python
def M_step(X, gamma):
    """
    Performs M-step on GMM model
    Each input is numpy array:
    X: (N x d), data points
    gamma: (N x C), distribution q(T)  
    
    Returns:
    pi: (C)
    mu: (C x d)
    sigma: (C x d x d)
    """
    N = X.shape[0] # number of objects
    C = gamma.shape[1] # number of clusters
    d = X.shape[1] # dimension of each object

    ### YOUR CODE HERE

    return pi, mu, sigma
```


```python
gamma = E_step(X, pi0, mu0, sigma0)
pi, mu, sigma = M_step(X, gamma)
grader.submit_m_step(pi, mu, sigma)
```

### Loss function

Finally, we need some function to track convergence. We will use variational lower bound $\mathcal{L}$ for this purpose. We will stop our EM iterations when $\mathcal{L}$ will saturate. Usually, you will need only about 10-20 iterations to converge. It is also useful to check that this function never decreases during training. If it does, you have a bug in your code.

<b>Task 3:</b> Implement a function that will compute $\mathcal{L}$ using template below.

$$\mathcal{L} = \sum_{i=1}^{N} \sum_{c=1}^{C} q(t_i =c) (\log \pi_c + \log f_{\!\mathcal{N}}(x_i \mid \mu_c, \Sigma_c)) - \sum_{i=1}^{N} \sum_{c=1}^{K} q(t_i =c) \log q(t_i =c)$$


```python
def compute_vlb(X, pi, mu, sigma, gamma):
    """
    Each input is numpy array:
    X: (N x d), data points
    gamma: (N x C), distribution q(T)  
    pi: (C)
    mu: (C x d)
    sigma: (C x d x d)
    
    Returns value of variational lower bound
    """
    N = X.shape[0] # number of objects
    C = gamma.shape[1] # number of clusters
    d = X.shape[1] # dimension of each object

    ### YOUR CODE HERE

    return loss
```


```python
pi, mu, sigma = pi0, mu0, sigma0
gamma = E_step(X, pi, mu, sigma)
pi, mu, sigma = M_step(X, gamma)
loss = compute_vlb(X, pi, mu, sigma, gamma)
grader.submit_VLB(loss)
```

### Bringing it all together

Now that we have E step, M step and VLB, we can implement the training loop. We will initialize values of $\pi$, $\mu$ and $\Sigma$ to some random numbers, train until $\mathcal{L}$ stops changing, and return the resulting points. We also know that the EM algorithm converges to local optima. To find a better local optima, we will restart the algorithm multiple times from different (random) starting positions. Each training trial should stop either when maximum number of iterations is reached or when relative improvement is smaller than given tolerance ($|\frac{\mathcal{L}_i-\mathcal{L}_{i-1}}{\mathcal{L}_{i-1}}| \le \text{rtol}$).

Remember, that initial (random) values of $\pi$ that you generate must be non-negative and sum up to 1. Also, $\Sigma$ matrices must be symmetric and positive semi-definite. If you don't know how to generate those matrices, you can use $\Sigma=I$ as initialization.

You will also sometimes get numerical errors because of component collapsing. The easiest way to deal with this problems is to restart the procedure.

<b>Task 4:</b> Implement training procedure


```python
def train_EM(X, C, rtol=1e-3, max_iter=100, restarts=10):
    '''
    Starts with random initialization *restarts* times
    Runs optimization until saturation with *rtol* reached
    or *max_iter* iterations were made.
    
    X: (N, d), data points
    C: int, number of clusters
    '''
    N = X.shape[0] # number of objects
    d = X.shape[1] # dimension of each object
    best_loss = None
    best_pi = None
    best_mu = None
    best_sigma = None

    for _ in range(restarts):
        try:
            ### YOUR CODE HERE

        except np.linalg.LinAlgError:
            print("Singular matrix: components collapsed")
            pass

    return best_loss, best_pi, best_mu, best_sigma
```


```python
best_loss, best_pi, best_mu, best_sigma = train_EM(X, 3)
grader.submit_EM(best_loss)
```

If you implemented all the steps correctly, your algorithm should converge in about 20 iterations. Let's plot the clusters to see it. We will assign a cluster label as the most probable cluster index. This can be found using a matrix $\gamma$ computed on last E-step. 


```python
gamma = E_step(X, best_pi, best_mu, best_sigma)
labels = gamma.argmax(axis=1)
colors = np.array([(31, 119, 180), (255, 127, 14), (44, 160, 44)]) / 255.
plt.scatter(X[:, 0], X[:, 1], c=colors[labels], s=30)
plt.axis('equal')
plt.show()
```

86a91f2450a0f53434c6d039ad4830763e584edf
ipynb
Jupyter Notebook
documents/atiam_ml/06b_expectation_maximization.ipynb
esling/esling.github.io
a50f29a2901713f2c69f6c331a2462c479e539c5
MIT
2020-10-28T22:49:38.000Z
2020-10-28T22:49:38.000Z
documents/atiam_ml/06b_expectation_maximization.ipynb
esling/esling.github.io
a50f29a2901713f2c69f6c331a2462c479e539c5
MIT
documents/atiam_ml/06b_expectation_maximization.ipynb
esling/esling.github.io
a50f29a2901713f2c69f6c331a2462c479e539c5
MIT
2016-10-19T16:01:59.000Z
2021-09-20T04:59:13.000Z
Qwen/Qwen-72B
1. YES
2. YES