Refs:¶

1. https://github.com/karpathy/micrograd/tree/master/micrograd

2. https://github.com/mattjj/autodidact

3. autograd/numpy/numpy_vjps.py

4. https://auto-ed.readthedocs.io/en/latest/mod2.html#ii-more-theory

5. docs/mod2.rst

Chain rule¶

Scalar single-variable chain rule¶

Recall the limit definition of derivative of a function,

From the limit definition you can find the value of $g(x+h)$ as

You can use this rule to find the chain rule of finding the chaining of two functions together,

Scalar two-variable chain rule¶

Consider a function of two variables $f(u(x),v(x))$. Find its derivative,

Now $f(u + \delta u, v + \delta v)$ should not be expanded in one step but in two steps. First keep $v + \delta v$ as it is, and expand with respect to $u + \delta u$

and then do the same with $v+\delta v$,

We use $\lim_{\delta v,\delta u \to 0} f'_u(u, v+\delta v)\delta u = \lim_{\delta v,\delta u \to 0} f'_u(u, v)\delta u + f''_{uv}(u, v)(\delta v)(\delta u) = \lim_{\delta u \to 0} f'_u(u, v)\delta u$ to get,

Going back to the chain rule,

Scalar valued vector function chain rule¶

Consider two functions $f(\bfg): \bbR^m \to \bbR$, $\bfg(x): \bbR \to \bbR^m$ that can be composed together $f(\bfg(x))$. We want to find the derivative of composition $f \circ g$ by chain rule.

Recall that the derivative (Jacobian) of $f(\bfg)$ is a row vector, $$\frac{\p f(\bfg)}{\p \bfg} =

$$

And the derivative (Jacobian) of $\bfg(x)$ is a column vector, $$\frac{\p \bfg(x)}{\p x} =

$$

Note that a vector function is a multi-variate scalar function

We can apply the multi-variate scalar function chain rule,

The derivatives of $\bfg$ can be separated from derivatives of $f$ as vector multiplication, $$ \frac{\p}{\p x} f(\bfg(x)) =

Hence the chain rule for vector derivatives works out for our definition of vector derivatives,

Note that the order of multiplication matters, specifically

This is a consequence of row-vector convention. If we chose a column-vector convention the result will be completely different.

General chain rule¶

Let the function be $\bff(\bfg): \bbR^m \to \bbR^n$ and $\bfg(\bfx): \bbR^p \to \bbR^m$, then the derivative (Jacobian) of their composition $\bff \circ \bfg$ is

Computational complexity of Forward vs Reverse mode differentiation¶

Consider three functions, $\bfh(\bfx): \bbR^m \to \bbR^n$, $\bfg(\bfh): \bbR^n \to \bbR^p$ and $\bff(\bfg): \bbR^p \to \bbR^q$ chained together for composition $\bff(\bfg(\bfh(\bfx))) : \bbR^m \to \bbR^q$. To find the derivative (Jacobian) the composite function, we use chain rule:

Computational complexity of matrix multiplication¶

Let’s say you multiply two matrices $A \in \bbR^{m \times n}$ and $B \in \bbR^{n \times p}$, total number of additions and multiplications (floating point operations) can be calculated by

where $\bfa_i^\top$ are the row-vectors of matrix $A$ and $\bfb_i$ are the column vectors of matrix $B$. Then matrix $C$ is written as

We note that $C$ matrix has $pm$ elments and each element requires computing dot product of size $n$ vectors,

Each dot product requires $n$ multiplications and $n-1$ additions. Hence matrix multiplication which has $pm$ dot products requires $pm(n + n -1)$ (floating point) operations.

Matrix multiplication has a computation complexity of $O(pmn)$ for matrices of size $m \times n$ and $n \times p$.

Computational complexity of forward-mode differentiation¶

In forward diff, we compute computaional complexity from input side to the output side.

The first two matrix multiplications $X_{p \times n} = \left( \frac{\p \bfg}{\p \bfh} \frac{\p \bfh}{\p \bfx} \right)$ are of the size $p \times m$ and $m \times n$, resulting in $O(pmn)$ complexity.

The second two matrix multiplications $\left(\frac{\p \bff}{\p \bfg} X_{p \times n} \right)$ are of the size $q \times p$ and $p \times n$, resulting in $O(qpn)$ complexity.

The total computational complexity of forward differentiation is $O(qpn + pmn) = O((qp+pm)n)$.

For a longer chain of functions of Jacobians of shape $q_i \times p_i$ with ($p_i = q_{i-1}$).

We get a computational complexity that looks like $O( (\sum_{i=1}^n q_i p_i) p_0)$. Note that the size of input $p_0$ is the only common factor for the entire chain.

Computational complexity of reverse-mode diff¶

In reverse-mode diff, we compute computaional complexity from input side to the output side.

The first two matrix multiplications $X_{q \times p} = \left( \frac{\p \bff}{\p \bfg} \frac{\p \bfg}{\p \bfh} \right)$ are of the size $q \times p$ and $p \times m$, resulting in $O(qpm)$ complexity.

