# Automatic Differentiation (AD)

In short, the promise of AD is

```julia
f(x) = 4x + x^2

df(x) = derivative(f, x)
```

such that

```julia
df(3) = 4 + 2*3 = 10
```

### What AD is not

**Symbolic rewriting:**
$$ f(x) = 4x + x^2 \quad \rightarrow \quad df(x) = 4 + 2x $$

**Numerical differentiation:**
$$ \frac{df}{dx} \approx \frac{f(x+h) - f(x)}{\Delta h} $$

## Forward mode AD

Key to AD is the application of the chain rule
$$\dfrac{d}{dx} f(g(x)) = \dfrac{df}{dg} \dfrac{dg}{dx}$$

Consider the function $f(a,b) = \ln(ab + \sin(a))$.


```julia
f(a,b) = log(a*b + sin(a))
```


```julia
f_derivative(a,b) = 1/(a*b + sin(a)) * (b + cos(a))
```


```julia
a = 3.1
b = 2.4
f_derivative(a,b)
```

Dividing the function into the elementary steps, it corresponds to the following "*computational graph*":




```julia
function f_graph(a,b)
    c1 = a*b
    c2 = sin(a)
    c3 = c1 + c2
    c4 = log(c3)
end
```


```julia
f(a,b) == f_graph(a,b)
```

To calculate $\frac{\partial f}{\partial a}$ we have to apply the chain rule multiple times.

$\dfrac{\partial f}{\partial a} = \dfrac{\partial f}{\partial c_4} \dfrac{\partial c_4}{\partial a} = \dfrac{\partial f}{\partial c_4} \left( \dfrac{\partial c_4}{\partial c_3} \dfrac{\partial c_3}{\partial a}  \right) = \dfrac{\partial f}{\partial c_4} \left( \dfrac{\partial c_4}{\partial c_3} \left( \dfrac{\partial c_3}{\partial c_2} \dfrac{\partial c_2}{\partial a} + \dfrac{\partial c_3}{\partial c_1} \dfrac{\partial c_1}{\partial a}\right)  \right)$


```julia
function f_graph_derivative(a,b)
    c1 = a*b
    c1_ϵ = b
    
    c2 = sin(a)
    c2_ϵ = cos(a)
    
    c3 = c1 + c2
    c3_ϵ = c1_ϵ + c2_ϵ
    
    c4 = log(c3)
    c4_ϵ = 1/c3 * c3_ϵ
    
    c4, c4_ϵ
end
```


```julia
f_graph_derivative(a,b)[2] == f_derivative(a,b)
```

**How can we automate this?**


```julia
# D for "dual number", invented by Clifford in 1873.
struct D <: Number
    x::Float64 # value
    ϵ::Float64 # derivative
end
```


```julia
import Base: +, *, /, -, sin, log, convert, promote_rule

a::D + b::D = D(a.x + b.x, a.ϵ + b.ϵ) # sum rule
a::D - b::D = D(a.x - b.x, a.ϵ - b.ϵ)
a::D * b::D = D(a.x * b.x, a.x * b.ϵ + a.ϵ * b.x) # product rule
a::D / b::D = D(a.x / b.x, (b.x * a.ϵ - a.x * b.ϵ)/b.x^2) # quotient rule

sin(a::D) = D(sin(a.x), cos(a.x) * a.ϵ)
log(a::D) = D(log(a.x), 1/a.x * a.ϵ)

Base.convert(::Type{D}, x::Real) = D(x, zero(x))
Base.promote_rule(::Type{D}, ::Type{<:Number}) = D
```


```julia
f(D(a,1), b)
```

Boom! That was easy!


```julia
f_derivative(a,b)
```


```julia
f(D(a,1), b).ϵ ≈ f_derivative(a,b)
```

**How does this work?!**

