SymbolicNeuralNetworks.jl

SymbolicNeuralNetworks builds a symbolic representation of a (small) neural network, lets you form arbitrary expressions from it — derivatives with respect to the input, derivatives with respect to the parameters, and combinations of the two — and compiles those expressions into ordinary Julia functions with RuntimeGeneratedFunctions.

It is built on AbstractNeuralNetworks and is mostly used together with GeometricMachineLearning and GeometricIntegrators.

When to reach for this package

The motivation is Zygote's difficulty with second-order derivatives: when a loss function itself contains a derivative of the network — as the losses of Hamiltonian and Lagrangian neural networks do — differentiating it again with reverse-mode AD is either slow or does not work at all. Computing those derivatives symbolically, once, ahead of time, side-steps the problem entirely.

The trade-off is that the whole network is unrolled into one expression, so this only makes sense for small networks. Code generation time grows with the size of the network, while evaluation is fast and allocates little.

Installation

using Pkg
Pkg.add("SymbolicNeuralNetworks")

Quickstart

Build a model with AbstractNeuralNetworks, wrap it in a SymbolicNeuralNetwork, and turn symbolic expressions into functions with build_nn_function:

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

c = Chain(Dense(2, 3, tanh), Dense(3, 1, tanh))
snn = SymbolicNeuralNetwork(c)

# the symbolic output of the network, an expression in `snn.input` and `params(snn)`
soutput = c(snn.input, params(snn))

\[ \begin{equation} \left[ \begin{array}{c} \tanh\left( \mathtt{W\_4_1} + \mathtt{W\_3_{1}ˏ_1} ~ \tanh\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\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\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} \]

forward = build_nn_function(soutput, snn)

nn = NeuralNetwork(c)             # the same model, with numeric parameters
forward([1.0, 2.0], params(nn))   # a single sample
1-element Vector{Float64}:
 -0.8881789235085311

The generated function also takes a whole batch, one sample per column:

forward(rand(2, 4), params(nn))
1×4 Matrix{Float64}:
 -0.970402  -0.950386  -0.96885  -0.958311

The derivative with respect to the input is a Jacobian, the derivative with respect to the parameters a Gradient, and both are built into functions the same way:

using SymbolicNeuralNetworks: Jacobian, derivative

jacobian = build_nn_function(derivative(Jacobian(snn)), snn)
jacobian([1.0, 2.0], params(nn))
1×2 Matrix{Float64}:
 0.0313863  0.0642525

For training, SymbolicPullback provides the derivative of a loss with respect to the parameters in the shape an optimizer expects:

using AbstractNeuralNetworks: FeedForwardLoss

pb = SymbolicPullback(snn, FeedForwardLoss())
loss_value, pullback = pb(params(nn), c, (rand(2, 4), rand(1, 4)))
pullback(1).L1.b
3-element Vector{Float64}:
 -0.6940588274660868
  0.37680977182921144
  0.14101037302717292

Where to go next