The second two matrix multiplications $\left(X_{q \times p}\frac{\p \bfh}{\p \bfx} \right)$ are of the size $q \times p$ and $p \times n$, resulting in $O(qpn)$ complexity.

The total computational complexity of forward differentiation is $O(qpm + qmn) = O(q(pm+mn))$.

For a longer chain of functions of Jacobians of shape $q_i \times p_i$ with ($p_i = q_{i-1}$).

We get a computational complexity that looks like $O( q_n (\sum_{i=0}^{n-1} q_i p_i))$. Note that the size of output $q_n$ is the only common factor for the entire chain.

Reverse-mode differentiation is called backpropagation¶

Reverse-mode differentiation is called backpropagation in neural networks. It is more popular because most of the times you compute the derivatives of the loss function which is a scalar function with output dimension as only 1. This makes reverse-mode differentiation clearly superior for loss function gradient.

Implementation of forward/reverse mode differentiation in Pytorch¶

Reverse mode¶

Let’s compute the derivatives of

Forward mode¶

Vector Jacobian product (vjp) for reverse-mode differentiation¶

Typical output of a neural network is a loss function. Loss function is always a scalar. Most neural network libraries implement reverse-mode differentiation only for a scalar output.

Hence, the first Jacobian on the output side of chain rule is a row-vector.

$\bfg(\bfx): \bbR^n \to \bbR^p$, $\bff(\bfg): \bbR^p \to \bbR^q$ and $l(\bff): \bbR^q \to \bbR$.

When you are writing a programmatic derivative function for reverse mode differentation, the function does two things:

1. Compute the local Jacobian of the function for example $\frac{\p \bff}{\p \bfg}$.

2. Left multiply the Jacobian with a row-vector of accumulated derivative so far. For example, $\frac{\p l}{\p \bff} \frac{\p \bff}{\p \bfg}$.

The template of the function is like this:

def g(arg1, arg2):
    # Compute g
    return g 
    
def g_vjp(arg1, arg2, dl_dg):
    # Compute vector Jacobian product with respect to each oargument
    return dl_arg1, dl_arg2

If you are given a function $\bfg(\bfx)$, and you want to implement vjp function for it. It is often easier to imagine a sclar loss function $l(\bfg(\bfx))$ whose accumulated gradient $\frac{\p l}{\p \bfg}$ is given as an input argument. The function vjp returns the derivative of the loss function with respect to the inputs,

which looks like a vector Jacobian product, but you are free to not compute the Jacobian separately. Sometimes it is computationally harder to compute the jacobian separately then multiply it by the vector.

Jacobian vector product (jvp) for forward-mode differentiation¶

It is also common to implment foward mode differentiation with only a scalar input assumption, say $t$.

Say $\bfh(\bfx): \bbR^m \to \bbR^n$, $\bfg(\bfh): \bbR^n \to \bbR^p$ and $\bff(\bfg): \bbR^p \to \bbR^q$

You can assume $\bfx$ to be function of scalar $t \in \bbR$, $\bfx(t)$. Then the chain rule is

You can compute the derivative with respect to one element of $\bfx$ at a time by setting that element’s derivative to be 1 and the rest to be zero. For example, if you want to compute the $\frac{\p \bff}{\p x_2}$ then set

.

For forward pass you typically implement a function called jvp which stands for Jacobian vector product:

1. The Jacobian is the local derivative. For example $\frac{\p \bfh}{\p \bfx}$

2. Multiplication of the jacobian with an incoming accumulated gradient which is a column-vector. For example, $\frac{\p \bfx}{\p t}$.

The template of the function is like this:

def g(arg1, arg2):
    # Compute g
    return g 
    
def g_jvp(arg1, arg2, darg1_dt, darg2_dt):
    # Compute Jacobian vector product with respect to t
    return dg_dt

If you are given a function $\bfg(\bfx)$, and you want to implement jvp function for it. It is often easier to imagine a sclar input variable $\bfg(\bfx(t)))$ whose accumulated gradient $\frac{\p \bfx}{\p t}$ are given as an input argument. The function jvp returns the derivative of the output with respect to the scalar input $t$,

which looks like a Jacobian vector product, but you are free to not compute the Jacobian separately. Sometimes it is computationally harder to compute the jacobian separately then multiply it by the vector.

Implementing numpy backpropagation for various operations¶

Vector Jacobian Product for addition¶

where $\bfa, \bfb, \bff \in \bbR^n$

Let $l(\bff(\bfa, \bfb)) \in \bbR$ be the eventual scalar output. We find $\frac{\p l}{\p \bfa}$ and $\frac{\p l}{\p \bfb}$ for Vector Jacobian product.

Similarly,

VJP for element-wise multiplication¶

where $\alpha, \beta, f \in \bbR$

Let $l(f(\alpha, \beta)) \in \bbR$ be the eventual scalar output. We find $\frac{\p l}{\p \alpha}$ and $\frac{\p l}{\p \beta}$ for Vector Jacobian product.

