Automatic Differentiation

Automatic Differentiation#

Automatic differentiation augments a computation with another computation that propagates derivatives through the same primitive operations. It applies the chain rule to the computational graph exactly, up to floating-point arithmetic.

Computational graphs#

A computational graph is a directed acyclic graph (or DAG) that encodes the computation of a multivariate function. We are going to use two types of nodes in the graph. Tensor nodes (i.e., numerical arrays) are represented by ellipses. Operation nodes (i.e., mathematical operations) are represented by shaded boxes. The mathematical operators we consider here are primitive operators: operations treated as indivisible at the chosen level of description. Which operators are primitive depends on where we build the computational graph. For example, if we are working directly at the level of the compiler of a programming language, the primitives are just the basic arithmetic operations (e.g., addition, multiplication). At a higher level, primitives can include matrix multiplication, trigonometric functions, exponentials, and power functions. Because we are working within Python, we will use the second approach. This is the level exposed by libraries such as PyTorch and JAX.

Here is an example. Consider the following function:

\[ f(x) = x^2 + x \]

Here is a good implementation of this function in Python:

def f(x):
    a = x ** 2
    b = a + x
    return b

The implementation names the intermediate results a and b. The node sqr below computes the square, and add computes the sum. When x is a tensor, libraries such as PyTorch and JAX record a computational graph of this form:

Show code cell source

Hide code cell source

import base64

import graphviz
from IPython.display import display


def new_graph(rankdir="LR", fontsize=14):
    return graphviz.Digraph(
        graph_attr={
            "rankdir": rankdir,
            "bgcolor": "transparent",
            "margin": "0.02",
            "ordering": "out",
            "nodesep": "0.12",
            "ranksep": "0.08",
        },
        node_attr={
            "fontname": "DejaVu Sans",
            "fontsize": str(fontsize),
            "color": "black",
            "fontcolor": "black",
            "penwidth": "1.0",
            "margin": "0",
        },
        edge_attr={
            "color": "black",
            "penwidth": "0.8",
            "arrowsize": "0.55",
        },
    )


tensor_node = {
    "shape": "ellipse",
    "style": "filled",
    "fillcolor": "white",
    "fixedsize": "true",
    "width": "0.48",
    "height": "0.30",
}
operation_node = {
    "shape": "box",
    "style": "filled",
    "fillcolor": "gray85",
    "fixedsize": "true",
    "width": "0.55",
    "height": "0.38",
}
derivative_tensor_node = {
    "shape": "ellipse",
    "style": "dashed,filled",
    "fillcolor": "gray90",
    "fixedsize": "true",
    "width": "0.58",
    "height": "0.30",
}
derivative_output_node = {
    **derivative_tensor_node,
    "width": "1.05",
}
derivative_operation_node = {
    "shape": "box",
    "style": "dashed,filled",
    "fillcolor": "gray75",
    "fixedsize": "true",
    "width": "0.72",
    "height": "0.40",
}
derivative_edge = {
    "style": "dashed",
    "color": "black",
    "penwidth": "1.0",
    "arrowsize": "0.55",
}


def display_graph(graph):
    """Display vector web output and print-compatible PDF/PNG fallbacks."""
    svg = graph.pipe(format="svg").decode("utf-8")
    pdf = graph.pipe(format="pdf")
    graph.graph_attr["dpi"] = "300"
    try:
        png = graph.pipe(format="png")
    finally:
        graph.graph_attr.pop("dpi", None)
    display(
        {
            "image/svg+xml": svg,
            "application/pdf": base64.b64encode(pdf).decode("ascii"),
            "image/png": base64.b64encode(png).decode("ascii"),
        },
        raw=True,
    )


g = new_graph(fontsize=9)

g.node('x', **tensor_node)
g.node('a', **tensor_node)
g.node('b', **tensor_node)
g.node('sqr', **operation_node)
g.node('add', **operation_node)

g.edge('x', 'sqr')
g.edge('sqr', 'a')
g.edge('a', 'add')
g.edge('x', 'add')
g.edge('add', 'b')

display_graph(g)

Computational graph for b equals x squared plus x, alternating tensor nodes with square and addition operation nodes.

The directed edges alternate between tensor values and primitive operations and follow the evaluation order from the input \(x\) to the output \(b\). An automatic-differentiation system treats each primitive as atomic and supplies local derivative rules for it. To differentiate \(b=f(x)\) with respect to \(x\), the system augments this graph with operations that propagate derivative information.

