Automatic Differentiation & Functional Operators in R

Symbolic differentation

Symbolic differentiation is basic calculus–manipulating function formulas to come up with other function formulas. Some software can do this as well, for example Wolfram Alpha can tell me that the \(\frac{d}{d x}(2 x + 3 x ^ 3)\) is \(2 + 9 x ^2\). This is pretty cool in itself, because the computer isn’t executing this code (it’s not multiplying any actual number by 2), it’s manipulating the formula as data.²

input: D[2 x + 3 x^3, x]
output: 2 + 9 x^2

However, because the functions used in deep learning are so huge and complex, this method isn’t used, as the formulas can grow exponentially in the size of the network.

Numerical differentation

This is the straightforward programming way to compute the derivative of a function at a point \(x\), or at least an approximation of it. Supposing f() is the function, and x = 2, we can evaluate the function at x and nearby:

x <- 2
xnear <- 2.0001
y <- f(x)
ynear <- f(xnear)
slope <- (ynear - y)/(xnear - x)

This doesn’t work well in practice: if xnear is very close to x then roundoff errors come into play, but if xnear is far from x then the approximation isn’t very good.

f() and f’() pairs

First, we need easy access to a function and its derivative function: they’re a matched set. I first thought to hold these in a small R list, starting with \sin() and \cos(), calling them mysin and mycos.

mycos <- list(f = cos, deriv = sin)
mysin <- list(f = sin, deriv = cos)

This worked to some degree, but to call the function or derivative I needed the awkward syntax mysin$f(x) or mysin$deriv(x). Worse, there’s no obvious way to take a second derivative: mysin$deriv(x) is a function in x, but I can’t get its derivative in an automated way.³

To work around the syntax issue, I made mycos() and mysin() callable functions:

mycos <- function(x) {
  return(cos(x))
}

mysin <- function(x) {
  return(sin(x))
}

Now, how can we attach the derivative of mysin to mysin itself, a function? Fortunately R allows us to assign “attributes” to any type of data, including functions.⁴ Ideally the derivative of mysin is mycos rather than just cos; this would allow for second derivatives and higher. My first attempt tried to assign the attributes directly.

attr(mycos, "getderiv") <- mysin
attr(mysin, "getderiv") <- mycos

However, order affects became apparent, because <- is assignment-by-copy. When we run attr(mycos, "getderiv") <- mysin it’s attaching a copy of mysin which as of yet doesn’t have a derivative attached, effectively breaking the ability to compute the second derivative of mycos. Then on attr(mysin, "getderiv") <- mycos, the mysin function gets its derivative as a copy of mycos that also has a workable first-level derivative. So for mysin we can get first and second derivatives, but not third. (Tl;dr: it didn’t work.)

Here’s the solution I landed on. In the "getderiv" attribute, we store an anonymous function whose only job is to return the derivative function. After setting the attribute, we can even call it with attr(mycos, "getderiv")() syntax.

attr(mycos, "getderiv") <- function() {return(mysin)}
attr(mysin, "getderiv") <- function() {return(mycos)}
deriv <- attr(mysin, "getderiv")()
print(deriv(c(-1, 0, 1)))

## [1] 0.5403023 1.0000000 0.5403023

print(cos(c(-1, 0, 1)))

## [1] 0.5403023 1.0000000 0.5403023

print(deriv)

## function(x) {
##   return(cos(x))
## }
## attr(,"getderiv")
## function () 
## {
##     return(mysin)
## }

Notice that deriv is a function that does indeed return cos(x), and furthermore, it itself has a "getderiv" function! This means we can get its derivative:

deriv_deriv <- attr(deriv, "getderiv")()
print(deriv_deriv(c(-1, 0, 1)))

## [1] -0.841471  0.000000  0.841471

print(deriv_deriv)

## function(x) {
##   return(sin(x))
## }
## attr(,"getderiv")
## function () 
## {
##     return(mycos)
## }

And we can continue to do so, because each function comes packaged with an anonymous function that returns the derivative. These anonymous functions are also closures: function() {return(mysin)} contains mysin bound to (referencing) the data in the calling scope, not a copy (see Appendix below for details). Note that mysin isn’t evaluated here until the anonymous function is run; otherwise we’d have an infinite-recursion issue.

Calling attr(mycos, "getderiv")() is really clunky, so let’s create a derivative function.

d <- function(f) {
  deriv <- attr(f, "getderiv")()
  return(deriv)
}

This is a higher-order function: it takes a function f (which is a function in x) and returns a function (the derivative). This makes good sense, as that’s what the derivative operator is: a function that maps functions to other functions. To call the derivative, we can now use d(mycos)(x), which is starting to look as much like mathematical notation as code.

x <- c(-1, 0, 1)
print(d(mysin)(x))    # calculate and compute derivative of cosin

## [1] 0.5403023 1.0000000 0.5403023

print(cos(x))       # should be the same

