Mathematics Basics - Multivariate Calculus (Multivariate Chain Rule)

Multivariate Chain Rule

We have seen an example previously about how to use chain rule to calculate total derivative of function f with respect to variable t through intermediate variables x, y and z as follows.

$$\begin{array}{rl}f& =\sin (x)e^{y{z}^2}\\ x& =t-1\\ y& =t^2\\ z& =\frac{1}{t}\\ \frac{df}{dt}& =\frac{\partial f}{\partial x}\cdot \frac{dx}{dt}+\frac{\partial f}{\partial y}\cdot \frac{dy}{dt}+\frac{\partial f}{\partial z}\cdot \frac{dz}{dt}\\ & =\cos (x)e^{y{z}^2}\cdot (1)+z^2\sin (x)e^{y{z}^2}\cdot (2t)+2y\sin (x)e^{y{z}^2}\cdot (-\frac{1}{t^2})\\ & =e^{y{z}^2}[\cos (x)+2t{z}^2\sin (x)-\frac{2yz}{t^2}\sin (x)]\end{array}$$

Now let’s formalize the definition of chain rule in multivariate calculus. We start with a function f with n input variables $x_1$, $x_2$,…, $x_n$.

$$f(x_1,x_2,\cdots ,x_n)=f(\mathbf{x})$$

To simplify our expression, we will denote function f by $f(\mathbf{x})$ with bold letter $\mathbf{x}$ representing a vector of n input variables.

If each of the input variable in $\mathbf{x}$ is also a function of another variable t, our function f can be re-expressed as $f(\mathbf{x}(t))$.

$$f(\mathbf{x}(t))=f(x_1(t),x_2(t),\cdots ,x_n(t))$$

To compute the derivative of function f with respect to t, we can write a dot product of two n-dimensional vectors.

$$\begin{array}{rl}\frac{df}{dt}& =\frac{\partial f}{\partial x_1}\cdot \frac{d{x}_1}{dt}+\frac{\partial f}{\partial x_2}\cdot \frac{d{x}_2}{dt}+\cdots +\frac{\partial f}{\partial x_n}\cdot \frac{d{x}_n}{dt}\\ & =\left(\begin{array}{cccc}\frac{\partial f}{\partial x_1}& \frac{\partial f}{\partial x_2}& \cdots & \frac{\partial f}{\partial x_n}\end{array}\right)\cdot \left(\begin{array}{c}\frac{d{x}_1}{dt}\\ \frac{d{x}_2}{dt}\\ ⋮\\ \frac{d{x}_n}{dt}\end{array}\right)\end{array}$$

What we get from the first vector is just the partial derivatives of function f with respect to all its input variables. Does that look familiar? Yes, that is our Jacobian vector. So our differentiation equation can be simplified as,

$$\frac{df}{dt}=(J_f)\cdot \frac{d\mathbf{x}}{dt}$$

We use $\frac{d\mathbf{x}}{dt}$ to compactly represent the vector of differentiation expression $\frac{d{x}_1}{dt},\frac{d{x}_2}{dt},\cdots ,\frac{d{x}_n}{dt}$.

This is multivariate chain rule for one intermediate layer. What if we have more than one intermediate layer? Let’s add one more layer in between $\mathbf{x}$ and t and denote the function as $f(\mathbf{x}(\mathbf{u}(t)))$. Now f is a function of a vector of n variables $\mathbf{x}$. Each variable of $\mathbf{x}$ is a function of a vector of m variables $\mathbf{u}$. Finally, each variable of $\mathbf{u}$ is a function of variable t. To differentiate function f with respect to t, we need

$$\begin{array}{rl}\frac{df}{dt}& =\left(\begin{array}{cccc}\frac{\partial f}{\partial x_1}& \frac{\partial f}{\partial x_2}& \cdots & \frac{\partial f}{\partial x_n}\end{array}\right)\cdot \left(\begin{array}{cccc}\frac{\partial x_1}{\partial u_1}& \frac{\partial x_1}{\partial u_2}& \cdots & \frac{\partial x_1}{\partial u_m}\\ \frac{\partial x_2}{\partial u_1}& \frac{\partial x_2}{\partial u_2}& \cdots & \frac{\partial x_2}{\partial u_m}\\ ⋮& ⋮& \ddots & ⋮\\ \frac{\partial x_n}{\partial u_1}& \frac{\partial x_n}{\partial u_2}& \cdots & \frac{\partial x_n}{\partial u_m}\end{array}\right)\cdot \left(\begin{array}{c}\frac{d{u}_1}{dt}\\ \frac{d{u}_2}{dt}\\ ⋮\\ \frac{d{u}_m}{dt}\end{array}\right)\\ & =(J_f)\cdot (J_{\mathbf{x}})\cdot \frac{d\mathbf{u}}{dt}\end{array}$$

Again, we have simplified our expression by two Jacobian terms. Note the middle Jacobian ($J_{\mathbf{x}}$) is a matrix instead of a vector.

With that, we can generalize the multivariate chain rule of n intermediate layers as

$$\frac{df}{dt}=(J_f)\cdot (J_{{\mathbf{x}}^{(1)}})\cdot (J_{{\mathbf{x}}^{(2)}})\cdots (J_{{\mathbf{x}}^{(n-1)}})\cdot \frac{d{\mathbf{x}}^{(n)}}{dt}$$

Here the superscript $(i)$ denotes the $i^{th}$ intermediate layer.

