A short introduction to automatic differentiation

Lecture notes, Optimization for Data Sciences

Samuel Vaiter

2026-09-16 · 2026-10-02 · #optimization #automatic-differentiation · PDF · Markdown

1 Introduction

Automatic differentiation libraries include, in Python, PyTorch, JAX and TensorFlow, and in Julia, Zygote, among others. They have been used in machine learning since roughly 2005–2010. The underlying ideas come from optimal control: Lions in the 1970s–1980s, then A. Walther, A. Griewank and others in the 1990s.

Automatic differentiation, as the name implies it, is a differentiation method. Given a smooth function , one seeks to determine the Jacobian of it (i.e., the matrix of partial derivatives) at some point . Note that typically in machine learning, , and what we look at is the gradient of the function . Nevertheless, let’s keep this generality for the moment.

1.1 Cheap Gradient Principle

Before diving into the core material, let us mention the cheap gradient principle from Baur–Strassen. Let be real rational functions and let .

Complexity.

Here is the minimum number of nonscalar multiplications and divisions needed to compute the rational functions ; additions, subtractions, and multiplication by real constants are free in this complexity model.

Theorem 1: Baur–Strassen

For ,

Backpropagation is exactly taking advantage of these “theoretical” result, preserving more or less the small constant in front of . This extends to algebraic functions (Griewank ’89) and to non-smooth functions (Bolte ’23). Keep in mind this comment when we are going to discuss ReLU activation.

1.2 Finite Differences

f(x) = at x =

Figure 1. Round-off error vs truncation error.

Apart from computing by hand if the analytical expression is known, a common way to approximate a derivative is to use a finite difference scheme. Let with derivative . By definition,

which suggests the approximation

The error decomposes into a round-off part and a truncation part of order . Between and , the error scales with , in a noiseless way. When becomes too big, then it seems to start to “diverge”: this is what we call the truncation error of first-order approximation. The less obvious thing is that, when starts to become smaller and smaller, the error grows again: this is what we call the rounding effect, and this time the error is noisy. The “optimal” step balancing the two is

where is the machine precision. The best achievable error is then of order .

So why bother doing something else? There are at least three reasons:

Computing a full gradient or Jacobian by forward differences takes one function evaluation per input coordinate, plus a shared evaluation at the base point.
It is highly sensitive to the choice .
Subtracting nearly equal values can amplify floating-point error when the step is small.

It is possible to improve the finite difference used above (called the forward finite difference) by using centered difference. Basically, it replaces the estimate by .

1.3 Complex Step Approximation

Can we avoid this nasty round-off error? Let be (complex) analytic with (for instance ). Given , the Taylor expansion at , evaluated at , gives

Taking real and imaginary parts,

hence with truncation error

Unlike finite differences, this involves no subtraction of nearby quantities, so it avoids subtractive cancellation, although floating-point round-off is still present. One of the issue is you need to have a stable complex arithmetic to be able to use this trick.

2 Forward differentiation over scalar functions

2.1 Dual Numbers

The complex-step method avoids subtractive cancellation. Can we now avoid the truncation error? It looks impossible, right? We would like a “number” such that it allows us

holds exactly, with no remainder.

The key idea is is to introduce such that and . The complex numbers can be constructed from the reals in many fashions: completed by an element such as , as with a specific addition and multiplication, or as the quotient space . Similarly, we can construct the dual numbers ring in several ways:

as completed by an element such that ;
as the set endowed with a specific law ;
or as the quotient .

Concretely, is the real plane equipped with the operations

Observe that essentially, going from to , we lose the term for the multiplication. People familiar with differential geometry would not be impressed, since these objects are quite “classical” in this context.

2.2 Linearize everything

So, can we linearize errors with this ring? Polynomials linearize: for ,

What kind of sorcery is behind this?

Extending this analysis to analytical function (over an open set of ), we can prove a similar statement, i.e.,

2.3 Forward-Mode Automatic Differentiation

It turns out that dual number represents one “mode” of automatic differentiation. For , the chain rule reads

and it works for composition:

Remark 2

In Python one implements a class Dual carrying self.real and self.dual, overloading __mul__ and __add__ to encode the rules above. In Julia the same behaviour is obtained through multiple dispatch. We call that operator overloading. Other possibilities include record tape, the original one by Wengert and used for instance by the eager mode of PyTorch, and source-to-source transformation (e.g., Tapenade for FORTRAN code).

3 Reminders on functions of several variables

We now recall basic definition of calculus of several variables, mainly as a pretex to define vector-Jacobian and Jacobian-vector products.

3.1 Gradient

For ,

A forward-difference approximation of the full gradient uses evaluations of : one at and one at for each .

3.2 Directional derivative

Here our function is from to and two objects could be of interest: the directional derivative of at for a given direction , or the full gradient of at . Note that the two are related – knowledge of the gradient is strictly superior to the directional derivative – by