## [1] 0.5403023 1.0000000 0.5403023

We can even do multiple derivatives:

x <- c(-1, 0, 1)
print(d(d(mycos))(x))    # calculate and compute 2nd derivative of cosin

## [1] 0.5403023 1.0000000 0.5403023

print(mycos(x))          # should be the same

## [1] 0.5403023 1.0000000 0.5403023

Let’s see if we can plot mysin(x) and d(mysin)(x) for a sanity check.

library(ggplot2)
library(tidyr)

# create columns based on function and derivative values
xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = mysin(xs), 
                 dfx = d(mysin)(xs))

# reshape for plotting
df_toplot <- gather(df, ytype, value, -x)

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

So cool.

Addition and Subtraction

Now that we’ve got the representation worked out, let’s try implementing addition of functions, and derivatives of additions. Mathematically, we want to represent \(\frac{d}{dx}(f(x) + g(x)) = \frac{d}{dx}f(x) + \frac{d}{dx}g(x)\). Here \(+\) is also acting as a higher-order function, taking two input functions \(f\) and \(g\) and returning the addition. It’s an “infix” function, with the function name in the middle of the two parameters.

R lets us create infix functions easily, by using % in the function name. If `%add%` <- function(a, b) { return(a + b)}, then we can call either `%add%`(2, 5) or 2 %add% 5 to get back 7. We just have to use the backticks when defining or using prefix notation.

Let’s define a %+% function for our functions. As above the result of %+% should also come paired with its derivative, the sum of the functions run through d().

`%+%` <- function(f, g) {
  func <- function(x) {
    return(f(x) + g(x))
  }
  
  attr(func, "getderiv") <- function() {return(d(f) %+% d(g))}
  return(func)
}

Some interesting stuff is happening here: true to form, %+% is a function mapping two input functions in x to a single output function in x. Defined inside, func calls the input functions on the given x (and these are bound via closure). Like mycos and mysin above, the "getderiv" attribute is a function that returns the derivative, and since we create that derivative via d() and the %+% operator, it will also have its own derivative attached. (We’re working now on the assumption that all higher-order functions we create will take and return such matched-pair functions.) Finally, just as above, d(f) %+% d(g) is not evaluated, just defined with variables bound, so there’s no infinite recursion. Basically a lazy list of functions.⁵

Let’s try it.

print(sin(0.5))

## [1] 0.4794255

print(cos(0.5))

## [1] 0.8775826

print(sin(0.5) + cos(0.5))

## [1] 1.357008

sumfunc <- mysin %+% mycos
print(sumfunc(0.5))

## [1] 1.357008

print(cos(0.5) + cos(0.5))

## [1] 1.755165

sumfunc <- mycos %+% d(mysin)    # cos + cos
print(sumfunc(0.5))

## [1] 1.755165

A plot:

xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = mycos(xs), 
                 gx = mysin(xs),
                 sumfxgx = (mycos %+% mysin)(xs))

df_toplot <- gather(df, ytype, value, -x)

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

Let’s get %-% out of the way while we’re at it.

`%-%` <- function(f, g) {
  func <- function(x) {
    return(f(x) - g(x))
  }
  
  attr(func, "getderiv") <- function() {return(d(f) %-% d(g))}
  return(func)
}

Multiplication and Division

Now it’s time to get brave, and implement %*% with the product rule, \(\frac{d}{dx}(f(x) g(x)) = g(x)\frac{d}{dx}f(x) + f(x)\frac{d}{dx}g(x)\), or more succinctly \((f \cdot g)' = f'\cdot g + g'\cdot f\). In this case the derivative uses all four of d(f), f, d(g) and g.

`%*%` <- function(f, g) {
  func <- function(x) {
    return(f(x) * g(x))
  }

  attr(func, "getderiv") <- function() {return((g %*% d(f)) %+% (f %*% d(g)))}
  return(func)
}

Crossing fingers…

xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = mysin(xs), 
                 gx = mycos(xs),
                 fxgx = (mycos %*% mysin)(xs))

df_toplot <- gather(df, ytype, value, -x)

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

And %/%, based on \((f/g)' = (f' g - g'f)/(g^2)\).

`%/%` <- function(f, g) {
  func <- function(x) {
    return(f(x)/g(x))
  }
  
  attr(func, "getderiv") <- function() {return( ((d(f) %*% g) %-% (d(g) %*% f)) %/% (g %*% g) )}
  return(func)
}

Constants and Variables

Now, given that we can handle multiplication, can we compute \(\frac{d}{dx}2x\) and get back just 2, or do we need a special case? It turns out we can use our %*% function, if we can treat 2 as a function. Specifically, we want it to be a function in x that always returns 2, no matter the input, since this will allow us to properly combine it with the others.

It will also need an attached derivative. The derivative of a constant is 0, also a constant, which we’ll represent in the same way. Let’s create a const() that returns such a function given the constant of interest; thus const(2) would return one of our functions, as would const(0).

const <- function(input) {
  func <- function(x) {
    return(input)
  }
  
  attr(func, "getderiv") <- function() { return(const(0)) }
  return(func)
}

Here again closures playing an integral role.⁶ I also like that a const() function’s derivative is just const(0), and d() maps const(0) to itself.

xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = (mysin %+% const(2))(xs), 
                 dfx = d(mysin %+% const(2))(xs))

