Study notes: automatic differentiation (part 3)

The side quest into automatic differentiation continues! In my previous post, I attempted a different design direction, and landed on something promising. Instead of representing functions using expressions, I introduced a type, V, which “augmented” floating point values to carry both their value, x, together with the evaluation of their derivative, dx.

After adding a battery of math functions to the type V, we could then do the following:

let f (x: V) = sin (x + 42.0)
let f' = diff f
f' 10.0
> val it: float = -0.1629907808

While this felt like an improvement, this has a major issue. We can get the value of f', which is great, but we can’t derive further. If we were for instance interested in the second order derivative, f'', we would be out of luck. So let’s see if we can fix that.

The root of the problem

The reason we can’t differentiate repeatedly boils down to this:

• to apply the diff function we need to pass in a function V -> V,
• the diff function returns a function float - float.

To get out of this pickle, we would need to transform the diff function into something like this:

diff: (V -> V) -> (V -> V)

So how do we do that?

My initial thought was to try a simple modification of my original implementation of diff. Instead of

// original implementation
let diff (f: V -> V) (x: float): float =
    f (V (x, 1.0))
    |> fun v -> v.DX

… I could try and keep things operating on the type V, perhaps like so:

// attempt: keep things "in V"
let diff (f: V -> V) (x: V): V =
    f (V (x.X, 1.0))
    |> fun v -> V (v.DX, ???)

However, this clearly won’t work. After the first derivation, we end up with a function f' that simply returns a constant. We lost all information about the original function itself at that point, and as a result there is nothing we can use to compute the next order derivative. It’s also not clear how to set the dx part of the result (marked ??? above) without computing the next order derivative and opening an infinite loop.

Looking into dfx

After much tinkering, I ended up taking a shortcut, and started looking into the dfx code base, a recent effort to re-invigorate the development of DiffSharp, the OG F# automatic differentiation library.

That code base is dense. However, after some digging, I ended up finding the piece I was missing, in the Tensor.fs file. DiffSharp / dfx is more ambitious than this little side-quest project of mine. It supports functions over tensors (Tensor -> Tensor), whereas I am at the moment just interested in very basic float -> float functions. So I proceeded to strip down the dfx code to fit my much simpler use case. After removing all the parts I did not need, this ended up translating into this modified version of V:

type V =
    | X of float
    | DX of V * V
    with
    member this.x =
        match this with
        | X x -> x
        | DX (x, _) -> x.x
    member this.dx =
        match this with
        | X _ -> X 0.0
        | DX (_, dx) -> dx

You can find the finished code of this version on Codeberg.

Instead of the simple struct we had previously, V is now a recursive discriminated union: a value V is either:

• a plain constant X, carrying its value a float,
• a value DX, carrying its value and derivative, both of type V.

In other words, V is now a tree-like structure, where each node is either a leaf X, or a branch DX with 2 branches, the value and the derivative. We will work through an example in a bit, but, at a high level, this is promising: a DX node now carries the full chain of calculations, which means that we should be able to carry on differentiating.

I added for convenience 2 methods, x, which returns the value of the node as a float, and dx, the V value of its derivative. Armed with that, we can proceed in a fashion similar to the previous episodes, and bolt on some math functions. Let’s start with addition:

type V =
    | X of float
    | DX of V * V
    with
    // omitted for brevity
    static member (+) (v1: V, v2: V) =
        match v1, v2 with
        | X x1, X x2 -> X(x1 + x2)
        | X x1, DX (x2, dx2) -> DX(X x1 + x2, dx2)
        | DX (x1, dx1), X x2 -> DX(x1 + X x2, dx1)
        | DX (x1, dx1), DX (x2, dx2) -> DX (x1 + x2, dx1 + dx2)

When adding 2 nodes, we need to distinguish 3 cases.

• If we have 2 constants X, the result is still a constant,
• If we have 2 DXs, both their values and derivatives are added,
• If we add a constant X to a DX, their values are added, and, because the derivative of the constant term is 0, we can ignore it.

