Automatic Differentiation

We can produce symbolic derivatives by using symbolic numbers if we’re willing to prepend:

(I’ve omitted type annotations in pursuit of minimalism, thus to avoid specialization to Double we disable the dreaded monomorphism restriction.) If done, then in GHCi, simply evaluate the derivative at the symbolic number x, and evaluate its linear map at 1:

We find the derivative of \(x e^x\) is:

Some elementary functions require Data.Complex and identities such as \(\arccos(z) = -i \ln (\sqrt{z^2 - 1} + z)\). It’s friendlier to define shortcuts:

Let’s do more examples. Suppose we wish to derive:

We convert to point-free form:

Then find their symbolic derivatives with:

Answers:

One more. Let’s confirm an example from Wikipedia:

We type:

and find the entries of the Jacobian matrix are indeed:

Twiddling each identifier is irksome. A more ergonomic solution that suits Haskell is to exploit typeclasses so something like:

represents the function \(x \mapsto x^2\), while:

represents its derivative, namely the function computing \(x \mapsto (\epsilon \mapsto 2 x \epsilon)\).

To achieve this, we’d overload mul so it’s our original mul as a Function but dMul as a Derivative. Similarly for the other building blocks such as (.) and id.

These last two deserve special mention. Principled overloading of (.) and id is known as category theory. This is a great boon, for we can stand on the shoulders of mathematical giants to view differentiation in a whole new light.

Instead of typeclasses, we persevere with renaming identifiers by hand. This gives us a cut-away view of category theory magic and also makes us appreciate typeclasses; absence makes the heart grow fonder.

For some applications, computing derivatives is most efficient using a method called reverse-mode automatic differentiation (RAD), which means we write our function so that all sequences of compositions are left-associative. For example:

This is fine for small examples, but grows tiresome for larger functions. Furthermore, we only care about left-folding compositions for the linear maps (derivatives), not for the function itself (the "primal").

Can we force left-associativity for the linear maps without explicit parentheses everywhere? Yes! Continuation-passing style (CPS) comes to our rescue.

Once again, with the right set of GHC extensions:

represents the derivative of \(x \mapsto x^2\) computed via RAD. And once again, we’re forgoing typeclasses, so we’re forced to employ tedious text substitution instead. Define:

To force our derivatives to be computed in reverse-mode, we apply rad to each building block and replace (<.) with (<<.). This leads to a problem with dCross, which does not expect arguments in CPS form; we fix this below.

The linear map is now represented as a continuation, so should be applied to id before use. For example, the derivative of \(f(x) = \sin (\cos x^2)\) via RAD is:

and we can compute \(f'(2)\) with:

Time to fix dCross. We need more machinery for its RAD variant:

For example:

In practice, the co-domain of each continuation is often the scalar field, that is, we wish to differentiate a function that utlimately spits out a single number.

For example, in neural networks, we define a cost function that has many parameters (the weights and biases) and outputs a single number. We differentiate to find the gradient, which tells us how best to tweak the parameters to reduce the cost. Namely, we perform gradient-based optimization. When computed via RAD, this is known as backpropagation.

It turns out we can simplify this special case because the linear maps from a vector space to a scalar field are isomorphic to that vector space; they are dual. We call this isomorphism and its inverse dot and undot for the 1-dimensional case. We also define 2-dimensional variants. With clever typeclasses, we could avoid inventing similar names and carefully matching the isomorphisms with various function types; the compiler would do all this for us.

Sandwiching a function of type ((b -> s) -> (a -> s)) between undot and dot yields a function of type b -> a.

Composition is the same as for the CPS form, so we can reuse (<<.) in our examples.

Ideally, compositions of inverses such has dot . undot would be optimized away. GHC can do this with rewrite rules.

In machine learning, a single neuron of a neural net is a bunch of weights and a bias that are typically represented with floats.

To compute the output of a neuron, we multiply each input with its corresponding weight, add the bias, sum it all up, then feed the result through some activation function f:

Training requires the derivative of this function with respect to each weight and bias. (We’ve simplified a little: we actually want the derivative of something like the square of the difference between the output of this function and the desired output.)

We could organize the weights into nested pairs to fit the function into our framework, but we’re better off supporting a vector of weights directly, because in practice there are many of them. For our toy demo we use lists. Since we’re shunning typeclasses, we invent yet more names for the list versions of jam and cross:

We can express firing a neuron as follows:

as well as its first derivative:

We can compute in reverse-mode with a few more mechanical code transformations:

For example:

computes the gradient at the input [0.2,0.7,0.1] for the weights [3,1,4] and bias -2:

This technique scales up for layers of neurons, but without typeclasses, we’d have to juggle closely related building blocks, painstakingly ensure the right type of block is in the right place.

How does this all work?

Our first lines of code are based on Conal Elliott’s blog post, who cites a paper by Jerzy Karczmarczuk. For the rest of our code, see Conal Elliott’s talk on automatic differentiation and category theory and his corresponding paper.

You may remember the CPS trick from computer science classics such as Reynolds' paper on definitional interpreters, but in fact the idea can be traced to an 1854 paper by Cayley.

In Haskell, sequences of (++) are fastest when evaluated in "forward-mode". Joachim Breitner explains how to transform to CPS to achieve this with an approach that resembles ours.

Our choice of (<.) was inspired by the flow package.

Conal Elliott laments that automatic differentiation is "typically presented in opposition to symbolic differentiation", which is clearly at odds with our examples above!

To avoid confusion, I propose renaming "symbolic differentiation" to "schoolbook differentiation" or perhaps "shambolic differentiation", because often it seems to refer to the centuries-old methods of differentiation taught in school, where expressions are small enough that sharing saves little, and has nothing to do with whether the answer is ultimately symbolic or numeric.

Less grating are the terms "forward-mode" and "reverse-mode" for what functional programmers might call right and left folds. The reverse function of Haskell98 is a left fold, a lovely coincidence stemming from the right-folded nature of Haskell lists.