The trick of forward mode AD is to let Julia implicitly perform the mapping `f -> f_graph_derivative` for you and then let the compiler optimize the resulting code structure (that's what compilers do!).


```julia
@code_typed f(D(a,1), b)
```

While this is somewhat hard to parse, plugging these operations manually into each other we find that this code equals

```julia
D.x = log(a.x*b + sin(a.x))
D.ϵ = 1/(a.x*b + sin(a.x)) * (a.x*0 + (a.ϵ*b) + cos(a.x)*a.ϵ)
```

which, if we drop `a.x*0`, set `a.ϵ = 1`, and rename `a.x` $\rightarrow$ `a`, reads

```julia
D.x = log(a*b + sin(a))
D.ϵ = 1/(a*b + sin(a)) * (b + cos(a)
```

This precisely matches our definitions from above:

```julia
f(a,b) = log(a*b + sin(a))

f_derivative(a,b) = 1/(a*b + sin(a)) * (b + cos(a))
```

Importantly, the compiler sees the entire "rewritten" code and can therefore apply optimizations. In this simple example, we find that the code produced by our simple Forward mode AD is essentially identical to the explicit implementation.


```julia
@code_llvm debuginfo=:none f_graph_derivative(a,b)
```


```julia
@code_llvm debuginfo=:none f(D(a,1), b)
```

Our AD is alreadly pretty powerful and general. Let's define the promised function `derivative`:


```julia
derivative(f::Function, x::Number) = f(D(x, one(x))).ϵ
```


```julia
g(x) = x + x^2
```


```julia
derivative(g, 3.0)
```

Anonymous function oft come in handy here:


```julia
derivative(x->3*x^2+4x+5, 2)
```


```julia
derivative(x->sin(x)*log(x), 3)
```

We can also define the partial derivative $\frac{df(a,b)}{da}$ from above:


```julia
df(x) = derivative(a->f(a,b),x)
```

Here, `b` is "wrapped into a closure".


```julia
df(1.23)
```

## Taking the derivative of *code*

> Repeat $t \leftarrow (t + x/2)/2$ until $t$ converges to $\sqrt{x}$.


```julia
@inline function Babylonian(x; N = 10)
    t = (1+x)/2
    for i = 2:N
        t = (t + x/t)/2
    end
    t
end
```


```julia
Babylonian(2)
```


```julia
sqrt(2)
```

Using our forward mode AD, that is our dual numbers, we can compute the derivative of `Babylonian` **with no rewrite at all**.


```julia
Babylonian(D(5, 1))
```


```julia
sqrt(5)
```


```julia
1 / (2*sqrt(5))
```

**It just works and is efficient!**


```julia
@code_native debuginfo=:none Babylonian(D(5, 1))
```

Recursion? Works as well...


```julia
function power(x, n)
    if n <= 0
        return 1
    else
        return x*power(x, n-1)
    end
end
```


```julia
4.0^3
```


```julia
derivative(x -> power(x,3), 4.0)
```


```julia
3*4.0^2 # 3*x^2
```

Deriving our Vandermonde matrix from yesterday?


```julia
function vander_generic(x::AbstractVector{T}) where T
    m = length(x)
    V = Matrix{T}(undef, m, m)
    for j = 1:m
        V[j,1] = one(x[j])
    end
    for i= 2:m
        for j = 1:m
            V[j,i] = x[j] * V[j,i-1]
            end
        end
    return V
end
```

\begin{align}V=\begin{bmatrix}1&a&a^{2} &a^3\\1&b&b^{2} &b^3\\1&c&c^{2} &c^3\\1&d&d^{2} &d^3\end{bmatrix}\end{align}

\begin{align}\frac{dV}{da}=\begin{bmatrix}0&1&2a &3a^2\\0&0&0 &0\\0&0&0 &0\\0&0&0 &0\end{bmatrix}\end{align}


```julia
a, b, c, d = 2, 3, 4, 5
V = vander_generic([D(a,1), D(b,0), D(c,0), D(d,0)])
```


```julia
[V[i,j].ϵ for i in axes(V,1), j in axes(V,2)]
```

## Symbolically (because we can)

The below is mathematically equivalent, **though not exactly what the computation is doing**. Our AD isn't performing symbolic manipulations.


```julia
using SymPy
```


```julia
@vars x
```


```julia
Babylonian(x; N=1)
```


```julia
diff(Babylonian(x; N=1))
```


```julia
simplify(Babylonian(x; N=5))
```


```julia
simplify(diff(simplify(Babylonian(x; N=5)), x))
```

## Don't reinvent the wheel: ForwardDiff.jl

Now that we have understood how forward AD works, we can use the more feature complete package [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl).


```julia
using ForwardDiff
```


```julia
ForwardDiff.derivative(Babylonian, 2)
```


```julia
@edit ForwardDiff.derivative(Babylonian, 2)
```

(Note: [DiffRules.jl](https://github.com/JuliaDiff/DiffRules.jl))

## If time permits: Reverse mode AD

Forward mode:
$\dfrac{\partial f}{\partial x} = \dfrac{\partial f}{\partial c_4} \dfrac{\partial c_4}{\partial x} = \dfrac{\partial f}{\partial c_4} \left( \dfrac{\partial c_4}{\partial c_3} \dfrac{\partial c_3}{\partial x}  \right) = \dfrac{\partial f}{\partial c_4} \left( \dfrac{\partial c_4}{\partial c_3} \left( \dfrac{\partial c_3}{\partial c_2} \dfrac{\partial c_2}{\partial x} + \dfrac{\partial c_3}{\partial c_1} \dfrac{\partial c_1}{\partial x}\right)  \right)$

Reverse mode:
$\dfrac{\partial f}{\partial x} = \dfrac{\partial f}{\partial c_4} \dfrac{\partial c_4}{\partial x} = \left( \dfrac{\partial f}{\partial c_3}\dfrac{\partial c_3}{\partial c_4}   \right) \dfrac{\partial c_4}{\partial x} = \left( \left( \dfrac{\partial f}{\partial c_2} \dfrac{\partial c_2}{\partial c_3} + \dfrac{\partial f}{\partial c_1} \dfrac{\partial c_1}{\partial c_3} \right) \dfrac{\partial c_3}{\partial c_4} \right) \dfrac{\partial c_4}{\partial x}$

Forward mode AD requires $n$ passes in order to compute an $n$-dimensional
gradient.

Reverse mode AD requires only a single run in order to compute a complete gradient but requires two passes through the graph: a forward pass during which necessary intermediate values are computed and a backward pass which computes the gradient.

*Rule of thumb:*

Forward mode is good for $\mathbb{R} \rightarrow \mathbb{R}^n$ while reverse mode is good for $\mathbb{R}^n \rightarrow \mathbb{R}$.

An efficient source-to-source reverse mode AD is implemented in [Zygote.jl](https://github.com/FluxML/Zygote.jl), the AD underlying [Flux.jl](https://fluxml.ai/) (since version 0.10).


```julia
using Zygote
```


```julia
f(x) = 5*x + 3
```


```julia
gradient(f, 5)
```


```julia
@code_llvm debuginfo=:none gradient(f,5)
```


```julia
@code_llvm debuginfo=:none derivative(f,5)
```

## Some nice reads

Papers:
* https://www.jmlr.org/papers/volume18/17-468/17-468.pdf

Lectures:


* https://mitmath.github.io/18337/lecture8/automatic_differentiation.html

Blog posts:

* ML in Julia: https://julialang.org/blog/2018/12/ml-language-compiler

* Nice example: https://fluxml.ai/2019/03/05/dp-vs-rl.html

* Nice interactive examples: https://fluxml.ai/experiments/

* Why Julia for ML? https://julialang.org/blog/2017/12/ml&pl

* Neural networks with differential equation layers: https://julialang.org/blog/2019/01/fluxdiffeq

* Implement Your Own Automatic Differentiation with Julia in ONE day : http://blog.rogerluo.me/2018/10/23/write-an-ad-in-one-day/

* Implement Your Own Source To Source AD in ONE day!: http://blog.rogerluo.me/2019/07/27/yassad/

Repositories:

* AD flavors, like forward and reverse mode AD: https://github.com/MikeInnes/diff-zoo (Mike is one of the smartest Julia ML heads)

Talks:

* AD is a compiler problem: https://juliacomputing.com/assets/pdf/CGO_C4ML_talk.pdf

44b2f002236f480247e4a09b76a84b31c82c4b84
ipynb
Jupyter Notebook
Day2/4a_automatic_differentiation.ipynb
tbdrstl/JuliaCologne21
0e2778cd8d0601291dcc8573c20469b3fb84ebe7
MIT
2021-03-09T19:05:00.000Z
2021-08-11T05:00:07.000Z
Day2/4a_automatic_differentiation.ipynb
tbdrstl/JuliaCologne21
0e2778cd8d0601291dcc8573c20469b3fb84ebe7
MIT
2021-03-18T09:27:40.000Z
2021-04-08T09:37:47.000Z
Day2/4a_automatic_differentiation.ipynb
tbdrstl/JuliaCologne21
0e2778cd8d0601291dcc8573c20469b3fb84ebe7
MIT
2021-03-18T09:30:02.000Z
2021-06-18T06:50:58.000Z
Qwen/Qwen-72B
1. YES
2. YES