通过图重构加速泰勒模式自动微分，提升科学机器学习中偏微分算子计算效率。
Collapsing Taylor Mode Automatic Differentiation
- 重构计算图，通过求和合并导数以压缩梯度传播路径
- 在典型PDE算子上提速，超越嵌套反向传播性能
- 兼容通用线性PDE与随机泰勒模式，编译器可自动处理
通过嵌套反向传播计算偏微分方程（PDE）算子成本高昂，虽广泛使用却严重限制了其在科学机器学习中的应用。近期方法如前向拉普拉斯算子与随机泰勒模式自动微分（AD）提出前向方案应对该问题。本文提出一种针对泰勒模式的优化技术，通过重写计算图‘坍缩’导数，实现对一般线性PDE算子及随机泰勒模式的有效应用。该方法仅需在计算图中向上传播累加和，可由机器学习编译器自动完成，无需暴露复杂性给用户。我们实现了该坍缩过程并在多个主流PDE算子上评估，证实其显著加速泰勒模式，并优于嵌套反向传播。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。