A Fully Annotated Autograd Tutorial

Backpropagation

Suppose instead we had an elaborate computation graph with thousands of variables and functions. Suppose also that \(z\) is some variable in the computation graph and that we'd like to calculate the derivative of \(z\) with respect to every single ancestor of \(z\). We'd need a method that systematically applies the chain rule to make these calculations. This is exactly what backpropagation accomplishes.

Think about the earlier statement, that said if we have already calculated \(\frac{dz}{dy}\) for all children \(y\) of \(x\), then the inductive chain rule can be invoked to calculate the \(\frac{dz}{dx}\). What we'd like then, is an ordering of all ancestors of \(z\) such that if we calculated the derivatives in that order, each ancestor would have its children's derivatives calculated before itself.

The algorithm that returns this ordering is called topological sort. I recommend learning about this algorithm if you are not familiar with it, there are many resources online. By triggering topological sort on \(z\), treating edges as dependencies, we obtain an ordering such that for any ancestor \(x\) of \(z\), \(x\) appears after in the ordering after its own children. Backpropagation then calculates derivatives in this order, and keeps track of all computed derivatives at each step in a dictionary. The pseudocode is shown below.

BACKPROPAGATION(Computation graph G, node z in G):

Let topo_order = TopologicalSort(G, z)

Let computed = {dz/dz : 1}

For each node x in topo_order (excluding z):

    Let dz/dx = 0

    # Invoke Inductive Chain Rule
    For each child y of x:
        Let partial_x = calculate(∂y/∂x)
        Let total_y = computed.get(dz/dy)
        dz/dx += total_y * partial_x

    insert dz/dx into computed

Returning to the computation graph in Figure 1, the ordering returned by topological sort might be \((z, y_2, y_1, x_2, x_1, x_3, x_4)\). Other orderings are possible too, as multiple orderings satisfy the dependencies in this case. From here, the loop on line 7 would proceed:

1. current node: \(y_2\)

   computed = \(\{\frac{dz}{dz} : 1\}\)

   calculate: \(\frac{\partial z}{\partial y_2} = 1\)

   invoke chain rule: \(\frac{dz}{dy_2} = \frac{dz}{dz} \cdot \frac{\partial z}{\partial y_2} = 1 \cdot 1 = 1\)

2. current node: \(y_1\)

   computed = \(\{\frac{dz}{dz} : 1, \frac{dz}{dy_2} : 1\}\)

   calculate: \(\frac{\partial z}{\partial y_1} = x_1\)

   invoke chain rule: \(\frac{dz}{dy_1} = \frac{dz}{dz} \cdot \frac{\partial z}{\partial y_1} = 1 \cdot x_1 = x_1\)

3. current node: \(x_2\)

   computed = \(\{\frac{dz}{dz} : 1, \frac{dz}{dy_2} : 1, \frac{dz}{dy_1} : x_1\}\)

   calculate: \(\frac{\partial y_1}{\partial x_2} = 1\)

   invoke chain rule: \(\frac{dz}{dx_2} = \frac{dz}{dy_1} \cdot \frac{\partial y_1}{\partial dx_2} = x_1 \cdot 1 = x_1\)

4. current node: \(x_1\)

   computed = \(\{\frac{dz}{dz} : 1, \frac{dz}{dy_2} : 1, \frac{dz}{dy_1} : x_1, \frac{dz}{dx_2} : x_1\}\)

   calculate: \(\frac{\partial z}{\partial x_1} = y_1\) and \(\frac{\partial y_1}{\partial x_1} = 1\)

   invoke chain rule: \(\frac{dz}{dx_1} = \frac{dz}{dz} \cdot \frac{\partial z}{\partial x_1} + \frac{dz}{dy_1} \cdot \frac{\partial y_1}{\partial x_1} = 1 \cdot y_1 + x_1 \cdot 1 = y_1 + x_1\)