which needs only evaluations.

3.3 Jacobian

For ,

For forward finite differences, the full Jacobian uses evaluations of the vector-valued function , sharing the evaluation at across all outputs. A Jacobian–vector product in one direction uses two evaluations of . The arithmetic cost of each evaluation still depends on and on the output dimension .

4 Chain Rule

Let .

If , then .
If and , then

5 JVP versus VJP

JVP (Jacobian–vector product).

The finite-difference approximation uses two evaluations of , at and ; forward-mode AD computes the JVP without a finite-difference step.

VJP (vector–Jacobian product).

For , the VJP is . Applying forward differences to the scalar function uses evaluations of ; reverse-mode AD computes the VJP in one backward sweep.

The dual part ( above) is a tangent vector, so forward mode JVP pushforward, while reverse mode VJP pullback. This is the JAX mental model (jvp/vjp):

JVP and VJP — **Figure 2.** JVP as pushforward and VJP as pullback.

6 Computational Graph

The key point to understand automatic differentiation is to view computer programs as evaluation of a computational graph. A computational graph represents the stack of evaluation performed by the program. Consider written as a composition

A chain of functions — **Figure 3.** A chain computational graph.

The intermediate variables are

By the chain rule,

Writing , , and , the order in which this product is evaluated determines the cost.

6.1 Forward mode

Since , computing the full gradient (case ) requires JVPs, with cost

For and , this is .

Algorithm 1. Forward mode

Input:
for do
return and

6.2 Reverse mode (backpropagation)

Since , computing a gradient (case ) requires a single VJP,

with cost

For and , this is .

Reverse mode stores the intermediate values needed by the backward pass, using memory proportional to the total size of those stored values. Forward mode can stream through a chain, retaining only the current primal and tangent states. Checkpointing is a way to alleviate this issue.

Algorithm 2. Reverse mode

Input:
, ()
for do
for to do
return and ()

7 Backpropagation on a General Graph

The graph above was a chain: every node had exactly one predecessor, so the Jacobian factorised as a single sequential product. A real program is not a chain but a directed acyclic graph (DAG): a variable may feed several operations (fan-out) and an operation may read several variables (fan-in). The running example is exactly of this kind – is consumed three times.

The graph.

Write for the DAG. Each non-source node computes an elementary operation of its parents,

where denotes the parents of . The sources are the inputs , and we assume a single sink (a scalar output; several outputs are handled by seeding several adjoints).

The running DAG — **Figure 4.** The running directed acyclic graph.

v1 = x1 + x2
v2 = x1 * x2
v3 = exp(x1)
v4 = v1 + v2
o  = v3 + v4

One assignment per line (+ - * /, ^, exp/sin/cos/log/sqrt/tanh). Above each node: value (forward). Below, in bold: adjoint ∂o/∂· (reverse).

Figure 5. Interactive computational graph: edit the “code”, build the DAG and its topological order, then play the reverse sweep watching each adjoint flow back to the input roots.

with

Topological order.

A topological order is a linear ordering of in which every node appears after all its parents; it exists if and only if is acyclic. After assigning the inputs to the source nodes, evaluating the remaining nodes in this order – the forward pass – computes each from already-available parents and reproduces .

The chain rule with fan-out.

Let the adjoint of be the (column) gradient of the output with respect to that node, . Applying the chain rule at every operation that consumes , the adjoint accumulates one contribution per child:

This sum is the entire difference with the chain case. When every node has a single child the sum has one term and we recover the sequential product of the linear case above; fan-out is precisely what a chain cannot express, and the gradient of a reused variable is the sum of the contributions of all its uses.

Reverse sweep.

Seed and visit the nodes in reverse topological order: this guarantees that when is reached, the adjoints of all its children are already final. In practice one organizes this as a push (or scatter): initialise every , and as each node is processed it adds to the accumulator of each parent . Each elementary operation then only needs its local VJP – how to turn an output cotangent into input cotangents – while the graph wiring and the reverse order handle the global accumulation. The adjoints of the inputs are the gradient, .

Algorithm 3. Reverse-mode AD on a DAG (backpropagation)

Input: DAG with inputs and sink
topologically sort as (sources first)
Assign the given input values to the source nodes
for do
- if is not a source then
  - (store )
for all ,
for do
- for each do
  - $\bar{x}_u \leftarrow \bar{x}_u
    - \big(\partial \phi_{v_t} / \partial x_u\big)^\top \bar{x}_{v_t}$
return

The example, backward.

The reverse sweep gives , then and . The reused input , with the three children , collects three terms,

which is exactly .

Cost.

Each edge is traversed once per pass, so the backward sweep costs a constant times the forward graph, the cheap-gradient principle on an arbitrary DAG, at the price of storing the intermediates (the tape). A finite loop can be unrolled into an acyclic computation and differentiated by backpropagation. For an output defined by an implicit relation, such as a fixed point, implicit differentiation is another option.