df_toplot <- gather(df, ytype, value, -x)

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

By the way, footnote readers know that we made a mistake in the definition of mycos earlier, in that d(mycos) should be const(-1) %*% mysin rather than just mysin. Now that we can work with constants like -1 we can finally fix it:

attr(mycos, "getderiv") <- function() {return(const(-1) %*% mysin)}

We still can’t quite compute \(\frac{d}{dx}2x\) because we don’t yet have a way to represent \(x\); all we have to work with are functions, and now constants (as functions). Much like our const() function though, we can treat \(x\) like a function: one that just returns its input, with a derivative of const(1).

xvar <- function(x) {
  return(x)
}

attr(xvar, "getderiv") <- function() {return(const(1))}

Now we can refer to the \(x\) variable as a function by xvar:

xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = (const(2) %*% xvar)(xs), 
                 dfx = d(const(2) %*% xvar)(xs))

df_toplot <- gather(df, ytype, value, -x)

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

Composition

What about something like \(\frac{d}{dx}\sin(\cos(x))\)? Mathematically this is derived with the chain rule \(\frac{d}{dx}f(g(x)) = f'(g(x))g'(x)\). But \(\sin(\cos(x))\) is difficult to represent in our scheme, since mysin(mycos) would nonsensically attempt to call mysin on the mycos function. What we need is the composition operator, \(\circ\), which is the infix version of calling \(f\) on the result of \(g\): \((f \circ g)(x) = f(g(x))\). We’ll use %o%:

`%o%` <- function(f, g) {
  func <- function(x) {
    return(f(g(x)))
  }
  
  attr(func, "getderiv") <- function() {return(  (d(f) %o% g) %*% d(g) )}
  return(func)
}

Let’s have fun with \(\sin(2\cos(x))\) and its first and second derivatives.

xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = (mysin %o% (const(2) %*% mycos))(xs), 
                 dfx = d(mysin %o% (const(2) %*% mycos))(xs),
                 ddfx = d(d(mysin %o% (const(2) %*% mycos)))(xs))

df_toplot <- gather(df, ytype, value, -x) 

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

Power Rules

Rather than implement the power rule, \(\frac{d}{dx}x^n = nx^{n-1}\), or even the general power rule, \(\frac{d}{dx}f(x)^n = nf(x)^{n-1}\frac{d}{dx}f(x)\), we’ll go straight for the functional generalized power rule: \((f^g)' = -f^g\left( f' \frac{g}{f} + g' \ln f \right)\). This version derivates a function raised to the power of another, as in \(\frac{d}{dx}\cos(x)^{\sin(s)}\).

First though, we need to implement the natural logarithm, whose derivative is \(\frac{d}{dx}\ln(x) = 1/x\). (R uses log() for natural logarithm.)

myln <- function(x) {
  return(log(x))
}
attr(myln, "getderiv") <- function() {return(const(1) %/% xvar)}


`%^%` <- function(g, h) {
  func <- function(x) {
    return(g(x) ^ h(x))
  }

  attr(func, "getderiv") <- function() {return(  
                                           const(-1) %*% (g %^% h) %*%
                                                   (  
                                                     (d(h) %*% (myln %o% g)) %+%
                                                     (h %*% (d(g) %/% g))
                                                   )
                                              )} 
  return(func)
}

This example only plots \(\cos^3(x)\) and \(\frac{d}{dx}\cos^3(x)\), but since const(3) is a function in x it fits in this scheme too.

xs <- seq(0, 3*pi, 0.1)
df <- data.frame(x = xs, 
                 fx = (mycos %^% const(3))(xs), 
                 dfx = d(mycos %^% const(3))(xs))

## Warning in log(x): NaNs produced

df_toplot <- gather(df, ytype, value, -x) 

p <- ggplot() +
  geom_line(data = df_toplot, aes(x = x, y = value, color = ytype)) 
plot(p)

## Warning: Removed 16 rows containing missing values (geom_path).

When \(f(x)\) is negative the derivative isn’t defined according to the generalized power rule, because the logarithm isn’t defined for negative numbers. The general power rule would work, but can only accept a const() as the power.

Speaking of logarithms and exponents, here’s \(e^x\), which is its own derivative:

# this one is too dang easy ;)
myexp <- function(x) {
  return(exp(x))
}
attr(myexp, "getderiv") <- function() {return(myexp)}

Acknowledgements

I’d like to thank Tarak Shah (@tarakc02) for helpful comments. Thanks Tarak!