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.
翻译:自动微分(AD)是基于物理信息的机器学习中的关键步骤,用于计算网络输出相对于配置点坐标的高阶导数。本文提出一种新颖轻量级的算法,用于实现基于物理信息算子学习中的自动微分,我们称之为"零坐标偏移"(ZCS)技巧。与传统方法将所有采样坐标作为叶变量不同,ZCS针对每个空间或时间维度仅引入一个标量叶变量,将所需导数从"多根多叶"简化为"单根多叶"结构,从而使反向模式AD可直接应用。这一方法通过避免沿函数(物理参数)维度复制计算图,实现了显著的性能提升。ZCS在现有深度学习库中易于实现;我们通过扩展DeepXDE包完成了具体实现。我们开展了全面的基准测试分析和多个案例研究,训练基于物理信息的DeepONet来求解无数据偏微分方程。结果表明,ZCS将GPU内存消耗和训练耗时持续降低一个数量级,且该缩减因子与函数数量成比例。作为一种底层优化技术,ZCS对数据、物理方程或网络架构均无限制,且不会在任一环节损害训练结果。