2.1 The Multivariable Chain Rule

Let us complicate our deep composite function a little. Suppose we had a function as follows:

\[ f(a) = f_m(f_{1k}(f_{1k-1}\cdots f_{11}(a)\cdots), f_{2k-1}(\cdots f_{21}(a)\cdots)) \]

How can we compute the derivative of something like that? Well, the first thing to note is that the inner composite functions are, from the point of view of \(f_m\), just functions. If we set \(f_{1} = f_{1k-1}(\cdots f_{11}(a)\cdots)\) and \(f_{2} = f_{2k-1}(\cdots f_{21}(a)\cdots)\), we can rewrite this as:

\[ f(a) = f_m(f_{1}(a), f_2(a)) \]

we can then apply the multivariable chain rule to obtain:

\[ \begin{equation} f^{\prime}(a) = f_m^{\prime}(f_1(a), f_2(a))f_1^\prime(a) + f_{m\prime}(f_1(a), f_2(a))f_2^\prime(a) \tag{2.1} \end{equation} \]

Assuming we have already done the forward pass, we know the values for \(f_1(a)\) and \(f_2(a)\), as well as the forms of the derivatives for \(f_m^\prime\) and \(f_{m^\prime}\). The values of \(f_m^{\prime}(f_1(a), f_2(a))\) and \(f_{m\prime}(f_1(a), f_2(a))\) are, therefore, immeditaly available, and we have the first step of reverse-mode autodiff. We do not have the derivatives for \(f_1^\prime(a))\) or \(f_2^\prime(a)\) yet but, as you can imagine, all we will need to do is unpack those with the chain rule too. Before we do that, however, we need to dive in a little deeper in the chain rule.

2.2 Unifying the Rules

So far, we have been dealing with composite functions on one input variable; i.e., our final function \(f(a)\) has always been a function of one single variable \(a\).

Suppose now that we had a two-variable function; i.e.: \[ f(x_1,x_2) = f_m(f_{1}(x_1, x_2)), f_2(x_1)) \]

And we wanted to determine the partial derivatives of \(f\) with respect to \(x_1\) and \(x_2\). Partial derivatives are essentially equivalent to taking a function and interpreting all but one of its variables - the one you are taking the partial derivative with respect to - as constants. This, by construction, means that any partial derivative we take reduces to a single-variable case; i.e., there is nothing special about multiple variables.

For a function \(f(x_1,x_2)\) for which we are computing \(f^\prime(x_1,x_2)\), we will write \(f^\prime(x_1, \cdot)\) to reinforce the idea that we do not care about about \(x_2\) in this case (and, conversely, we will write \(f_\prime(\cdot,x_2)\) when we derive with respect to \(x_2\)).

We now take the partial derivative of \(f\) with respect to \(x_1\), which reduces to the multivariable chain rule:

