Physics-Driven Deep Neural Networks for Solving the Optimal Transport Problem Associated With the Monge-Ampère Equation

Abstract

Monge-Ampère equations (MAEs) are fully nonlinear second-order partial differential equations (PDEs), which are closely related to various fields including optimal transport (OT) theory, geometrical optics and affine geometry. Despite their sig-nificance, MAEs are extremely challenging to solve. Although some classical numerical approaches can solve MAEs, their computational efficiency deteriorates significantly on fine grids, with convergence often heavily dependent on the quality of initial estimate. Research on deep learning methods for solving MAEs is still in its early stages, which predominantly addresses simple formulations with basic Dirichlet boundary conditions. Here, we propose a deep learning method based on physics- driven deep neural networks, enabling the solution of both simple and generalised MAEs with transport boundary condi-tions. In this method, we deal with two first-order sub-equations separated from MAE instead of solving the single MAE directly, which facilitates the imposition of transport boundary conditions and simplifies the training of neural networks. Moreover, we constrain the convexity of solution using the Lagrange multiplier method and maintain the optimisation process differentiable with bilinear interpolation. We provide three progressively complex examples ranging from a simple MAE with an analytical solution to a highly nonlinear variant arising in phase retrieval to validate the effectiveness of our method. For comparison, we benchmark against state-of-the-art deep learning approaches that have been systematically adapted to accommodate the specific requirements of each example.

Keywords

AI for Science / artificial inteligence / deep learning / neural nets / unsupervised learning

Acknowledgements

The work is funded by the CAAI-Huawei MindSpore Open Fund (CAAIXSJLJJ-2022-010A).

Conflicts of Interest

The authors declare no confiicts of interest.

3 Data Availability Statement

Data underlying the results presented in this paper are not publicly available at this time but may be obtained from the authors upon reasonable request.

References

[1] G. Monge, “Mémoire Sur la Théorie Des Déblais Et Des Remblais,” Histoire de l’Acadmie Royale des Sciences (1781): 666-704.

[2] X. Tian, et al., “Optimal Transport-Based Graph Matching for 3D Retinal OCT Image Registration,” in 2022 IEEE International Conference on Image Processing (ICIP) 2022), 2791-2795.

[3] Z. Su, Y. Wang, R. Shi, et al., “Optimal Mass Transport for Shape Matching and Comparison,” IEEE Transactions on Pattern Analysis and Machine Intelligence 37, no. 11 (2015): 2246-2259, https://doi.org/10.1109/tpami.2015.2408346.

[4] R. Wu, L. Xu, P. Liu, et al., “Freeform Illumination Design: A Nonlinear Boundary Problem for the Elliptic Monge-Ampère Equation,” Optics Letters 38, no. 2 (2013): 229-231, https://doi.org/10.1364/ol.38.000229.

[5] N. Lei, Y. Guo, K. Su, S. Liu, Z. Luo, et al., “A Geometric Under-standing of Deep Learning,” Engineering 6, no. 3 (2020): 361-374, https://doi.org/10.1016/j.eng.2019.09.010.

[6] J.-D. Benamou B. D. Froese and A. M. Oberman, “Two Numerical Methods for the Elliptic Monge-Ampère Equation,” ESAIM: Mathe-matical Modelling and Numerical Analysis 44, no. 4 (2010): 737-758, https://doi.org/10.1051/m2an/2010017.

[7] A. M. Oberman, “Wide Stencil Finite Difference Schemes for the Elliptic Monge-Ampère Equation and Functions of the Eigenvalues of the Hessian,” Discrete and Continuous Dynamical Systems—B 10, no. 1 (2008): 221-238, https://doi.org/10.3934/dcdsb.2008.10.221.

[8] A. Caboussat, R. Glowinski, and D. Gourzoulidis, “A Least-Squares/ Relaxation Method for the Numerical Solution of the Three- Dimensional Elliptic Monge-Ampère Equation,” Journal of Scientific Computing 77, no. 1 (2018): 53-78, https://doi.org/10.1007/s10915-018-0698-6.

[9] J. Liu, B. D. Froese, A. M. Oberman, and M. Xiao, “A Multigrid Scheme for 3D Monge-Ampère Equations,” International Journal of Computer Mathematics 94, no. 9 (2017): 1850-1866, https://doi.org/10.1080/00207160.2016.1247443.

[10] G. Loeper and F. Rapetti, “Numerical Solution of the Monge-Ampère Equation by a Newton’s Algorithm,” Comptes Rendus Mathe-matique 340, no. 4 (2005): 319-324, https://doi.org/10.1016/j.crma.2004.12.018.

[11] Z. Feng, D. Cheng, and Y. Wang, “Transferring Freeform Lens Design Into Phase Retrieval Through Intermediate Irradiance Trans-port,” Optics Letters 44, no. 22 (2019): 5501-5504, https://doi.org/10.1364/ol.44.005501.

[12] M. Raissi, P. Perdikaris, and G. E. Karniadakis, “Physics-Informed Neural Networks: A Deep Learning Framework for Solving Forward and Inverse Problems Involving Nonlinear Partial Differential Equa-tions,” Journal of Computational Physics 378 (2019): 686-707, https://doi.org/10.1016/j.jcp.2018.10.045.

[13] H. Eivazi, M. Tahani, P. Schlatter, and R. Vinuesa, “Physics- Informed Neural Networks for Solving Reynolds-Averaged Navier-Stokes Equations,” Physics of Fluids 34, no. 7 (2022): 075117, https://doi.org/10.1063/5.0095270.

