Let’s say we want to calculate the partial derivative of \(y\) with respect to \(x_{1}\), with \(x_{1} = 1.5\) and \(x_{2} = 0.5\). As we mentioned above, we can try do it one intermediate variable at a time. Note that we are only calculating the numerical value of the derivative. For each \(v_{i}\), we calculate \(\dot{v}_{i} = \frac{\partial v_{i}}{\partial x_{1}}\). Let’s try a few variables to see how it goes:

For these calculations, we are simply applying chain rules with basic derivatives. Note that how \(\dot{v}_{i}\) only depends on the derivatives and values of the earlier variables. We can now augment Table 2 to include the derivatives.

The values of the intermediate variables are sometimes called the primal trace, and the derivative values the tangent trace.

This generalizes nicely to a generic vector-valued function \(f: \mathbb{R}^{n} \rightarrow \mathbb{R}^{m}\). Assume we are trying to evaluate the function at \(\mathbf{x} = \mathbf{a}, \mathbf{a} \in \mathbb{R}^{n}\). In that case, the Jacobian matrix is in the form of

and each column consists of

which are the partial derivatives of \(y_{j}\) with respect to each \(x_{i}\). So we can obtain the columns one by one by setting the corresponding \(\dot{x}_{i} = 1\) and setting the other entries zero. This opens some interesting techniques. For example, if we want to compute the Jacobian-vector product with \(\mathbf{r}\), we can simply set \(\dot{\mathbf{x}}=\mathbf{r}\) instead of a unit vector:

In terms of complexity, forward mode is on the order of \(O(n)\), where \(n\) is the dimension of the input vector. Hence, it will be a great choice for \(f: \mathbb{R} \rightarrow \mathbb{R}^{m}\), while performing poorly for \(f: \mathbb{R}^{n} \rightarrow \mathbb{R}\).

One popular way to implement forward mode AD is to use dual numbers. Dual numbers were introduced by William Clifford in 1873 through his paper Preliminary Sketch of Biquaternions. Clifford was studying how to extend Hamilton’s vectors to consider rotations around not only the origin, but any arbitrary lines in the 3-dimensional space. He called his extension on vectors rotors, and the sum of such rotors motors. He then introduces a symbol \(\omega\) to convert motors into vectors, and that \(\omega^{2} = 0\). His reasoning was purely geometric:

The symbol \(\omega\), applied to any motor, changes it into a vector parallel to its axis and proportional to the rotor part of it. … and if made to operate directly on a vector, reduces it to zero.

However, it turns out that dual numbers can be conveniently to represent differentiation.

Formally, we represent dual numbers as \(a + b\epsilon\). They follow the usual component-wise addition rule:

And their multiplications work like this (similar to complex numbers):

with the rule that \(\epsilon^{2} = 0\).

The connection between dual numbers and differentiation becomes clear once we look at the Taylor series expansion. Given an arbitrary real function \(f : \mathbb{R} \rightarrow \mathbb{R}\), we can express its Taylor series expansion at \(x_{0}\) as

as long as the function is \(n+1\) times differentiable and has bounded \(n+1\) derivative. If we instead extend them to using dual numbers as inputs, it follows that

So if we set \(b=1\), we can get the derivative at \(x=a\) by taking the coefficient of \(\epsilon\) after evaluation:

where we expanded the Taylor series at \(a\), and since \(\epsilon^{2} = 0\), all terms involving \(\epsilon^{2}\) and higher vanish. Under this formulation, addition works:

Multiplication works:

And chain rule works as expected:

This essentially gives us the way to conduct forward mode AD: by using dual numbers, we can get the primal and tangent trace simultaneously.

So, how do we take this to higher dimensions? We simply add an \(\epsilon\) for each component. Assume a function \(f: \mathbb{R}^{n} \rightarrow \mathbb{R}\). We define a vector \(\boldsymbol{\epsilon} \in \mathbb{R}^{n}\) where \(\boldsymbol{\epsilon}_{i}^{2} = \boldsymbol{\epsilon}_{i} \boldsymbol{\epsilon}_{j} = 0\). If you write out the component-wise computation following Eq. \ref{eq:dual-single-v}, it follows that

\(\nabla f(\boldsymbol{a}) \cdot \boldsymbol{\epsilon}\) can be interpreted as the directional derivative of \(f\) in the direction of \(\boldsymbol{\epsilon}\). Addition, multiplication and chain rule work as expected:

For \(f : \mathbb{R}^{n} \rightarrow \mathbb{R}^{m}\), instead of using a vector of \(\epsilon\), we can use a matrix of \(\epsilon\). And for each row of that matrix, we follow the exact same rule as Eq. \ref{eq:dual-multi-v}.

Reverse mode automatic differentiation, also known as adjoint mode, calculates the derivative by going from the end of the evaluation trace to the beginning. The intuition comes from the chain rule. Consider a function \(y = f(x(t))\). From the chain rule, it follows that

