Cracking the Code of Automatic Differentiation in Haskell

Forward Mode AD

This is the simplest way of calculating the derivative of a multivariable function. Consider the following function, which takes two independent variables x1, x2, and has one dependent variable y=f(x1,x2).

f(x1,x2) = x1 * x2 + x1²

The forward mode AD calculates the derivative by traversing the graph for each of the independent variables by calculating derivatives of each vi partial function by considering the following for a given independent variable x1:

v̇i = ∂vi/∂x 1

For example, to calculate the derivative x1 at point x1=2 and x2=3, the algorithm traverses the graph and evaluates the function and its derivative:

Forward Automatic Differentiation is inappropriate for machine learning optimization (check this article) due to its inefficiency in calculating the gradient for scalar functions with multiple independent variables. In such cases, multiple graph traversals are necessary, making Reverse mode Automatic Differentiation a more suitable choice.

There is another way to compute derivatives on Forward mode by using a special number called Dual Number. However, the use of a Dual Number is out of the scope of this post.

Reverse mode AD

The reverse mode automatic differentiation is calculated in two phases: The first phase calculates each of the intermediate values vi and records them in the computational graph. The second phase, referred to as backpropagation, calculates the derivatives by considering the adjoints of vi from the outputs to the inputs. The adjoints are defined as follows:

V̄i =∂y/∂v i

So, we're looking at how sensitive the current output is to changes in each vi. Adjoints show how much each vi contributes to the change in the output of the function f(X). During backpropagation, we start with the derivative of the function with respect to itself, which is 1.

The following table shows an example of reverse mode AD, by considering the same function as before.

Take a look at V̄1. Going back to the computational graph we see that v1 contributes to both v3 and v4 and that's why we need to sum their respective derivatives with respect to v1. Now if we evaluate ∂y/∂x1 withx1=2 and x2=3 we get the value of 7, the same value we calculated with forward mode AD.

In summary, Reverse mode AD can be defined with the following steps:

• The forward phase evaluation to get the function evaluation for the given inputs.
• The reverse phase starts with ȳ=1 and visits each of the parent nodes.
• For each parent node vi: calculate the contribution V̄i of the change in each vi to the change in the output y.
• At the end of the reverse mode you get the partial derivatives of the function with respect to each independent variable.

The AD Haskell library

Everything above was the theory behind Automatic differentiation. We now switch to Haskell and a library that implements this paradigm. The name is just ad, and at the time of writing this post, the latest release is 4.5.6 and has 376 stars on GitHub. This library offers a set of modules to compute derivatives (gradients, jacobians, directional derivatives and Hessians) by using either Reverse or Forward mode.

The core of this library is that each mode is represented by types and typeclasses; for example, AD and its forward mode have the following definitions:

newtype AD s a = AD { runAD :: a }
  deriving (Eq,Ord,Show,Read,Bounded,Num,Real,Fractional,Floating,Enum,RealFrac,RealFloat,Erf,InvErf,Typeable)

data Forward a
  = Forward !a a
  | Lift !a
  | Zero
  deriving (Show, Data, Typeable)

Take a look at the first constructor, (Forward !a a). The first argument is the value of the function evaluation and the second is its accumulative derivative (da).

The reverse mode is defined as follows:

data Reverse s a where
  Zero :: Reverse s a
  Lift :: a -> Reverse s a
  Reverse :: {-# UNPACK #-} !Int -> a -> Reverse s a
  deriving Show

And each mode is represented by a typeclass:

class (Num t, Num (Scalar t)) => Mode t where
  type Scalar t
  type Scalar t = t

  -- | allowed to return False for items with a zero derivative, but we'll give more NaNs than strictly necessary
  isKnownConstant :: t -> Bool
  isKnownConstant _ = False

  asKnownConstant :: t -> Maybe (Scalar t)
  asKnownConstant _ = Nothing

  -- | allowed to return False for zero, but we give more NaN's than strictly necessary
  isKnownZero :: t -> Bool
  isKnownZero _ = False

  -- | Embed a constant
  auto  :: Scalar t -> t
  default auto :: (Scalar t ~ t) => Scalar t -> t
  auto = id

  -- | Scalar-vector multiplication
  (*^) :: Scalar t -> t -> t
  a *^ b = auto a * b

  -- | Vector-scalar multiplication
  (^*) :: t -> Scalar t -> t
  a ^* b = a * auto b

  -- | Scalar division
  (^/) :: Fractional (Scalar t) => t -> Scalar t -> t
  a ^/ b = a ^* recip b

  -- |
  -- @'zero' = 'lift' 0@
  zero :: t
  zero = auto 0

Computing gradients with the AD library

Given the above function that we used for our manual examples, we can compute the gradient for the same inputs by using the grad function from the library. The following code snippet, shows the type signature and an example of computing the gradient with Forward mode:

ghci> :t grad
grad
  :: (Traversable f, Num a) =>
    (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f a -> f a
ghci> import Numeric.AD.Mode.Forward
ghci> grad (\[x1,x2]-> x1*x2 + x1*x1) [2,3]
[7,2]
ghci>

It takes:

• The function that we want to compute the gradient for any given mathematical function:

(forall s. f (AD s (Forward a)) -> AD s (Forward a)).

• The values where the gradient would be evaluated: f a

And returns the gradient in the context f a: [7,2] which is the gradient evaluated at [2,3].

There is also a grad function that implements the reverse mode. In the following code snippet, we can see that it has almost the same type signature as the forward mode version:

ghci> import Numeric.AD.Mode.Reverse
ghci> :t grad
grad
  :: (Traversable f, Num a) =>
    (forall s.(Reifies s Tape,Typeable s) =>f (Reverse s a) -> Reverse s a)
    -> f a 
-> f a
ghci> grad (\[x1,x2]-> x1*x2 + x1*x1) [2,3]
[7,2]
ghci>