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

We appreciate your insightful comments and suggestions. Here are our responses to your questions:

Q1: It would strengthen the effectiveness of HEAP if some other constraints were demonstrated

A1: The additional experiment results are reported in the rebuttal supplementary material (RSM), available at https://anonymous.4open.science/r/HEAP/RSM.pdf. We provide four additional experiments, two of which are Black-Scholes with no-arbitrage constraints (Eq. 1-5 of RSM), and the other two are burgers and reaction-diffusion with rate-of-change constraints. The results are shown in Tables 1, 2, 6, 7 of RSM.

---

Q2: Does HEAP have some boundaries for applications? Are there any constraints that HEAP cannot handle?

A2: HEAP is based on a PDHG algorithm variant, designed for convex constraints. Therefore, HEAP can handle any convex constraints in principle. For more general non-convex constraints, the convergence of HEAP is not guaranteed.

We will add this discussion and fix the typos in the final version once the paper update is allowed.

Thank you for the insightful question. The additional experiment results are reported in the rebuttal supplementary material (RSM), available at https://anonymous.4open.science/r/HEAP/RSM.pdf. Regarding your question, we provide the following explanations:

Q1: What does constraint mean here, if PDE has a fixed solution?

A1: The purpose of constraints, both equality and inequality, is to regularize the solution space and select the meaningful PDE solutions when multiple solutions exist.

Firstly, in a general sense, the IC/BCs are also constraints of PDE, albeit equality constraints.

Secondly, in practice, proper IC/BCs are sometimes not easy to obtain, while inequality is usually more accessible. For example, under which IC/BCs the Navier-Stokes equation has a unique solution is still an open problem (the Millennium Prize Problems) [1]. In this case, the inequality constraints are added to ensure a physically reasonable solution, e.g., the upper-bounded total energy (Eq. 7 in [1]). Constraints do affect the solution of PDEs, as we show in Fig. 2 of RSM, where the solution of the Black-Scholes (BS) equation with different financial constraints is compared. The BS equation is a fundamental PDE in quantitative finance (Nobel prize 1997) [2], which does not have a unique solution in practice.

Thirdly, the constrained PDE problems have been formulated and studied in previous AI4PDE works [3, 4]. However, the existing methods are not scalable to high-dimensional problems, which is the focus of our work.

---

Q2: Discuss the Evolutional deep neural network (EDNN) paper

A2: The EDNN also formulates the evolution of surrogate model parameters as an optimization problem, but it 1) solves the problem by a numerical solver instead of a neural network, 2) considers only equality constraints, and 3) is not directly scalable to high-dimensional problems. We will add a comparison with EDNN in the final version.

---

Q3: Is there any more complex problem? Are all the testing cases with exact or trusted numerical solutions?

A3: We provide four additional experiments, two of which are Black-Scholes with no-arbitrage constraints (Eq. 1-5 of RSM), and the other two are burgers and reaction-diffusion with rate-of-change constraints. The results are shown in Tab. 1,2,6,7 of RSM. All the additional testing cases are without exact solutions or trusted numerical solutions.

We have added all the above discussions and RSM to the draft, but we are not able to provide the updated draft here due to the rebuttal rules this year. We hope this response addresses your concerns, and please let us know if you have any further questions.

Reference:

[1] Fefferman, Charles L. "Existence and smoothness of the Navier-Stokes equation." The millennium prize problems 57.67 (2006): 22.

[2] Hull, John C., and Sankarshan Basu. Options, futures, and other derivatives. Pearson Education India, 2016.

[3] Hoshisashi, Kentaro, Carolyn E. Phelan, and Paolo Barucca. "Physics-Informed Neural Networks for Derivative-Constrained PDEs." ICML 2024 AI for Science Workshop.

[4] Moro, Viggo, and Luiz FO Chamon. "Solving Differential Equations with Constrained Learning." The Thirteenth International Conference on Learning Representations.

Thank you for the insightful questions and suggestions. The additional experiment results are reported in the rebuttal supplementary material (RSM), available at https://anonymous.4open.science/r/HEAP/RSM.pdf. Here we provide brief responses to your questions point by point:

Q1: Computational efficiency: memory usage and training/inference speed for d<=20 and higher dimensions scalability.

A1: We provide results for d=5, 10, 15, 20, 50, 100 in the Tab. 4 of RSM. The HEAP costs 1.6x GPU memory and about 2x training/inference time compared to the baseline CSO. The time and memory usage are nearly linearly increasing with the dimension, thus HEAP is scalable to higher dimensions in terms of computation.

---

Q2: Experimental design: parameters in the Gaussian distribution, hyperparameters for CSO, penalty coefficient in CSO

A2: The parameters in the Gaussian distribution are mean=0, std=$\sqrt{0.5}$ for training and mean=0 std=1 for testing. The penalty coefficients are PDE residual loss (with weight $w_1$), the constraint violation loss ($w_2$), and the numerical constraint loss ($w_3$). To balance their magnitudes during training, the weights are set as $w_1 = 10^{-2}$, $w_2 = 1.0$, and $w_3 = 10$. The remaining hyperparameters are displayed in Tab. 3 of RSM.

---

Q3: Hyperparameter Ablation: iteration count K, extended dimension in the hypernetwork.

A3: As shown in Tab. 5 of RSM, both iteration count K = 1,2,3,4 and extended dimension = 5,10,20 in the hypernetwork have a significant impact on the final PDE accuracy and constraint satisfaction. However, K=3 is a relatively robust choice across different PDE families.

---

Q4: Handling more general, non-quadratic, or nonlinear constraints