\[ \begin{equation} f^\prime(x_1,x_2) = f_m^\prime(f_1(x_1,\cdot), f_2(x_1))\times f_1^\prime(x_1,\cdot) + f_{m\prime}(f_2(x_1,\cdot), f_2(x_1)\times f_{2}^\prime(x_1) \tag{2.2} \end{equation} \] and the partial derivative with respect to \(x_2\), which reduces to the simple chain rule:

\[ \begin{equation} f_\prime(x_1,x_2) = f_{m^\prime}(f_1(\cdot,x_2), f_2(\cdot)) \times f_{1^\prime}(\cdot, x_2) \tag{2.3} \end{equation} \]

Now note that the chain rules from Eq. (2.2) and (2.3) share a certain motif. Namely, if the outer function \(f_m(x_1, \cdots, x_n)\) is a function on \(n\) parameters, and the \(k^{th}\) parameter of \(f_m\) contains a function \(f_j(x_i,\cdots)\) that involves an input variable \(x_i\), then the partial derivative of \(f\) (our composite function which involves \(f_m\)) with respect to \(x_j\) will be a summation in which one of the terms will be of the form:

\[ \begin{equation} \frac{\partial}{\partial x_j}f_m(\cdots, f_j(x_i, \cdots), \cdots) \times \frac{\partial}{\partial x_j} f_j(x_i, \cdots) \tag{2.4} \end{equation} \]

Indeed, the partial derivative with respect to \(x_j\) will be a sum containing one such term for every parameter of \(f_m\) that involves \(x_j\), either directly; if, say, \(f_j(x_i) = x_i\), or indirectly; if, say, \(f_j(x_i)\) is itself a composite function like \(f_j(x_i) = f_k(f_s(x_i))\).

More formally, suppose we had \(k\) possible input variables \(\{x_1, \cdots, x_k\}\) and a function \(f\) defined on a subset of these variables, i.e., suppose we had:

\[ f\left(f_1\left(x_{\pi(1)}\right), \cdots, f_n\left(x_{\pi(n)}\right)\right) \]

where \(1 \leq \pi(i) \leq k\), \(1 \leq i \leq n\), and the \(\pi(i)\) can take on the same value for different values of \(i\); i.e., it is possible that \(i \neq j\) but \(\pi(i) = \pi(j)\) for \(1 \leq j \leq n\).

If we wanted to compute the partial derivative of \(f\) with respect to \(x_j\), we would have, by the (multivariable) chain rule:

\[ \frac{\partial}{\partial x_j}f\left(f_1\left(x_{\pi(1)}\right), \cdots, f_n\left(x_{\pi(n)}\right)\right) = \sum_{\{i~\mid~\pi(i) = j\}} \frac{\partial}{\partial x_i}f\left(f_1\left(x_{\pi(1)}\right), \cdots, f_n\left(x_{\pi(n)}\right)\right) \times \frac{\partial}{\partial x_j} f_i\left(x_{\pi(i)}\right) \]

If this is still looking too abstract, we will work through an example from beginning to end in the next section.

2.3 Expression Graphs

In taking the partial derivative of the composite \(f\) with respect to a variable \(x_i\), we had to:

1. identify which parameters of \(f\) contained functions which referred to \(x_i\);
2. apply the chain rule and add a term of the form given by Eq. (2.4) for each parameter that did;
3. sum those terms together to obtain the final result.

In particular, we have so far relied on visually inspecting the function to do step (1), but this is not practical for complex functions. It is much more convenient (and the reason will hopefully become clear as we proceed), instead, to represent our expression as a graph (a DAG, really) where nodes are either: i) intermediate functions or ii) input variables (Fig. 2.1).

Let’s look again at a two-variable example:

example <- quote(f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1)))

\[ f(x_1,x_2) = f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1)) \]

But this time let us assign concrete functions to them:

• \(f_1(x, y) = \sin{\left(\pi x \right)} + \cos{\left(\pi y \right)}\)
• \(f_2(x, y) = x^{y}\)
• \(f_3(x) = \log{\left(x \right)}\)

Now, supposing we wanted to compute \(\frac{\partial f}{\partial x_1}\), we would get, by the (multivariable) chain rule:

\[ \begin{align} \frac{\partial}{\partial x_1}f(x_1, x_2) =~ &\frac{\partial}{\partial x}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1))\times \frac{\partial}{\partial x_1}f_2(x_1, f_2(x_2, x_1))) + \frac{\partial}{\partial y}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1)) \times \frac{\partial }{\partial x_1}f_3(x_1) = \\ =~&\frac{\partial}{\partial x}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1))\times \left[\frac{\partial}{\partial x}f_2(x_1, f_2(x_2, x_1))) + \frac{\partial}{\partial y}f_2(x_1, f_2(x_2, x_1)))\times \frac{\partial}{\partial x_1}f_2(x_2, x_1)\right] +\\ &+ \frac{\partial}{\partial y}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1)) \times \frac{\partial }{\partial x_1}f_3(x_1) \end{align} \]