VJP for matrix-matrix, matrix-vector and vector-vector multiplication¶

Case 1: VJP for vector-vector multiplication¶

where $f \in \bbR$, and $\bfb, \bfa \in \bbR^n$

Let $l(f(\bfa, \bfb)) \in \bbR$ be the eventual scalar output. We find $\frac{\p l}{\p \bfa}$ and $\frac{\p l}{\p \bfb}$ for Vector Jacobian product.

Similarly,

Case 2: VJP for matrix-vector multiplication¶

Let

where $\bff \in \bbR^m$, $\bfb \in \bbR^n$, and $\bfA \in \bbR^{m \times n}$

Let $l(\bff(\bfA, \bfb)) \in \bbR$ be the eventual scalar output. We want to findfind $\frac{\p l}{\p \bfA}$ and $\frac{\p l}{\p \bfb}$ for Vector Jacobian product.

Let

Define matrix derivative of scalar to be:

Note that

Since $\bfa_i^\top \bfb$ is a scalar, it is easier to find its derivative with respect to the matrix $\bfA$.

Let

Then \begin{align}\frac{\p l}{\p \bff}\frac{\p }{\p \bfA}\bfa_i^\top \bfb =

= \frac{\p l}{\p f_i}\bfb^\top \in \bbR^{1 \times n} \end{align}

Returning to our original quest for

Note that

We can group the terms inside a single transpose.

Which results in

The derivative with respect to $\bfb$ is simpler:

Case 3: VJP for matrix-matrix multiplication¶

Let

where $\bfF \in \bbR^{m \times p}$, $\bfB \in \bbR^{n \times p}$, and $\bfA \in \bbR^{m \times n}$

Let $l(\bfF(\bfA, \bfB)) \in \bbR$ be the eventual scalar output. We want to find $\frac{\p l}{\p \bfA}$ and $\frac{\p l}{\p \bfB}$ for Vector Jacobian product.

Note that a matrix-matrix multiplication can be written in terms horizontal stacking of matrix-vector multiplications. Specifically, write $\bfF$ and $\bfB$ in terms of their column vectors:

Then for all $i$

From the VJP of matrix-vector multiplication, we can write

and for all $i \ne j$

Instead of writing $l(\bfF)$, we can also write $l(\bff_1, \bff_2, \dots, \bff_p)$, then by chain rule of functions with multiple arguments, we have,

It turns out that some of outer products can be compactly written as matrix-matrix multiplication:

Hence,

The vector Jacobian product for $\bfB$ can be found by applying the above rule to $\bfF_2(\bfA, \bfC) = \bfF^\top(\bfA, \bfB) = \bfB^\top \bfA^\top = \bfC \bfA^\top$ where $\bfC = \bfB^\top$ and $\bfF_2 = \bfF^\top$.

Take transpose of both sides

Put back, $\bfC = \bfB^\top$ and $\bfF_2 = \bfF^\top$,

Customizing backward step (vector-Jacobian product) in PyTorch¶

Consider the derivative of Sigmoid activation function

The above derivative is computed by chain rule. However, there is much simpler expression that can avoid unnecessary computations,

Whats the deal with the unbroadcast function¶

To understand unbroadcast function, first it is important to understand numpy broadcasting. Once you understand that then we can proceed to understand unbroadcast function

Once you understand broadcasting, we can appreciate the unbroadcast function. Suppose you have A.shape = (100, 1) and B.shape = (1, 100) and you add them together

F = A + B

Then F.shape = (100, 100). Note that as if we created copies of A and B along the relevant dimensions.

Let’s think about backpropagation now. The incoming grad will have a dimension of dl__dF.shape = (100, 100) but the returned dl__dA must have shape dl__dA.shape = (100, 1) and dl__dB.shape = (1, 100).

From the chain rule of addition we only know that dF_{ij}/dA_{ij} = 1 and dF_{ij}/dB_{ij} = 1. What happens when you create copies of a variable and they contribute to a function? Let’s consider a very simple function, of two variables:

What if the $y₁$ and $y₂$ were copies of the same variable?

We can use the chain rule to find out the result

If we know that $y₁$ and $y₂$ are just copies of $x$, then $∂ₓy₁ = 1$ and $∂ₓy₂ = 1$ .

In other words, if input variables are copied multiple times and then the function is computed using the copies, then the derivative is the sum of the derivatives with respect to the copies.

That’s what the unbroadcast function does. Because broadcasting creates the copies of the inputs along the mismatching axis, unbroadcasting sums the derivatives along the mismatching axis.

Coming back to the F = A + B example. We get dl__dF.shape = (100, 100). A.shape = (100, 1) was copied 100 times along the axis=1 (2nd dimension). So dl__dA = dl__dF.sum(axis=1) . Same for the B as well. dl__dB = dl__dF.sum(axis=0).

That’s what unbroadcast function does but in a more generalizable way