关于机器学习:含有PDE约束的双层深度学习优化方法

简介

对于具备简单束缚 - 偏微分方程的问题，传统的办法次要依赖自随同算子，随着深度学习办法的倒退，许多基于神经网络的办法也开始呈现，这些办法的外围是应用神经网络进行推理的减速，以实现几十倍甚至几百倍的提速，这是相较于传统办法的劣势，劣势就是精度不够。深度学习就是这种取舍，通过深度学习来克服传统办法的一些有余，然而其本身也会有诸多问题，可解释性、计算复杂度和推理速度等，从目前来看基于物理机器学习的办法在最优控制问题上具备极强的后劲。

本文介绍的办法是来源于下面两种办法的延长，也能够是看作算子学习实践的拓展，联合了双层优化实践，这算是一个翻新，成果相交于目前的办法还是有长足的提高

Hao, Z., Ying, C., Su, H., Zhu, J., Song, J., & Cheng, Z. (2022). Bi-level Physics-Informed Neural Networks for PDE Constrained Optimization using Broyden’s Hypergradients. arXiv preprint arXiv:2209.07075.

传统的自随同办法能够参考以下的文章：

Xu, S., Li, Y., Huang, X., & Wang, D. (2021). Robust Newton–Krylov Adjoint Solver for the Sensitivity Analysis of Turbomachinery Aerodynamics. AIAA Journal, 59(10), 4014-4030.
Xu, S., Radford, D., Meyer, M., & Müller, J. D. (2015). Stabilisation of discrete steady adjoint solvers. Journal of computational physics, 299, 175-195.
Xu, S., & Timme, S. (2017). Robust and efficient adjoint solver for complex flow conditions. Computers & Fluids, 148, 26-38.

目前解决偏微分方程束缚的最优化问题能够参考一下三篇论文:
这一篇的思路是间接暴力优化，成果较差：

Mowlavi, S., & Nabi, S. (2021). Optimal control of PDEs using physics-informed neural networks. arXiv preprint arXiv:2111.09880.

这两篇的思路就是将问题转化为无约束问题并借助梯度降落办法进行求解：

Hwang, R., Lee, J. Y., Shin, J. Y., & Hwang, H. J. (2022, June). Solving pde-constrained control problems using operator learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 36, No. 4, pp. 4504-4512).
Wang, S., Bhouri, M. A., & Perdikaris, P. (2021). Fast PDE-constrained optimization via self-supervised operator learning. arXiv preprint arXiv:2110.13297.

问题形容

首先咱们思考以下的问题：

$$
\begin{matrix}
\min_{y\in Y_{\mathrm{ad}},u\in U_{\mathrm{ad}}}& \mathcal{J} (y,u)\\
\,\,\mathrm{s}.\mathrm{t}.& \begin{array}{c}
\mathcal{F} (y,u)(x)=0,\forall x\in \Omega\\
\mathcal{B} (y,u)(x)=0,\forall x\in \partial \Omega\\
\end{array}\\
\end{matrix}
$$

这个时候算子学习就开始发挥作用了，咱们别离应用两个神经网络对管制输出 \(u\) 和对应的偏微分方程束缚的解 \(y\)，别离是 \(u_\theta\) 和 \(y_w\)。这里引入了两层优化的思维，也就是 \(u_\theta\) 和 \(y_w\) 是相互影响的，然而一起优化又会造成性能进化。因为损失函数的大小并不是在一个量级，会造成对算子学习或者优化过程至多一个会失败，因而，能够将问题投影为以下的模式：

$$
\begin{array}{ll}
\min _\theta & \mathcal{J}\left(w^\star, \theta\right) \\
\text {s.t.} & w^\star=\arg \min _w \mathcal{E}(w, \theta) .
\end{array}
$$

其中 \(\mathcal{E}=\int_{\Omega}\left|\mathcal{F}\left(y_w, u_\theta\right)(x)\right|^2 \mathrm{~d} x+\int_{\partial \Omega}\left|\mathcal{B}\left(y_w, u_\theta\right)(x)\right|^2 \mathrm{~d} x\)。
这个优化的步骤如下：

1. 外层优化过程：固定 \(\theta\)，用梯度降落办法更新 \(w\)
2. 内层优化过程：固定 \(w\)，用梯度降落办法更新 \(\theta\)

这个过程能够看作是博弈的过程，能够参考双层优化的综述

Liu, R., Gao, J., Zhang, J., Meng, D., & Lin, Z. (2021). Investigating bi-level optimization for learning and vision from a unified perspective: A survey and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence.

算法细节

算法的具体过程如下图所示：

须要指出的是，在算子学习的过程，也即是内层优化 \(w\) 时，梯度降落与参数 \(\theta\) 无关。然而在外层优化 \(\theta\) 时，梯度降落同时依赖两个参数，上面就要介绍为什么是这个模式。
首先思考这个优化过程，上面我用一张图来形容

