Nested AD
Nested automatic differentiation means applying automatic differentiation inside another automatic differentiation computation.
Nested AD
Nested automatic differentiation means applying automatic differentiation inside another automatic differentiation computation.
Conceptually, this means differentiating derivative computations.
Examples include:
| Expression | Meaning |
|---|---|
| $\nabla(\nabla f)$ | Hessian |
| $D(Df)$ | second directional derivative |
| $\nabla(Df[v])$ | Hessian-vector product |
| $D(\nabla f)$ | Jacobian of gradient |
| $\nabla(\text{optimization step})$ | meta-gradient |
Nested AD is essential for higher-order derivatives, implicit differentiation, meta-learning, differentiable optimization, and differentiable programming systems.
The mathematics is straightforward. The implementation is subtle.
First-Order AD as a Transformation
An AD transform maps a program into another program.
If
$$ f : X \to Y, $$
then forward mode produces a transformed program:
$$ Df : TX \to TY, $$
where $T$ represents tangent information.
Reverse mode produces a program that computes pullbacks:
$$ f^\ast : T^\ast Y \to T^\ast X. $$
Nested AD applies these transformations repeatedly.
For example:
$$ D(Df) $$
means applying forward-mode transformation twice.
Similarly,
$$ \nabla(\nabla f) $$
means applying reverse mode to a reverse-mode derivative computation.
The resulting program has multiple derivative levels active simultaneously.
Example: Forward-over-Forward
Consider
$$ f(x) = x^3. $$
First forward mode computes:
$$ f(x + \epsilon v) = x^3 + 3x^2v,\epsilon. $$
The tangent coefficient is:
$$ Df(x)[v] = 3x^2v. $$
Apply forward mode again with another perturbation:
$$ x + \epsilon_1 v + \epsilon_2 w + \epsilon_1\epsilon_2 u. $$
Now mixed infinitesimal terms appear:
$$ f(x + \cdots) = x^3 + 3x^2(\cdots) + 6xvw,\epsilon_1\epsilon_2 + \cdots $$
The coefficient of
$$ \epsilon_1\epsilon_2 $$
contains second-order information.
Nested forward mode therefore naturally produces higher derivatives.
Dual Number Nesting
Forward mode is commonly implemented using dual numbers:
$$ a + b\epsilon, \quad \epsilon^2 = 0. $$
Nested forward mode uses nested dual structures:
$$ (a + b\epsilon_1) + (c + d\epsilon_1)\epsilon_2. $$
Expanding gives:
$$ a + b\epsilon_1 + c\epsilon_2 + d\epsilon_1\epsilon_2. $$
Each infinitesimal direction corresponds to a derivative level.
The mixed term
$$ d\epsilon_1\epsilon_2 $$
encodes second-order interaction.
This construction generalizes to arbitrary derivative order, but the algebra grows rapidly.
Reverse-over-Reverse Nesting
Nested reverse mode is more complicated.
A reverse-mode computation creates adjoints and backward passes. Differentiating that computation introduces adjoints of adjoints.
Suppose:
g = grad(f)
H = grad(g)
The second gradient differentiates the reverse pass used by the first gradient.
This means:
- The original primal computation must remain differentiable.
- The backward pass must itself behave like differentiable code.
- Adjoint accumulation logic becomes part of the differentiated computation.
The system must distinguish derivative levels carefully.
Mixed-Mode Nesting
Many practical systems use mixed-mode nesting.
Examples include:
| Nesting | Purpose |
|---|---|
| forward-over-reverse | Hessian-vector products |
| reverse-over-forward | directional derivative gradients |
| reverse-over-reverse | higher-order scalar derivatives |
| forward-over-forward | Taylor coefficients |
| reverse-over-forward-over-reverse | advanced implicit differentiation |
Mixed mode is often preferable because different AD modes have complementary strengths.
For scalar-output functions:
| Mode | Efficient dimension |
|---|---|
| forward mode | small input dimension |
| reverse mode | small output dimension |
Nested systems combine these strengths.
Perturbation Confusion
The most famous nested-AD failure mode is perturbation confusion.
Suppose two forward-mode computations accidentally share the same infinitesimal symbol:
$$ \epsilon. $$
Then derivative information from different levels mixes incorrectly.
For example:
$$ (a + b\epsilon) + (c + d\epsilon) = (a+c) + (b+d)\epsilon. $$
If these infinitesimals were meant to represent different derivative levels, the result is wrong.
The problem becomes severe when derivative computations are passed through higher-order functions or closures.
Tagging Derivative Levels
Correct nested forward mode assigns a unique tag to each perturbation level.
Instead of one universal infinitesimal:
$$ \epsilon, $$
the system uses:
$$ \epsilon_1, \epsilon_2, \epsilon_3, \ldots $$
with independent nilpotent behavior.
Then:
$$ \epsilon_i \epsilon_j \ne 0 \quad \text{for } i \ne j, $$
while
$$ \epsilon_i^2 = 0. $$
Each derivative transform introduces a fresh perturbation identity.
Implementation-wise, systems usually represent this with:
| Technique | Description |
|---|---|
| unique IDs | each AD transform gets fresh perturbation tag |
| lexical scoping | perturbations scoped to derivative level |
| type-level tagging | derivative levels encoded in types |
| runtime tagging | tags stored dynamically |
| staged transforms | derivative levels separated during compilation |
Without tagging, nested forward mode is unreliable.
Cotangent Level Separation
Reverse mode has a similar issue.
Adjoints from different derivative levels must remain separate.
Suppose:
outer gradient
inner gradient
The inner reverse pass should not accidentally consume cotangents from the outer reverse pass.
A correct nested reverse system tracks cotangent levels explicitly.
Conceptually:
| Level | Meaning |
|---|---|
| primal level | original computation |
| tangent/cotangent level 1 | first derivative |
| tangent/cotangent level 2 | second derivative |
| tangent/cotangent level k | kth derivative |
This separation is part of the semantics, not merely debugging metadata.
Closures and Higher-Order Functions
Nested AD becomes harder in languages with closures and higher-order functions.
Example:
def outer(x):
def inner(y):
return x * y
return grad(inner)
The derivative transform captures free variables from the outer scope.
The AD system must decide:
| Question | Requirement |
|---|---|
| which values are primal | original computation |
| which values are tangent | derivative level |
| which values are cotangent | reverse accumulation |
| which closures capture derivative state | nested differentiation |
| which tapes belong to which level | nesting correctness |
Functional languages often model this more cleanly because closures are explicit semantic objects.
Dynamic Computation Graphs
Nested AD interacts strongly with dynamic graph systems.
If the computation graph depends on runtime control flow, the graph structure itself may differ between derivative levels.
For example:
if grad(f)(x) > 0:
...
The derivative computation influences the primal control flow.
A nested AD system must specify:
- Which computations are traced.
- Which branches are differentiated.
- Whether graph structure is static or dynamic.
- Whether nested derivative traces are compositional.
Different frameworks make different choices.
Differentiating Optimizers
Nested AD is central in differentiable optimization.
Suppose gradient descent performs:
$$ x_{t+1} = x_t - \eta \nabla f(x_t). $$
Now suppose we want derivatives with respect to hyperparameters:
$$ \frac{\partial x_T}{\partial \eta}. $$
The optimizer itself becomes part of the differentiated program.
This requires differentiating through:
| Object | Derivative target |
|---|---|
| gradients | first-order derivatives |
| update rules | optimization dynamics |
| momentum accumulators | optimizer state |
| learning rate schedules | hyperparameters |
| inner training loops | meta-learning |
Nested AD enables this.
Implicit Differentiation
Nested AD also appears in implicit differentiation.
Suppose:
$$ g(x, y(x)) = 0. $$
We may avoid differentiating every iteration of a solver by differentiating the fixed-point condition itself.
Still, the resulting derivative computations often involve nested linearizations and reverse passes.
This is especially common in:
| Area | Example |
|---|---|
| meta-learning | differentiating equilibrium states |
| optimization layers | differentiating argmin operators |
| physics simulation | differentiating steady states |
| probabilistic inference | differentiating fixed-point solvers |
Nested AD provides the underlying machinery.
Tape Nesting
In tape-based reverse mode, nested AD introduces nested tapes.
Possible models include:
| Model | Description |
|---|---|
| independent tapes | each derivative level has separate tape |
| hierarchical tapes | tapes reference parent tapes |
| reentrant tapes | backward passes may themselves record operations |
| staged tapes | derivative levels separated during compilation |
Incorrect tape interaction can cause:
| Failure | Meaning |
|---|---|
| tape corruption | nested passes overwrite state |
| missing gradients | tape lifetime ends too early |
| duplicated gradients | nested replay occurs twice |
| memory explosion | all nested levels retained simultaneously |
Robust nested systems need explicit tape ownership semantics.
Compiler Perspective
A compiler-based AD system often handles nesting better than runtime operator overloading systems.
Instead of dynamically stacking derivative objects, the compiler transforms intermediate representations explicitly:
program
→ linearized program
→ transposed program
→ differentiated transposed program
This exposes derivative levels structurally.
Compiler IRs can annotate:
| IR annotation | Purpose |
|---|---|
| primal variable | original value |
| tangent variable | forward derivative |
| cotangent variable | reverse derivative |
| residual | saved intermediate |
| derivative level | nesting separation |
This structure reduces ambiguity.
Complexity Explosion
Higher-order derivatives grow rapidly.
For dimension $n$:
| Derivative order | Tensor size |
|---|---|
| gradient | $n$ |
| Hessian | $n^2$ |
| third derivative tensor | $n^3$ |
| fourth derivative tensor | $n^4$ |
Nested AD can therefore create exponential growth in storage and computation.
Most practical systems avoid materializing full higher-order tensors.
Instead they compute:
| Operation | Scalable form |
|---|---|
| Hessian-vector product | $Hv$ |
| Jacobian-vector product | $Jv$ |
| vector-Hessian-vector | $v^\top Hv$ |
| directional kth derivative | scalar directional expansion |
Operator forms scale better than explicit tensors.
Practical Design Principles
A robust nested AD system should:
| Principle | Reason |
|---|---|
| separate derivative levels | avoid perturbation confusion |
| represent derivatives explicitly | improve correctness |
| isolate tapes per level | prevent state corruption |
| distinguish primal/tangent/cotangent data | preserve semantics |
| expose operator APIs | avoid tensor explosion |
| support compositional transforms | enable higher-order programming |
Nested AD is fundamentally about composing derivative transformations safely and predictably.
Conceptual View
Automatic differentiation is often introduced as a technique for computing gradients.
Nested AD reveals a deeper interpretation.
Differentiation becomes a compositional program transform:
$$ \mathcal{D}(\mathcal{D}(f)), \quad \mathcal{D}(\mathcal{D}(\mathcal{D}(f))), \quad \ldots $$
The challenge is no longer merely computing derivatives. The challenge is preserving semantic structure across multiple interacting derivative levels while controlling memory, complexity, and numerical stability.