Automatic Differentiation

Automatic differentiation(AD) is one of the core techniques in modern deep learning frameworks. There are two main methods: forward & reverse mode AD.

The Chain Rule

Chain rule is the core principle for both forward & reverse AD. The working environment for AD is a calculation DAG, with each node being a tensor $\mathbb{R}^{n}$. Take two node $x$ & $y$, given that there exists a path from $x$ to $y$, then

\[\begin{align} \frac{\partial y}{\partial x} = \sum_{i\in\mathrm{pa}(y)} \frac{\partial y}{\partial i}\cdot \frac{\partial i}{\partial x} \\ \frac{\partial y}{\partial x} = \sum_{i\in\mathrm{ch}(x)} \frac{\partial y}{\partial i}\cdot \frac{\partial i}{\partial x} \end{align}\]

where $\mathrm{pa}, \mathrm{ch}$ represent direct predecessors & successors of the given node, and $\partial \cdot / \partial x$ represents the Jacobian matrix, derivative in the context of tensor calculus.

Next we may apply the principle on the DAG. Take reverse mode as an example:

1. Start by calculating $\partial j / \partial i$ for every edge $<i,j>$;
2. Note this would’ve had all $\partial z / \partial i,~i\in \mathrm{pa}(z)$ calculated.
3. Now for each node $x$ with its $\partial z / \partial x$ not calculated, and $\partial z / \partial y$ calculated for all its children, apply the chain rule. This could be carried out in a (reverse) topological order.

Optimizing: Vector-Jacobian Product

Carrying out matrix product on each step with the DAG is, obviously, temporal & spatial unfriendly.

The first step toward optimization is the observation that our final node $z$ is a scalar, which means every $\partial z / \partial x$ is a vector, not a matrix. In fact, they’re just gradients $\nabla_{x}z$. Notice this brings convenience to the reverse mode but not the forward mode, for $z$’s appearance in the recurrence. Therefore, rewrite the chain rule as

\[\nabla_{y}z = \sum_{i\in\mathrm{ch}(y)} \nabla_{i}z \cdot \frac{\partial i}{\partial y}\]

For each node, only the storage of its $\nabla z$ is required, which costs $O(n_{i})$, comparing to the forward mode $\partial i / \partial x$ which costs $O(n_{i}\times n_{x})$ (if materialized densely).

Where do the Jacobians come from? In fact, our automatic differentiation system is implemented on only a limited set of primitive operations, for each of which we already know how to construct its Jacobian. For example, given $x_{1} + x_{2} = y$ (regarding broadcasting for simplicity), we know for each edge $<x_{i}, y>$ the Jacobian is the identical matrix $I^{n\times n}$, where $x_{i}, y \in \mathbb{R}^n$.

Furthermore, should we really build that identical matrix and carry out matrix multiplication? If we rewrite, again, the chain rule in a somewhat more expressive manner:

\[\nabla_{y}z = g_{G}\left(\left\{ \nabla_{i}z ~|~ i\in C \subseteq \mathrm{ch}(y) \right\}\right)\]

It’s just a total mapping of gradients defined on the DAG, composed by single mappings represented with Jacobians. So if we know the rules, why not write its equivalent, e.g. for addition:

\[x + \dots = y \implies \nabla_{x}z = g_{G}(\nabla_{y}z) := \nabla_{y}z\]

Generally, for operation $y_{1},\dots,y_{m}=f(x_{1},\dots,x_{n})$, there are $n$ gradient mappings, or vector-Jacobian product defined upon:

\[g_{i}: (\nabla_{y_{1}}z ,\dots,\nabla_{y_{m}}z) \mapsto \nabla_{x_{i}}z, ~i=1,\dots,n\]