HEAP: Hyper Extended A-PDHG Operator for Constrained High-dim PDEs

Neural operators have emerged as a promising approach for solving high-dimensional partial differential equations (PDEs). However, existing neural operators often have difficulty in dealing with constrained PDEs, where the solution must satisfy additional equality or inequality constraints beyond the governing equations. To close this gap, we propose a novel neural operator, Hyper Extended Adaptive PDHG (HEAP) for constrained high-dim PDEs, where the learned operator evolves in the parameter space of PDEs. We first show that the evolution operator learning can be formulated as a quadratic programming (QP) problem, then unroll the adaptive primal-dual hybrid gradient (APDHG) algorithm as the QP-solver into the neural operator architecture. This allows us to improve efficiency while retaining theoretical guarantees of the constrained optimization. Empirical results on a variety of high-dim PDEs show that HEAP outperforms the state-of-the-art neural operator model.

Partial differential equations (PDEs) describe many complex processes in science, engineering, and finance. Solving high-dimensional PDEs—especially with practical constraints like temperature limits or financial rules—is challenging and often computationally intensive. We introduce HEAP, an innovative machine learning method that efficiently solves these constrained PDEs by translating the problem into an optimization task. HEAP quickly finds accurate solutions while strictly meeting necessary constraints, outperforming current methods in accuracy and efficiency. This advancement enables faster and more reliable solutions for practical problems, such as calculating heat distributions and pricing financial derivatives, significantly broadening the applications of machine learning in crucial fields.