Full Jacobians#

We begin with a chain of four differentiable primitive functions, \(f_1\), \(f_2\), \(f_3\), and \(f_4\). Generic DAGs are discussed later. Their composition defines \(f\) by

\[ f(x) = f_4(f_3(f_2(f_1(x)))). \]

Let their domains and codomains be

\[\begin{split} \begin{aligned} f_1&: \mathbb{R}^n \rightarrow \mathbb{R}^p, \\ f_2&: \mathbb{R}^p \rightarrow \mathbb{R}^q, \\ f_3&: \mathbb{R}^q \rightarrow \mathbb{R}^r, \\ f_4&: \mathbb{R}^r \rightarrow \mathbb{R}^m, \end{aligned} \end{split}\]

so that

\[ f: \mathbb{R}^n \rightarrow \mathbb{R}^m. \]

Using functional notation, we can write \(f\) as the composition of the four functions:

\[ f = f_4 \circ f_3 \circ f_2 \circ f_1. \]

For two stages, let \(f_{1,k}\) denote component \(k\) of \(f_1\) and \(f_{2,i}\) component \(i\) of \(f_2\). The chain rule gives, for \(i=1,\ldots,q\) and \(j=1,\ldots,n\),

\[ \partial_j (f_2 \circ f_1)_i(x) = \sum_{k=1}^p \partial_k f_{2,i}(f_1(x)) \partial_j f_{1,k}(x) = \partial_k f_{2,i}(f_1(x)) \partial_j f_{1,k}(x), \]

where \(\partial_j\) denotes differentiation with respect to input coordinate \(j\). The last expression uses the Einstein convention: the repeated index \(k=1,\ldots,p\) is summed. We write \(J_f(x)\) for the Jacobian matrix with entries \([J_f(x)]_{ij}=\partial_j f_i(x)\). In matrix form,

\[ J_{f_2 \circ f_1}(x) = J_{f_2}(f_1(x)) J_{f_1}(x). \]

In other words, the Jacobian of \(f_2\circ f_1\) is a \(q\times n\) matrix, and it is the product of the Jacobian of \(f_2\) and the Jacobian of \(f_1\). Let’s now add the rest of the functions:

\[ J_f(x) = J_{f_4}(f_3(f_2(f_1(x)))) J_{f_3}(f_2(f_1(x))) J_{f_2}(f_1(x)) J_{f_1}(x). \]

The corresponding code names each intermediate value:

def f(x):
    a = f1(x)
    b = f2(a)
    c = f3(b)
    d = f4(c)
    return d

Thus,

\[\begin{split} \begin{aligned} a &= f_1(x), \\ b &= f_2(a), \\ c &= f_3(b), \\ d &= f_4(c). \end{aligned} \end{split}\]

The resulting computational graph is:

A chain of four primitive functions maps x through intermediate values a, b, and c to output d.

Using the intermediate variables, the chain rule becomes

\[ J_f(x) = J_{f_4}(c) J_{f_3}(b) J_{f_2}(a) J_{f_1}(x). \]

The desired Jacobian is therefore a product of the primitive Jacobians. A direct algorithm would

Run the computational graph forward to compute the intermediate variables.
Compute the Jacobian of each primitive function.
Multiply the Jacobians in chain-rule order.

The intermediate Jacobians can be prohibitively large in machine-learning applications. Automatic differentiation avoids forming them explicitly.

Forward-mode automatic differentiation#

For a direction \(v\in\mathbb{R}^n\), forward mode computes the Jacobian-vector product (JVP) \(J_f(x)v\):

\[ J_f(x)v = J_{f_4}(c)\Big(J_{f_3}(b)\big(J_{f_2}(a)(J_{f_1}(x)v)\big)\Big). \]

This calculation requires only the local JVP rule for each primitive. We write primitive.jvp(z, w) for the primitive’s Jacobian at \(z\) applied to the incoming tangent vector \(w\). The vectors v1, v2, and v3 below are the propagated tangents:

a = f1(x)
v1 = f1.jvp(x, v)
b = f2(a)
v2 = f2.jvp(a, v1)
c = f3(b)
v3 = f3.jvp(b, v2)
d = f4(c)
Jfv = f4.jvp(c, v3)

The computational graph below overlays this tangent calculation on the ordinary value computation, often called the primal computation. Solid nodes and edges show the primal computation; dashed nodes and edges show the JVP calculation.

