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$?