Collapsing Taylor Mode Automatic Differentiation

TL;DR

This paper introduces a graph rewriting optimization for Taylor mode automatic differentiation that accelerates PDE operator computations, outperforming traditional nested backpropagation methods and enabling more efficient scientific machine learning applications.

Contribution

The paper proposes a novel graph rewriting technique to 'collapse' derivatives in Taylor mode AD, improving computational efficiency for PDE operators and integrating seamlessly with machine learning compilers.

Findings

Accelerates Taylor mode AD for PDE operators

Outperforms nested backpropagation in speed

Applicable to general linear PDE operators

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.