Gaussian Mixture M-Step

Write the M-step of the EM algorithm for a Gaussian mixture model.

The E-step gave you responsibilities resp[i, c], the soft share of point i that belongs to component c. The M-step refits each component to its share of the data, as if those shares were weights. Let N_c = sum_i resp[i, c]. Then

weights[c] = N_c / n
means[c]   = sum_i resp[i, c] * x_i / N_c
covs[c]    = sum_i resp[i, c] * (x_i - means[c]) (x_i - means[c])^T / N_c

Alternating the E-step and this M-step never lowers the likelihood of the data.

Function signature

def gmm_m_step(X, resp):

Example

Input

X = np.array([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0], [2.0, 2.0]])
resp = np.array([[1.0, 0.0], [0.5, 0.5], [0.5, 0.5], [0.0, 1.0]])
gmm_m_step(X, resp)

Output

weights = array([0.5, 0.5])
means   = array([[0.5, 0.5], [1.5, 1.5]])
covs    = array([[[ 0.75, -0.25], [-0.25,  0.75]],
                 [[ 0.75, -0.25], [-0.25,  0.75]]])

Reasoning

Each component holds a total share of 2, so both weights are 2 / 4. The first mean is (1 * (0, 0) + 0.5 * (2, 0) + 0.5 * (0, 2)) / 2 = (0.5, 0.5). The covariance uses the same shares on the outer products of the centred points.

Constraints and notes

• X has shape (n, d) and resp has shape (n, k).
• Return (weights, means, covs) with shapes (k,), (k, d) and (k, d, d).
• Divide the covariance by N_c, not by N_c - 1.