Continuing from my last post on implementing forward-mode automatic differentiation (AD) using C# operator overloading, this is just a quick follow-up showing how easy reverse mode is to achieve, and why it's important.
Why Reverse Mode Automatic Differentiation?
As explained in the last post, the vector representation of forward-mode AD can compute the derivatives of all parameter simultaneously, but it does so with considerable space cost: each operation creates a vector computing the derivative of each parameter. So N parameters with M operations would allocation O(N*M) space. It turns out, this is unnecessary!
Reverse mode AD allocates only O(N+M) space to compute the derivatives of N parameters across M operations. In general, forward mode AD is best suited to differentiating functions of type:
R → RN
That is, functions of 1 parameter that compute multiple outputs. Reverse mode AD is suited to the dual scenario:
RN → R
That is, functions of many parameters that return a single real number. A lot of problems are better suited to reverse mode AD, and some modern machine learning frameworks now employ reverse mode AD internally (thousands of parameters, single output that's compared to a goal).
How does Reverse Mode Work?
The identities I described in the other article still apply since they're simply the chain rule, but reverse mode computes derivatives backwards. Forward-mode AD is easy to implement using dual numbers in which the evaluation order matches C#'s normal evaluation order: just compute a second number corresponding to the derivative along side the normal computation. Since reverse mode runs backwards, we have to do the computational dual: build a (restricted) continuation!
This is a rough sketch showcasing both forward mode and reverse mode and how they're duals. Forward mode AD using dual numbers will look something like this:
public readonly struct Fwd
{
public readonly double Magnitude;
public readonly double Derivative;
public Fwd(double mag, double deriv)
{
this.Magnitude = mag;
this.Derivative = deriv;
}
public Fwd Pow(int k) =>
new Fwd(Math.Pow(Magnitude, k), k * Math.Pow(Magnitude, k - 1) * Derivative);
public static Fwd operator +(Fwd lhs, Fwd rhs) =>
new Fwd(lhs.Magnitude + rhs.Magnitude, lhs.Derivative + rhs.Derivative);
public static Fwd operator *(Fwd lhs, Fwd rhs) =>
new Fwd(lhs.Magnitude + rhs.Magnitude,
lhs.Derivative * rhs.Magnitude + rhs.Derivative * lhs.Magnitude);
public static Func<double, Fwd> Differentiate(Func<Fwd, Fwd> f) =>
x => f(new Fwd(x, 1));
public static Func<double, double, Fwd> DifferentiateX0(Func<Fwd, Fwd, Fwd> f) =>
(x0, x1) => f(new Fwd(x0, 1), new Fwd(x1, 0));
public static Func<double, double, Fwd> DifferentiateX1(Func<Fwd, Fwd, Fwd> f) =>
(x0, x1) => f(new Fwd(x0, 0), new Fwd(x1, 1));
}
Translating this into reverse mode entails replacing Fwd.Derivative with a continuation like so:
public readonly struct Rev
{
public readonly double Magnitude;
readonly Action<double> Derivative;
public Rev(double y, Action<double> dy)
{
this.Magnitude = y;
this.Derivative = dy;
}
public Rev Pow(int e)
{
var x = Magnitude;
var k = Derivative;
return new Rev(Math.Pow(Magnitude, e),
dx => k(e * Math.Pow(x, e - 1) * dx));
}
public static Rev operator +(Rev lhs, Rev rhs) =>
new Rev(lhs.Magnitude + rhs.Magnitude, dx =>
{
lhs.Derivative(dx);
rhs.Derivative(dx);
});
public static Rev operator *(Rev lhs, Rev rhs) =>
new Rev(lhs.Magnitude * rhs.Magnitude,
dx =>
{
lhs.Derivative(dx * rhs.Magnitude);
rhs.Derivative(dx * lhs.Magnitude);
});
public static Func<double, (double, double)> Differentiate(Func<Rev, Rev> f) =>
x =>
{
double dx = 0;
var y = f(new Rev(x, dy => dx += dy));
y.Derivative(1);
return (y.Magnitude, dx);
};
public static Func<double, double, (double, double, double)> Differentiate(Func<Rev, Rev, Rev> f) =>
(x0, x1) =>
{
double dx0 = 0, dx1 = 0;
var y = f(new Rev(x0, dy => dx0 += dy),
new Rev(x1, dy => dx1 += dy));
y.Derivative(1);
return (y.Magnitude, dx0, dx1);
};
public static Func<double, double, double, (double, double, double, double)> Differentiate(Func<Rev, Rev, Rev, Rev> f) =>
(x0, x1, x2) =>
{
double dx0 = 0, dx1 = 0, dx2 = 0;
var y = f(new Rev(x0, dy => dx0 += dy),
new Rev(x1, dy => dx1 += dy),
new Rev(x2, dy => dx2 += dy));
y.Derivative(1);
return (y.Magnitude, dx0, dx1, dx2);
};
}
As I mentioned in my last post, my goal here isn't the most efficient implementation for reverse mode AD, but to distill its essence to make it direct and understandable. This representation builds a whole new continuation on every invocation of the function being differentiated. More efficient representations would only compute this continuation once for any number of invocations, and there are plenty of other optimizations that can be applied to both forward and reverse mode representations.
Comments
I'm not really concerned about the performance of this particular incarnation anyway, it's really for clarity that I'll use to test against more sophisticated implementations.
For instance, it should be possible to lift the trace delegate outside of the inner delegate returned by Differentiate by restricting loops from differentiable functions, ie. replace loops with an explicit Sum() function.
Once you do that, it should then be possible to replace delegates with LINQ expressions and get a very efficient compiled equivalent with no virtual dispatch or allocation.
I don't think I've come across a dual representation for dual numbers, have you?
I thought differentiating this function would be a pathological case:
Rev f(Rev x){
for(int i=0; i<100; i++) x = x + x;
return x;
}
I think you'd need a more sophisticated representation with lookahead to solve this fully.
A simple extension for partial lookahead could add a special type for multiplication/division by constants, which accumulates operations until it hits another Codual, then it performs the operation with two Coduals, and then apply the accumulated constants.
So the specific accumulation form is novel, but not the general idea of using CPS to backpropagate. Thanks for the reference!
By contrast, forward diff can differentiate R -> R^n in the same time and memory complexity as the original function.
Solved! The identities I linked before + lifting multiplicaton/division by constants to a second field, like dual numbers, ensures that the identities are as aggressive as possible in culling the expression trace. This is definitely a novel representation. The expression loops we've looked at so far all seem to be addressed. Can you see any cases where this still falls down?
I think the only problem might be a slight change in floating precision because I change the order of operations during division/multiplication by constants.
internal Codual(double x)
{
this.Magnitude = x;
this.Derivative = 0;
}
public Codual Sin()
{
var lhs = this;
var y = new Codual(Math.Sin(Magnitude));
trace.Add(() => lhs.Derivative += y.Derivative * Math.Cos(lhs.Magnitude));
return y;
}
Where trace is a global list of callbacks. When you compute the derivative of a function, you first run the function, and then you iterate over the trace backwards and call all the callbacks. After that, the Derivative field of your variables will be set to the correct value.
The trick is that instead of accumulating the derivative only at the variables with +=, we do so that at every node. Because we iterate over the trace backwards, the y.Derivative has the correct final value before the callback in the code above gets executed.
internal Codual(double x, Action callback)
{
this.Magnitude = x;
this.Derivative = 0;
val self = this;
trace.Add(() => callback(self));
}
public Codual Sin()
{
return new Codual(Math.Sin(Magnitude), Propagate(this, Math.Cos(this.Magnitude)));
}
Where
Propagate(x, d) = (y) => y.Derivative += x.Derivative * d.
For operations with two arguments you'd need Propagate2(x1,d1,x2,d2).
Then you can use exactly the same code for the operations on dual and codual numbers, changing only the constructor and the definition of Propagate. E.g. for a single variable:
internal Dual(double x, double dx)
{
this.Magnitude = x;
this.Derivative = dx;
}
public Dual Sin()
{
return new Dual(Math.Sin(Magnitude), Propagate(this, Math.Cos(this.Magnitude)));
}
Where Propagate(x, d) = x.Derivative * d.
By the way, just like you can make forward mode handle multiple inputs by changing Derivative to be an array, you can make reverse mode handle multiple outputs by making Derivative an array.
That's an interesting breakdown of Dual/Codual via Propagators though! If only more languages had better partial evaluation support so that more general code like that could be more common.
There must be a purely functional structure that can merge the parallel branches without just reimplementing a state monad for a tape. The more I think about it, the more convinced I am that this will probably just end up as some specialization of Conal's approach, so maybe I should just start there.