The two terms on the right-hand side can be seen as going backwards: \(\frac{\partial y}{\partial x}\) can be determined once we calculate \(y\) from \(x\), and \(\frac{\partial x}{\partial t}\) can be calculated once we calculate \(x\) from \(t\). Extending this to multivariate functions, we have the multivariate chain rule:

The proof of it can be found in Appendix A4 of Vector Calculus, Linear Algebra, and Differential Forms: a Unified Approach by John H. Hubbard and Barbara Burke Hubbard. For example, for a function \(F(t) = f(g(t)) = f(x(t), y(t))\) where \(f : \mathbf{R}^{2} \rightarrow \mathbf{R}\) and \(g : \mathbf{R} \rightarrow \mathbf{R}^{2}\), we have

So it follows

Applying this idea to automatic differentiation, we have the reverse mode. To calculate derivatives in this mode, we need to conduct two passes. First, we need to do a forward pass, where we obtain the primal trace (Table 2). We then propagate the partials backward to obtain the desired derivatives (following the chain rule). Let’s go back to our example before to see reverse mode in action. Consider again the function

Following the same way of assigning intermediate variables as in Table 2 and in Fig. 3, we can assemble an adjoint trace. An adjoint \(\bar{v}_{i}\) is defined as

which is equivalent to the product of all the partials from \(y_{j}\) up until \(v_{i}\). We start from the last variable \(v_{6} = y\):

And repeat the same process for all the intermediate variables, we have Table 4 below.

Note how we are able to obtain the adjoints of the input variables \(x\) and \(y\) (which are equivalent to the partial derivatives of \(f\) with respect to \(x\) and \(y\)) at the same time. For a function \(f : \mathbf{R}^{n} \rightarrow \mathbf{R}\), it takes only one application of reverse mode to compute the entire gradient. In general, if the dimension of the outputs is significantly smaller than that of inputs, reverse mode is a better choice.

To implement forward mode in Rust, I opt to create a structure to represent dual numbers:

pub struct DualScalar {
    pub v: f64,
    pub dv: f64,
}

The operations are then implemented on DualScalar, following standard traits such as std::ops::{Add, Div, Mul, Neg, Sub}. For example,

impl Add for DualScalar {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        DualScalar {
            v: self.v + rhs.v,
            dv: self.dv + rhs.dv,
        }
    }
}

pub fn deriv(&self) -> f64 {
    self.dv.clone()
}

where the addition rule is implemented for the add(self, rhs: Self) function. To access the derivative value, we have And evaluating a function to get the derivative is done through the derivative<F>(func: F, x0: f64 ) function:

pub fn derivative<F>(func: F, x0: f64 ) -> f64
    where F : Fn(DualScalar) -> DualScalar,
{
    func(DualScalar{ v: x0, dv: 1.0}).deriv()
}

For multivariate functions, I simply represent the inputs as a Vec holding DualScalar.

pub fn gradient<F>(func: F, x0: &[f64]) -> Vec<f64>
    where F: Fn(&[DualScalar]) -> DualScalar,
{
    // To get all the partials, we set each var to have dv=1
    // and the others dv=0, and pass them through the function
    let mut inputs: Vec<DualScalar> = x0.iter().map(|&v| DualScalar { v: v, dv: 0. }).collect();
    (0..x0.len()).map(
        |i| {
            inputs[i].dv = 1.;
            let partial = func(&inputs).deriv();
            inputs[i].dv = 0.;
            partial
        }
    ).collect()
}

Note that we have to iterate through all the inputs by setting dv to be 1, then collect the partials. For Jacobians, we obtain the columns one by one:

pub fn jacobian<F, const N: usize, const M: usize>(func: F, x0: &[f64]) -> SMatrix<f64, M, N>
    where F: Fn(&[DualScalar]) -> Vec<DualScalar>,
{
    // To get all the partials, we set each var to have dv=1
    // and the others dv=0, and pass them through the function
    let mut jacobian: SMatrix<f64, M, N> = SMatrix::zeros();
    let mut inputs: Vec<DualScalar> = x0.iter().map(|&v| DualScalar { v: v, dv: 0. }).collect();

    // every time we call the func we can get one column of partials
    for (i, mut col) in jacobian.column_iter_mut().enumerate() {
        inputs[i].dv = 1.;
        let col_result = func(&inputs);
        for j in 0..M {
            col[j] = col_result[j].dv;
        }
        inputs[i].dv = 0.;
    }

    return jacobian;
}

The API of these functions are pretty straightforward:

// API for forward mode
//
// derivative(f, x)
let f1_test = |x: fdiff::DualScalar| x * x;
let f1_result: f64 = fdiff::derivative(f1_test, 2.0);
println!("Derivative of f(x) = x^2 at {} is {}", 2.0, f1_result);