Note that while X carries a float, the DX branches are of type V, so when mixing X and DX terms, we re-wrap the value from the X term inside the result, like so: | X x1, DX (x2, dx2) -> DX((X x1) + x2, dx2).

We’ll skip the other operators, which follow the same pattern, and take a look at the sin function. Because its derivative is the cos function, we’ll need that as well, and do them both at once:

type V =
    | X of float
    | DX of V * V
    with
    // omitted for brevity
    static member Cos(v: V) =
        match v with
        | X x -> X (System.Math.Cos x)
        | DX (x, dx) -> DX (V.Cos(x), -1.0 * V.Sin(x) * dx)
    static member Sin(v: V) =
        match v with
        | X x -> X (System.Math.Sin x)
        | DX (x, dx) -> DX (V.Sin(x), V.Cos(x) * dx)

Again we need to distinguish 2 cases.

• The sin of a constant X is still a constant,
• The sin of a DX value has 2 parts: the value, and its derivative, computed by applying the chain rule.

Note how, while the X case uses System.Math.Sin to compute the sin of the floating value constant, the DX case uses V.Sin and V.Cos, so the results are still of type V.

Working through an example

Let’s see how this works out on the same function we used in previous episodes, f (x: V) = sin (x + 42.0). When evaluating the derivative at some value x, we will push the value x forward, with a derivative of 1, and let it bubble through the computation graph. Step by step, this is how it works out:

The derivative part of the final node is a value V which carries the full computation graph, so we should be able to differentiate that again.

Does it work?

At that point we are more or less done. All we need is to rewrite our diff function, passing in a derivative of 1:

let diff (f: V -> V) =
    fun (x: V) ->
        DX(x, X 1.0)
        |> f
        |> fun fx -> fx.dx

So… Does it work?

let f (x: V) = sin (x + 42.0)
f (X 10.0)

> val it: V = X 0.986627592

let f' = diff f
f' (X 10.0)

> val it: V = X -0.1629907808

So far, so good, same results as last time. Now for the real test:

let f'' = diff f'
f'' (X 10.0)

> val it: V = X -0.986627592

Victory!

Parting thoughts

First, I am very excited to see dfx bring new energy to DiffSharp. DiffSharp is beautiful, but had been dormant for a while, I hope dfx will take that foundation and build on it!

I had a lot of fun digging into the dfx codebase. This is something I had intended to do for a while, but I always felt a bit intimidated. Building a toy project like Galet first helped me a lot. By working through a similar problem on my own, unassisted, I gained a decent mental picture of the domain, which in turn helped my understanding of what I was looking at in dfx, and where to even look. We all have our own ways of working through problems, and personally I have found multiple times that trying to implement my own naive solution first is a great way to gain a better appreciation for the design choices of existing libraries.

I struggled quite a bit to figure out how to modify the V type to handle repeated differentiation. I had been dancing around the right solution for a
couple of days when I finally decided to “cheat” and look into how DiffSharp / dfx did it. In hindsight, the solution felt rather obvious. I had the right questions (I knew I was missing the computation graph, and being forced to pass a value with a “fake” derivative to call a function felt off), but couldn’t get unstuck. This might be because the change needed was a relatively small modification. Often, I find it easier to try out a complete rewrite, whereas I am more “shy” making gradual changes to an existing design.

One thing that caught me by surprise is that this version is slightly slower than my first take, where I represented functions as expressions and performed symbolic derivation. Granted, I haven’t spent any effort optimizing either of them, but still. I suspect that part of the speed improvements I saw in the previous post were due to using a struct, which I can’t do now that I have a recursive data structure. I am vaguely wondering if I could somehow rewrite that tree, flattening it as a list.

Anyways, this feels like real progress! I’ll stop here for now. Next step: partial differentiation, and / or functions of multiple variables.

You can find the finished code of this version on Codeberg.