[14] J. Lim and D. Psaltis, “Maxwellnet: Physics-Driven Deep Neural Network Training Based on Maxwell’s Equations,” APL Photonics 7, no. 1 (2022): 011301, https://doi.org/10.1063/5.0071616.

[15] S. Savović, M. Ivanović and R. Min, “A Comparative Study of the Explicit Finite Difference Method and Physics-Informed Neural Net-works for Solving the Burgers’Equation,” Axioms 12, no. 10 (2023): 982, https://doi.org/10.3390/axioms12100982.

[16] G. Wu, F. Wang, and L. Qiu, “Physics-Informed Neural Network for Solving Hausdorff Derivative Poisson Equations,” Fractals 31, no. 6 (2023): 2340103, https://doi.org/10.1142/s0218348x23401035.

[17] B. Zhang, G. Wu, Y. Gu, X. Wang, and F. Wang, “Multi-Domain Physics-Informed Neural Network for Solving Forward and Inverse Problems of Steady-State Heat Conduction in Multilayer Media,” Physics of Fluids 34, no. 11 (2022): 116116, https://doi.org/10.1063/5.0116038.

[18] B. Zhang, F. Wang, and L. Qiu, “Multi-Domain Physics-Informed Neural Networks for Solving Transient Heat Conduction Problems in Multilayer Materials,” Journal of Applied Physics 133, no. 24 (2023): 245103, https://doi.org/10.1063/5.0153705.

[19] K. Nyström and M. Vestberg, “Solving the Dirichlet Problem for the Monge-Ampère Equation Using Neural Networks,” Journal of Computational Mathematics and Data Science 8 (2023): 100080, https://doi.org/10.1016/j.jcmds.2023.100080.

[20] A. Singh, M. Bauer, and S. Joshi, “Physics Informed Convex Arti-ficial Neural Networks (Picanns) for Optimal Transport Based Density Estimation,” arXiv preprint arXiv:2104. 01194 (2021).

[21] B. Amos, L. Xu, and J. Z. Kolter, “Input Convex Neural Networks ” in Proceedings of Machine Learning Research (PMLR), (2017), 146-155.

[22] C. Bösel and H. Gross, “Single Freeform Surface Design for Pre-scribed Input Wavefront and Target Irradiance,” Journal of the Optical Society of America A 34, no. 9 (2017): 1490-1499, https://doi.org/10.1364/josaa.34.001490.

[23] K. He, X. Zhang, S. Ren, et al., “Deep Residual Learning for Image Recognition,” in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), (2016), 770-778.

[24] P. Ramachandran, B. Zoph, and Q. V. Le, “Searching for Activation Functions ” arXiv preprint arXiv:1710.05941 (2017).

[25] S. Joshi, G. Ramakrishnan, and S. Chandra, “Using Sequential Unconstrained Minimization Techniques to Simplify SVM Solvers,” Neurocomputing 77, no. 1 (2012): 253-260, https://doi.org/10.1016/j.neucom.2011.07.010.

[26] B. D. Froese, Numerical Methods for the Elliptic Monge-Ampère Equation and Optimal Transport (Simon Fraser University, 2012).

[27] Y. Brenier, “Polar Factorization and Monotone Rearrangement of vector-valued Functions,” Communications on Pure and Applied Mathe-matics 44, no. 4 (1991): 375-417, https://doi.org/10.1002/cpa.3160440402.

[28] A. Makkuva, A. Taghvaei, S. Oh, “Optimal Transport Mapping via Input Convex Neural Networks,” in Proceedings of Machine Learning Research (PMLR), (2020), 6672-6681.

[29] L. Lyu, Z. Zhang, M. Chen, and J. Chen, “MIM: A Deep Mixed Residual Method for Solving High-Order Partial Differential Equations,” Journal of Computational Physics 452 (2022): 110930, https://doi.org/10.1016/j.jcp.2021.110930.

[30] A. G. Baydin, B. A. Pearlmutter, “Automatic Differentiation in Machine Learning: A Survey,” Journal of Machine Learning Research 18 (2018): 1-43, https://doi.org/10.48550/arXiv.1502.05767.

[31] J. Berg and K. Nyström, “A Unified Deep Artificial Neural Network Approach to Partial Differential Equations in Complex Geometries,” Neurocomputing 317 (2018): 28-41, https://doi.org/10.1016/j.neucom.2018.06.056.

[32] Kingma D. P., Ba J., “Adam: A Method for Stochastic Optimiza-tion,” arXiv preprint arXiv:1412.6980. (2014).

[33] L. Chen, Deep Learning and Practice with Mindspore (Springer Na-ture, 2021).

[34] K. He, X. Zhang, S. Ren, J. Sun, “Delving Deep INTO Rectifiers: Surpassing Human-Level Performance on Imagenet Classification,” in Proceedings of the IEEE International Conference on Computer Vision (ICCV, 2015), 1026-1034.

[35] A. D. Jagtap, K. Kawaguchi, and G. E. Karniadakis, “Adaptive Activation Functions Accelerate Convergence in Deep and Physics- Informed Neural Networks,” Journal of Computational Physics 404 (2020): 109136, https://doi.org/10.1016/j.jcp.2019.109136.

[36] R. Timofte, E. Agustsson, L. Van Gool, et al., “Ntire 2017 Challenge on Single Image Super-Resolution: Methods and Results,” in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), (2017), 114-125.

[37] E. Agustsson and R. Timofte,“Ntire 2017 Challenge on Single Image Super-Resolution: Dataset and Study,” in Proceedings of the IEEE Con-ference on Computer Vision and Pattern Recognition (CVPR), (2017), 126-135.

Funding

CAAI-Huawei MindSpore Open Fund(CAAIXSJLJJ-2022-010A)