/prelude/ad

Definitions related to automatic differentiation.

Futhark supports a fairly complete facility for computing derivatives of functions through automatic differentiation (AD). The following does not constitute a full introduction to AD, which is a substantial topic by itself, but rather introduces enough concepts to describe Futhark's functional AD interface. AD is a program transformation that is implemented in the compiler itself, but the user interface is a handful of fairly simple functions.

AD is useful when optimising cost functions, and for any other purpose where you might need derivatives, such as for example computing surface normals for signed distance functions.

Futhark's AD support includes the following:

• Differential operators for forward-mode (jvp) and reverse-mode (vjp).

• Almost arbitrary control flow in differentiable code (some limitations apply when using GPU backends, see below).

• Higher order derivatives by nesting differentiation operators, including arbitrary mixing of forward- and reverse mode (although using multiple rounds of reverse mode is rarely useful and often slow).

• Custom derivatives (with_vjp).

• Vector AD (mjp, jmp), sometimes also known as "batched" or "multi-directional" AD.

• Checkpointing of sequential loops.

↑val jvp2 'a 'b: (f: a -> b) -> (x: a) -> (x': a) -> (b, b)

Jacobian-Vector Product ("forward mode"), producing also the primal result as the first element of the result tuple.

↑val vjp2 'a 'b: (f: a -> b) -> (x: a) -> (y': b) -> (b, a)

Vector-Jacobian Product ("reverse mode"), producing also the primal result as the first element of the result tuple.

↑val jmp2 'a 'b [n]: (f: a -> b) -> (x: a) -> (x': [n]a) -> (b, [n]b)

Jacobian-Matrix Product, returning also the primal result. As jvp2, but accepts an array of seed vectors (hence "matrix", although transposed). Semantically equivalent to mapping, but may be more efficient. If used with #[unroll], tangent calculations are unrolled when possible.

↑val mjp2 'a 'b [n]: (f: a -> b) -> (x: a) -> (y': [n]b) -> (b, [n]a)

Matrix-Jacobian Product, returning also the primal result. As vjp2, but accepts an array of seed vectors (hence "matrix"). Semantically equivalent to mapping, but may be more efficient. If used with #[unroll], adjoint calculations are unrolled when possible.

↑val jvp 'a 'b: (f: a -> b) -> (x: a) -> (x': a) -> b

Jacobian-Vector Product ("forward mode").

↑val vjp 'a 'b: (f: a -> b) -> (x: a) -> (y': b) -> a

Vector-Jacobian Product ("reverse mode").

↑val jmp 'a 'b [n]: (f: a -> b) -> (x: a) -> (x': [n]a) -> [n]b

Jacobian-Matrix Product. As jvp, but accepts a vector of seed values. Semantically equivalent to mapping, but may be more efficient.

↑val mjp 'a 'b [n]: (f: a -> b) -> (x: a) -> (y': [n]b) -> [n]a

Matrix-Jacobian product. As vjp, but accepts a vector of seed values. Semantically equivalent to mapping, but may be more efficient.

↑val with_vjp 'a 'b: (f: a -> b) -> (f': (res: b) -> (b_adj: b) -> a) -> (x: a) -> b

Provide custom reverse-mode adjoint code for a given function. This is useful when the adjoint synthesised by AD is not as good as one that is known analytically.

The function f returns a result of type b. In the return sweep, the function f' is invoked first with the result of f and second with the cotangents of the result (be careful not to mix up the order), and must return the sensitivity with respect to the input.

A common pattern is that b is a tuple where some part is the intended primal result of with_vjp, and some part is only used in f'.

Beware: if f uses any free variables, these will not be taken into account when computing the adjoint. Make these part of the argument instead.

↑val with_vjp_tape 'a 'b 'c: (f: a -> (c, b)) -> (f': (c, b) -> a) -> (x: a) -> b

A variant of with_vjp where the intermediate result necessary for the adjoint (c) is explicitly separated from the primal result (b).

Jacobians

For a differentiable function f whose input comprise n scalars and whose output comprises m scalars, the Jacobian for a given input point is an m by n matrix of scalars that each represent a partial derivative. Intuitively, position (i,j) of the Jacobian describes how sensitive output i is to input j. The notion of Jacobian generalises to functions that accept or produce compound structures such as arrays, records, sums, and so on, simply by "flattening out" the values and considering only their constituent scalars.

Computing the full Jacobian is usually costly and sometimes not necessary, and it is not part of the AD facility provided by Futhark. Instead it is possible to compute parts of the Jacobian, which semantically (but not operationally) can be seen as multiplying the Jacobian with a vector, producing a vector. However, it is important to understand that the full Jacobian is not constructed as an intermediate step.

We can take the product of an m by n Jacobian with an n-element tangent vector to produce an m-element vector (Jacobian-vector product). Such a product can be computed in a single (augmented) execution of the function f. This is provided by the function jvp.

We can also take the product of an m-element vector cotangent vector with the m by n Jacobian to produce an n-element vector (vector-Jacobian product). This too can be computed in a single execution of f, with vjp.

A tangent has the same structure as the input and represents a direction in input space. A cotangent has the same structure as the output and represents sensitivities flowing backwards through the computation.

Using an elementary (co-)tangent vector, we can use the jvp function to produce a column of the full Jacobian, and vjp to produce a row, with the nonzero element of the vector identifying which column or row is extracted. Which is superior for a given situation depends on whether the function has more inputs or outputs.

We can freely nest vjp and jvp to compute higher-order derivatives.

Efficiency

Both jvp and vjp work by transforming the program to carry along extra information associated with each scalar value.

In the case of jvp ("forward mode", or "tangent mode"), this extra information takes the form of an additional scalar representing the tangent, which is then propagated in each scalar computation using essentially the chain rule. Therefore, jvp has a memory overhead of approximately 2x, and a computational overhead of slightly more, but usually less than 4x.

In the case of vjp ("reverse mode" or "adjoint mode"), since our starting point is a cotangent, the function is essentially first run forward, then backwards (the return sweep) to propagate the cotangent. During the return sweep, all intermediate results computed during the forward sweep must still be available, and must therefore be stored in memory during the forward sweep

• this is called "the tape". This means that the memory usage of vjp is proportional to the number of sequential steps of the original function (essentially turning time into space). The compiler does a nontrivial amount of optimisation to ameliorate this overhead (see AD for an Array Language with Nested Parallelism), but it can still be substantial for programs with deep sequential loops.

Nesting vjp, understood as applying vjp to the result of vjp, is usually a bad idea, as the code structure produced by vjp is fairly complicated, due to the tape management. Passing the output of jvp to vjp, or the other way, is however fine. As a rule of thumb, whenever you stack multiple differential operators, make sure only one of them is vjp or related ones.

When using vector AD (mjp/jmp), each scalar is associated with a vector of tangents or cotangents, and the space overhead for storing these is therefore multiplied with the vector size. However, in the case of vjp, the intermediate results are only stored once. It varies on a case-by-case basis whether vector AD is faster than using map on top of vjp/jvp. Vector AD essentially converts propagation of (co-)tangents from scalar to array operations, which can have a significant impact on memory accesses, depending on how the compiler manages to optimise the resulting code. It is hard to predict whether this offsets the reduction in primal work. If the vector size is a constant, and the #[unroll] attribute is put on the AD operator, then the vectors become unrolled (turned into tuples, essentially), although this should only be done when the vector size is quite small, as the increase in code size is substantial.

Differentiable functions

AD only gives meaningful results for differentiable functions. The Futhark type system does not distinguish differentiable from non-differentiable operations. As a rule of thumb, a function is differentiable if its results are computed using a composition of primitive floating-point operations, without ever converting to or from integers. Most functions will also have discontinuities around values that influence control flow.

Note that a function whose input or output is a sum type with more than one constructor is not differentiable (or at least the sum-typed part is not). This is because the choice of constructor is not a continuous quantity.

Limitations

jvp is expected to work in all cases. vjp has limitations when using the GPU backends similar to those for irregular flattening. Specifically, you should avoid structures with variant sizes, such as loops that carry an array that changes size through the execution of the loop.