Automatic Differentiation: From Forward to Reverse in Small Steps
Tags: algorithmicdifferentiation, differentiableprogramming, rust, deeplearning
Published: 2022-09-24 08:48:19


---
title: Automatic Differentiation: From Forward to Reverse in Small Steps
published: true
date: 2022-09-23 12:40:43 UTC
tags: algorithmicdifferentiation,differentiableprogramming,rust,deeplearning
canonical_url: https://getcode.substack.com/p/automatic-differentiation-from-forward
---
**An in-depth explanation of differentiable programming, for programmers**

I associate derivatives with practicing differentiation rules on increasingly implausible looking functions. If you have similar memories, you'll remember it's a mechanical, rather boring process. Boring for humans is good for computers though, because that means it's amenable to automation. So automate it we will, using a technique called Automatic Differentiation (AD). What's surprising about AD is that it generalizes so well: AD not only calculates the derivative of mathematical functions, but also of programs using constructs like conditionals, loops, recursion and higher-order functions. To highlight this, people have coined _differentiable programming_ as an umbrella term for the various uses and applications AD enables.

![Impressive looking blue neuron squid thing generated by stable diffusion](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/jd5egrxfzmyypxyiafjg.png)

_What a neuron looks like according to Stable Diffusion._

AD has become a widely used technique, as it plays a pivotal role in deep learning. Gradient descent-based optimization - how the billions of weights in neural networks are trained - needs derivatives, and AD calculates them extremely efficiently. It's up to 6 or 7 orders of magnitude faster than naïve approaches. That's the difference between one day to train a model and 2,738 years! With the current explosion of interest in deep learning, AD is a hot research topic as well. That said, it has many other applications, notably computer graphics, and risk analysis in finance.

This post explains how AD works in-depth, both its forward and somewhat baffling reverse mode. It is aimed at programmers - I'll assume you remember some notions of what differentiation is and some rules, but will link to further information to explore at your leisure. Examples are in Rust, so some experience with typed, functional languages is helpful but not necessary. The implementation does not rely on Rust-specific features.

