Collapsing Taylor Mode Automatic Differentiation

TL;DR

This work tackles the computational burden of evaluating PDE operators with neural networks by introducing Collapsed Taylor Mode AD, a graph-rewrite optimization that leverages linearity to propagate the sum of highest-order Taylor coefficients instead of aggregating results post-hoc. The method generalizes forward schemes like the forward Laplacian and extends to randomized Taylor mode and arbitrary linear differential operators via Faà di Bruno expansions and Griewank interpolation. Empirical results on Laplacian, weighted Laplacian, and biharmonic operators show substantial speedups and memory reductions over standard Taylor mode and outperform nested first-order AD, with potential integration into ML compilers to automate these rewrites. The approach is implemented in PyTorch using torch.fx and demonstrates practical gains for scientific ML tasks such as PINNs and VMC, suggesting broad applicability and a route toward broader adoption of Taylor mode in PDE contexts.

Abstract

Computing partial differential equation (PDE) operators via nested backpropagation is expensive, yet popular, and severely restricts their utility for scientific machine learning. Recent advances, like the forward Laplacian and randomizing Taylor mode automatic differentiation (AD), propose forward schemes to address this. We introduce an optimization technique for Taylor mode that 'collapses' derivatives by rewriting the computational graph, and demonstrate how to apply it to general linear PDE operators, and randomized Taylor mode. The modifications simply require propagating a sum up the computational graph, which could -- or should -- be done by a machine learning compiler, without exposing complexity to users. We implement our collapsing procedure and evaluate it on popular PDE operators, confirming it accelerates Taylor mode and outperforms nested backpropagation.

Figures (13)

• Figure 1: $\blacktriangle$ Vanilla Taylor mode is not enough to beat nested 1st-order AD. Illustrated for computing the Laplacian of a $\mathrm{tanh}$-activated $50 \!\to\! 768 \!\to\! 768 \!\to\! 512 \!\to\! 512 \!\to\! 1$ MLP with JAX (+ jit) on GPU (details in \ref{['sec:jax-benchmark']}). We show how to automatically obtain the specialized forward Laplacian through simple graph transformations that "collapse" vanilla Taylor mode.
• Figure 2: $\blacktriangleright$ Collapsed Taylor mode directly propagates the sum of highest degree coefficients. Visualized for pushing 4 $K$-jets through a ${\mathbb{R}}^5 \!\to\! {\mathbb{R}}^3 \!\to\! {\mathbb{R}}$ function ($K=2$ yields the forward Laplacian).
• Figure 3: Taylor mode propagates Taylor coefficients of a path in input space. This results in the function-transformed path's Taylor coefficients. The Taylor expansion of degree $K$ is called a $K$-jet; hence Taylor mode propagates the input $K$-jet to the output $K$-jet.
• Figure 4: Illustration of \ref{['eq:ttc-general']} for the biharmonic operator, i.e., the 5 values of ${\bm{j}}$ with $\lVert {\bm{j}} \rVert_1 = 4$ and their coefficients $\gamma_{{\bm{i}}, {\bm{j}}}$ to interpolate the desired mixed partials.
• Figure 5: Collapsed Taylor mode accelerates standard Taylor mode and outperforms nested 1st-order AD. Exact computation varies the batch size, stochastic computation fixes a batch size and varies the samples such that $S < D$ (Laplacians), and $2 + 3S < 9/2 D^2 - 3/2D + 4$ (biharmonic operator); we could compute exactly otherwise. Opaque markers are non-differentiable computations.