Collapsing Taylor Mode Automatic Differentiation

Felix Dangel, Tim Siebert, Marius Zeinhofer, Andrea Walther

Keywords: Automatic differentiation, Taylor mode, Differential operators, Higher-order derivatives, PINNs, VMC, Laplacian

TL;DR: We accelerate Taylor mode for practically relevant differential operators by collapsing Taylor coefficients; this can be done automatically with compute graph simplifications

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.

Primary Area: Machine learning for sciences (e.g. climate, health, life sciences, physics, social sciences)

Submission Number: 8671