5. current node: \(x_3\)

   computed = \(\{\frac{dz}{dz} : 1, \frac{dz}{dy_2} : 1, \frac{dz}{dy_1} : x_1, \frac{dz}{dx_2} : x_1, \frac{dz}{dx_1} : y_1 + x_1\}\)

   calculate:\(\frac{\partial y_1}{\partial x_3} = 1\)

   invoke chain rule: \(\frac{dz}{dx_3} = \frac{dz}{dy_1} \cdot \frac{\partial y_1}{\partial x_3} = x_1 \cdot 1 = x_1\)

6. current node: \(x_4\)

   computed = \(\{\frac{dz}{dz} : 1, \frac{dz}{dy_2} : 1, \frac{dz}{dy_1} : x_1, \frac{dz}{dx_2} : x_1, \frac{dz}{dx_1} : y_1 + x_1, \frac{dz}{dx_3} : x_1\}\)

   calculate: \(\frac{\partial y_1}{\partial x_4} = 1\) and \(\frac{\partial y_2}{\partial x_4} = 2x_4\)

   invoke chain rule: \(\frac{dz}{dx_4} = \frac{dz}{dy_1} \cdot \frac{\partial y_1}{\partial x_4} + \frac{dz}{dy_2} \cdot \frac{\partial y_2}{\partial x_4} = x_1 \cdot 1 + 1 \cdot 2x_4 = x_1 + 2x_4\)

7. Final output: \(\{\frac{dz}{dz} : 1, \frac{dz}{dy_2} : 1, \frac{dz}{dy_1} : x_1, \frac{dz}{dx_2} : x_1, \frac{dz}{dx_1} : y_1 + x_1, \frac{dz}{dx_3} : x_1, \frac{dz}{dx_4} : x_1 + 2x_4 \}\)

At each iteration of backpropagation, for the current node \(x\), we invoke the chain rule using the previously calculated total derivatives of \(z\) with respect to children of \(x\), combined with the newly calculated partial derivatives of children of \(x\) with respect to \(x\) itself. At each iteration, the dictionary of computed total derivatives grows by one.

Extending Backpropagation to Tensors

We just saw the inductive chain rule and backpropagation in the context of scalar variables and functions. We will now extend it to tensor valued variables and functions. This will allow us to express variables and compositions of functions in a more compact way, and it will allow us to invoke the chain rule "in bulk". We'll start with an example.

Let \(A,B \in \mathbb{R}^{2 \times 2}\) be two matrices, and suppose we let \(z = \text{sum}(A \times B)\), meaning \(z\) is a scalar resulting from summing the entries of the matrix given by \(A \times B\). It makes perfect sense to for example be intersted in the derivative of \(z\) with respect to the entry \(a_{11}\) of \(A\). In fact, \(z = \text{sum}(A \times B)\) can secretly be expressed as a computation graph with scalar valued functions and variables, like we saw before:

\[ A = \begin{bmatrix} a_{11} & a_{12} \\ a_{21} & a_{22} \end{bmatrix} = \begin{bmatrix} 3 & 7 \\ 2 & 5 \end{bmatrix} \] \[ B = \begin{bmatrix} b_{11} & b_{12} \\ b_{21} & b_{22} \end{bmatrix} = \begin{bmatrix} 2 & 0 \\ 0 & 4 \end{bmatrix} \] \[ C = A \times B = \begin{bmatrix} c_{11} & c_{12} \\ c_{21} & c_{22} \end{bmatrix} = \begin{bmatrix} 3*2 + 7*0 & 3*0+7*4 \\ 2*2+5*0 & 2*0 + 5*4 \end{bmatrix} = \begin{bmatrix} 6 & 28 \\ 4 & 20 \end{bmatrix} \]

where the computation graph is given by:

1. \(c_{11} = a_{11}*b_{11} + a_{12}*b_{21}\)
2. \(c_{12} = a_{11}*b_{12} + a_{12}*b_{22}\)
3. \(c_{21} = a_{21}*b_{11} + a_{22}*b_{21}\)
4. \(c_{22} = a_{21}*b_{12} + a_{22}*b_{22}\)
5. \(z = c_{11} + c_{12} + c_{21} + c_{22}\)