Forward-mode differentiation graph carrying tangent v through four Jacobian-vector product operations to produce the product of the full Jacobian with v.

Forward mode carries one tangent vector alongside each primal value and never forms a full intermediate Jacobian.

Choosing \(v=e_j\), the \(j\)-th standard basis vector of \(\mathbb{R}^n\), extracts one Jacobian column:

\[ J_f(x)e_j = \partial_j f(x). \]

Computing every column requires \(n\) forward sweeps. For a scalar loss, \(n\) is the number of model parameters while \(m=1\), so this cost scales with \(n\), as coordinatewise finite differences do. Unlike finite differences, automatic differentiation has no step size and introduces no truncation error. Reverse mode exploits the small output dimension.

Reverse-mode automatic differentiation#

For a cotangent \(u\in\mathbb{R}^m\), a vector of sensitivities attached to the output, reverse mode computes the vector-Jacobian product (VJP) \(u^{\mathsf T} J_f(x)\):

\[ u^{\mathsf T} J_f(x) = \Big(\big((u^{\mathsf T} J_{f_4}(c))J_{f_3}(b)\big)J_{f_2}(a)\Big)J_{f_1}(x). \]

Suppose each primitive provides a local VJP rule. The vectors u1, u2, and u3 below are the cotangents propagated from the output back toward the input:

# Forward pass
a = f1(x)
b = f2(a)
c = f3(b)
d = f4(c)

# Backward pass
u1 = f4.vjp(c, u)
u2 = f3.vjp(b, u1)
u3 = f2.vjp(a, u2)
uTJf = f1.vjp(x, u3)

In the graph below, solid edges supply stored primal values to the VJP operations, while dashed edges carry cotangents backward.

Reverse-mode differentiation graph carrying cotangent u backward through four vector-Jacobian product operations to produce u transpose times the full Jacobian.

Reverse mode normally stores the intermediate primal values and traverses the graph once forward and once backward. Computing the entire Jacobian requires one forward evaluation followed by \(m\) reverse sweeps, using \(u=e_i\) for \(i=1,\ldots,m\). For a scalar loss, \(m=1\) and \(u=1\), so one reverse sweep after the forward evaluation computes the gradient with respect to all \(n\) parameters. In machine learning, this reverse-mode calculation is called backpropagation.

Generic computational graphs#

Branches require one additional local rule. Consider

\[ f(x) = x + x, \]

implemented as

def f(x):
    a = x + x
    return a

The primitive \(\operatorname{add}:\mathbb{R}^n\times\mathbb{R}^n\to\mathbb{R}^n\) maps a pair of tensors to their sum. The map \(\operatorname{fanout}:\mathbb{R}^n\to\mathbb{R}^n\times\mathbb{R}^n\) duplicates its input: \(\operatorname{fanout}(x)=(x,x)\). Therefore,

\[ f = \operatorname{add} \circ \operatorname{fanout}. \]

After identifying \(\mathbb{R}^n\times\mathbb{R}^n\) with \(\mathbb{R}^{2n}\), the fanout Jacobian is the \(2n\times n\) matrix

\[\begin{split} J_{\operatorname{fanout}}(x) = \begin{bmatrix} I_n \\ I_n \end{bmatrix}. \end{split}\]

Its JVP duplicates an incoming tangent. In reverse mode, cotangent contributions arriving through the two outgoing paths are added before propagation continues toward the input.

The same construction applies to

\[ f(x) = x^2 + x, \]

implemented as:

def f(x):
    a = x * x
    b = a + x
    return b

Here the square and the addition act componentwise. Let \(\operatorname{sqr}(x)=x\odot x\), where \(\odot\) denotes componentwise multiplication, and let \(\operatorname{id}\) be the identity map. Then

\[ f = \operatorname{add} \circ (\operatorname{sqr}\times\operatorname{id}) \circ \operatorname{fanout}. \]

The product map \(\operatorname{sqr}\times\operatorname{id}:\mathbb{R}^n\times\mathbb{R}^n\to\mathbb{R}^n\times\mathbb{R}^n\) sends \((x_1,x_2)\) to \((\operatorname{sqr}(x_1),\operatorname{id}(x_2))\). Traversing operations in dependency order, applying local JVP/VJP rules, and adding cotangent contributions at branches extend automatic differentiation from chains to arbitrary DAGs.