This is part 2 of a series of blog posts about implementing automatic differentiation. You can read part 1 here. In this post, we extend our automatic differentiation system to support higher order derivatives.

Like the previous post, some knowledge of calculus is required and Racket-y stuff will be explained as we go.

In part 1, we implemented first order automatic differentiation. Our derivative function took in a DNumber representing the output and a DNumber representing the input and returned a plain number representing the derivative of the output with respect to the input. Recall that a DNumber represents the result of a differentiable computation and stores a mapping from each input to the numerical value of its derivative.

1  Higher Order Derivatives

At the end of part 1, I teased higher order derivatives. What are higher order derivatives? In part one, we only took the first derivatives. In math, we can compute the second derivative by taking the first derivative twice. For example, the second derivative of \(x^2\) with respect to \(x\) is \(2\), which can be found by taking the first derivative, getting \(2x\), and taking the first derivative again, getting \(2\).

More precisely, to compute the \(n\)th derivative of \(y\) with respect to \(x\),

\[\frac{d^0y}{dx^0} = y\] \[\frac{d^ny}{dx^n} = \frac{d}{dx}\frac{d^{n-1}y}{dx}\]

where \(n \ge 0\).

Great! We have our derivative function from part 1, so we can just call it \(n\) times, right? No. The signature doesn’t line up. Our derivative function returns a plain number, but the first argument has to be a DNumber. Ok, can we just make derivative return a DNumber then?

Let’s take a step back and think about what would mean. A DNumber is the result of a differentiable computation. A DNumber is created by storing the computation’s inputs and the partial derivative of the output with respect to each input. An operator like multiplication is responsible for recording these inputs and computing these partial derivatives.

Currently, these partial derivatives are stored as plain numbers. And when we compute the derivative, we are adding and multiplying plain numbers together with Racket’s built-in arithmetic operators, not our custom differentiable ones. In order for the result of derivative to be a DNumber, the computation of the derivative itself must be differentiable. This means DNumbers must store partial derivatives as DNumbers instead of plain numbers and derivative must use our differentiable addition and multiplication to combine partial derivatives. This also means that our operators, which construct a DNumber for the result, will also have to construct a DNumber to compute each partial derivative, rather than using plain Racket numbers.

Let’s revise our data definitions:

All we changed was the data type of the derivative field of dchild. It used to be a plain number, and now it’s a DNumber.

Ok, so far so good. Let’s rewrite multiplication:

Let’s compare this to the previous implementation:

Previously, the partial derivative with respect to one input was the numerical value of the other input. Now, we just use the other input directly! This is actually simpler and more straightforward than the old implementation.

Let’s do addition now:

The only change we made was creating dummy DNumbers for the 1s.

2  The Problem

Things start to break down when we try to do this to an operator like \(\frac{1}{x}\):

\[ \frac{d}{dx} \frac{1}{x} = -\frac{1}{x^2} \]

The derivative of the reciprocal operation contains another reciprocal operation. To see why this is a problem, let’s look at the old and new implementations:

Previously, we just use Racket’s built-in / function to compute the derivative. But now, since the computation of the derivative must be differentiable, we have to use our own operators to produce a DNumber. For add and mul, the partial derivatives are simple and don’t involve other operations. But for reciprocal, the partial derivative is another call to reciprocal. This recursion never terminates. This makes sense because you can keep taking the derivative of the reciprocal function and the exponent in the denominator will continue growing. Since the computation of the partial derivatives must be differentiable, they will have partial derivatives as well. If the function can be differentiated arbitrarily many times and continue producing different expressions with more partial derivatives, the computation graph will be infinite and its construction will never terminate.

2.1  Laziness

To avoid this problem, we can be lazy and construct the computation graph on demand, rather than all at once. Remember, to compute first-order derivatives, all we needed was the numerical values of the partial derivatives. Similarly, to compute the second derivative, we only need the numerical value of the partial derivatives’ partial derivatives. We can build the computation graph just as deeply as we need it, and no more.

To make this idea of laziness more concrete, let’s explore a common example of laziness: streams.

A stream is like a list, but it can be infinite, and it can be constructed "on demand". For example, let’s construct a stream containing all integers, starting at a given one:

A brief Racket aside: Racket has linked lists. The function cons takes in an element and a list and returns a list with that element at the beginning of the list. For example:

null is the empty list, also written as (list) and '(). And in the output, we see lists like '(1 2). This is the same as (list 1 2).

The list function we’ve been working with is shorthand for uses of cons and null.

stream-cons is similar to cons for lists, but delays the evaluation of the rest of the stream (the second argument) until its value is needed. stream-take takes in a stream and a number n and returns another stream that produces the first n elements of the given stream. stream->list forces the evaluation of all elements of the stream and converts the stream to a list. If we did (stream->list (integers-from 1)), it would never terminate.

Rather than store the entire stream in memory like a list, a stream stores a function that computes the rest of the stream. However, once part of the stream is computed, it is remembered:

The begin expression prints "hello!" and then returns 2.

The (displayln "hello!") only runs once because the first stream->list forces the evaluation of the whole stream. Once a stream is forced, it remembers the values of its elements, so when you ask for them again, it doesn’t re-compute the begin with the print. This is nice because if, instead of printing, there was a big calculation to compute one of the elements, you wouldn’t want to run that more than once.