$$
u_{\theta}\rightarrow y_w\rightarrow \mathcal{E} \\
u_{\theta}\rightarrow y_{w^\star}\rightarrow \mathcal{J} \left(w,\theta \right)
$$

依照这个链式法则就能清晰了，双层优化的内环或者外环过程中，咱们固定了一个参数，只用梯度降落更新另一个参数。对于 \(w\) 的优化，因为链式法则终止于 \(w\) 本身，那么是不必思考 \(\theta\)，然而 \(\theta\) 的优化就会波及到 \(w\)，那么就必须思考到 \(w\) 在链式法则中的作用，因而，上面介绍了对 \(\theta\) 的梯度优化。

具体而言，外层优化指标 \(\mathcal{J}\left(w^\star, \theta\right) \) 内层优化的最优参数 \(w^\star\)

$$
\frac{\mathrm{d} \mathcal{J}}{\mathrm{d} \theta}=\frac{\partial \mathcal{J}}{\partial \theta}+\frac{\partial \mathcal{J}}{\partial w^\star} \frac{\partial w^\star}{\partial \theta}
$$

其中 \(\frac{\partial \mathcal{J}}{\partial \theta}\) 和 \(\frac{\partial \mathcal{J}}{\partial w^\star} \) 都是能够间接利用 Pytorch 或者 Tensorflow 中的 API 间接计算失去，然而后一项 \(\frac{\partial w^\star}{\partial \theta} \) 是无奈间接计算失去的，也是本文的次要奉献。

基于 Cauchy 隐性函数定理能够失去：如果对于一些 使得低层次优化问题有解，也就是 \(\left.\frac{\partial \mathcal{E}}{\partial w}\right|_{\left(w^{\prime}, \theta^{\prime}\right)}=0\) 和 \(\left(\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right)^{-1} \) 是可逆的，那么存在一个函数 \(w^\star=w^\star(\theta)\) 在 \(\left(w^{\prime}, \theta^{\prime}\right) \text {s.t.}\left.\frac{\partial \mathcal{E}}{\partial w}\right|_{\left(w^\star\left(\theta^{\prime}\right), \theta^{\prime}\right)}=0\)，咱们有

$$
\left.\frac{\partial w^\star}{\partial \theta}\right|_{\theta^{\prime}}=-\left.\left[\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right]^{-1} \cdot \frac{\partial^2 \mathcal{E}}{\partial w \partial \theta^{\top}}\right|_{\left(w^\star\left(\theta^{\prime}\right), \theta^{\prime}\right)}
$$

因而，该问题能够被求解为

$$
\frac{\mathrm{d} \mathcal{J}}{\mathrm{d} \theta}=\frac{\partial \mathcal{J}}{\partial \theta}-\left.\frac{\partial \mathcal{J}}{\partial w^\star} \cdot\left(\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right)^{-1} \cdot \frac{\partial^2 \mathcal{E}}{\partial w \partial \theta^{\top}}\right|_{\left(w^\star(\theta), \theta\right)}
$$

然而，对于神经网络的参数来说，计算 Hessian 矩阵的逆值是难以实现的。为了应答这一挑战，咱们能够首先计算 \(z^\star \triangleq \frac{\partial \mathcal{J}}{\partial w^\star} \cdot\left(\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right)^{-1}\)，这也被称为反向量 - 黑森积。求解 \(z^\star \) 能够通过求解上面的线性方程：

$$
g_w(z)=\frac{\partial}{\partial w}\left(\frac{\partial \mathcal{E}}{\partial w^{\top}} \cdot z\right)-\left.\frac{\partial \mathcal{J}}{\partial w}\right|_{w=w^\star}=0
$$

咱们只须要以较低的计算成本计算两个 Jacobian 向量乘积，这不须要创立一个微小的 Hessian 矩阵实例。具体来说，咱们应用低秩的 Broyden 办法来迭代近似解：
首先，近似将 \(\left(\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right)^{-1} \) 示意为：

$$
\left(\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right)^{-1} \approx B_i=-I+\sum_{k=1}^i u_k v_k^{\top}
$$

随后更新参数

$$
\begin{aligned}
z_{i+1} &=z_i-\alpha \cdot B_i g_i(z) \\
u_{i+1} &=\frac{\Delta z_{i+1}-B_i \Delta g_{i+1}}{\left(\Delta z_{i+1}\right)^{\top} B_i \Delta g_{i+1}} \\
v_{i+1} &=B_i \Delta z_{i+1}
\end{aligned}
$$

其中 \(\Delta z_{i+1}=z_{i+1}-z_i, \Delta g_{i+1}=g_{i+1}-g_i\)，\(\alpha\) 是学习率。
因而，逆 Hessian 矩阵能够通过以下形式更新

