[论文解读] Improved architectures and training algorithms for deep operator networks

Sifan Wang, Hanwen Wang+1 more|arXiv (Cornell University)|Oct 4, 2021

Model Reduction and Neural Networks被引用 8

一句话总结

本文提出改进的架构与训练算法，以提升深度算子网络（DeepONets）在学习参数化PDE解算子方面的性能。通过神经正切核（NTK）理论分析训练动态，作者识别出反向传播中梯度的幅值偏差，并引入自适应损失加权与一种新型架构，缓解梯度消失问题，使预测精度提升10–50倍，尤其在无成对数据的自监督设置下表现显著。

ABSTRACT

Operator learning techniques have recently emerged as a powerful tool for learning maps between infinite-dimensional Banach spaces. Trained under appropriate constraints, they can also be effective in learning the solution operator of partial differential equations (PDEs) in an entirely self-supervised manner. In this work we analyze the training dynamics of deep operator networks (DeepONets) through the lens of Neural Tangent Kernel (NTK) theory, and reveal a bias that favors the approximation of functions with larger magnitudes. To correct this bias we propose to adaptively re-weight the importance of each training example, and demonstrate how this procedure can effectively balance the magnitude of back-propagated gradients during training via gradient descent. We also propose a novel network architecture that is more resilient to vanishing gradient pathologies. Taken together, our developments provide new insights into the training of DeepONets and consistently improve their predictive accuracy by a factor of 10-50x, demonstrated in the challenging setting of learning PDE solution operators in the absence of paired input-output observations. All code and data accompanying this manuscript are publicly available at \url{https://github.com/PredictiveIntelligenceLab/ImprovedDeepONets.}

研究动机与目标

• 为解决DeepONet训练中的幅值偏差问题，即模型因梯度不平衡而更倾向于近似高幅值函数。
• 开发一种无需成对输入-输出数据的自监督训练框架，用于学习PDE解算子。
• 基于NTK分析，提出一种理论基础扎实的自适应损失加权方案，以平衡训练样本间梯度幅值。
• 设计一种新型DeepONet架构，以抵抗梯度消失病态并改善信号在各层间的传播。
• 在多种PDE基准测试中展示显著的预测精度提升，尤其在具有挑战性的自监督设置下。

提出的方法

• 作者采用神经正切核（NTK）理论分析DeepONets的训练动态，并识别出梯度更新中幅值偏差的根本原因。
• 基于Hessian矩阵最大对角元与对角元之比，推导出自适应损失加权方案，定义为 $ \lambda = \sqrt{\frac{\|\bm{H}\|_{\infty}}{\text{diag}(\bm{H})}} $，以在反向传播中平衡梯度幅值。
• 提出一种新型DeepONet架构，通过改进残差连接与归一化策略，提升信号传播能力并减少梯度消失效应。
• 训练算法将自适应加权集成至随机梯度下降中，实现梯度平衡与收敛速率优化之间的连续插值。
• 在无标签输入-输出对、仅依赖PDE约束的物理信息自监督设置下评估该方法。
• 在Burgers'方程、Allen-Cahn方程和Stokes方程等基准PDE上进行实证验证，实现对未见初始/边界条件与几何结构的完整泛化。

实验结果

研究问题

• RQ1 为何DeepONets在训练过程中会偏向近似幅值较大的函数？
• RQ2 能否利用神经正切核（NTK）理论推导出一种原则性、自适应的损失加权方案，以平衡训练样本间的梯度幅值？
• RQ3 如何通过改进的DeepONet架构提升信号传播能力并减少梯度消失问题？
• RQ4 自适应加权与架构改进在多大程度上可提升自监督PDE解算子学习中的预测精度？
• RQ5 在缺乏成对训练数据的情况下，该方法能否在复杂、高维PDE上实现显著性能提升，尤其是对高维PDE？

主要发现

• 基于NTK分析推导出的自适应损失加权方案，能有效平衡反向传播梯度的幅值，显著降低模型对高幅值函数的偏向。
• 新型DeepONet架构在所有基准测试中均持续优于传统DeepONets，展现出更强的梯度消失病态抵抗能力。
• 在Burgers'方程的自监督学习中，测试误差从20.49%降至4.19%，平均提升5倍。
• 对于Allen-Cahn方程，误差从30.87%降至10.86%，实现2.8倍提升，且无需任何标注数据。
• 在具有挑战性的Stokes方程基准中，传统物理信息DeepONet完全失效（压力误差达99.63%），而改进模型仅达2.41%误差，提升40倍。
• 总体而言，该方法在所有测试案例中实现10–50倍的预测精度提升，且在泛化能力与鲁棒性方面均有稳定增益。