Build Your Own Automatic Differentiation Program | by Jonathan Kernes | Towards Data Science

For anyone who’s completely lost at how graphs can be used to compute derivatives, or just wants to know how TensorFlow works at a fundamental level, this is your guide.
Deep learning is so popular and so widespread these days, it’s easy to ask the question, “Why did it take so long for people to figure this out?” There are a number of reasons why, and getting into the history of deep learning is really not the focus of this post (the introduction to Reference [1] does a nice job), but one of the key advances that makes this revolution roll is automatic differentiation (AutoDiff).
If you want to get right to technical details, skip to the section “Building an AutoDiff”, if you just want the code, skip to the section “Implementing AutoDiff in Python”
For a companion Jupyter notebook with the code used in this article, please see the following link: https://github.com/Jmkernes/Automatic-Differentiation
What’s so great about automatic differentiation you ask? Imagine you’ve got some function you’re playing around with, and you really want to know its derivative. If it’s a relatively simple function, you can probably just rattle off the answer. If f(x)=x3, then we can quickly conclude df/dx=3x2. What if it’s more complicated and you don’t just know the answer? Well, you don’t really have a lot of options here.
Option 1: if you want the numerical answer, you can go back to basics and straight up compute finite differences, using formulas like
and their more sophisticated and more accurate cousins [2]. But that’s not that satisfying. Maybe we want the symbolic answer, in terms of x’s and y’s and stuff, in which case a numerical answer just isn’t going to cut it. Or, maybe our differentiation variable x is actually a large multi-dimensional tensor, and computing the numerical difference one-by-one for every entry of that thing is going to be horribly inefficient and time-intensive.
Option 2: derive the derivative. This is supposed to be meta. What does it mean to derive a derivative? How did we know the answer to df/dx=3x2 before? Well, it turns out there’s really only two ways you figure out derivatives. If it’s a totally new function, you just have to derive the answer via its limit representation shown above. In reality though, most functions aren’t new, they’re actually just compositions of other functions. In this case, we can figure out new derivatives via the chain-rule:
One of the first physics projects I worked on involved the use some special functions called generalized hypergeometric functions. If you don’t know what those are, don’t worry, that’s the point. Like every special function, they’re miserable to deal with, and contain enough identities and special cases to fill up a Harry Potter spell book. So, disheartened, I went to my advisor and directly asked him, “Hey what in the world is this thing?”
His response was pretty simple — “It’s just some function that we know some stuff about.” That seems like a useless answer, but it’s actually kind of deep. He then asked me “you know what a cosine is, right?” To which I snarkily responded, “yeah of course I know what a cosine is.” Without hesitation, he snapped right back, “But how do you know what a cosine is? How is knowing what a hypergeometric function any different that knowing what a cosine is? At the end of the day, you know what a cosine is because you can name some identities and you can make a pretty plot of it. That’s it.”
At this point, I was stumped. I was having a mathematical existential crisis, and re-evaluating everything I had ever learned. “I know what a cosine is because I know it’s derivative is -sin(x)!”, I thought. “And, I know what sine is! Because it’s derivative is cosi...“, oh no. An infinite loop.
After reviving from my existential crisis, I came to the realization that there are two types of functions in the world: things we are familiar with like x2 or x+5, and things we are unfamiliar with, but can be composed of familiar things — like (x+5)2. Automatic differentiation is centered around this latter concept. We can frame its mission statement as:
Given a collection of elementary functions, things like e^x, cos(x), or x2, then using the rules of calculus, it is possible to determine the derivative of any function that is composed of these elementary functions.
Since we can represent functions in all sorts of ways, like through power series, Fourier series, etc. that’s a pretty powerful statement. It all works via the chain rule. The awesome thing about AutoDiff, is that it’s the mathematical equivalent of journeying around the world with zero planning. You can just keep composing and stacking functions as your heart desires, going deeper and deeper, always assured that AutoDiff is going to be able follow the breadcrumbs and compute a derivative for you.
For the rest of the post, we will assume that every function we want the derivative of can be broken up into elementary functions (if you want the derivative of a modified Bessel function of the second kind, you’re going to have to look elsewhere)
By the end of this post, you will be able to generate graphs like the following:
But with as much complexity as you like. Let’s break down what this graph means. Suppose that we are given the function
and we would like to compute the partial derivatives of f with respect to its variables x, y, and z. We can decompose this function into a composition of very basic functions, namely, a function add(x,y)=x+y that takes two number and adds them, and a function multiply(x,y)=x*y which takes two numbers and multiplies them. Basic, right? In this new form we have
It’s more a more wordy representation, but... it’s not wrong.
This form is useful, because we immediately know the derivatives of add and multiply. Furthermore, we have rewritten this as a composition of two elementary functions, so we can apply the chain rule. We obtain
Using basic calculus, we get 1*y=y, which is obviously the right answer. Though you may not notice it, whenever you compute derivatives you implicitly perform the above operation. Now, let’s make a computer do that.
In order to specify the procedure, we need to do two things: 1) tell the computer what the fundamental building blocks are 2) tell it how to put those blocks together to compute a derivative.
Let’s focus on part one first. There are four fundamental types of blocks, which we will call nodes, making use of the graph-theoretic language. These include
Variables — mutable objects (tensors) of which we want to know the gradients of.Constants — immutable objects (tensors) whose value, once specified, does not change.Placeholders — these are basically empty containers specifying where a future Variable or Constant will go, once we decide what that value should be.Operators — these are the conceptually hardest to picture, but also the most important.
Variables — mutable objects (tensors) of which we want to know the gradients of.
Constants — immutable objects (tensors) whose value, once specified, does not change.
Placeholders — these are basically empty containers specifying where a future Variable or Constant will go, once we decide what that value should be.
Operators — these are the conceptually hardest to picture, but also the most important.
The first three should be relatively straightforward. What is slightly confusing about computation graphs that takes some time to get used to, is that the nodes of the graph are actually operators and NOT variables, constants, or placeholders. Intuitively, we like to think of functions as composed of their values. When we wrote f(x, y, z) we used x, y, z. We did not write f(add, multiply). However, that’s exactly what our computer will be looking for. In general, an operator node takes inputs and returns outputs. We assume that the node emits one output only; though a tensor contains many entries, it can be considered a single output
We can now understand the graph at the beginning of the section. We only used numbers, but those could also be considered to be scalar tensors. The graph thus shows how tensors flow through computation. Get it? Tensors...Flow...TensorFlow. Now let’s consider the opposite direction, computing the gradients.
To compute the actual value of the output f(x,y,z) we run the graph forward, in the direction of the arrows, beginning with the variables x, y, z, and flowing through the Operators. Let’s now take our equation and think about it graphically. We will add one additional, trivial computation
which will help us initialize our back-propagation procedure. Inspecting our equation for the partial derivatives of f, an important thing to note is that a derivative is also a function! This means that the chain-rule derivative is also a composition of functions. Remember, composition of functions = graph. That’s the rule. What’s really nice, is that our graph structure is identical, but run backwards. You can see the structure below:
The red arrows signify the flow of derivatives from the final output to the start as a reversed computation graph. It can be computed exactly the same way, where we supply the first node with a derivative of 1, using the trivial identity df/df=1. Our goal should now be clear:
Specify all variables, placeholders, and constants in our graph
Compose these together using operators to form the final, scalar output
Compute a forward-pass which inputs values and flows through the operators to get the output. Store each node’s output value in that node.
Compute a backward-pass to compute gradients, beginning from a gradient of 1 at the final node (the output).
Making this work is the goal of the remainder of this article.
We want warn the reader up front that this involves global variables. Global variables are taboo, but the reason we use them should hopefully become clear as we go along. The graph is our global object (ok not a global variable actually a global object, sorry for the lie), which we will be adding nodes to. As far as I can tell, there isn’t an easy way to explicitly scope objects, and I believe TensorFlow, at least in V1, also used a global graph variable. If anyone has a solution to avoid global variables please post it in the comments!
Ok, warning aside, let’s get started. We first need to define our default graph class. This will be a basic structure that just keeps track of the current nodes and their types.
A few things: everything after the __init__ method you can just forget about. It doesn’t affect the code and is just there to try to make it slightly easier to deal with creating localized graphs. You can also just delete it, and make sure to run Graph()to initialize the global default graph, and del _g to erase it before and after defining a bunch of nodes. Also, we don’t really need numpy, it’s there for helping us do matrix multiplication later.
Great, we’ve got a working graph. Now let’s define each of those four classes we discussed earlier. They’re more or less all the same, with each containing a couple of unique features
The Operator class needs an inputs attribute, the Constant class needs to protect it’s values, and the rest are kind of similar. We put in some default names, like Var1, Var2, ... which are about the minimal level of descriptiveness. The name field is important for later retrieving nodes from the graph, since we won’t have any other way to identify them
Also, it’s very important we keep track of the names of the Placeholders, otherwise we really won’t have any way of hunting them down in the computational graph.
We put in some other things like a __repr__ function to automatically display nodes as a string in a format like “Operator: name:mul/2, value:3.875”. That should help us know right away what node we’re dealing with.
Now, to actually make our graph do something, we need to give it some real operators. These are going to precisely the elementary operations we’ve been talking about this whole time. Things like addition, subtraction, multiplication. We will define some of these, and to make things much easier to code, we are going to overload all of these operators (note: I couldn’t think of a better way to do this that doesn’t either repeat a ton of code or require adding methods to the base Node class in a weird way. Open to better suggestions). The one tricky operator is __neg__. Remember x -> -x is actually a multiplication operator, so the graph needs to know about that.
We’ve now got all of our infrastructure in place. We can easily build the graph by just defining Variables, Constants, or Placeholders, and performing computation per usual. The operators will take care of adding nodes to the graph and everything. We can create a graph g and add some operations via the following syntax
with Graph() as g:  x = Variable(1.3)  y = Variable(0.9)  z = x*y+5g.variables # Will output {Variable: name:Var/0, value:1.3, Variable: name:Var/1, value:0.9}g.operators # Will output {Operator: name:add/0, Operator: name:mul/0}g.constants # Will output {Constant: name:Const/0, value:5}
As you can see, the graph correctly created a constant, and also kept track of the two operations, x*y multiplication as well as addition of the constant 5. If we run with Graph() as g again, it should erase the current default graph and begin building anew. Cool, everything works. Now, we have to use a little bit of algorithms knowledge to compute our forward and backward passes.
As it stands, our default graph just keeps track of the current nodes. It has no structure. Ultimately, we want to run our graph forward, meaning we feed it all the Variables, Constants, and Placeholders, then these flow through the operators to our output. The order that we do this in is a classical algorithms problem. We need to perform a topological sort (the CLRS book [3] is always good, or just use wikipedia). Our graph is a directed acyclic graph, meaning that each node has a beginning and end point. There are no loops. The topological sort guarantees that if we run every node in order, we will never be lacking information to compute a node. It’s like ordering class prerequisites; you adjust your schedule to to take Calc I before Calc III.
After a topological sort, our graph has truly become a graph. It finally has some solid structure! Even if this structure is just an array of ordered nodes. To complete our program, we perform the passes. The forward pass is really simple; just compute the value of each node in order, and store its value in that node so future nodes can access it if need be.
The backward pass is slightly more difficult, but not much so. Initialize the root node with a value of 1 (recall the discussion of using df/df=1). Then, only traverse the graph via its Operators. In other words, compute the gradients for all inputs of the current operator, then recursively do the same for any of those inputs that were also an operator. Here’s the code:
That’s it! We’re done with the main part of the program. What remains is to play around with this and analyze it a bit. Here’s an example of how you might calculate the gradients for a simple function
which depends on two variables and one constant. We can easily compute its derivative by hand and check that our code gets it right. First the forward pass:
val1, val2, val3 = 0.9, 0.4, 1.3with Graph() as g:  x = Variable(val1, name='x')  y = Variable(val2, name='y')  c = Constant(val3, name='c')  z = (x*y+c)*c + x  order = topological_sort(z)  res = forward_pass(order)  grads = backward_pass(order)  print("Node ordering:")  for node in order:      print(node)   print('-'*10)  print(f"Forward pass expected: {(val1*val2+val3)*val3+val1}")  print(f"Forward pass computed: {res}")---------------------------------------------------""" Output:Node ordering:Variable: name:x, value:0.9Variable: name:y, value:0.4Operator: name:mul/6Constant: name:c, value:1.3Operator: name:add/6Operator: name:mul/7Operator: name:add/7----------Forward pass expected: 3.0580000000000003Forward pass computed: 3.0580000000000003"""
And for the backward pass:
dzdx_node = [a for a in order if a.name=='x'][0]dzdy_node = [a for a in order if a.name=='y'][0]dzdc_node = [a for a in order if a.name=='c'][0]print(f"dz/dx expected = {val3*val2+1}")print(f"dz/dx computed = {dzdx_node.gradient}")print(f"dz/dy expected = {val1*val3}")print(f"dz/dy computed = {dzdy_node.gradient}")print(f"dz/dc expected = {val1*val2+2*val3}")print(f"dz/dc computed = {dzdc_node.gradient}")---------------------------------------------------""" Output:dz/dx expected = 1.52dz/dx computed = 1.52dz/dy expected = 1.1700000000000002dz/dy computed = 1.1700000000000002dz/dc expected = 2.96dz/dc computed = 2.96"""
Ok, so I promised you would be able to make pretty pictures by the end of this. Here’s me keeping that promise. We will use the python library Graphviz. You may have to pip install graphviz which is why I left this off until the end. To build the graph, we just create a node for every node in the graph, then we draw edges using operators. The operators dictate the flow of the graph, so just like when we did the backward_pass, we just need the operator edges. By using the following code
We can make pretty pictures like the following
Which represents the graph of the test code we ran earlier.
If you’ve made it this far, thank you for reading! To summarize, we’ve shown that automatic differentiation is built on two ideas
All functions we are usually interested in computing the derivatives of can be represented as compositions of elementary functionsThese elementary functions are simple, or at the very least have known derivatives with which we can supply our AutoDiff program
All functions we are usually interested in computing the derivatives of can be represented as compositions of elementary functions
These elementary functions are simple, or at the very least have known derivatives with which we can supply our AutoDiff program
From there, we can create a graph, topologically sort, then run a forward-backward pass to get gradients. I hope this has been useful, and good introduction to how AutoDiff systems like TensorFlow work. Sometimes, if you just think a little bit harder about something you already know pretty well, you might just find something interesting!
[1] Goodfellow, Ian, et al. Deep learning. Vol. 1. No2. Cambridge: MIT press, 2016.
[2] https://en.wikipedia.org/wiki/Five-point_stencil
[3] Cormen, Thomas H., et al. Introduction to algorithms. MIT press, 2009.