# Parameterized Representations via Implicit Stochastic Modulation for High-Dimensional and High-Order Neural PDE Solvers

arXiv: https://arxiv.org/abs/2606.22150 · HTML: https://mlsanity.com/paper/2606.22150/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 33 of its 56 references that are in the map.
- First posted on arXiv 20 Jun 2026.

## Builds on (most cited first, 30 of 33)

- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Physics-informed neural networks (PINNs) for fluid mechanics: A review](https://mlsanity.com/paper/2105.09506/) (2021)
- [Hidden Physics Models: Machine Learning of Nonlinear Partial Differential Equations](https://mlsanity.com/paper/1708.00588/) (2017)
- [Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations](https://mlsanity.com/paper/1706.04702/) (2017)
- [Adversarial Uncertainty Quantification in Physics-Informed Neural Networks](https://mlsanity.com/paper/1811.04026/) (2018)
- [Estimates on the generalization error of Physics Informed Neural Networks (PINNs) for approximating PDEs](https://mlsanity.com/paper/2006.16144/) (2020)
- [On the Almost Sure Convergence of Stochastic Gradient Descent in Non-Convex Problems](https://mlsanity.com/paper/2006.11144/) (2020)
- [Convergence rates for the stochastic gradient descent method for non-convex objective functions](https://mlsanity.com/paper/1904.01517/) (2019)
- [Machine learning approximation algorithms for high-dimensional fully nonlinear partial differential equations and second-order backward stochastic differential equations](https://mlsanity.com/paper/1709.05963/) (2017)
- [When Do Extended Physics-Informed Neural Networks (XPINNs) Improve Generalization?](https://mlsanity.com/paper/2109.09444/) (2021)
- [Forward-Backward Stochastic Neural Networks: Deep Learning of High-dimensional Partial Differential Equations](https://mlsanity.com/paper/1804.07010/) (2018)
- [Deep splitting method for parabolic PDEs](https://mlsanity.com/paper/1907.03452/) (2019)
- [Stochastic Gradient Descent for Nonconvex Learning without Bounded Gradient Assumptions](https://mlsanity.com/paper/1902.00908/) (2019)
- [Augmented Physics-Informed Neural Networks (APINNs): A gating network-based soft domain decomposition methodology](https://mlsanity.com/paper/2211.08939/) (2022)
- [Two-Layer Neural Networks for Partial Differential Equations: Optimization and Generalization Theory](https://mlsanity.com/paper/2006.15733/) (2020)
- [A novel meta-learning initialization method for physics-informed neural networks](https://mlsanity.com/paper/2107.10991/) (2021)
- [Learning stiff chemical kinetics using extended deep neural operators](https://mlsanity.com/paper/2302.12645/) (2023)
- [Machine Learning for semi linear PDEs](https://mlsanity.com/paper/1809.07609/) (2018)
- [Learning Physics-Informed Neural Networks without Stacked Back-propagation](https://mlsanity.com/paper/2202.09340/) (2022)
- [Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks](https://mlsanity.com/paper/2312.14499/) (2023)
- [Solving high-dimensional optimal stopping problems using deep learning](https://mlsanity.com/paper/1908.01602/) (2019)
- [Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators](https://mlsanity.com/paper/2412.00088/) (2024)
- [Deep neural operators can serve as accurate surrogates for shape optimization: A case study for airfoils](https://mlsanity.com/paper/2302.00807/) (2023)
- [A Priori Generalization Analysis of the Deep Ritz Method for Solving High Dimensional Elliptic Equations](https://mlsanity.com/paper/2101.01708/) (2021)

## Abstract (arXiv)

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.