Derivatives

There are two directions in which a neural network can be differentiated, and this package has a struct for each:

Jacobian differentiates with respect to the input,
Gradient differentiates with respect to the parameters.

Both store a symbolic expression, which derivative returns and build_nn_function compiles. Because the result of one is again a symbolic expression, they compose freely — see Double Derivatives.

using SymbolicNeuralNetworks
using SymbolicNeuralNetworks: Jacobian, Gradient, derivative
using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
import Random
Random.seed!(123)

c = Chain(Dense(2, 3, tanh), Dense(3, 2, tanh))
snn = SymbolicNeuralNetwork(c)
nn = NeuralNetwork(c)
ps = params(nn)

Jacobian

Jacobian differentiates an expression with respect to nn.input. Without an expression it takes the output of the network:

j = Jacobian(snn)
size(derivative(j))

(2, 2)

The convention is $\square_{ij} = \partial{}f_i/\partial{}x_j$, so the result is $\mathrm{output\_dim}\times\mathrm{input\_dim}$ — the same convention Zygote and ForwardDiff use:

import ForwardDiff

input = rand(2)
jacobian = build_nn_function(derivative(j), snn)
jacobian(input, ps) ≈ ForwardDiff.jacobian(x -> c(x, ps), input)

true

An expression can also be given explicitly, in which case it is flattened with vec first, so the rows of the result are indexed by vec(f):

Jacobian(c(snn.input, params(snn)) .^ 2, snn) |> derivative |> size

(2, 2)

Gradient differentiates with respect to params(snn). Its result has the shape of the parameters: for each entry of the differentiated expression there is one full parameter set holding the derivative with respect to each parameter.

g = Gradient(snn)
derivative(g)[1].L1.b

\[ \begin{equation} \left[ \begin{array}{c} \mathtt{W\_3_{1}ˏ_1} ~ \left( 1 - \tanh^{2}\left( \mathtt{W\_4_1} + \mathtt{W\_3_{1}ˏ_1} ~ \tanh^{2}\left( \mathtt{W\_2_1} + \mathtt{W\_1_{1}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{1}ˏ_2} ~ \mathtt{x_2} \right) + \mathtt{W\_3_{1}ˏ_2} ~ \tanh^{2}\left( \mathtt{W\_2_2} + \mathtt{W\_1_{2}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{2}ˏ_2} ~ \mathtt{x_2} \right) + \mathtt{W\_3_{1}ˏ_3} ~ \tanh^{2}\left( \mathtt{W\_2_3} + \mathtt{W\_1_{3}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{3}ˏ_2} ~ \mathtt{x_2} \right) \right) \right) ~ \left( 1 - \tanh^{2}\left( \mathtt{W\_2_1} + \mathtt{W\_1_{1}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{1}ˏ_2} ~ \mathtt{x_2} \right) \right) \\ \mathtt{W\_3_{1}ˏ_2} ~ \left( 1 - \tanh^{2}\left( \mathtt{W\_2_2} + \mathtt{W\_1_{2}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{2}ˏ_2} ~ \mathtt{x_2} \right) \right) ~ \left( 1 - \tanh^{2}\left( \mathtt{W\_4_1} + \mathtt{W\_3_{1}ˏ_1} ~ \tanh^{2}\left( \mathtt{W\_2_1} + \mathtt{W\_1_{1}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{1}ˏ_2} ~ \mathtt{x_2} \right) + \mathtt{W\_3_{1}ˏ_2} ~ \tanh^{2}\left( \mathtt{W\_2_2} + \mathtt{W\_1_{2}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{2}ˏ_2} ~ \mathtt{x_2} \right) + \mathtt{W\_3_{1}ˏ_3} ~ \tanh^{2}\left( \mathtt{W\_2_3} + \mathtt{W\_1_{3}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{3}ˏ_2} ~ \mathtt{x_2} \right) \right) \right) \\ \mathtt{W\_3_{1}ˏ_3} ~ \left( 1 - \tanh^{2}\left( \mathtt{W\_4_1} + \mathtt{W\_3_{1}ˏ_1} ~ \tanh^{2}\left( \mathtt{W\_2_1} + \mathtt{W\_1_{1}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{1}ˏ_2} ~ \mathtt{x_2} \right) + \mathtt{W\_3_{1}ˏ_2} ~ \tanh^{2}\left( \mathtt{W\_2_2} + \mathtt{W\_1_{2}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{2}ˏ_2} ~ \mathtt{x_2} \right) + \mathtt{W\_3_{1}ˏ_3} ~ \tanh^{2}\left( \mathtt{W\_2_3} + \mathtt{W\_1_{3}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{3}ˏ_2} ~ \mathtt{x_2} \right) \right) \right) ~ \left( 1 - \tanh^{2}\left( \mathtt{W\_2_3} + \mathtt{W\_1_{3}ˏ_1} ~ \mathtt{x_1} + \mathtt{W\_1_{3}ˏ_2} ~ \mathtt{x_2} \right) \right) \\ \end{array} \right] \end{equation} \]

Terminology

The name Gradient is not used in the usual sense here. A gradient normally collects the partial derivatives of a scalar function; Gradient differentiates every entry of an array with respect to every parameter, so the gradient of a matrix is a matrix of parameter sets:

\[\mathtt{Gradient}\left( \begin{pmatrix} m_{11} & \cdots & m_{1m} \\ \vdots & \vdots & \vdots \\ m_{n1} & \cdots & m_{nm} \end{pmatrix} \right) = \begin{pmatrix} \nabla_{\mathbb{P}}m_{11} & \cdots & \nabla_{\mathbb{P}}m_{1m} \\ \vdots & \vdots & \vdots \\ \nabla_{\mathbb{P}}m_{n1} & \cdots & \nabla_{\mathbb{P}}m_{nm} \end{pmatrix},\]

where $\mathbb{P}$ are the parameters of the network.

The underlying function is symbolic_parameter_gradient, which can also be used directly. For a scalar expression it returns a single parameter set rather than an array of them:

using SymbolicNeuralNetworks: symbolic_parameter_gradient
using LinearAlgebra: norm

gradient = symbolic_parameter_gradient(norm(c(snn.input, params(snn))) ^ 2, snn)
keys(gradient)

(:L1, :L2)

Such a parameter-shaped expression is an equation set, which build_nn_function builds as a single function whose result has the same nesting. Summing over the batch with reduce = + gives the gradient of the summed per-sample expression:

f = build_nn_function(gradient, params(snn), snn.input; reduce = +)
f(rand(2, 8), ps).L1.b

3-element Vector{Float64}:
 -2.406760476493319
  1.5060435320698913
 -3.6789741121788766

SymbolicPullback

SymbolicPullback packages that up for training: it differentiates an AbstractNeuralNetworks.NetworkLoss with respect to the parameters and presents the result the way an optimizer expects a pullback to look.

using AbstractNeuralNetworks: FeedForwardLoss

pb = SymbolicPullback(snn, FeedForwardLoss())

input, output = rand(2, 8), rand(2, 8)
loss_value, pullback = pb(params(nn), c, (input, output))
loss_value

0.5547550239357568

The second entry is a function of the output sensitivities. A NetworkLoss is scalar-valued, so those are just 1:

pullback(1).L1.b

3-element Vector{Float64}:
 -0.0018979669434240481
 -1.5858177322012235
  1.6985287943810916

Constructing the pullback is where most of the time goes; evaluating it is fast and allocates once per call. See Training a Symbolic Neural Network for a complete training run, and Limitations for the assumption SymbolicPullback makes about the loss.