Double Derivatives

SymbolicNeuralNetworks can compute derivatives of arbitrary order, by feeding the symbolic result of one derivative into the next. This is the case Zygote-based AD struggles with, and the reason this package exists.

The two building blocks are Jacobian (differentiate with respect to the input) and Gradient (differentiate with respect to the parameters); see Derivatives for each on its own.

Jacobian of a neural network

using AbstractNeuralNetworks
using SymbolicNeuralNetworks
using SymbolicNeuralNetworks: Jacobian, Gradient, derivative
using Latexify: latexify

c = Chain(Dense(2, 1, tanh; use_bias = false))
nn = SymbolicNeuralNetwork(c)
□ = Jacobian(nn)
derivative(□) |> latexify

The output of nn is one-dimensional and the convention is

\[\square_{ij} = [\mathrm{jacobian}_{x}f]_{ij} = \frac{\partial}{\partial{}x_j}f_i,\]

so the result has shape $\mathrm{output\_dim}\times\mathrm{input\_dim} = 1\times2$:

size(derivative(□))

(1, 2)

Gradient of a neural network

Gradient differentiates every element of an array-valued expression with respect to the network parameters:

g = Gradient(nn)

derivative(g)[1].L1.W |> latexify

Combining the two

Feeding the Jacobian into a Gradient differentiates the network twice — first with respect to its input, then with respect to its parameters:

g = Gradient(derivative(□), nn)

The result is a matrix (of the shape of the Jacobian) of parameter sets. To read the derivative of the first Jacobian entry with respect to the weight W of the first layer:

matrix_index = (1, 1)
layer = :L1
weight = :W
derivative(g)[matrix_index...][layer][weight] |> latexify

build_nn_function turns the whole thing into an executable function. We evaluate it at

\[x = \begin{pmatrix} 1 \\ 0 \end{pmatrix}, \quad W = \begin{bmatrix} 1 & 0 \end{bmatrix}\]

using AbstractNeuralNetworks: params
using NeuralNetworkParameters: NetworkParameters

built_function = build_nn_function(derivative(g), params(nn), nn.input)

x = [1.0, 0.0]
ps = NetworkParameters((L1 = (W = [1.0 0.0],),))
built_function(x, ps)[matrix_index...][layer][weight]

1×2 Matrix{Float64}:
 -0.219726  -0.0