$$
\left(\frac{\partial^2 \mathcal{E}}{\partial w \partial w^{\top}}\right)^{-1} \approx B_{i+1}=B_i+\frac{\Delta z_{i+1}-B_i \Delta g_{i+1}}{\left(\Delta z_{i+1}\right)^{\top} B_i \Delta g_{i+1}} \Delta z_{i+1}^{\top} B_i
$$

最初，通过 m 次迭代失去的解 \(z_m\) 就能够应用去计算梯度 \(\frac{\mathrm{d} \mathcal{J}}{\mathrm{d} \theta} \)，这样就实现了外层优化过程。

作者随后还给出了一个实践剖析，也就是说内层优化过程是收敛的，也就是说，

$$
\left\| w_t-w \right\| _2\leqslant p_t\left\| w \right\| _2
$$

其中 \(p_t<1\)，\(p_t\rightarrow 0,t\rightarrow \infty \)。

定理：如果一些假如成立且对于 \(z\) 线性方程能够用 Broyden 办法求解，那么存在一个常数 \(M>0 \)，使以下不等式成立：

$$
\left\|\mathrm{d}_2 J_t-\mathrm{d}_2 J\right\|_2 \leqslant\left(L\left(1+\kappa+\frac{\rho M}{\mu^2}\right)+\frac{\tau M}{\mu}\right) p_t\|w\|_2+M \kappa\left(1-\frac{1}{m \kappa}\right)^{\frac{m^2}{2}} .
$$

其中，\(\mathrm{d}_2 J_t \) 和 \(\mathrm{d}_2 J\) 别离是计算出的和实在的对于 \(\theta\) 的超阶梯。

试验后果

这篇文章次要和 hPINN、PINN-LS、PI-DeepONet、Truncated Unrolled Differentiation(TRMD) 和 Neumann Series (Neumann) 这几种办法进行比拟，次要在在泊松方程、热方程、伯格斯方程和纳维亚 - 斯托克斯方程下面进行了试验，具体成果如下：

$$
\scriptsize
\begin{array}{c|cccccc}
\hline \text {Objective}(\mathcal{J}) & \text {Poisson’s 2d CG} & \text {Heat 2d} & \text {Burgers 1d} & \text {NS Shape} & \text {NS 2inlets} & \text {NS backstep} \\
\hline \text {Initial guess} & 0.400 & 0.142 & 0.0974 & 1.78 & 0.0830 & 0.138 \\
\text {hPINN-P} & 0.373 & 0.132 & 0.0941 & – & 0.0732 & 0.0586 \\
\text {hPINN-A} & 0.352 & 0.138 & 0.0895 & – & 0.0867 & 0.0696 \\
\text {PINN-LS} & 0.239 & 0.112 & 0.0702 & – & 0.0634 & 0.0831 \\
\text {PI-DeepONet} & 0.392 & 0.0501 & 0.0710 & \mathbf{1 . 2 7} & 0.0850 & 0.0670 \\
\hline \text {Ours} & \mathbf{0 . 1 6 0} & \mathbf{0 . 0 3 7 9} & \mathbf{0 . 0 6 6 2} & \mathbf{1 . 2 7} & \mathbf{0 . 0 3 6 9} & \mathbf{0 . 0 3 6 5} \\
\text {Reference values} & 0.159 & 0.0378 & 0.0634 & 1.27 & 0.0392 & 0.0382 \\
\hline
\end{array}\\
\begin{array}{c|cccccc}
\hline \text {Objective}(\mathcal{J}) & \text {Poisson’s 2d CG} & \text {Heat 2d} & \text {Burgers 1d} & \text {NS Shape} & \text {NS 2inlets} & \text {NS backstep} \\
\hline \text {Initial guess} & 0.400 & 0.142 & 0.0974 & 1.77 & 0.0830 & 0.138 \\
\text {TRMD} & 0.224 & 0.0735 & 0.0671 & – & 0.0985 & 0.0391 \\
T 1-T 2 & 0.425 & 0.0392 & 0.0865 & 1.68 & 0.115 & 0.0879 \\
\text {Neumann} & 0.225 & 0.0782 & 0.0703 & 1.31 & 0.0962 & 0.0522 \\
\hline \text {Broyden(Ours) } & \mathbf{0 . 1 6 0} & \mathbf{0 . 0 3 7 9} & \mathbf{0 . 0 6 6 2} & \mathbf{1 . 2 7} & \mathbf{0 . 0 3 6 9} & \mathbf{0 . 0 3 6 5} \\
\hline
\end{array}
$$

总结

本文介绍的这篇文章算是偏微分方程束缚的最优化问题的 SOTA 后果，作者引入了双层优化的办法，防止了 loss 相加导致的学习失败的情景，是一篇十分好的文章。将来的钻研方向能够是引入成果更好的优化办法、算子学习实践的优化以及生成管制输出的网络结构设计，这些都属于尚未摸索的畛域，置信这些会为 Physics Informed Machine Learning 带来全新的视角。