Calculating the Jacobian of a neural network

I'm trying to calculate confidence intervals for a neural network (rather than prediction intervals). I'm following this paper, which treats them in the same framework as any parametric (parameter-involving?) nonlinear model (which a neural net basically is). Calculating these CI's involves computing the Jacobian -- the matrix of partial derivatives of $\hat y$ with respect to the parameters.

In the one-outcome regression case, a neural net with two hidden layers (and no biases, WLOG) is $$y = \sigma(\sigma(\mathbf{X\alpha)\beta)\gamma} + \epsilon$$ where

• $y$ is the $N\times 1$ outcome
• $\mathbf{X}$ is the $N\times P_x$ data
• $\alpha$ is the $P_x \times P_\beta$ first set of weights
• $\beta$ is the $P_\beta \times P_\gamma$ second set of weights
• $\gamma$ is the $P_\gamma \times 1$ final set of weights, linking to $y$
• $\epsilon$ is the mean-zero error
• $\sigma()$ is the activation function

The chain rule gives me the following partial derivatives of the fitted model: $$\begin{array}{rcl} \displaystyle\frac{\partial\hat y}{\partial \hat\gamma} & = & \sigma(\sigma(\mathbf{X\alpha})\beta) \\ \displaystyle\frac{\partial\hat y}{\partial \hat\beta} & = & \sigma'(\sigma(\mathbf{X\alpha})\beta)\gamma\sigma(\mathbf{X}\alpha) \\ \displaystyle\frac{\partial\hat y}{\partial \hat\alpha} & = & \sigma'(\sigma(\mathbf{X\alpha})\beta)\gamma\sigma'(\mathbf{X}\alpha)\beta \mathbf{X} \end{array}$$

Now, the matrix dimensions on the above expressions are not conformable, which suggests that vectorizing the Jacobian calculation is not possible. Is this true?

Next, if I need to calculate the partial derivatives parameter-by-parameter, what is the appropriate way to deal with the network structure when evaluating each derivative?

For example, $\alpha_1$ links the $\mathbf{X}_1$ to $\sigma(\mathbf{X\alpha})_1$. That node in turn links to several different upper nodes. Let's say that there are $g$ hidden units in the topmost layer. Should the Jacobian column for $\alpha_1$ include $\displaystyle\sum_g \sigma(\mathbf{X\alpha})_1 \beta_g$ (i.e.: row-wise sum)? Or would the combination of the $\beta$'s take some other form than a sum?

Likewise, $\hat\alpha_1$ influences $\hat y$ through all of the top layer hidden-units. So $\hat\alpha_1$'s Jacobian column should also include $\sigma'(\sigma(\mathbf{X\alpha})\beta)\gamma$, treated as a matrix, which evaluates to a $N \times P_\gamma \times P_\gamma \times 1 = N \times 1$ matrix?

$\hat\alpha_1$'s Jacobian column would thus be $$\frac{\partial\hat y}{\partial \hat\alpha_1} = \sigma'(\sigma(\mathbf{X\alpha})\beta)\gamma\circ\displaystyle\sum_g\sigma'(\mathbf{X}_1\alpha_1)\beta_g \circ\mathbf{X}_1$$

Is this correct?
Is there a more efficient way to do this? Can a general statement be made about how much better the analytical Jacobian would be than one calculated by numerical approximation, for example in the numDeriv package in R?

• Is there a specific reason you do not use backpropagation? – bayerj Nov 15 '16 at 20:40
• I train the model with backprop. Is there some way to get a linear taylor expansion of the fitted model from backprop? – generic_user Nov 15 '16 at 20:54
• No, but backprop essentially calculates the Jacobian. So if you just backprop using the identity as a loss function, you should get the Jacobian of the outputs wrt the weights. – bayerj Nov 18 '16 at 21:04
• @bayerj That's a good idea. I haven't written my code such that the loss function is a drop-in argument, but if I ever implement such a setup I might do as you suggest. – generic_user Nov 24 '16 at 16:37