We can invoke backpropagation on this computation graph just like before, and arrive at the derivatives (try and calculate these for yourself):

\(\frac{dz}{dc_{11}} = \frac{dz}{dc_{12}} = \frac{dz}{dc_{21}} = \frac{dz}{dc_{22}} = 1, \frac{dz}{da_{11}} = \frac{dz}{da_{21}} = b_{11} + b_{12}, \frac{dz}{da_{12}} = \frac{dz}{da_{22}} = b_{21} + b_{22}, \) \(\frac{dz}{db_{11}} = \frac{dz}{db_{12}} = a_{11} + a_{21}, \text{ and } \frac{dz}{db_{21}} = \frac{dz}{db_{22}} = a_{12} + a_{22}.\)

It gets cumbersome to keep track of this many individual derivatives, thus introducing the need for a new kind of notation. Let's define the derivative of a scalar \(z\) with respect to a whole tensor \(T\) as \(\frac{dz}{dT}\), where \(\frac{dz}{dT}\) has the same shape as \(T\) itself, and where each entry in \(\frac{dz}{dT}\) is the derivative of \(z\) with respect to the corresponding entry in \(T\). So all the above derivatives can now be expressed more compactly:

\[ \begin{align} \frac{dz}{dA} &= \begin{bmatrix} \frac{dz}{da_{11}} & \frac{dz}{da_{12}} \\ \frac{dz}{da_{21}} & \frac{dz}{da_{22}} \end{bmatrix} = \begin{bmatrix} b_{11} + b_{12} & b_{21} + b_{22} \\ b_{11} + b_{12} & b_{21} + b_{22} \end{bmatrix},\\ \frac{dz}{dB} &= \begin{bmatrix} \frac{dz}{db_{11}} & \frac{dz}{db_{12}} \\ \frac{dz}{db_{21}} & \frac{dz}{db_{22}} \end{bmatrix} = \begin{bmatrix} a_{11} + a_{21} & a_{11} + a_{21} \\ a_{12} + a_{22} & a_{12} + a_{22} \end{bmatrix},\\ \frac{dz}{dC} &= \begin{bmatrix} \frac{dz}{dc_{11}} & \frac{dz}{dc_{12}} \\ \frac{dz}{dc_{21}} & \frac{dz}{dc_{22}} \end{bmatrix} = \begin{bmatrix} 1 & 1 \\ 1 & 1 \end{bmatrix}.\\ \end{align} \]

Now you may have noticed that these derivatives closely follow a pattern. Whenever the entries of a matrix have a summation pattern, your instinct should be that it is the result of multiplying two matrices, which is in fact the case here.

\[ \begin{align} \frac{dz}{dA} &= \begin{bmatrix} \frac{dz}{da_{11}} & \frac{dz}{da_{12}} \\ \frac{dz}{da_{21}} & \frac{dz}{da_{22}} \end{bmatrix}\\ &= \begin{bmatrix} b_{11} + b_{12} & b_{21} + b_{22} \\ b_{11} + b_{12} & b_{21} + b_{22} \end{bmatrix}\\ &= \begin{bmatrix} 1 * b_{11} + 1 * b_{12} & 1 * b_{21} + 1 * b_{22} \\ 1 * b_{11} + 1 * b_{12} & 1 * b_{21} + 1 * b_{22} \end{bmatrix}\\ &= \begin{bmatrix} 1 & 1 \\ 1 & 1 \end{bmatrix} \times \begin{bmatrix} b_{11} & b_{21} \\ b_{12} & b_{22} \end{bmatrix}\\ &= \frac{dz}{dC} \times B^\mathsf{T} \end{align} \]

and similarly \(\frac{dz}{dB} = A^\mathsf{T} \times \frac{dz}{dC}\). I won't prove this here, but this result is no coincidence. It is a special case of a more general result that compactly summarizes the derivative of a scalar with respect to a matrix that was multiplied with another matrix.