This is the kind of behavior we want for our new computation graphs. Like a stream, our computation graph might be infinite, but we don’t need the whole thing at once. And we don’t want to re-run potentially expensive numerical computations if we don’t have to.

We will achieve this laziness using Racket’s promises. From the Racket documentation:

A promise encapsulates an expression to be evaluated on demand via force. After a promise has been forced, every later force of the promise produces the same result.

These are not like JavaScript’s promises, which are used for sequencing asynchronous computations. Racket’s promises are used for delayed computations.

For our purposes, we will be using delay, which creates a simple promise. If you want to learn more about promises, see my blog post on composable promises!

Here is an example showing how delay and force work:

Like a stream’s elements, the value of the promise is computed only when forced, and the result is remembered so subsequent forces don’t re-compute the result.

For our computation graphs, what should we wrap in a promise? In other words, what should only be computed on demand? Sometimes, we don’t care about a DNumber’s derivatives at all. For example, if we just want a DNumber’s value, we don’t need any of its derivatives. But when we are computing the derivative of a DNumber with respect to another, we will look at all of its derivatives. Considering this, it makes sense to wrap the children list of a DNumber in a promise. This means that the only thing that will be eagerly computed when constructing a DNumber is its value. Derivatives will only be computed on-demand. And when they are forced, they’ll all be forced at once, and only their values will be computed eagerly. If we want the first derivative, we won’t even compute partial derivatives of partial derivatives. We’ll only compute the immediate values of the immediate partial derivatives.

3  Lazy Computation Graph

Here are our new data definitions:

Rather than a list of DChildren, we store a promise of a list of DChildren. This achieves the desired laziness. Let’s rewrite multiplication and addition again:

It’s the same as before, except we wrap the children lists in a delay. Now let’s write reciprocal:

Again, we just wrap the children lists in a delay. Importantly, in this case, the delay protects us from an infinite loop. The evaluation of the list of children is delayed until it is forced, so the function terminates without actually making a recursive call. Like the infinite stream of integers, this infinite computation graph is computed on-demand.

Now let’s re-write derivative:

There are a few differences between the old derivative function and this new one: Instead of returning 1, we create a fresh constant DNumber for it. And since dnumber-children is a promise, we must force it to demand the values of the inputs. Additionally, since we’re working with DNumbers instead of plain numbers, we can’t use Racket’s for/sum to add up the partial derivatives. Instead, we use for/fold, which accumulates the sum variable as we loop over the inputs.

Here is an example of using for/fold:

The last difference is that we use our custom add operator to compute the rolling sum and we use our mul operator to compute the chain rule since we’re working with DNumbers now.

Let’s run the same tests as before:

The results are the same, but derivative returns DNumbers now instead of plain numbers, so we use dnumber-value to get the numerical value of the derivative. Now let’s compute some higher order derivatives:

Nice! One cool thing about promises is that when we print them out, we can see whether they’re forced or not:

If the promise is forced, we see its value.

We can use this to see how much of our computation graph has been forced:

We can see that, initially, no derivatives have been computed. But as we take higher order derivatives, more of the compuation graph gets forced as needed.

Now let’s make some more operators, starting with \(e^x\):

Since \(\frac{d}{dx}e^x = e^x\), we can store the result itself as the partial derivative of the result with respect to x. Not only is this pretty cool, but it also has the benefit that even though our computation can be differentiated arbitrarily many times, the computation graph will stay the same size. This is not the case for reciprocal, which grows exponentially with successive derivatives.

Another interesting thing is that our computation graph now has a cycle, but it is actually fine since it will never cause the computation of a derivative to infinitely loop. In particular, our assumption that no computation is an input to itself is not violated since the cycle is in the derivative, not the input itself.

We can see this cycle in the way e^4 gets printed after its children have been forced. That #0= and #0# stuff is what happens when racket prints a value which contains cycles. The #0# is where the result is stored in itself.

Now let’s do exponentiation. Recall the partial derivatives of the exponential:

\[ \frac{\partial a^b}{\partial a} = b a^{b-1} \]

\[ \frac{\partial a^b}{\partial b} = \ln(a) a^b \]

Since the computation of our partial derivatives must be differentiable, we must also have a differentiable operator for \(\ln\).

\[ \frac{\partial \ln(x)}{\partial x} = \frac{1}{x} \]

Again, we use a cycle to minimize the size of the computation graph of the exponential. However, our computation graph still grows with successive derivatives.

Just for fun, let’s implement sine and cosine too.

The definitions are mutually recursive in their derivatives.

Now our little automatic differentiation library supports higher order derivatives.

Let’s recap:

We had automatic differentiation for first order dervatives, but we couln’t do higher order derivatives since the result of a derivative was a plain number with no computation graph. In order to make the derivative return a DNumber, we had to make the computation of the derivative itself differentiable. This led to infinite computation graphs, so we used laziness to compute the computation graph on demand.

With higher order derivatives, you can do some crazy stuff. For example, you can do meta-gradient descent, where you use gradient descent to optimize the hyperparameters of gradient descent. This implementation is a little too inefficient for that to be practical, but it is possible.

The source code for the full "library" can be found here. In there, I have an example of using automatic differentiation for machine learning.