// gradient(f, x)
let f2_test = |x: &[fdiff::DualScalar]| x[0] + x[1];
let f2_result: Vec<f64> = fdiff::gradient(f2_test, vec![1., 2.].as_slice());
println!("Gradient of f(x,y) = x + y at ({}, {}) is {:?}", 1.0, 2.0, f2_result);

// jacobian(f, x)
// f(1) = x^2 * y
// f(2) = x + y
let f3_test = |x: &[fdiff::DualScalar]| {
    vec![x[0] * x[0] * x[1], x[0] + x[1]]
};
let f3_result: SMatrix<f64, 2, 2> = fdiff::jacobian(f3_test, vec![1., 2.].as_slice());
println!("Jacobian of f(x,y) = [x^2 * y , x + y] at ({}, {}) is {:?}", 1.0, 2.0, f3_result);

Reverse mode turns out to be a bit harder in Rust. Initially, I struggled with the ownership model as I wanted to dynamically represent the computation graph as linked nodes. Then I stumbled across this excellent post by Rufflewind, which goes through a tape-based implementation. Essentially, we need to separate the computation graph with the expression, so that users can treat variables as float-like, because we are simply storing indices to the locations on an array.

We first define a Tape and Node:

pub struct Tape {
    pub nodes: RefCell<Vec<Node>>,
}

#[derive(Clone, Copy)]
pub struct Node {
    pub partials: [f64; 2],
    pub parents: [usize; 2],
}

The parents hold the node’s dependencies in the computation graph, and partials are the respective partials to the parents. We then define a few operations on the Tape:

impl Tape {
    pub fn var(&self, value: f64) -> Var {
        let len = self.nodes.borrow().len();
        self.nodes.borrow_mut().push(
            Node {
                partials: [0.0, 0.0],
                // for a single (input) variable, we point the parents to itself
                parents: [len, len],
            }
        );
        Var {
            tape: self,
            index: len,
            v: value,
        }
    }

    pub fn binary_op(&self, lhs_partial: f64, rhs_partial: f64,
                     lhs_index: usize, rhs_index: usize, new_value: f64) -> Var {
        let len = self.nodes.borrow().len();
        self.nodes.borrow_mut().push(
            Node {
                partials: [lhs_partial, rhs_partial],
                // for a single (input) variable, we point the parents to itself
                parents: [lhs_index, rhs_index],
            }
        );
        Var {
            tape: self,
            index: len,
            v: new_value,
        }
    }
}

The var function is called when we create a new variable on the tape, and binary_op is a helper function for creating intermediate variables as results of mathematical operations (such as +). As for the users, we implement a Var struct (which has a one-to-one correspondence to a Node on the tape)

#[derive(Clone, Copy)]
pub struct Var<'t> {
    pub tape: &'t Tape,
    pub index: usize,
    pub v: f64,
}

where tape points to the tape we currently have, index is the location of the corresponding Node on the tape, and v is the actual value. The reason why we need a separate Var struct is that users don’t need to deal with the ownership of the Node by the tape, and that Var can be copied just like a float. The actual reverse pass (back-propagation) is easy:

impl Var<'_> {
    /// Perform back propagation
    pub fn backprop(&self) -> Grad {
        // vector storing the gradients
        let tape_len = self.tape.nodes.borrow().len();
        let mut grad = vec![0.0; tape_len];
        grad[self.index] = 1.0;

        for i in (0..tape_len).rev() {
            let node = self.tape.nodes.borrow()[i];
            // increment gradient contribution to the left parent
            let lhs_dep = node.parents[0];
            let lhs_partial = node.partials[0];
            grad[lhs_dep] += lhs_partial * grad[i];

            // increment gradient contribution to the right parent
            // note that in cases of unary operations, because
            // partial was set to zero, it won't affect the computation
            let rhs_dep = node.parents[1];
            let rhs_partial = node.partials[1];
            grad[rhs_dep] += rhs_partial * grad[i];
        }

        Grad { grad }
    }
}

where Grad is simply a struct holding the gradients, with a wrt() function that gets the specific partial derivatives for a specific Var. To implement specific differentiation rules, we implement the corresponding traits on Var:

impl<'t> Add for Var<'t> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        self.tape.binary_op(1.0, 1.0,
                            self.index, rhs.index, self.v + rhs.v)
    }
}

For the users, they can obtain partial derivatives easily:

let tape = rdiff::Tape::new();
let x = tape.var(1.0);
let y = tape.var(1.0);
let z = -2.0 * x + x * x * x * y + 2.0 * y;
let grad = z.backprop();
println!("dz/dx of z = -2x + x^3 * y + 2y at x=1.0, y=1.0 is {}", grad.wrt(x));
println!("dz/dy of z = -2x + x^3 * y + 2y at x=1.0, y=1.0 is {}", grad.wrt(y));