What is the Hessian of the Gaussian likelihood

Question

I am trying to learn the fine differences between different methods of Kronecker factoring for approximate curvature (like [1], and [2]) which require taking the Hessian of the pre-activations of the last layer. [2] says that the pre-activation Hessian of the squared loss $\frac{(\hat{\mathbf{y}} - \mathbf{y})^2}{2} = I$ which makes sense because the first partial w.r.t. multiple independent output variables $\hat{\mathbf{y}}$ would be a vector of $\mathbf{\hat{y}} - \mathbf{y}$ and then expanding to the Hessian w.r.t. each element of $\hat{\mathbf{y}}$ would be $\frac{\partial^2}{\partial \hat{y}_i \hat{y}_j}$ would be equal to $I$.

I would like to do the same thing for the Gaussian loss used in regression tasks which output a Gaussian likelihood, but I am unsure if I am doing it correctly.

Given the log Gaussian likelihood below parameters $(\mathbf{\mu}, \mathbf{\sigma}) = \tau$, what are the Jacobian and Hessian? (assuming, as in the first case, $\mu, \sigma$ represent multiple outputs).

$$ -\log\Big( \frac{1}{\sigma \sqrt{2\pi}}e^{-\frac{1}{2}(\frac{x - \mu}{\sigma})^2}\Big) = \log \sigma + \frac{1}{2} \log 2\pi + \frac{1}{2\sigma^2}(x - \mu)^2 $$

The Jacobian would be the first partial derivatives of the negative log density with respect to the parameters of the distribution, \begin{aligned} \frac{\partial}{\partial\tau}\log \sigma + \frac{1}{2} \log 2\pi + \frac{1}{2\sigma^2}(x - \mu)^2 &= \begin{bmatrix}\frac{\partial\mathcal{N}}{\partial \mu}, & \frac{\partial\mathcal{N}}{\partial \sigma}\end{bmatrix} \\ &= \begin{bmatrix} -\frac{(x - \mu)}{\sigma^2}, & \frac{1}{\sigma} - \frac{(x - \mu)^2}{\sigma^3}\end{bmatrix} \end{aligned}

And the Hessian would be the second partial derivatives w.r.t each entry in the Jacobian, \begin{aligned} \frac{\partial^2 \mathcal{N}}{\partial^2 \tau} &= \begin{bmatrix}\frac{\partial^2\mathcal{N}}{\partial \mu\mu}, & \frac{\partial^2\mathcal{N}}{\partial \mu\sigma} \\ \frac{\partial^2\mathcal{N}}{\partial \sigma\mu}, & \frac{\partial^2\mathcal{N}}{\partial \sigma\sigma}\end{bmatrix} \\ &= \begin{bmatrix} \frac{1}{\sigma^2}, & \frac{2(x - \mu)}{\sigma^3} \\ \frac{2(x - \mu)}{\sigma^3}, & -\frac{1}{\sigma^2} + 3\frac{(x - \mu)^2}{\sigma^4} \end{bmatrix} \end{aligned}

Does this look correct? What I cannot justify is why I do not have a 3 dimensional Hessian. The Jacobian of the squared loss (single output variable, multiple outputs) introduced above is a vector and then the Hessian is a 2 dimensional object of size $n \times n$.

In the Gaussian case we have two output variables, and also multiple outputs, so then would the Hessian of this Gaussian be $n \times n \times 4$, or $n \times n \times 2$?

[1] https://openreview.net/pdf?id=Skdvd2xAZ

[2] http://proceedings.mlr.press/v70/botev17a/botev17a.pdf

Answer 1

I think your confusion is that you are directly applying a probabilistic result to a statistical sample.

That is, if your vector of observations $[Y_1, Y_2, \ldots, Y_n] \sim_{iid} \mathcal{N}(\mu, \sigma^2)$ then I believe you are considering that the "Jacobian of the sample" would be a tensor, i.e. the $n \times 2 \times 2$ vector of $\partial ^2 [Y_1, Y_2, \ldots, Y_n] / \partial ^2 \tau$. While technically true, in practice this doesn't make sense. Rather, we use the random sample to estimate the functionals of the Jacobian and gradient for a single observation. Your result for the gradient and Jacobian are correct for a single observation, i.e. for a probability case. The missing piece is how we use the random sample to estimate those functionals. The most straightforward application is to use the Mann Wold theorem. This gives you that if you have consistent estimators of $\mu$ and $\sigma^2$, you can just plug them into the functional and boom there's your estimate.

So using a random sample of size $n$, I would have as the Jacobian:

\begin{aligned} \widehat{\frac{\partial^2 \mathcal{N}}{\partial^2 \tau}} &= \begin{bmatrix} \frac{1}{\hat{\sigma}^2}, & \frac{2(x - \hat{\mu})}{\hat{\sigma^3}} \\ \frac{2(x - \hat{\mu})}{\hat{\sigma^3}}, & -\frac{1}{\hat{\sigma^2}} + 3\frac{(x - \hat{\mu})^2}{\hat{\sigma}^4} \end{bmatrix} \end{aligned}

This is what some first year probability students (ahem myself in the past) would have called "putting hats on everything".