All code is [on GitHub](https://github.com/kurtschelfthout/minidiff).

## What are derivatives, really?

Derivatives are typically explained in visual terms using graphs of functions. For a scalar function ℝ→ℝ, the derivative at a particular point is represented by the tangent of the function at that point. This interpretation builds useful intuition, and I recommend [3blue1brown's excellent visual series](https://www.3blue1brown.com/topics/calculus) if you're unfamiliar or need a refresher. The visual interpretation works for functions that are two or three dimensional - for example a function 𝑓:ℝ²→ℝ describes the height of an imagined terrain as a function of (x,y) coordinates. The derivative of such a function has two components, one for each input: ∂𝑓/∂𝑥 and ∂𝑓/∂𝑦. These are called the _partial derivatives_, and together describe the slope of the terrain at a particular (x,y) coordinate.

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/u2f4oe5ppqoi823zc8bl.png)

_A tangent line at x=3_

However, no one can imagine a four-dimensional space, let alone [GPT-3](https://www.wikiwand.com/en/GPT-3)'s 175 billion-dimensional space of weights. To deal with that, another perspective is that the derivative is a measure of the sensitivity of a function to its inputs. (I'll use the terms inputs and outputs of functions rather than parameters or domain, because it's more familiar to programmers.)

This view is common in risk analysis. Say you have a function that, based on various observable values, calculates today's value of an investment. The observables could be interest rates, inflation rates, foreign exchange rates, anything that affects the value of the investment. To gauge how risky the investment is, we'd like to know how its value changes if the observables change. If the value of the investment changes significantly for a small change in say interest rate, then that indicates it is particularly sensitive to interest rates. If we don't like such a high exposure to interest rates, we may want to hedge this risk.

It turns out that the derivative is exactly a measure for the sensitivity of the output(s) of a function to each of its inputs. More accurately, for a function 𝑓: ℝ³→ℝ, i.e. a function with three inputs and one output, the derivative function 𝑑𝑓: ℝ³→ℝ³ has three partial derivatives at each input value. The partial derivative is a measure of how the value of 𝑓 changes, if the corresponding input value changes - in other words, how sensitive 𝑓 is to changes in the value of each of its inputs. What's neat about partial derivatives is that you can linearly combine them to get the sensitivity to several inputs as well (for well-behaved functions, at least). The partial derivatives are exactly the quantities you need to figure out everything there is to know about the sensitivity of the function to input changes.

But how is this measure interpreted and where does it break down? It turns out that the derivative is just an approximation. To be exact, the derivative is [the best linear approximation](https://math.stackexchange.com/questions/3551202/proof-that-derivative-is-the-best-linear-approximation) of the function.

In practice linear approximations of non-linear functions are only good close to the point where they are evaluated. If our investment is worth $1000 today, at an interest rate of 1%, and sensitivity to interest rate is $10 per 0.01% _at 1% interest rate_, then a reasonable estimate of the value of the investment at 1.01% interest rate is $1010. However it is just an estimate, and the further away from the interest rate of 1%, the more inaccurate the estimate will be. The exception to this is if the function is itself linear - in that case a linear approximation is exact.

To change tack, having rid ourselves of the necessity to imagine 15-dimensional landscapes, we can now also understand how deep learning networks are trained through an optimization process called gradient descent. A neural network can be understood as a machine that has potentially billions of parameters, called weights, which need to be adjusted so that the machine produces the desired outputs. During training, we give the machine an input and it produces an output. We compute the error - the difference between the machine's output, and the desired output. We'd like to make this error as small as possible. Because the error is a function of the weights of the machine, we can calculate how sensitive the error is to small changes in each of the weights by calculating the partial derivative of the error function to all the (potentially billions) of weights. Once we know how the error changes with the weights, we can adjust the weights to make the error smaller - if the sensitivity is positive, we'll make that weight a bit smaller, if it's negative, we'll make it a bit bigger. Lather, rinse, repeat, and before you know it you can [generate some amazing images](https://lexica.art/).

## What differentiable programming means

Let's start with a simple example - a function that takes and returns a single floating point number.

```rust
fn f(t: f64) -> f64 {
    t.powi(2) + t + 1.0
}
```

Let's say we're interested in having a function `df` which represents the derivative of `f`.

How can we figure out `df` from `f` automatically? Besides automatic differentiation, there are two other approaches: symbolic and numeric differentiation. I'll briefly discuss them to better understand what makes AD so effective.

### Symbolic differentiation

An exhaustive [set of rules](https://www.wikiwand.com/en/Differentiation_rules) exists to calculate `df` from `f`, so by applying these rules or [asking something smarter than me](https://www.wolframalpha.com/input?i=derivative+of+t%5E2+%2B+t+%2B+1), we can write `df` directly:

```rust
fn df_sym(t: f64) -> f64 {
    2.0 * t + 1.0
}
```

This approach is called symbolic differentiation. It's easy to do for this simple example, but with common programming constructs like shared variables, iteration, recursion and conditionals, it can be very hard to come up with an efficient closed formula for the derivative.

### Numeric differentiation

An alternative is numeric approximation, which follows from the [definition](https://www.wikiwand.com/en/Derivative#/Definition) of derivative:

```rust
fn df_num(t: f64, h: f64) -> f64 {
    (f(t + h) - f(t)) / h
}
```

For some small number `h`. Numeric differentiation, while simple to implement, is not without problems: choosing an `h` that is just right, and doesn't lead to numerical instabilities is tricky. There are approaches that solve some of these problems, but simplicity is lost. Also this needs two evaluations of `f`, which is not ideal.

Comparing symbolic and numeric differentiation:

```rust
let t = 5.0;
println!("Symbolic: {}", df_sym(t));

let h = 0.000001;
println!("Numeric: {}", df_num(t, h))
```

Results in:

```
Symbolic: 11
Numeric: 11.000001002514637
```

### Automatic differentiation

The idea of automatic differentiation is to compute the derivative at the same time as the function itself, by applying differentiation rules one operation at a time. Let's do that manually first.

We're aiming for a function `f_ad(t_dual: (f64, f64)) -> (f64, f64)`. The argument `t_dual` is a tuple with the value of `t` as the first element, and its derivative as the second element. The combination of these two is typically called a _dual number_. We need `t_dual` to be a dual number because the input to `f_ad` might be the result of another automatically differentiated function. The result of `f_ad` is also a dual number: it returns the value `f` would normally compute, as well as its derivative.

Here's the implementation in AD style:

```rust
fn f_ad(t_dual: (f64, f64)) -> (f64, f64) {
    let (primalt, tangentt) = t_dual;
    let (primal0, tangent0) = (primalt.powi(2), 2.0 * primalt * tangentt);
    let (primal1, tangent1) = (primal0 + primalt, tangent0 + tangentt);
    let (primal2, tangent2) = (primal1 + 1.0, tangent1);
    (primal2, tangent2)
}
```

Each operation in the original `f` is on a separate line. On each line, the regular value, the _primal_, is updated, along with the _tangent_, which represents the derivative. The calculation of the tangent is the result of applying differentiation rules, most importantly the chain rule. It's worth dwelling on the [the chain rule](https://www.wikiwand.com/en/Differentiation_rules#/The_chain_rule) for a bit, because it's essential to AD.

The chain rule tells us how to calculate the derivative of function composition - a fancy name for passing the result of a function `e` into a function `f`, i.e. `f(e(t))`. This is written as 𝑓 ∘ 𝑒, where ∘ denotes function composition. (I used `e` here instead of the more typical `g`, because e comes before f). The chain rule is used when calculating `tangent0`.

The first component of `tangent0` is `2.0 * primalt`, because of the [power rule](https://www.wikiwand.com/en/Differentiation_rules#/The_polynomial_or_elementary_power_rule). Then that's multiplied by `tangentt` because of the chain rule, which comes into play because `t_dual` may be the result of applying another function `e`. `t_dual`'s tangent is the derivative of `e(t)`. `tangent0` then needs to determine the derivative of `e(t).powi(2)`, in other words the composition of `e` and `powi`. Hence the chain rule tells us to derive `powi`, and then multiply that by the derivative of `e`, which is `tangentt`.

Note that calculating the primal at each step only needs previously calculated primal values - this seems obvious but is not true for the tangent, which needs both primal and previous tangent values. `tangent0` for example uses the value `2.0` from its own primal, as well as `primalt` and `tangentt`.

Finally, here's how to call `f_ad`:

```rust
println!("Automatic: {}", f_ad((t, 1.0)).1);
```

Why is the initial tangent 1.0? Because that's the derivative of the identity function 𝑖𝑑(𝑥)=𝑥.

This outputs:

```
Automatic: 11
```

Have a go yourself at the [Rust playground](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=3575c2c0ce68498e6f1528855c364e7c).

### Towards automatic differentiation

The next sections develop a minimal implementation of AD. It proceeds in steps: first, we make the manual implementation of AD actually automatic. Then, we'll investigate the efficiency of the implementation, and iteratively improve it. Each step reveals an insight into how and why AD works.

Generally, AD can be implemented in two ways: source to source transformation, or via operator overloading. Source to source transformation analyses the source of the original program `f` and generates the code for `f_ad`. Depending on the constructs the transformation supports, it is entirely transparent to the programmer. However, it needs a lot of machinery for parsing and producing efficient code, which is not relevant to the core ideas of AD. I'll use operator overloading here, but the ideas are transferable.

## Actually automatic through operator overloading

First let's refactor the tuples into a proper type of dual numbers: `Dual`. The idea is to define operators like sum and product on `Dual` using operator overloading, so that we can write programs that work with dual numbers as if we're dealing with floating point numbers. Except now functions compute the value and derivative at the same time! This process is somewhat similar to introducing a library for complex or rational numbers. Here's the type for dual numbers:

```rust
struct Dual {
    primal: f64,
    tangent: f64,
}
```

Primitive `Dual` values are `constant` and `var`:

```rust
impl Dual {
    fn constant(value: f64) -> Self {
        Dual {
            primal: value,
            tangent: 0.0,
        }
    }

    fn var(value: f64) -> Self {
        Dual {
            primal: value,
            tangent: 1.0,
        }
    }
}
```

Both `constant` and `var` have a user-defined primal value, and only differ in their derivative. `constant` represents the constant function f(x) = C. [Constant functions have a derivative that's 0 everywhere](https://www.wikiwand.com/en/Differentiation_rules#/Constant_Term_Rule). `var` represents the identity function f(x) = x, with [a derivative that's 1 everywhere](https://www.wikiwand.com/en/Differentiation_rules#/The_polynomial_or_elementary_power_rule). (Why call it `var` and not `id`? Because typically it's used to define new variables we'd like to calculate the derivative of.)

Next is overloading addition and multiplication for `Dual` by using the [sum](https://www.wikiwand.com/en/Differentiation_rules#/Differentiation_is_linear) and [product](https://www.wikiwand.com/en/Differentiation_rules#/The_product_rule) rules in `tangent`:

```rust
// Add trait omitted
fn add(self, rhs: Self) -> Self::Output {
    Dual {
        primal: self.primal + rhs.primal,
        tangent: self.tangent + rhs.tangent,
    }
}

// Mul trait omitted
fn mul(self, rhs: Self) -> Self::Output {
    Dual {
        primal: self.primal * rhs.primal,
        tangent: rhs.tangent * self.primal + self.tangent * rhs.primal,
    }
}
```

And also define `powi` for `Dual`:

```rust
impl Dual {
    fn powi(self, exp: i32) -> Self {
        Dual {
            primal: self.primal.powi(exp),
            tangent: f64::from(exp) * self.primal.powi(exp - 1) * self.tangent,
        }
    }
}
```

The example function `f` now becomes:

```rust
fn f_ad_dual(t: Dual) -> Dual {
    t.powi(2) + t + Dual::constant(1.0)
}

// call site:
f_ad_dual(Dual::var(3.0))
```

It looks like the original function `f`, with a bit of syntactic noise. Some languages allow defining new literals, or implicit conversions, which allows getting rid of the call to `Dual::constant`. Another way to achieve that would be to define overloads for all operators for `Dual` and `f64` in all possible combinations. In [unityped](https://existentialtype.wordpress.com/2011/03/19/dynamic-languages-are-static-languages/) languages with operator overloading, there isn't even be a type signature to change, so `f` may not need to change at all.

I'll be using graphs throughout this post to illustrate the calculation of the tangent, which is straightforward now but will become more involved. Here's a graph representing the calculation of `df(3)`, i.e. calculation of `Dual.tangent` at t=3:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/py6dpybnr4vh8963yg5r.png)

On the left is the tangent of the input `t`, shown as Δ𝑡. Edges scale the tangent with a factor, in this case because of the chain rule. If no factor appears on the edge, it is equal to one. At each `+` node, tangents are added.

Here's how the calculation proceeds - from left to right, a detail that will become important.

{% embed https://youtu.be/JguTRH_O1Fo %}

### Intermezzo: Dual numbers

Mathematically, [Dual numbers](https://www.wikiwand.com/en/Dual_number) are expressions of the form a+bϵ where a and b are real numbers, and ϵ is a "small quantity" such that ϵ²=0. Think of complex numbers but ϵ instead of 𝑖.

What's interesting about dual numbers is that we can think of the first part of the dual number as the primal, and the second part as the derivative, and end up with rules for differentiation that are the same as the rules for symbolic differentiation. For sums and products this needs only simple arithmetic:

- Sum: (a+bϵ) + (c+dϵ) = (a+b) + (c+d)ϵ.

- Product: (a+bϵ)(c+dϵ) = ac + adϵ + bcϵ + bdϵ² = ac + (ad + bc)ϵ since ϵ²=0.

Compare this with the implementation of `Dual` for sum and product: the first part of the dual number corresponds to the primal, the second part to the tangent. A [chain rule with a similar shape](https://www.wikiwand.com/en/Dual_number#/Differentiation) can be derived as well.

## Generalizing dual numbers

We left the previous section with a `Dual` type that consists of two floating point numbers. However, the derivative of a function with multiple inputs has multiple components: the partial derivative with respect to each input. This means that a single float may not be enough for the `tangent` field. To address that, let's refactor `Dual`. At this point some of this might seem abstraction for the sake of abstraction, but it will come in handy later.

First redefine `Dual` as follows:

```rust
struct Dual<T> {
    primal: f64,
    tangent: T,
}
```

The `primal` is still a `f64` value, but `tangent` is left generic. We can recover our old `Dual` by instantiating `T` as `f64`, so this is strictly a generalization. However, just any `T` won't do - we must be able to apply differentiation rules to it.

It turns out that to calculate derivatives, `T` only needs a zero element, an addition and scale operation - which corresponds to `T` forming a [vector space](https://www.wikiwand.com/en/Vector_space). I don't know how to explain intuitively why this is true! But you can check the differentiation rules to see that it is. An accurate statement and proof can be found on the [Linearity of differentiation](https://www.wikiwand.com/en/Linearity_of_differentiation) Wikipedia page.

Anyway, trusting the math hasn't failed me yet, so we can represent this requirement on `T` by a Rust trait:

```rust
trait VectorSpace {
    fn zero() -> Self;
    fn add(&self, rhs: &Self) -> Self;
    fn scale(&self, factor: f64) -> Self;
}
```

And implement `VectorSpace` for floats:

```rust
impl VectorSpace for f64 {
    fn zero() -> Self {
        0.0
    }

    fn add(&self, rhs: &Self) -> Self {
        self + rhs
    }

    fn scale(&self, factor: f64) -> Self {
        self * factor
    }
}
```

Now we can overload operators on `Dual<T>` for any `T` which is a `VectorSpace`:

```rust
impl<T: VectorSpace> Dual<T> {
    fn constant(primal: f64) -> Self {
        Dual {
            primal,
            tangent: T::zero(),
        }
    }

    fn add_impl(&self, rhs: &Dual<T>) -> Dual<T> {
        Dual::new(self.primal + rhs.primal, self.tangent.add(&rhs.tangent))
    }

    fn mul_impl(&self, rhs: &Dual<T>) -> Dual<T> {
        Dual::new(
            self.primal * rhs.primal,
            rhs.tangent
                .scale(self.primal)
                .add(&self.tangent.scale(rhs.primal)),
        )
    }
}
```

I'm omitting the ceremonial implementations of the `Add` and `Mul` traits. It's now possible to write `f` with our new `Dual<T>` type:

```rust
fn f<T: VectorSpace>(x: &Dual<T>) -> Dual<T> {
    t.powi(2) + t + Dual::constant(1.0)
}

fn df(t: f64) -> f64 {
    let res = f(&Dual::new(t, 1.0));
    res.tangent
}
```

`Dual<T>` can now be used like `Dual` could, and we seem to be back where we started. We have gained some power though - let's actually leverage it next.

## Forward mode and Reverse mode

AD comes in two flavours, called modes: _forward mode_ and _reverse mode_. These two modes can also be mixed. So far we've implemented forward mode, which is the easiest to understand. It's called forward mode because the computation of the derivative flows in the same "forward direction" as the evaluation of the main program. In reverse mode, part of the calculation will happen in the opposite direction - we'll build up to how exactly that works. For the moment, let's consider our current implementation more in depth, to understand what its limits are.

For functions that have a single input, but many outputs, forward mode works efficiently. More concretely, calculating the derivative only adds a constant factor overhead to each operation - a few more additions and multiplications. Here's an example of a function with multiple outputs, and how we can get its derivatives:

```rust
fn f_2out<T: VectorSpace>(x: &Dual<T>) -> (Dual<T>, Dual<T>) {
    let ref y = x.sin() * x.sin();
    (y + x * Dual::constant(10.0), x + y * Dual::constant(20.0))
}

fn df_2out<T: VectorSpace>(x: f64) -> (f64, f64) {
    let (dx1, dx2) = f_2out(&Dual::new(x, 1.0));
    (dx1.tangent, dx2.tangent)
}
```

`df_2out` returns the derivative of each of its outputs with respect to `x`, and needed just one call to `f_2out` to achieve that. There are two derivative values now, but forward mode AD can calculate both of them with a constant factor overhead.

This is great news, however a function with multiple inputs tells a different story.

```rust
fn f<T: VectorSpace>(x: &Dual<T>, y: &Dual<T>) -> Dual<T> {
    x * y
}

fn df_v1(x: f64, y: f64) -> (f64, f64) {
    (
        f(&Dual::new(x, 1.0), &Dual::constant(y)).tangent,
        f(&Dual::constant(x), &Dual::new(y, 1.0)).tangent,
    )
}
```

Because we need to calculate the partial derivative to each of the two inputs separately, we need two calls to `f`. This is bad for efficiency: for gradient optimization problems in deep learning with billions of inputs, we'd need billions of calls to the function! This approach will certainly not do.

Let's visualize what's going on. Here's a graph representing the calculation of the derivative of `df(3, 2)`, i.e. calculation of `Dual.tangent` at x=3 and y=2:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/pbqjz5f0vqw88dc98fro.png)

On the left are the tangents of `x` and `y`, shown as Δ𝑥 and Δ𝑦. The factors on the edges from Δ𝑥 and Δ𝑦 are because of the product rule: (𝑥𝑦)′ = 𝑥𝑦′ + 𝑦𝑥′. This shows visually why we had to evaluate twice, once with Δ𝑥=1 and once with Δ𝑦=1: a single evaluation of `f` with both Δ𝑥=1 _and_ Δ𝑦=1 would give the _total derivative_. The total derivative is the sensitivity of the function when all its inputs change at once; but we're not interested in that. We need the partial derivatives, so we get the sensitivity to each input individually.

## First fix: many calls to one call

If you've ever dealt with a function that has significant overhead on each call - say, it accesses a database - then you'll have learnt that it's not a good idea to call that function multiple times in a loop. It's much better to pass all the inputs at once, and re-implement the function to deal with multiple values more efficiently by accessing the database just once.

We can apply a similar idea here. We're calling `f` multiple times, each time with different tangent values. Instead we can make the tangent `T` an array, to pass them all at once into a single call.

This is where the `Dual<T>` abstraction pays off. No need to change the `Dual` type - just implement `VectorSpace` for arrays:

```rust
impl<T: Copy + VectorSpace, const N: usize> VectorSpace for [T; N] {
    fn zero() -> Self {
        [T::zero(); N]
    }

    fn add(&self, rhs: &Self) -> Self {
        let mut result = [T::zero(); N];
        for (i, (l, r)) in self.iter().zip(rhs).enumerate() {
            result[i] = l.add(r);
        }
        result
    }

    fn scale(&self, factor: f64) -> Self {
        let mut result = [T::zero(); N];
        for (i, v) in self.iter().enumerate() {
            result[i] = v.scale(factor);
        }
        result
    }
}
```

Arrays are added element by element (point-wise), and scaling applies to each element.

Keeping `f` the same, but changing `df`:

```rust
fn df_v2(x: f64, y: f64) -> [f64; 2] {
    f(&Dual::new(x, [1.0, 0.0]), &Dual::new(y, [0.0, 1.0])).tangent
}
```

This works out better: the function is only called once. The following animation visualizes the computation now:

{% embed https://youtu.be/jDdkNsYG0CU %}

Evaluation still proceeds forwards through the graph, as before. The array representation allows us to evaluate the function only once, and still get all the partial derivatives.

Alas: memory usage for the arrays is quadratic in the number of inputs! That's workable for a low number of inputs, but is intractable for billions.

## Second fix: introduce an intermediate representation

Those arrays certainly look sparse. Perhaps a sparse encoding is a way forward? Unfortunately, that won't work, because as the computation proceeds they don't remain sparse for long, as the visualization shows. Instead, let's use a related idea, and replace the arrays with an intermediate representation. This splits the calculation in two passes: the first builds an intermediate representation, the second uses the intermediate representation to compute the derivatives.

For the intermediate representation, we introduce a new `Delta` type:

```rust
enum Delta {
    Zero,
    Var(usize),
    Scale(f64, Rc<Delta>),
    Add(Rc<Delta>, Rc<Delta>),
}
```

If you're not familiar with Rust: `Delta` is a union type with four possible cases. `Rc` is a reference counted smart pointer - for this post it's safe to just squint and pretend the `Rc`s are not there.

`Delta` is a reified representation of `VectorSpace` operations - and captures the structure of the derivative computation as a directed acyclic graph. `Var` represents an input variable like x or y; the argument is a unique number to identify the variable.

To use `Dual<Delta>`, we need to make `Delta` a `VectorSpace`:

```rust
impl VectorSpace for Rc<Delta> {
    fn zero() -> Self {
        Rc::new(Delta::Zero)
    }

    fn add(&self, rhs: &Self) -> Self {
        Rc::new(Delta::Add(Rc::clone(self), Rc::clone(rhs)))
    }

    fn scale(&self, factor: f64) -> Self {
        Rc::new(Delta::Scale(factor, Rc::clone(self)))
    }
}
```

The last piece of the puzzle is evaluating the intermediate representation to the partial derivatives.

```rust
pub fn eval_delta(scale_acc: f64, delta: &Delta, result: &mut Vec<f64>) {
    match *delta {
        Delta::Zero => (),
        Delta::Var(i) => result[i] += scale_acc,
        Delta::Scale(factor, ref d2) => eval_delta(scale_acc * factor, d2, result),
        Delta::Add(ref l, ref r) => {
            eval_delta(scale_acc, l, result);
            eval_delta(scale_acc, r, result);
        }
    }
}
```

`eval_delta` is memory-efficient: it updates a single array in place. It iterates depth-first through the graph described by `Delta`, and accumulates scale factors along each path from the root to the `Var` or `Zero` leaves. When it encounters a `Var` case, which is necessarily a leaf, it just adds the accumulated scale factor to the right place in the result array.

Putting all of that together, the third version of `df` becomes:

```rust
fn df_v3(x: f64, y: f64) -> Vec<f64> {
    let dual_delta = f(
        &Dual::new(x, Rc::new(Delta::Var(0))),
        &Dual::new(y, Rc::new(Delta::Var(1))),
    );
    let mut result = vec![0.0, 0.0];
    eval_delta(1.0, &dual_delta.tangent, &mut result);
    result
}
```

This looks much better! The first pass, executing `f`, returns a `Delta` intermediate representation, and then in a second pass, `eval_delta` accumulates the scale factors "in reverse" with respect to the usual evaluation order of the program. This reversal is visible as the `scale_acc` parameter of `eval_delta`. The evaluation process can be visualized as follows, as before for `df(3,2)`. Note that we now have an explicit representation of this graph as a result of the forward pass, for this example `Add(Scale(2.0, Var(0)), Scale(3.0, Var(1)))`.

{% embed https://youtu.be/FLeO5StO3n0 %}

In other words, this is a "beta version" of reverse mode AD - there is still an ingredient missing, as we'll see - but it explains where the name comes from. Using a single depth-first traversal of the `Delta` graph, it's possible to get the partial derivatives with respect to _all_ the inputs in one reverse sweep. And importantly, we don't need an amount of memory that's the square of the number of inputs anymore.

## Finally efficient: dealing with sharing

Now there is one remaining problem - if the original program contains shared variables, `eval_delta` does too much work. Here's a problematic example:

```rust
fn f_sharing_bad<T: VectorSpace>(x: &Dual<T>, y: &Dual<T>) -> Dual<T> {
    let ref s = x * y;
    s + s
}
```

The `s` variable is used twice in the final result, and as expected the primal of `s` is only calculated once and then reused in the addition. But, the calculation of the tangent of `s` is done twice in `eval_delta`. Here's the `Delta` graph for x=3 and y=2:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/u2e1bsos19mgdn9ubo5r.png)

We want our derived program to have asymptotically the same complexity as the main program, and not respecting sharing can easily blow up the complexity.

This smells like a dynamic programming problem: parts of the graph are shared, but they're being evaluated repeatedly. To achieve the same asymptotic performance as the main program, we have to make sure that each node in the `Delta` graph is evaluated only once. As usual in dynamic programming, part of the answer is memoization: as we calculate scale factors, we'll need to store for each node what the factor is up to that point. To guarantee we visit every node only once, and we have accumulated the necessary factors, we have to visit the graph in reverse topological order. Luckily, we already have such an order: it's just the reverse order of the forward pass!

To accomplish all this, we need to make two changes:

1. Introduce a `Trace` type to record the evaluation order of the operations in the forward pass; and

2. Use `Delta::Var` to represent intermediate variables explicitly. Since we can't detect which variables are shared and which ones are not, we assign the result of each operation to a temporary `Var` in the trace. Put another way, each node in the graph is represented by a `Var`.

Here's the salient parts of the implementation. First, the `Trace` type.

```rust
struct Trace {
    trace: RefCell<Vec<Rc<Delta>>>,
}
```

Ignoring the `RefCell` and `Rc`, this is just a resizable array (`Vec` in Rust) of `Delta` values. We're going to treat this `Trace` as a stack: in the forward sweep, each tangent operation gets pushed on it, and in the reverse sweep, we'll be popping operations to calculate derivatives. The stack ensures we are visiting each node just once and in reverse topological order.

Then the `Dual` type needs to be augmented with a `Trace`, so we can record each operation in it:

```rust
struct DualTrace<'a> {
    trace: &'a Trace,
    dual: Dual<Rc<Delta>>,
}
```

This means we need to redo the addition and multiplication overloads we did earlier for `DualTrace`. For each operation we can delegate to the overload for `Dual`, but need to do two extra things:

1. Push the operation on the `Trace`;

2. Return a new `Var` to represent this operation. This intermediate variable will then used to store and share the result, so it is only evaluated once.

Remember way back in the introduction, we rewrote `f` to look like this:

```rust
fn f_ad(t_dual: (f64, f64)) -> (f64, f64) {
    let (primalt, tangentt) = t_dual;
    let (primal0, tangent0) = (primalt.powi(2), 2.0 * primalt * tangentt);
    let (primal1, tangent1) = (primal0 + primalt, tangent0 + tangentt);
    let (primal2, tangent2) = (primal1 + 1.0, tangent1);
    (primal2, tangent2)
}
```

That's very similar to how we're using `Var` now: each operation becomes an implicit `let Var_i = Add(...)` in the trace.

Here's what that looks like for addition:

```rust
impl Trace {
    fn push(&self, op: Dual<Rc<Delta>>) -> Dual<Rc<Delta>> {
        let mut trace = self.trace.borrow_mut();
        let var = Dual {
            primal: op.primal,
            tangent: Rc::new(Delta::Var(trace.len())),
        };
        trace.push(op.tangent);
        var
    }
}
```

`Trace.push` adds a new operation (`Scale`, `Add` or `Zero`) to the end of the trace, and returns a new `Delta::Var` to represent the result of that operation. Now addition:

```rust
impl<'a> DualTrace<'a> {

    fn delta_push(&self, op: Dual<Rc<Delta>>) -> DualTrace<'a> {
        let dual = self.trace.push(op);
        DualTrace {
            trace: self.trace,
            dual,
        }
    }

    fn add_impl(&self, rhs: &DualTrace<'a>) -> DualTrace<'a> {
        let op = &self.dual + &rhs.dual;
        self.delta_push(op)
    }
}
```

And the same for other operations. This maintains the invariant that the result of any operation always has a `Var` as tangent. We'll leverage this fact when doing the reverse sweep, as it allows us to detect sharing explicitly.

Here's the latest and greatest `eval`. It uses `eval_delta` which is unchanged from before.

```rust
pub fn eval(inputs: usize, dual_trace: &DualTrace) -> Vec<f64> {
    let mut result = vec![0.0; dual_trace.trace_len()];
    let mut trace = dual_trace.trace_mut();

    eval_delta(1.0, &dual_trace.dual.tangent, &mut result);

    while trace.len() > inputs {
        let deltavar = trace.pop().unwrap();
        let idx = trace.len();
        if result[idx] != 0.0 {
            eval_delta(result.pop().unwrap(), &deltavar, &mut result);
        }
    }
    result
}
```

Let's see how to use all this, then dig into some details.

```rust
fn f_sharing_fixed<'a>(x: &DualTrace<'a>, y: &DualTrace<'a>) -> DualTrace<'a> {
    let ref s = x * y;
    s + s
}

fn df_sharing_fixed(x: f64, y: f64) -> Vec<f64> {
    let trace = Trace::new();
    let x = &trace.var(x);
    let y = &trace.var(y);
    let dual_trace = f_sharing_fixed(x, y);
    eval(2, &dual_trace)
}
```

For this example with x=3 and y=2, the trace after the forward pass is an array of 5 elements:

```rust
[ Var 0, // x
  Var 1, // y
  Add(Scale(2.0, Var 0), Scale(3.0, Var 1)), //let s = 2x + 3y
  Add(Var 2, Var 2) // let r = s + s
  Var 3 ] // r
```

(I've omitted some parentheses to make it easier to read)

The trace makes sharing explicit. Each `Var`'s id is a valid index in the array, which refers to where that "variable" is defined. In `eval`, a result array is allocated for each element of the trace, and this array is where the derivative for the corresponding variable is accumulated. The trace is also used to make sure each node is visited only once. Here `eval`'s reverse sweep is animated, always for x=3 and y=2. The result array is shown at the top.

{% embed https://youtu.be/lXMfCMXXvjM %}

And there you have it! From forward to efficient, reverse mode AD in four steps.

If you read this far, you may as well consider subscribing to [my newsletter](https://getcode.substack.com/), and/or follow me on [Twitter](https://twitter.com/kurt2001).

## Acknowledgements and further references

The main inspiration and origin of this post is the paper _"Provably Correct, Asymptotically Efficient, Higher-Order Reverse-Mode Automatic Differentiation" by Faustyna Krawiec, Neel Krishnaswami, Simon Peyton Jones, Tom Ellis, Andrew Fitzgibbon, and Richard Eisenberg_, along with Simon Peyton Jones' presentation at Haskell Exchange 2021. You can find paper and presentation [here](https://simon.peytonjones.org/provably-correct/).

[Conal Elliot](http://conal.net/)'s papers were also very helpful and influential: [Beautiful differentiation](http://conal.net/papers/beautiful-differentiation/) and [The simple essence of automatic differentiation](http://conal.net/papers/essence-of-ad/).

It helps having things explained from different perspectives. Here are a few other good links on the topic:

- [Calculus on Computational Graphs: Backpropagation](https://colah.github.io/posts/2015-08-Backprop/) by Christopher Olah has more of a deep learning focus.

- [Differentiable Programming from Scratch](https://thenumb.at/Autodiff/) by Max Slater has a more mathematical focus, has inline JavaScript graphs you can play with, and shows how AD is used to de-blur images.

Any errors and omissions are entirely my responsibility.