Automatic Differentiation

Motivation

• Need to calculate the derivative of some arbitrary function when calculating the function
• Example: minimizing the loss function for gradient descent for machine learning

Aside: gradient descent

• You have some neural network, that has an associated loss function, f(x, y, ...)
• Loss function represents the error in your model
• Goal: minimize the loss function by varying the parameters (x, y, ...)
• But how do we decide which direction will reduce fastest?

Aside: gradient descent

• Calculus tells us that this occurs along the direction of the negative gradient
• This means we need to compute the value of the derivative of the loss function
• One such way (which is used in TensorFlow and PyTorch) is automatic differentiation

Alternatives

• Numerical Differentiation
• Symbolic Differentiation

Numerical Differentiation

$$\frac{d}{{dx}}f\left( x \right) = \mathop {\lim }\limits_{h \to 0} \frac{{f\left( {x + h } \right) - f\left( x \right)}}{h }$$

• Follows from definition of a derivative
• Choose small h to approximate the derivative

Numerical Differentiation

• Problem: rounding errors due to inexactness of floating point math
• Problem: can be inefficient

Numerical Differentiation

$$f(x) = x^2$$

$$f'(2) = \frac{f(2 + 0.000001) - f(2)}{0.000001}$$

$$f'(2) = \frac{2.000001 ^ 2 - 2 ^ 2}{0.000001}$$

Alternatives

• Numerical Differentiation
• Symbolic Differentiation

Symbolic Differentiation

• Manipulate the mathematical expression itself, and apply differentiation rules

Symbolic Differentiation

• Problem: Need to convert the code into a single expression, in a form the differentiator can manipulate
• Problem: Control flow like if and while statements, and intermediate variables are difficult to handle

Solution: Automatic Differentiation!

Demo

https://repl.it/@NeilParikh1/Automatic-Differentiation

from ad import AD, sin

def f(x):
    return (x + x) ** 3

def g(x):
    return sin(((x + x) * x) ** 3)

print(f(AD(2))) # AD(x = 64, dx = 96)
print(g(AD(2))) # AD(x = 0.07951849401287635, dx = -1531.1360883428383)

How does this work?

Core idea: chain rule

def g(x):
    return sin(((x + x) * x) ** 3)

def g(x):
    a = x + x
    return sin((a * x) ** 3)

def g(x):
    a = x + x
    b = a * x
    return sin(b ** 3)

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

Derivative is:

dg/dx = cos(c) * ...

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

Derivative is:

dg/dx = cos(c) * dc/dx

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

Derivative is:

dc/dx = 3 * b**2 * db/dx
dg/dx = cos(c) * dc/dx

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

Derivative is:

db/dx = a * dx/dx + x * da/dx
dc/dx = 3 * b**2 * db/dx
dg/dx = cos(c) * dc/dx

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

Derivative is:

da/dx = dx/dx + dx/dx
db/dx = a * dx/dx + x * da/dx
dc/dx = 3 * b**2 * db/dx
dg/dx = cos(c) * dc/dx

def g(x):
    a = x + x
    b = a * x
    c = b ** 3
    return sin(c)

g(AD(2))

def g(x):
    a = AD(2) + AD(2)
    b = a * AD(2)
    c = b ** 3
    return sin(c)

class AD:
    def __init__(self, x, dx = 1):
        self.x = x
        self.dx = dx

def g(x):
    a = AD(2, 1) + AD(2, 1)
    b = a * AD(2, 1)
    c = b ** 3
    return sin(c)

def g(x):
    a = AD.add(AD(2, 1), AD(2, 1))
    b = AD.mult(a, AD(2, 1))
    c = AD.pow(b, 3)
    return AD.sin(c)

• When doing computations, combine both the value and derivative using chain rule

def add(a, b):
  x = a.x + b.x
  dx = a.dx + b.dx
  return AD(x, dx)

def g(x):
    a = add(AD(2, 1), AD(2, 1))
    b = mult(a, AD(2, 1))
    c = pow(b, 3)
    return sin(c)

def g(x):
    a = AD(2 + 2, 1 + 1) # add(AD(2, 1), AD(2, 1))
    b = mult(a, AD(2, 1))
    c = pow(b, 3)
    return sin(c)

def g(x):
    a = AD(4, 2)
    b = mult(a, AD(2, 1))
    c = pow(b, 3)
    return sin(c)

def g(x):
    b = mult(AD(4, 2), AD(2, 1))
    c = pow(b, 3)
    return sin(c)

def mult(a, b):
  x = a.x * b.x
  dx = a.x * b.dx + b.x * a.dx
  return AD(x, dx)

def g(x):
    b = mult(AD(4, 2), AD(2, 1))
    c = pow(b, 3)
    return sin(c)

def g(x):
    b = AD(4 * 2, 4 * 1 + 2 * 2) # mult(AD(4, 2), AD(2, 1))
    c = pow(b, 3)
    return sin(c)

def g(x):
    b = AD(8, 8)
    c = pow(b, 3)
    return sin(c)

def g(x):
    c = pow(AD(8, 8), 3)
    return sin(c)

def pow(a, b):
  x = a.x ** b
  dx = b * a.x ** (b - 1) * a.dx
  return AD(x, dx)

def g(x):
    c = pow(AD(8, 8), 3)
    return sin(c)

def g(x):
    c = AD(8 ** 3, 3 * 8 ** 2 * 8) # pow(AD(8, 8), 3)
    return sin(c)

def g(x):
    c = AD(512, 1536)
    return sin(c)

def g(x):
    return sin(AD(512, 1536))

def sin(a):
    x = math.sin(a.x)
    dx = a.dx * math.cos(a.x)
    return AD(x, dx)

def g(x):
    return sin(AD(512, 1536))

def g(x):
    return AD(math.sin(512), 1536 * math.cos(512)) # sin(AD(512, 1536))

def g(x):
    return AD(0.07951849401287635, -1531.1360883428383)

Other notes

• This technique is used in PyTorch, through the autograd package
• Implemented in TensorFlow as tf.GradientTape
• Calculating f'(x) is just as fast (asymptotically) as calculating f(x)!

Questions?