S2-PINN: Stochastic Separable Physics-Informed Neural Networks
Abstract
Uncertainty quantification (UQ) for random partial differential equations (PDEs) is ubiquitous in computational science and engineering. However, classical spectral solvers for this class of problems face the curse of dimensionality, and existing neural solvers often ignore the stochastic structure that makes moments and calibration tractable. We introduce a stochastic separable physics-informed neural network, dubbed S2-PINN, that represents the solution u(t,x,Z) of a random PDE with a learnable Gaussian spatial dictionary, Fourier temporal features, and a generalized polynomial chaos (gPC) stochastic basis, coupled by a low-rank Canonical Polyadic (CP) tensor decomposition core. The method is trained with a hybrid strong-form and gPC-projected residual loss. Our theoretical analysis establishes that the separable class is dense in L2 under mild conditions, and the projected residual corresponds exactly to a stochastic Galerkin constraint. Furthermore, we show that mini-batch projection coefficients are logarithmically dependent on the number of gPC modes, and that the orthogonality penalty controls the conditioning of the learned spatial dictionary. Using four manufactured random PDE benchmarks, we show that S2-PINN outperforms nine baselines in terms of mean and variance accuracy, as well as calibration, while using significantly fewer parameters. Further evaluations on non-manufactured Poisson and Darcy problems, a stochastic Navier--Stokes problem, a diffusion scaling study of higher random dimensions, and two stochastic inverse problems reveal the generalization capabilities of the proposed structure. Together, these results support stochastic separability as an effective design principle for physics-informed neural UQ. The code for the experiments can be found in https://github.com/DMax1314/s2pinn
Related
- Lazy training of quantum physics informed neural networks
- AKAPINN: Adaptive Kolmogorov-Arnold Physics-Informed Neural Networks for approximating solutions to quasilinear partial differential equations
- Solving Fredholm Integral Equations with Physics-Informed Neural Networks: Single-and double-network approaches
- Error Estimates in Physics-Informed Neural Networks to Heat-Wave Fluid-Structure Interactive System
- Lie Symmetry-Guided Adaptive Sampling for Physics-Informed Neural Networks
BibTeX
@misc{li2026s2pinn,
title = {{S$^2$-PINN: Stochastic Separable Physics-Informed Neural Networks}},
author = {Zhendong Li and Akwum Onwunta},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.03303},
eprint = {2610.03303},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.03303},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - S^2-PINN: Stochastic Separable Physics-Informed Neural Networks
AU - Zhendong Li
AU - Akwum Onwunta
PY - 2026
DA - 2026/10/02
PB - arXiv
DO - 10.48550/arxiv.2610.03303
UR - https://arxiv.org/abs/2610.03303
AB - Uncertainty quantification (UQ) for random partial differential equations (PDEs) is ubiquitous in computational science and engineering. However, classical spectral solvers for this class of problems face the curse of dimensionality, and existing neural solvers often ignore the stochastic structure that makes moments and calibration tractable. We introduce a stochastic separable physics-informed neural network, dubbed S^2-PINN, that represents the solution u(t,x,Z) of a random PDE with a learnable Gaussian spatial dictionary, Fourier temporal features, and a generalized polynomial chaos (gPC) stochastic basis, coupled by a low-rank Canonical Polyadic (CP) tensor decomposition core. The method is trained with a hybrid strong-form and gPC-projected residual loss. Our theoretical analysis establishes that the separable class is dense in L^2 under mild conditions, and the projected residual corresponds exactly to a stochastic Galerkin constraint. Furthermore, we show that mini-batch projection coefficients are logarithmically dependent on the number of gPC modes, and that the orthogonality penalty controls the conditioning of the learned spatial dictionary. Using four manufactured random PDE benchmarks, we show that S^2-PINN outperforms nine baselines in terms of mean and variance accuracy, as well as calibration, while using significantly fewer parameters. Further evaluations on non-manufactured Poisson and Darcy problems, a stochastic Navier--Stokes problem, a diffusion scaling study of higher random dimensions, and two stochastic inverse problems reveal the generalization capabilities of the proposed structure. Together, these results support stochastic separability as an effective design principle for physics-informed neural UQ. The code for the experiments can be found in https://github.com/DMax1314/s2pinn
ER -