LightweightAutoDiff

LAD: Lightweight Automatic Differentiation - a small autodiff tool. Support reverse mode throught a Taper that records instruction an then performs back-propagation WRT a given variable.

Forward mode

Forward mode automatic differenciation computes $\nabla f(x) dS$. In LAD forward mode is implemented with dual numbers. A Dual number is similar to a complex number and can be written as $x = a + \epsilon b$, where $a$ is the so-called real part of $x$, $b$ is the infinitesimal part of $x$ and where $\epsilon^2 = 0$. Thus:

• $x_1 + x_2 = a_1 + a_2 + \epsilon (b_1 + b_2)$
• $x_1 x_2 = a_1 a_2 + \epsilon (a_1 b_2 + b_1 b_2)$
• $\cos(a + \epsilon b) = \cos(a) - \epsilon b\sin(a)$
• ... This is called forward mode automatic differentiation because the value of the partial derivative (often called the seed) is set on the leaves of the expression tree and propagated to the root. Forward mode can thus compute a directional derivative with one function evaluation.

Here is how dual number are used with LAD: First, we define a function.

template<typename T> T f(const T& x, const T& y) { return x*x * sin(y); }

The in our main.cpp

constexpr LAD::Dual<double> eps(0., 1.);

const LAD::Dual<double> x(1.);
const LAD::Dual<double> y(0.25*std::numbers::pi_v<double>);

// let us compute ∂f/∂x at (x=1, y=pi/4)
{
	const LAD::Dual<double> fx = f(x + eps, y);

	fmt::println("f(x,y) = {}", fx.getValue());
	fmt::println("analytical ∂f/∂x(x,y) = {}", 2.*std::sin(y.getValue()));
	fmt::println("estimated  ∂f/∂x(x,y) = {}", fx.getDiff());
}
// now we compute ∂f/∂y at (x=1, y=pi/4)
{
	const LAD::Dual<double> fx = f(x, y + eps);

	fmt::println("f(x,y) = {}", fx.getValue());
	fmt::println("analytical ∂f/∂x(x,y) = {}", std::cos(y.getValue()));
	fmt::println("estimated  ∂f/∂x(x,y) = {}", fx.getDiff());
}
// now we compute -2 ∂f/∂x + 3 ∂f/∂y at (x=1, y=pi/4)
{
	const LAD::Dual<double> fx = f(x - 2.*eps, y + 3.*eps);

	fmt::println("f(x,y) = {}", fx.getValue());
	fmt::println("analytical directional derivatives = {}", -4.*std::sin(y.getValue()) + 3.*std::cos(y.getValue()));
	fmt::println("estimated  directional derivatives = {}", fx.getDiff());
}

Forward mode can compute a directional derivative in one function evaluation at relatively low cost. However, evaluating the full gradient requires $n$ functuion evaluations.

Reverse mode

Reverse mode automatic differentiation follows the opposite path and assumes the seed is set at the root of the call tree and then uses the chain rule to compute the derivatives at the leaves. More precisely, assume $\frac{\partial J}{\partial a}$, is known, where $J$ is the function we want to differentiate and where $a$ is a binary operator. Then

$$\begin{align} \frac{\partial J}{\partial u} &= \frac{\partial J}{\partial a}\frac{\partial a}{\partial u}\\\ \frac{\partial J}{\partial v} &= \frac{\partial J}{\partial a}\frac{\partial a}{\partial v} \end{align}$$

With, off course, $\frac{\partial J}{\partial J} = 1$. The automatic differentiation works in two passes:

1. The function is evaluated and all the operations and intermediate variables of the expression tree are stored into a "Tape"
2. The expression tree is then traversed from root to leaves propagating the partial derivatives.

Here is how reverse mode is used in LAD:

// we create a "tape" that records the operations 
// carried out upon the evaluation of our function.
LAD::Tape<double> tape;
// we create a scope that will be used by all subsequent variables.
// the scope grant access to our tape to all the temporaries 
// created with our function. 
LAD::TapeScope<double> scope(tape);

LAD::Variable<double> x(1.);
LAD::Variable<double> y(0.25*std::numbers::pi_v<double>);

const LAD::Variable<double> fxy = f(x, y);

// Reverse mode computes the gradient WRT all the variables, including temporaries
// involved in the computation of fxy.
// getDiffWRT return a std::range::view (a subset) of the gradient that corresponds 
// to the specified variables.
// if another `getDiffWRT` is called or if the tape is destroyeed the grad will be invalidated.
// One can store the gradient into a container by calling:
//  std::ranges::copy(grad, std::ranges::begin(dst))
const auto grad = fxy.getDiffWRT(x, y); 