A4: The effectiveness of HEAP largely depends on its underlying APDHG algorithm and the structure of constraints. For convex but non-linear constraints, modifications for constraint projections may still preserve convergence guarantees, thus HEAP can be easily extended. For non-convex constraints, however, theoretical convergence becomes challenging, though heuristic adaptations like penalty methods or relaxed Lagrangian formulations might work in practice.

---

Q5: Convergences of learnable step sizes (tau, sigma, etc.)

A5: As shown in Fig. 3 of RSM, the learnable step sizes converge to stable values during training after 60 batches.

---

Q6: Rationale for using ResNet

A6: We follow the backbone choice of previous work, CSO, where ResNet is chosen for its simplicity and effectiveness in training deep networks. The skip connections in ResNet help to alleviate the vanishing gradient problem. Other architectures like Transformer or LSTM are also possible for HEAP, but we leave them for future work.

---

Q7: Add related work STDE.

A7: STDE is a state-of-the-art method for efficiently estimating extra-high-dimensional gradients, which can be integrated into HEAP as a gradient estimator. We will add a discussion and cite to STDE in the final version.

---

Q8: Validity of the approximation by truncating the Taylor expansion of the constraints (Eqn. 4).

A8: The original constraint is about the function (infinite dimension), which must be approximated by some finite-dimensional numerical scheme for computation. The truncating Taylor expansion is actually an implicit Euler method, in the sense that the constraint is enforced at the next time step. The implicit Euler method is a widely used stable scheme for PDEs, and the error can be controlled by the square of the step size.

Thanks for your suggestions. The additional experiment results and visualizations are reported in the rebuttal supplementary material (RSM), available at https://anonymous.4open.science/r/HEAP/RSM.pdf, and will be added to the draft once allowed. Here we provide responses to your questions point by point:

Q1: A more detailed explanation of the experiment problems and possible visualization.

A1: Here, we briefly explain two of the existing experiment problems and a newly added problem.

1. Heat equation: The heat equation serves as a fundamental prototype PDEs in physics and mathematics. As a paradigmatic parabolic PDE, it governs ubiquitous physical phenomena ranging from thermal diffusion in materials to probability distribution evolution in stochastic processes. However, solving high-dim instances presents formidable challenges: both traditional mesh-based discretization techniques and neural network methods suffer from the "curse of dimensionality", with storage and computational costs growing exponentially as dimension increases.

2. Heat equation with constraints: The incorporation of monotonically decreasing temperature fields introduces a physically grounded regularization to the classical heat equation, elevating both its modeling fidelity and computational complexity. This constraint encodes the thermodynamic irreversibility of cooling processes in materials with latent heat barriers. Traditional numerical time-stepping schemes may violate the monotonicity condition unless rigorously coupled with projection operators or barrier methods.

3. Black-Sholes (newly added, with visualization): The Black-Scholes equation (Nobel prize 1997) is the cornerstone of modern quantitative finance, providing a rigorous framework for pricing financial derivatives under idealized market conditions. In practice, the no-arbitrage principle ensures the absence of trivial risk-free profit opportunities in the financial market. The no-arbitrage constraint arises naturally when modeling structured products with embedded downside protections or regulatory circuit breakers. The visualizations are Fig. 1 and 2 of RSM.

---

Q2: A clearer explanation of how QP formulation and HEAP algorithm overcome limitations in existing approaches.

A2:

1. QP formulation: the existing approaches of high-dim PDEs formulate the parameter evolution either as a least square problem or a black-box optimization, without explicitly considering the potential constraints. The QP formulation allows the incorporation of linear constraints into the parameter evolution, extending the applicability of NN-based solvers to a broader range of high-dim PDEs.

2. HEAP algorithm: the existing methods for solving large-scale QP problem are either iterative solvers like APDHG or vanilla neural operators. The former requires a large number (100+) of iterations to converge, while the latter suffers from the generalization issues. HEAP combines the best of both: it unrolls the APDHG algorithm into a neural network, which achieves better accuracy than vanilla neural operator due to algorithm priors and solves within 3 iterations, much fewer iterations than the APDHG due to learned knowledge from data.

---

Q3: Intuitive explanations for concepts.

A3: Intuitions of key components:

1. Evolving operator: the solution of a high-dim PDE is surrogated by a subnetwork, whose parameters evolve over time to approximate PDE. When constraints are added, the parameter evolution is formulated as a QP problem, where the objective is to minimize the discrepancy between the left and right-hand side of PDE, and the linear constraints represent the given constraints.

2. APDHG algorithm: the APDHG algorithm is a primal-dual iteration algorithm for QP, where the primal variable x is the solution variable and the dual variable y is the Lagrange multiplier or the penalty of constraints. The primal step updates x by a gradient descent of the objective function and penalty term, while the dual step updates y by a gradient ascent of the penalty term. Both updates are projected to the feasible set of the constraints.

3. HEAP network: HEAP unrolls a fixed number of iterations of the APDHG algorithm into a neural network, where the primal and dual variables are extended from a vector to a matrix, i.e. parallelized over a new dimension. The extension is parameterized by the output of the hypernetwork, which takes the information of the problem state as input. The HEAP network can be trained end-to-end by backpropagation.

4. Theorems: Theorem 3.1 shows HEAP network is truly an extended version of APDHG algorithm, since the HEAP network can be reduced to the APDHG algorithm when the parameters are set to ones. Theorem 3.2 shows that HEAP convergence of solving QP problems in terms of the required number of parameters. The proof naturally inherits from the convergence rate of the APDHG algorithm.