Maximum Likelihood with Gradient Descent or Coordinate Descent blows up

Question

Context

The maximum likelihood estimators for a Normal distribution with unknown mean and unknown variance are $$ \widehat{\mu} = \frac{1}{n}\sum_{i=1}^n x_i \qquad \text{and} \qquad \widehat{\sigma}^2 = \frac{1}{n}\sum_{i=1}^n (x_i - \mu)^2 $$ These can be found (for example) by taking derivatives of the average log-likelihood $$ \frac{1}{n}\sum_{i=1}^n \log p(x_i) = -\frac{1}{2}\log(2\pi) - \frac{1}{2n\sigma^2}\sum^n_{i=1} (x^{(i)} - \mu)^2 - \log \sigma $$

Question: What if I want to use a gradient-based method?

Yes, I know I can just use the estimators found above. However, I want to find such estimators using a gradient-based method such as coordinate descent or gradient descent. These are the gradients with respect to $\mu$ and with respect to $\sigma$ (which you can set equal to zero to find the estimators above)

$$ \begin{align} \frac{\partial}{\partial \mu} \frac{1}{n} \sum^n_{i=1} \log p(x^{(i)}) &= \frac{\overline{x}}{\sigma^2} - \frac{\mu}{\sigma^2} \\ \frac{\partial}{\partial \sigma} \frac{1}{n}\sum^n_{i=1} \log p(x^{(i)}) &= \frac{1}{n\sigma^3}\sum^n_{i=1}(x^{(i)} - \mu)^2 - \frac{1}{\sigma} \end{align} $$ I tried using them in gradient descent $$ \begin{align} \mu_{t+1} &\longleftarrow \mu_t + \gamma \left(\frac{\overline{x}}{\sigma^2_t} - \frac{\mu_t}{\sigma^2_t}\right) \\ \sigma_{t+1} &\longleftarrow \sigma_t + \gamma\left(\frac{1}{n\sigma^3_t}\sum^n_{i=1}(x^{(i)} - \mu_{t+1})^2 - \frac{1}{\sigma_t}\right) \end{align} $$ or in coordinate ascent (where I would keep, say $\sigma_t$ fixed and optimize $\mu_t$ for $n_{\text{inner}}$ times and then switch: keep $\mu_t$ fixed and optimize $\sigma_t$ for $n_{\text{inner}}$ times. All this for $n_{\text{outer}}$ times. However it seems to blow up for some reason and not give me the obvious answer. You can run the code here.

What am I doing wrong?

Answer 1

You correctly computed the gradients, but the optimization step is incorrect. Your gradient is computed at point $(\mu,\sigma)$, and, therefore, it should be used to update $\mu$ and $\sigma$ simultaneously. Correct update rule is $$ \begin{align} \mu_{t+1} &\longleftarrow \mu_t + \gamma \left(\frac{\overline{x}}{\sigma^2_t} - \frac{\mu_t}{\sigma^2_t}\right) \\ \sigma_{t+1} &\longleftarrow \sigma_t + \gamma\left(\frac{1}{n\sigma^3_t}\sum^n_{i=1}(x^{(i)} - \mu_{t})^2 - \frac{1}{\sigma_t}\right) \end{align} $$ The formula that you have in the question implements the alternating optimization which may or may not converge (it is guaranteed to converge if the gradient norm along the optimization path is at most one).

Therefore, I corrected your code as follows:

# Loop through and update mu and sigma
mus = [mu]
sigmas = [sigma]
for i in range(n_iter):
    # Compute gradients
    mu_grad = (np.mean(x) - mu) / (sigma**2)
    sigma_grad = (np.mean((x - mu)**2)/sigma**3 - 1 / sigma)
    # Update mu and sigma
    mu, sigma = mu + gamma_mu*mu_grad, sigma + gamma_sigma*sigma_grad
    # Store mu and sigma
    mus.append(mu)
    sigmas.append(sigma)

The result converges quickly.