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

Authors: Mingquan Feng, Weixin Liao, Yixin Huang, Yifan Fu, Qifu Zheng, Junchi Yan

ICML 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental Empirical results on a variety of high-dim PDEs show that HEAP outperforms the state-of-the-art neural operator model. All experiments were conducted on a machine with 1TB memory, 144 cores Intel Xeon Platinum 8352V CPU, and 8 GPUs (NV Ge Force RTX 4090). We benchmark HEAP on four prototype equations, each tested in both an unconstrained form (suffix -UC) and a constrained form (suffix -C). Tables 2, 3, 4, 5, and 6 present experimental results including PDE residual, physical constraint violation, time and space usage, and ablation studies.
Researcher Affiliation Academia 1Sch. of Computer Science & Sch. of Artificial Intelligence, Shanghai Jiao Tong University, China.
Pseudocode Yes Algorithm 1 A-PDHG embodiment for solving QP in Eq. 3. Algorithm 2 Hyper-Extended A-PDHG (i.e. HEAP).
Open Source Code Yes Code available at GitHub.
Open Datasets No The paper defines the PDEs and initial conditions for generating data, e.g., "The training set consists of 150K randomly generated θ from normal distribution, and the test set is 250 randomly generated θ from another normal distribution with slightly different mean and variance." This describes how to generate the data rather than providing concrete access information to a pre-existing publicly available dataset.
Dataset Splits Yes The training set consists of 150K randomly generated θ from normal distribution, and the test set is 250 randomly generated θ from another normal distribution with slightly different mean and variance. Unless otherwise noted, the initial ansatz parameters θ0 are drawn i.i.d. from a Gaussian distribution with mean 0 and standard deviation 0.5 for training, and from a wider N(0, 1) for test; this encourages robustness to distribution shift while keeping the training residual numerically stable.
Hardware Specification Yes All experiments were conducted on a machine with 1TB memory, 144 cores Intel Xeon Platinum 8352V CPU, and 8 GPUs (NV Ge Force RTX 4090).
Software Dependencies No The paper mentions "Adam optimizer" and "Runge-Kutta" for integration, but does not provide specific version numbers for any software libraries or dependencies used.
Experiment Setup Yes The extended dimension of HEAP is 20; the iteration is 3, and the hypernetwork Nw is a 5-layer Res Net with 1000 hidden units. The HEAP is trained with Adam optimizer, default batch size 64, and 100 epochs. E.g., Table 1 summarises the settings used for CSO on our Burgers type equation experiments; HEAP employs the same configuration. The ODE is integrated with the Runge-Kutta (Butcher, 1996) with a time step 1/10 of the time horizon. To balance their magnitudes during optimisation we set λpde = 10 2, λpc = 1, and λnc = 10, respectively, for all experiments.