Parameterized Representations via Implicit Stochastic Modulation for High-Dimensional and High-Order Neural PDE Solvers
About
Solving high-dimensional and high-order PDEs is challenged by the coupled growth of spatial dimensionality and derivative order. Recent stochastic derivative estimators reduce this cost by replacing full derivative tensors with randomized dimension or Taylor estimators, but they are mostly designed for fixed physical parameters and require retraining for each new parameter. We show that direct conditional parameterization of such solvers entangles physical parameters with the high-order automatic differentiation graph, causing extra memory growth and parameter-induced variance amplification. We propose Parameterized Representations via Implicit Stochastic Modulation (PRISM), a plug-and-play framework for parameterized high-dimensional and high-order stochastic neural PDE solvers. PRISM uses a hyper-generator to map physical parameters to affine modulators that scale and shift a purely spatial latent manifold, while keeping parameter branches value-connected but spatial-tangent-disconnected. This design preserves unbiased stochastic dimension and Taylor estimators, removes the parameter encoder from high-order spatial AD, and provides a variance-aware Lipschitz envelope over the parameter space. We prove parameterized unbiasedness, estimation-error bounds, and convergence under bounded stochastic variance. Experiments with PRISM-STDE and PRISM-SDGD on nonlinear parameterized PDEs show stable zero-shot generalization, reduced memory usage, and scalability up to 100,000 dimensions on a single GPU, with efficient low-rank SVD adaptation for unseen parameters.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Solving partial differential equations | Parametrized Poisson continuous parameter space (test) | Relative L2 Error4.91e-4 | 13 | |
| Solving partial differential equations | Allen-Cahn continuous parameter space (test) | Memory (MB)1.33e+3 | 10 | |
| Solving partial differential equations | Sine-Gordon continuous parameter space (test) | Memory Usage (MB)1.33e+3 | 10 | |
| Solving Hamilton-Jacobi-Bellman (HJB) Equations | HJB-Log | Computation Time (s)210 | 8 | |
| Solving Hamilton-Jacobi-Bellman (HJB) Equations | HJB-Rosenbrock | Computation Time (minutes)210 | 8 | |
| Solving parameterized PDEs | Parameterized Allen-Cahn PDE | Relative L2 Error0.0707 | 5 | |
| Solving parameterized PDEs | Parameterized Sine-Gordon PDE | Relative L2 Error5.40e-4 | 5 | |
| Solving Parameterized Partial Differential Equations | Parameterized Time-Dependent Allen-Cahn Equation 1K D | Relative L1 Error0.126 | 4 | |
| Solving Parameterized Partial Differential Equations | Parameterized Time-Dependent Allen-Cahn Equation 10K D | Relative L1 Error0.0379 | 4 | |
| Solving Parameterized Partial Differential Equations | Parameterized Time-Dependent Semilinear Heat Equation 10D | Relative L1 Error6.99e-5 | 4 |