Maximum Likelihood for Normal Distribution with Unknown Variance - Gradient Descent not working

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

Given the likelihood function $P(\mathcal{D};\mu,\sigma^2)$, where $\mathcal{D}$ is the dataset, $\mu$ is the mean and $\sigma^2$ is the variance, gradient descent first involves combining the parameters $\mu$ and $\sigma^2$ into a parameter vector $\theta=[\mu,\sigma^2]^T$ such that:

$$P(\mathcal{D};\mu,\sigma^2) = P(\mathcal{D};\theta)$$

Then, the parameter vector $\theta$ is updated as follows:

$$\theta_{t+1} \leftarrow \theta_t + \epsilon \nabla_{\theta} P(\mathcal{D};\theta)$$

Where $\epsilon$ is the learning rate, and $\nabla_{\theta} P(\mathcal{D};\theta)$ is the gradient of the likelihood function $P(\mathcal{D};\theta)$ with respect to $\theta$. Both coordinate descent and gradient descent are equivalent for a single parameter, but for multiple parameters, updating the parameters individually corresponds to coordinate descent, while updating them simultaneously after combining them into a parameter vector is equivalent to gradient descent. Please see this Jupyter notebook that I made that explains the process of estimating these parameters both analytically and using gradient descent. Try to open the notebook locally as GitHub does not render notebooks well.

Hope this helps.