Automatic differentiation (AD) is a critical step in physics-informed machine learning, required for computing the high-order derivatives of network output w.r.t. coordinates of collocation points. In this paper, we present a novel and lightweight algorithm to conduct AD for physics-informed operator learning, which we call the trick of Zero Coordinate Shift (ZCS). Instead of making all sampled coordinates as leaf variables, ZCS introduces only one scalar-valued leaf variable for each spatial or temporal dimension, simplifying the wanted derivatives from "many-roots-many-leaves" to "one-root-many-leaves" whereby reverse-mode AD becomes directly utilisable. It has led to an outstanding performance leap by avoiding the duplication of the computational graph along the dimension of functions (physical parameters). ZCS is easy to implement with current deep learning libraries; our own implementation is achieved by extending the DeepXDE package. We carry out a comprehensive benchmark analysis and several case studies, training physics-informed DeepONets to solve partial differential equations (PDEs) without data. The results show that ZCS has persistently reduced GPU memory consumption and wall time for training by an order of magnitude, and such reduction factor scales with the number of functions. As a low-level optimisation technique, ZCS imposes no restrictions on data, physics (PDE) or network architecture and does not compromise training results from any aspect.
翻译:自动微分是物理信息机器学习中的关键步骤,用于计算网络输出相对于配置点坐标的高阶导数。本文提出一种新颖且轻量级的算法来实现物理信息算子学习的自动微分,我们称之为零坐标偏移技巧。与将所有采样坐标设为叶变量不同,ZCS针对每个空间或时间维度仅引入一个标量叶变量,将所需导数从"多根多叶"简化为"单根多叶"结构,从而使得反向模式自动微分可直接使用。通过避免沿函数(物理参数)维度重复计算图,该方法实现了显著的性能跃升。ZCS易于在现有深度学习库中实现,我们通过扩展DeepXDE包完成了具体实现。我们进行了全面的基准测试分析和多项案例研究,训练物理信息DeepONet在无数据条件下求解偏微分方程。结果表明,ZCS持续将GPU内存消耗和训练耗时降低一个数量级,且该缩减系数与函数数量呈正比。作为底层优化技术,ZCS不对数据、物理场(PDE)或网络架构施加任何限制,且不会从任何方面影响训练结果。