fmt::println("f(x,y) = {}", fxy.getValue());
fmt::println("analytical ∂f/∂x(x,y) = {}", 2.*std::sin(y.getValue()));
fmt::println("estimated  ∂f/∂x(x,y) = {}", grad[0]);
fmt::println("analytical ∂f/∂y(x,y) = {}", std::cos(y.getValue()));
fmt::println("estimated  ∂f/∂y(x,y) = {}", grad[1]);

Forward Over Reverse

One can compute the product between the Hessain of a function at a given point $H$ and a given direction $d$. We introduce the function:

$$g(x,y) = \langle \nabla f(x,y), d\rangle$$

Whose gradient is $\nabla g = H d$. We can thus compute $g(x,y)$ with the forward mode autodiff and then compute the gradient of $g$ with the reverse mode autodiff.

While LAD do not support this directly, one can implement it as follows:

LAD::Tape<double> tape;
LAD::TapeScope<double> scope(tape);

// we want to differentiate g(x,y) = -2 ∂f/∂x (x,y) + 3 ∂f/∂y (x,y)
// AKA H_f*[-2; 3]

const LAD::DualVariable<double> eps(0., 1.);

LAD::DualVariable<double> x(1.);
LAD::DualVariable<double> y(0.25*std::numbers::pi_v<double>);

// The forward mode gives us the directional derivative of J
const DirectionalDiff gradJ_d = J.getDiff();
// The reverse mode allows us to compute the gradient of the directional derivative of J
// i.e. the product between the Hessian of J and a direction d.
const auto hessProd = gradJ_d.getDiffWRT(x, y);

fmt::println("f(x,y) = {}", J.getValue());
fmt::println("analytical -2 ∂²f/∂x²2 (x,y) + 3 ∂²f/∂x∂y (x,y) = {}", -4.*std::sin(y.getValue()) + 6.*x.getValue()*std::cos(y.getValue()));
fmt::println("estimated  -2 ∂²f/∂x²2 (x,y) + 3 ∂²f/∂x∂y (x,y) = {}", hessProd[0]);
fmt::println("analytical  -2 ∂²f/∂x∂y (x,y) + 3 ∂²f/∂y² (x,y) = {}", -4.*x.getValue()*std::cos(y.getValue()) - 3.*x.getValue()*x.getValue()*std::sin(y.getValue()));
fmt::println("estimated   -2 ∂²f/∂x∂y (x,y) + 3 ∂²f/∂y² (x,y) = {}", hessProd[1]);

Practical example: Rosenbrock function

We define Rosenbrock function:

template<typename T>
inline T rosenbrock(const std::span<const T> x)
{
	T J{};
	
	for (size_t i=0; i!=size_t(x.size()-1); ++i)
	{
		J += 100.*(x[i+1] - x[i]*x[i])*(x[i+1] - x[i]*x[i]) + (x[i] - 1.)*(x[i] - 1.);
	}
	
	return J;
}

Then we compute it's gradient at $x=(0, \dots, 0)$:

const std::vector<LAD::Variable<double>> x(10, LAD::Variable<double>());
		
const LAD::Variable<double> J = rosenbrock(std::span(x));
		
const auto grad = J.getDiffWRT(x);
		
fmt::println("J({}) = {}", fmt::join(x, ", "), J.getValue()); // J(0, 0, 0, 0, 0, 0, 0, 0, 0, 0) = 9
fmt::println("grad({}) = {}", fmt::join(x, ", "), grad);      // grad(0, 0, 0, 0, 0, 0, 0, 0, 0, 0) = [-18, -18, -18, -18, -18, -18, -18, -18, -18, -18]

Now let use compute its gradient at $J$'s minimum, $x=(1, \dots, 1)$:

const std::vector<LAD::Variable<double>> x(10, LAD::Variable<double>(1.));
		
const LAD::Variable<double> J = rosenbrock(std::span(x));
		
const auto grad = J.getDiffWRT(x);
		
fmt::println("J({}) = {}", fmt::join(x, ", "), J.getValue()); // J(1, 1, 1, 1, 1, 1, 1, 1, 1, 1) = 0
fmt::println("grad({}) = {}", fmt::join(x, ", "), grad);      // grad(1, 1, 1, 1, 1, 1, 1, 1, 1, 1) = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]