This is a very exciting result, and hints at how autodiff systems actually make computations. The pseudocode backpropagation algorithm I showed earlier is theoretically sound, and it is the most general form of backpropagation where every edge is explicitly present in the computation graph, but it is not how backpropagation is implemented in practice. Instead, autodiff uses formulas like those shown above to apply the scalar chain rule "in bulk". The computation graph then does not need to keep track of every edge and scalar entry in every tensor. Our autodiff system will express the example in this section like so:

\[ \begin{array}{ccccccccc} A& \searrow & & &&\\ B & \rightarrow & \text{matmul} & \rightarrow & C & \rightarrow & \text{sum} & \rightarrow & z\\ \end{array} \]

Then knowing \(z = \text{sum}(C)\) it will reason that \(\frac{dz}{dC} = \begin{bmatrix} 1 & 1 \\ 1 & 1 \end{bmatrix}\) , and finally knowing \(C = \text{matmul}(A,B)\), it will apply Equations (3) thus obtaining \(\frac{dz}{dA}\) and \(\frac{dz}{dB}\).

More generally from the point of view of our autodiff system, any function \(f\) takes a list of tensors \(X_1, X_2, \ldots X_k\) as input and produces a tensor \(Y\) as output, like so:

\[f(X_1, X_2, \ldots X_k) = Y.\]

If autodiff knows the derivative of a scalar \(z\) with respect to \(f\)'s output (also known as the upstream gradient, denoted \(\frac{dz}{dY}\)), then autodiff uses what's called the "backward function" of \(f\) to compute the derivative of \(z\) with respect to each of \(f\)'s inputs. Every function \(f\) in the system must have a backward function implemented. This backward function is typically a clever result that implicitly invokes the inductive chain rule upon all scalars in the input tensors. In the example above, the backwards function receives \(\frac{dz}{dC}\) as the upstream gradient and invokes Equations (3) to calculate the derivative of \(z\) with respect to both \(A\) and \(B\).

needs more explanations here before transitioning

This approach to autodiff introduces a problem. The issue is that tensors can be used as input to more than just one function. Suppose we augmented the example by letting \( D = \begin{bmatrix} 5 & 3 \\ 1 & 9 \end{bmatrix} , E = B + D, W = C + E,\) and finally \(z = \text{sum}(W)\) . The computation graph would change to

\[ \begin{array}{ccccccccc} A& \searrow & & &&\\ B & \rightarrow & \text{matmul} & \rightarrow & C & \searrow & \\ & \searrow&&&&&\text{add} & \rightarrow & W & \rightarrow& \text{sum} & \rightarrow & z\\ D & \rightarrow &\text{add} & \rightarrow &E &\nearrow \\ \end{array} \]

Now \(B\) is an input to both matmul and add. So autodiff as described so far, given upstream gradients \(\frac{dz}{dC}\) and \(\frac{dz}{dE}\), would calculate two different values for \(\frac{dz}{dB}\), coming from matmul backwards and add backwards. Let's call these \(G_{BC}\) and \(G_{BE}\), respectively. What should autodiff do now with these values?

While neither \(G_{BC} = \frac{dz}{dB}\) nor \(G_{BE} = \frac{dz}{dB}\), it actually turns out that \(\frac{dz}{dB} = G_{BC} + G_{BE}\), which is a result that follows from the inductive chain rule. Let's see why. Remember, the inductive chain rule dictates that for any scalar \(b_{ij} \in B\),

\[ \frac{dz}{db_{ij}} = \sum_{y \in c(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}}, \]

which we can break down further by letting \(c_{C}(b_{ij})\) denote the children of \(b_{ij}\) that are in \(C\), and similarly defining \(c_{E}(b_{ij})\),

\[ \frac{dz}{db_{ij}} = \sum_{y \in c_{C}(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}} + \sum_{y \in c_{E}(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}}. \]

