Automatic Differentiation
Derivatives are everywhere in computing — optimisation, machine learning, physics. But how does a computer find a derivative? Not by doing calculus on paper, and (it turns out) not by the approximate finite differences you may have seen. The real workhorse is automatic differentiation (AD), and the forward-mode version is delightfully simple: a tiny new kind of number.
The MIT lecture uses ForwardDiff.jl. WebAssembly can't run that package, so here we build forward-mode AD from scratch in a dozen lines — and it computes the exact derivative, live in your browser.
begin
using PlutoUI, WasmMakie
endFirst, a reminder: the derivative
The derivative $f'(x)$ is the slope of $f$ at $x$. The definition you learned is a limit of a finite difference:
$$f'(x) \approx \frac{f(x+h) - f(x)}{h}, \qquad h \to 0.$$
Let's watch that for $f = \sin$ at $x = 1$, where the true answer is $\cos(1) \approx 0.5403$. Shrink the step $h$ and the estimate creeps toward it — but never exactly, and if you make $h$ too small, floating-point round-off ruins it. AD will fix both problems.
step size h =
fd_estimate = (sin(1.0 + h) - sin(1.0)) / hfinite-difference estimate $f'(1) \approx$ 0.536085981011869
exact $\cos(1) =$ 0.5403023058681398
The gap is the error — try the smallest steps and watch round-off creep back in.
Forward-mode AD: dual numbers
Here is the trick. Invent a new symbol $\varepsilon$ with the rule $\varepsilon^2 = 0$ (it is not a real number — think of it as "infinitely small"). A dual number is $a + b\,\varepsilon$. Now feed $x + \varepsilon$ into any function and expand:
$$f(x + \varepsilon) = f(x) + f'(x)\,\varepsilon + \tfrac12 f''(x)\,\varepsilon^2 + \dots = f(x) + f'(x)\,\varepsilon$$
— every term with $\varepsilon^2$ vanishes! So the value lands in the first slot and the exact derivative lands in the $\varepsilon$ slot, for free. We just need a number type that carries both slots and knows the arithmetic rules.
begin
# A dual number carries a value `v` = f(x) and a derivative `d` = f'(x).
struct Dual
v::Float64
d::Float64
end
# Arithmetic: each rule is just the sum / product / quotient rule of calculus.
Base.:+(a::Dual, b::Dual) = Dual(a.v + b.v, a.d + b.d)
Base.:-(a::Dual, b::Dual) = Dual(a.v - b.v, a.d - b.d)
Base.:*(a::Dual, b::Dual) = Dual(a.v * b.v, a.v * b.d + a.d * b.v) # product rule
Base.:/(a::Dual, b::Dual) = Dual(a.v / b.v, (a.d * b.v - a.v * b.d) / (b.v * b.v))
# Mixing a dual with a plain number (a constant has derivative 0).
Base.:+(a::Dual, b::Real) = Dual(a.v + b, a.d)
Base.:+(a::Real, b::Dual) = Dual(a + b.v, b.d)
Base.:-(a::Dual, b::Real) = Dual(a.v - b, a.d)
Base.:-(a::Real, b::Dual) = Dual(a - b.v, -b.d)
Base.:-(a::Dual) = Dual(-a.v, -a.d)
Base.:*(a::Dual, b::Real) = Dual(a.v * b, a.d * b)
Base.:*(a::Real, b::Dual) = Dual(a * b.v, a * b.d)
Base.:/(a::Dual, b::Real) = Dual(a.v / b, a.d / b)
# Elementary functions via the chain rule.
Base.sin(a::Dual) = Dual(sin(a.v), cos(a.v) * a.d)
Base.cos(a::Dual) = Dual(cos(a.v), -sin(a.v) * a.d)
Base.exp(a::Dual) = Dual(exp(a.v), exp(a.v) * a.d)
end# The whole of forward-mode AD: evaluate f at the dual number x + 1·ε and read off
# the ε-slot. No finite differences, no round-off — the exact slope.
ad_derivative(f, x::Float64) = f(Dual(x, 1.0)).d# Example functions. Note: NO type annotation on `x`, so the SAME code runs on plain
# Float64 (to draw the curve) and on Dual numbers (to get the derivative).
function example(x, which::Int64)
if which == 1
return x * x - 2.0 # f(x) = x² − 2 , f′ = 2x
elseif which == 2
return sin(x) # f(x) = sin x , f′ = cos x
else
return x * x * x - 2.0 * x + 1.0 # f(x) = x³ − 2x + 1 , f′ = 3x² − 2
end
endSee it work
Pick a function and a point. The number is the derivative our dual numbers compute; the orange line is the tangent with that slope — drag $x_0$ and it stays glued to the curve.
example f =
point x₀ =
slope_ad = ad_derivative(u -> example(u, which), x0)f′(x₀) by dual numbers = 2.0 — the exact slope of the tangent line below.
let
# sample the curve with an integer loop (StepRangeLen iteration is not
# wasm-compilable yet), exactly like the Newton notebook
lo = -2.5
hi = 2.5
steps = 240
xs = Float64[]
ys = Float64[]
for k in 0:steps
t = lo + (hi - lo) * k / steps
push!(xs, t)
push!(ys, example(t, which))
end
# the tangent line through (x0, f(x0)) with the AD-computed slope
y0 = example(x0, which)
m = ad_derivative(u -> example(u, which), x0)
fig = Figure(size = (480, 320))
ax = Axis(fig[1, 1])
lines!(ax, xs, ys) # the curve f
lines!(ax, [lo, hi], [0.0, 0.0]) # the x-axis
lines!(ax, [lo, hi], [y0 + m * (lo - x0), y0 + m * (hi - x0)]) # tangent from AD
lines!(ax, [x0, x0], [y0 - 0.4, y0 + 0.4]) # mark the point x0
fig
endWhy this is a big deal
We never wrote a single derivative formula for example — we wrote the function once, and the dual-number arithmetic carried the derivative through automatically. That is exactly how ForwardDiff.jl and the autodiff engines inside PyTorch, JAX and Flux work (just with more bookkeeping and more elementary functions defined).
The original MIT lecture pushes this further: derivatives of multivariate functions (the gradient) and of vector-valued functions that transform images (the Jacobian). The idea is identical — carry an $\varepsilon$ in every input direction — only the bookkeeping grows.