TaylorDiff.jl: Taking Higher-order Derivatives Seriously

Songchen Tan

Introduction

## Understanding higher-order AD

First-order AD is all about finding an efficient and maintainable way for executing the **chain rule formula**

$$
\frac{\mathrm df(g(x))}{\mathrm dx}=f'(g(x))g'(x)
$$

in which forward mode AD (based on operator overloading) achieve that by using **dual number arithmetic**.
---
Higher-order AD is, by analogy, finding an efficient and maintainable way for executing the **Faa di Bruno’s formula** ($\pi=(k_1,\cdots,k_n)$ is a partition of integer $n$, such that $\sum_iik_i=n$)

$$
\frac{\mathrm d^nf(g(x))}{\mathrm dx^n}=\sum_{\pi \in \Pi} f^{(|\pi|)}(g(x)) \cdot \prod_{B \in \pi} g^{(B)}(x)
$$

in which Taylor mode AD is a generalization of operator overloading forward mode AD, using **polynomial arithmetic**.

Faa di Bruno's formula itself didn't give much performance gain, since there are exponentially many partitions. However, it gives insights on how we could reuse computations across different order.
---
## Current solutions

### TaylorSeries.jl

- Capable for calculating the full derivative tensor for arbitrary order
- Operator-overloaded approach on one- and multi-variable Taylor polynomials
- For $n$-variable derivatives up to $k$ order, TaylorSeries.jl forces you to compute the full derivative, which takes $O(kn^k)$ time; however, **full derivative is rarely needed**
- Hard-coded higher-order rules, only a limited amount of primitives
- Dynamically-sized and mutating array, efficiency and compatibility problem

---
### Diffractor.jl

- Next-generation SCT-based forward-mode and reverse-mode AD, designed with support for higher-order derivatives in mind
- Has a direct implementation of Faa di Bruno's formula
- Relies on writing specific rules for efficiency
- Type inference issues

---
### `jax.experimental.jet`

- Python implementation of Taylor mode AD in JAX
- Can only compute higher-order directional derivatives
- Benefits from vectorized code
- Hard-coded higher-order rules (again)
---

## Design goals and non-goals

Let's combine the best part of them:

- Focus on directional derivative (that's where Taylor mode shines) as in `jax.jet`
    - Cross-derivatives achieved by nesting this AD (like $\partial^5 u/\partial x^2\partial y^3$)
- Order as a type parameter, as in Diffractor.jl
    - No performance penalty in lower-order
- Intuitive interface, as in TaylorSeries.jl
- Smarter fallbacks than generic Faa di Bruno's formula
    - By lifting ChainRules's `frule` to higher-order
- Composable with other ADs

Theory

## Order as a type parameter

A Taylor polynomial is a order `N` polynomial represented by a wrapper over `NTuple{N + 1, P}` for arbitrary primal type `P`.

- Fully unrolled and non-allocating for performance
- No more mutating arrays issue

(You might be concerned about the fact that tangent space can be different from primal space, but that can be fixed by treating it separately like `P` and `NTuple{N, T}`. The mapping from primal space to tangent space is likely *idempotent*, so we won't have very different spaces for each order.)

---
## Taylor rules

For each function definition `f(x...) = Ω`, we make a overload `f(Tx...) = TΩ`, where

- `Tx` could be the primal type (optionally promoted to the polynomial type) or the polynomial type
- `TΩ` could be the primal type or the polynomial type

Overloading is per order for more aggressive optimization.

---
## Generation of Taylor rules from ChainRules `frule`s

How to obtain the Taylor rules automatically, so that we don't need to write them one-by-one? It turns out that we can make use of the higher-order product rule or division rule,

$$
f(g(x))^{(n)}=[f'(g(x))g'(x)]^{(n-1)}=\sum_{i=0}^{n-1}[f'(g(x))]^{(i)}g^{(n-i)}
$$

- Obtain $f'$ (often closely related to $f$ and $g$) from ChainRules
- Recurse to get the $n - 1$-th order derivative
- Raise the order via product rule or division rule

This process can be automated with SymbolicsUtils.jl.

---
## Complexity analysis

Let's assume that $f'(x)$ can be reconstructed from $f$ and $x$ via primitive operations ($+, -, \times, \div$) and finite mutual recursion.

- Each order of reconstruction is $O(n)$
- Each raising is $O(n)$
- Computing the Taylor polynomial is therefore at most $O(n^2)$
- For linear functions, it is in addition $O(n)$

---
Let $t = (t_0,\cdots,t_n)$ be a polynomial of order $n$, and $f=(f_0,\cdots,f_n)$ is the output polynomial of a function.

### Case study: $f = \tan t$

Let $s=(\tan t)'=1+\tan^2t=1+f^2$, compute in the following order:

$$
f_0, s_0\to f_1\to s_1\to f_2\to s_2\to\cdots \to f_n
$$

### Case study: $f = \arctan t$

Let $u=1/(\arctan t)'=1+t^2$, compute in the following order:

$$
u_0\to u_1\to\cdots\to u_{n-1}\to (f_0,\cdots,f_n)
$$

Implementation

Checkout the code at TaylorDiff.jl!

## Type definition and primitives

We define a new data structure `TaylorScalar` as described (it doesn't `<: Number` for technical reasons);

```julia
struct TaylorScalar{T,N}
    values::NTuple{N,T}
end
```

Taylor rules are unrolled by generated functions;
    
```julia
@generated function *(a::TS{T, N}, b::TS{T, N}) where {T, N}
    return quote
        va, vb = value(a), value(b)
        @inbounds TaylorScalar($([:(+($([:(va[$j] * vb[$(i + 1 - j)])
         for j in 1:i]...))) for i in 1:N]...))
    end
end
```
---
## Code generation

```julia
@eval @generated function (op::$F)(t::TaylorScalar{T, N}) where {T, N}
    der_expr = $(QuoteNode(toexpr(term)))
    f = $func
    quote
        $(Expr(:meta, :inline))
        t₁ = TaylorScalar{T, N - 1}(t)
        df = $der_expr
        $$raiser($f(value(t)[1]), df, t)
    end
end
```

---
## API

```julia
@inline function derivative(f, x::T, ::Val{N}) where {T, N}
    t = TaylorScalar{T, N}(x, one(x))
    return extract_derivative(f(t), N)
end

@inline function derivative(f, x::Vector{T}, l::Vector{T},
                            vN::Val{N}) where {T, N}
    t = map(TaylorScalar{T, N}, x, l)
    return extract_derivative(f(t), N)
end
```

---
## Downstream

Essentially some custom `rrule`s to avoid mutating code, e.g.

```julia
function rrule(::typeof(*), A::AbstractMatrix,
               t::AbstractVector{TaylorScalar{T, N}}) where {N, T}
    project_A = ProjectTo(A)
    function gemv_pullback(x̄)
        NoTangent(), @thunk(project_A(contract.(x̄, transpose(t)))), 
        @thunk(transpose(A)*x̄)
    end
    return A * t, gemv_pullback
end
```

Results

Scenario 1: Scaling comparison

We compare the efficiency of TaylorDiff.jl to ForwardDiff.jl (nesting first-order derivative), with scalar function $f(x)=\sin(x)$ up to order 10 and a simple MLP $f(x)=W_2\cdot\exp.(W_1\cdot x+b_1)+b_2$ up to order 7:

TaylorDiff.jl: Taking Higher-order Derivatives Seriously

Introduction

Theory

Implementation

Results

Scenario 1: Scaling comparison

Scenario 2: Taylor expansion

Scenario 3: Physical model and optimizations

Conclusion