提出高效计算高阶微分算子的新方法,显著提升大规模物理神经网络求解速度。
Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators
- 通过构造输入切线实现多变量函数任意阶导数的高效收缩
- 在单张A100 GPU上8分钟完成百万维偏微分方程求解
- 适合需要高阶微分算子的大规模科学计算任务
使用包含高维、高阶微分算子的损失函数优化神经网络时,传统反向传播因导数张量规模 $\mathcal{O}(d^{k})$ 及计算图复杂度 $\mathcal{O}(2^{k-1}L)$ 而效率低下,其中 $d$ 为定义域维度,$L$ 为前向计算图中操作数,$k$ 为导数阶数。以往工作通过随机化对 $d$ 的多项式增长进行缓解,而 $k$ 的指数增长在单变量情形下通过高阶自动微分(AD)解决。本文提出一种新方法:通过合理构造输入切线,将多变量函数的任意阶导数张量收缩转化为单变量高阶AD问题,可高效随机化任意微分算子。应用于物理信息神经网络(PINNs)时,相比一阶随机化AD,速度提升超1000倍,内存减少超30倍,并可在单张NVIDIA A100 GPU上8分钟内求解100万维偏微分方程,使高阶微分算子在大规模问题中的应用成为可能。
原文摘要 · Abstract (English)
Optimizing neural networks with loss that contain high-dimensional and high-order differential operators is expensive to evaluate with back-propagation due to $\mathcal{O}(d^{k})$ scaling of the derivative tensor size and the $\mathcal{O}(2^{k-1}L)$ scaling in the computation graph, where $d$ is the dimension of the domain, $L$ is the number of ops in the forward computation graph, and $k$ is the derivative order. In previous works, the polynomial scaling in $d$ was addressed by amortizing the computation over the optimization process via randomization. Separately, the exponential scaling in $k$ for univariate functions ($d=1$) was addressed with high-order auto-differentiation (AD). In this work, we show how to efficiently perform arbitrary contraction of the derivative tensor of arbitrary order for multivariate functions, by properly constructing the input tangents to univariate high-order AD, which can be used to efficiently randomize any differential operator. When applied to Physics-Informed Neural Networks (PINNs), our method provides >1000$\times$ speed-up and >30$\times$ memory reduction over randomization with first-order AD, and we can now solve \emph{1-million-dimensional PDEs in 8 minutes on a single NVIDIA A100 GPU}. This work opens the possibility of using high-order differential operators in large-scale problems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。