This is a powerful rule. As long as we can clearly define the variable expressions from one layer to the next, we can always differentiate function $f$ with respect to the input variable at any layer. We will see this rule being used a lot in practice.

Neural Networks

Multivariate chain rule is one essential component of neural networks. Each layer of neurons in the network are controlled by a set of parameters. These parameters are optimized by individually differentiating the cost function. Moreover, the differentiation result is propagated from one layer to the next, just like how our chain rule works.

Let’s use a simple neural network setup to explain the notations we are going to use.

We have circles representing nodes and lines representing connections between each node in a network. The above network can be written mathematically as

$$\begin{array}{rl}{a}_0^{(1)}& =\sigma (w_{00}{a}_0^{(0)}+w_{10}{a}_1^{(0)}+w_{20}{a}_2^{(0)}+b_0)\\ a_1^{(1)}& =\sigma (w_{01}{a}_0^{(0)}+w_{11}{a}_1^{(0)}+w_{21}{a}_2^{(0)}+b_1)\end{array}$$

Here we use a to represent value of a node (it can be an input or a output node). The superscript of a indicates which layer the node is at and the subscript of a indicates the nth node of that layer. w and b are the weight and bias respectively applied to the input node. Finally, σ is an activation function that transforms a number of input nodes to the output node.

We can further simply our network expression by vector notation.

$${\mathbf{a}}^{(1)}=\sigma (\mathbf{w}\cdot {\mathbf{a}}^{(0)}+\mathbf{b})$$

Now ${\mathbf{a}}^{(0)}$ is the vector of input nodes at layer 0. ${\mathbf{a}}^{(1)}$ is the vector of output nodes at layer 1. $\mathbf{w}$ is the weight matrix relating each input node to each output node. $\mathbf{b}$ is the bias vector for each output node.

If we extend the same idea to any layer of a neural network, we can get following relation between input layer $L-1$ and output layer $L$.

$${\mathbf{a}}^{(L)}=\sigma ({\mathbf{w}}^{(L)}\cdot {\mathbf{a}}^{(L-1)}+\mathbf{b})$$

You can imagine we stack up these intermediate layers of nodes between the input and output layers for a complex neural network. Each intermediate layer is simultaneously the output from previous layer and the input to the next layer. In solving such a neural network, we are interested to find the right parameters $\mathbf{w}$ and $\mathbf{b}$ at each layer that successfully map the model input to the final output.

In order to do that, we need to know how much our current network output is deviated from the actual output. This deviation is defined by cost function.

$$C=\sum _i(a_i^{(L)}-y_i{)}^2$$

We take the square of difference between the model output and the actual output and sum over all possible outputs. The goal for our neural network becomes solving for parameters $\mathbf{w}$ and $\mathbf{b}$ at each layer such that C is minimized. This requires us to perform differentiation of C with respect to every variable $\mathbf{w}$ and $\mathbf{b}$. Since $\mathbf{w}$ and $\mathbf{b}$ are connected to C via layers of nodes, we can apply multivariate chain rule here.

Let’s start simple with just one layer of network with single input and output node.

$$a^{(1)}=\sigma (w{a}^0+b)$$

We will add one assist variable z to represent the output value before activation. So the network model is connected to cost function as:

$$\begin{array}{rl}{z}^{(1)}& =w{a}^0+b\\ a^{(1)}& =\sigma (z^{(1)})\\ C& =(a^{(1)}-y{)}^2\end{array}$$

We can then find the derivative of $C$ with respect to $w$ and $b$ as

$$\begin{array}{rl}\frac{dC}{dw}& =\frac{\partial C}{\partial a^{(1)}}\cdot \frac{\partial a^{(1)}}{\partial z^{(1)}}\cdot \frac{d{z}^{(1)}}{dw}=2(a^{(1)}-y)\cdot {\sigma }^{-1}(z^{(1)})\cdot (a^0)\\ \frac{dC}{db}& =\frac{\partial C}{\partial a^{(1)}}\cdot \frac{\partial a^{(1)}}{\partial z^{(1)}}\cdot \frac{d{z}^{(1)}}{db}=2(a^{(1)}-y)\cdot {\sigma }^{-1}(z^{(1)})\cdot (1)\end{array}$$

This differentiation equation will be evaluated to a scalar, which is the step we take to move $w$ and $b$ from their current values. We then recalculate the differentiation result and update parameters w and b iteratively until their derivatives are zero. That is where the cost is minimized.

When we add more nodes and more layers to the network, it will get more complicated to differentiate $C$ with respect to each of $\mathbf{w}$ and $\mathbf{b}$. Nonetheless, we shall still apply the multivariate chain rule to iteratively differentiate and evaluate each layer of the network till all the parameters are optimized. This is a very laboratory process, so I will not demonstrate the details here. Fortunately, most of the machine learning solver has been equipped with the differentiation process. You do not need to spell out the expression yourself. What really matters is the principle behind all the calculation. It could help you better design your network structure and/or choose the parameters.

We have completed the study of multivariate chain rule. It is a very powerful tool when we are dealing with complex optimization problems such as neural networks. We can traverse from one set of variables to the next and calculate their influence on the final cost function with chained partial derivatives. Hope you have enjoyed reading so far. We will dive into more interesting topics in multivariate calculus in the coming articles.

---

(Inspired by Mathematics for Machine Learning lecture series from Imperial College London)