\[ \begin{align} A^\mathsf{T} \times \frac{dz}{dC} &= \begin{bmatrix} a_{11} & a_{21} \\ a_{12} & a_{22} \end{bmatrix} \times \begin{bmatrix} \frac{dz}{dc_{11}} & \frac{dz}{dc_{12}} \\ \frac{dz}{dc_{21}} & \frac{dz}{dc_{22}} \end{bmatrix}\\ &= \begin{bmatrix} \frac{dz}{dc_{11}} \cdot a_{11} + \frac{dz}{dc_{21}} \cdot a_{21} & \frac{dz}{dc_{12}} \cdot a_{11} + \frac{dz}{dc_{22}} \cdot a_{21} \\ \frac{dz}{dc_{11}} \cdot a_{12} + \frac{dz}{dc_{21}} \cdot a_{22} & \frac{dz}{dc_{12}} \cdot a_{12} + \frac{dz}{dc_{22}} \cdot a_{22} \end{bmatrix}\\ &= \begin{bmatrix} \frac{dz}{dc_{11}} \cdot \frac{\partial c_{11}}{\partial b_{11}} + \frac{dz}{dc_{21}} \cdot \frac{\partial c_{21}}{\partial b_{11}} & \frac{dz}{dc_{12}} \cdot \frac{\partial c_{12}}{\partial b_{12}} + \frac{dz}{dc_{22}} \cdot \frac{\partial c_{22}}{\partial b_{12}} \\ \frac{dz}{dc_{11}} \cdot \frac{\partial c_{11}}{\partial b_{21}} + \frac{dz}{dc_{21}} \cdot \frac{\partial c_{21}}{\partial b_{21}} & \frac{dz}{dc_{12}} \cdot \frac{\partial c_{12}}{\partial b_{22}} + \frac{dz}{dc_{22}} \cdot \frac{\partial c_{22}}{\partial b_{22}} \end{bmatrix}\\ &= \begin{bmatrix} \sum_{y \in c_{C}(b_{11})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{11}} & \sum_{y \in c_{C}(b_{12})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{12}} \\ \sum_{y \in c_{C}(b_{21})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{21}} & \sum_{y \in c_{C}(b_{22})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{22}} \\ \end{bmatrix}\\ \end{align} \]

and so the \(ij\) entry of \(G_{BC}\) is

\[ (G_{BC})_{ij} = \sum_{y \in c_{C}(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}} . \]

So we can see that \((G_{BC})_{ij}\) is exactly the part of the sum in \(\frac{dz}{db_{ij}}\) that corresponds to the children of \(b_{ij}\) in \(C\). Similarly, you could conclude that

\[ (G_{BE})_{ij} = \sum_{y \in c_{E}(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}}, \]

and finally

\[ (G_{BC})_{ij} + (G_{BE})_{ij} = \sum_{y \in c_{C}(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}} + \sum_{y \in c_{E}(b_{ij})} \frac{dz}{dy} \cdot \frac{\partial y}{\partial b_{ij}} = \frac{z}{db_{ij}}, \]

thus establishing \(\frac{dz}{dB} = G_{BC} + G_{BE}\).

This result generalizes beyond just this example. Suppose some tensor \(X\) is the input to several functions, so

\[ \begin{align} f_1(\ldots, X, \ldots) &= Y_1,\\ f_k(\ldots, X, \ldots) &= Y_2,\\ \vdots\\ f_k(\ldots, X,\ldots) &= Y_k,\\ \end{align} \]

and \(\frac{dz}{dY_1}, \frac{dz}{dY_2}, \ldots, \frac{dz}{dY_k}\) have already been calculated. Then autodiff will invoke the backwards function for each of \(f_1, f_2, \ldots, f_k\), thus obtaining \(G_{XY_1}, G_{XY_2}, \ldots, G_{XY_k}\), and finally set \(\frac{dz}{dX} = G_{XY_1} + G_{XY_2} + \ldots + G_{XY_k}\). Remember, \((G_{XY_n})_{ij}\) is the portion of the summation of \(\frac{dz}{dX_{ij}}\) corresponding to the children of \(X\) present in \(Y_n\).

Backpropagation via Autodiff System

Finally, we have all the theory needed for the automatic differentiation system that we will implement.