And this expands into the sum of three products:

\[ \begin{align} \frac{\partial}{\partial x_1}f(x_1, x_2) =&\frac{\partial}{\partial x}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1))\times\frac{\partial}{\partial x}f_2(x_1, f_2(x_2, x_1)))~+\\ + &\frac{\partial}{\partial x}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1))\times\frac{\partial}{\partial y}f_2(x_1, f_2(x_2, x_1)))\times \frac{\partial}{\partial x_1}f_2(x_2, x_1)~+ \\ + &\frac{\partial}{\partial y}f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1)) \times \frac{\partial }{\partial x_1}f_3(x_1) \tag{2.5} \end{align} \]

Let us now bridge this to the expression graph of \(f\), shown in Fig. 2.1.

express(!!example) %>% render_graph()

Looking at Fig. 2.1, if we assign, to each outlink in the the DAG, the partial derivative of the function represented by the node from which the edge departs with respect to each of its parameters, it is not hard to see that that the final partial derivative \(\frac{\partial f}{\partial x_1}\) is given by the sum of the products of these edge expressions along each path that connects \(f_1\) to \(x_1\). Indeed, the three products in Eq. (2.5) correspond to the three paths; \((f_1, f_2); (f_1, f_2, f_2); (f_1, f_3)\), that connect \(f_1\) to \(x_1\) in the graph.

Let’s show that this is indeed true by using this method to compute the partial derivatives of \(f\) at the point \((2,3)\), and then comparing the results with the full symbolic derivative obtained by expanding Eq. (2.5).

We first need our building blocks; i.e., the partial derivatives of each of the composing functions:

\[ \frac{\partial}{\partial x} f_1(x, y) = \pi \cos{\left(\pi x \right)}\\ \frac{\partial}{\partial y} f_1(x, y) = - \pi \sin{\left(\pi y \right)}\\ \frac{\partial}{\partial x} f_2(x, y) = \frac{x^{y} y}{x}\\ \frac{\partial}{\partial y} f_2(x, y) = x^{y} \log{\left(x \right)}\\ \frac{\partial}{\partial x} f_3(x) = \frac{1}{x} \]

And we also need the results of the forward propagation step:

\[ \begin{align} x_1 &= 2\\ x_2 &= 3\\ f_2(x_2, x_1) &= 9\\ f_2(x_1, f_2(x_2, x_1)) &= 512\\ f_3(x_1) &= 0.693147180559945 \end{align} \]

Finally, we can put these together by computing the products along the paths from \(f_1\) to \(x_1\):

\[ \begin{align} (f_1, f_2, x_1) &= 7238.22947387088\\ (f_1, f_2, f_2, x_1) &= 11023.8236395328\\ (f_1, f_3) &= -1.29038221381431 \end{align} \]

So that the sum yields 18260.7627311899. As for the symbolic derivative, we get that the formula is:

x_1, x_2 = symbols('x_1 x_2')
df_x1 = f_1(f_2(x_1, f_2(x_2, x_1)), f_3(x_1)).diff(x_1)
print('$$ \\frac{\\partial}{\\partial x_1} f(x_1, x_2) = ' + latex(df_x1) + '$$')

\[ \frac{\partial}{\partial x_1} f(x_1, x_2) = \pi x_{1}^{x_{2}^{x_{1}}} \left(x_{2}^{x_{1}} \log{\left(x_{1} \right)} \log{\left(x_{2} \right)} + \frac{x_{2}^{x_{1}}}{x_{1}}\right) \cos{\left(\pi x_{1}^{x_{2}^{x_{1}}} \right)} - \frac{\pi \sin{\left(\pi \log{\left(x_{1} \right)} \right)}}{x_{1}}\]

Which, computing for \((2, 3)\), yields:

print(df_x1.subs({'x_1': 2, 'x_2': 3}).evalf())

## 18260.7627311899

as expected.