🔢 mathematics

S $^{2}$ -PINN: Stochastic Separable Physics-Informed Neural Networks

This paper introduces S $^{2}$ -PINN, a stochastic separable physics-informed neural network that leverages a low-rank tensor decomposition of learnable spatial, temporal, and stochastic bases to efficiently and accurately quantify uncertainty in random partial differential equations while overcoming the curse of dimensionality and outperforming existing baselines.

Original authors: Zhendong Li, Akwum Onwunta

Published 2026-10-05

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhendong Li, Akwum Onwunta

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of engineering and climate science, predicting the future is rarely a matter of simple certainty. When scientists model how heat spreads through a metal beam, how water filters through soil, or how air swirls around an airplane wing, they must account for the fact that the inputs are never perfectly known. A material might be slightly rougher in one spot, or the wind might blow a fraction harder than expected. These small, random variations ripple through the equations, creating a vast cloud of possible outcomes rather than a single, fixed answer. To understand the risk, engineers need to know not just the most likely result, but the entire range of possibilities and how likely each one is. This field, known as uncertainty quantification, is essential for designing safe bridges, reliable medical devices, and accurate climate models. However, calculating these probabilities has historically been a monumental task. Traditional methods often require running the same complex simulation thousands of times, each with slightly different random inputs, to build a picture of the whole. As the number of uncertain factors grows, the computing power required explodes, making many real-world problems too expensive to solve with the tools currently available.

A team of researchers at Lehigh University has introduced a new approach that tackles this computational bottleneck by changing how the solution is built from the ground up. They developed a system called S2-PINN, which stands for Stochastic Separable Physics-Informed Neural Network. Instead of treating the problem as a massive, tangled web where space, time, and randomness are all mixed together, this method breaks the solution into three distinct, manageable pieces. Imagine the solution as a tapestry; rather than weaving every thread of color, time, and location into a single, dense block, the researchers separate the threads into three clear bundles: one for the physical space, one for the passage of time, and one for the random variations. These bundles are then woven back together using a specific, efficient pattern that keeps the total number of threads low. This separation allows the computer to handle the randomness without getting bogged down by the sheer volume of calculations usually required.

The core of this new system is a clever combination of three specialized tools working in harmony. First, it uses a flexible set of building blocks to describe the shape of the solution in space, adapting to the geometry of the problem. Second, it employs a rhythmic set of features to track how the solution changes over time, similar to how a musician might use a set of notes to describe a melody. Third, and most crucially for the uncertainty aspect, it uses a mathematical library of standard shapes to represent the random variations. These three sets of tools are linked together by a compact, low-rank connector that ensures they work in sync without requiring a separate, massive calculation for every single combination of space, time, and randomness. The system is trained by a dual process: it checks that the solution obeys the fundamental laws of physics at every point, and it also ensures that the statistical patterns of the randomness match the expected behavior. This dual check acts like a rigorous quality control, forcing the model to be both physically accurate and statistically sound.

When the researchers tested this new method against a wide variety of challenging problems, the results were striking. They ran simulations on four different types of physical systems, including models for heat diffusion, fluid flow, and chemical reactions, all with random inputs. In every case, the new method produced more accurate predictions of the average outcome and the spread of possible outcomes than nine other leading methods. It achieved this high level of precision while using significantly fewer computational resources; in one comparison, the new model used roughly thirty times fewer adjustable parameters than a standard, traditional neural network. This efficiency is vital because it means the method can scale to problems with many more random factors without crashing the computer. The researchers pushed the system to its limits, testing it with up to fifty different random variables, a level of complexity that would typically overwhelm standard approaches. Even in these extreme tests, the method remained stable and accurate.

The study also explored whether this approach could work on problems where the answer is not known in advance, such as predicting water flow through a rock with unpredictable permeability. Here, the method again outperformed traditional techniques, providing reliable estimates of both the average flow and the uncertainty around it. Furthermore, the researchers demonstrated that the system could work in reverse. Instead of just predicting the outcome of a known set of conditions, the system could look at observed data and work backward to figure out the hidden, random properties of the material itself. This ability to solve "inverse problems" opens the door to using the method for identifying unknown material properties in real-world scenarios, such as detecting flaws in infrastructure or characterizing underground reservoirs.

One of the most significant findings was how the method handles the relationship between the randomness and the physics. The researchers discovered that when the random variations follow a smooth, predictable pattern, the statistical part of the model acts as a powerful guide, sharpening the accuracy of the prediction. However, they also found that if the randomness is extremely erratic or follows a pattern that does not fit the standard shapes the model uses, the statistical guide can sometimes introduce a slight bias. This is not a failure of the system, but rather a clear signal that the nature of the randomness needs to be adjusted or that the model's library of shapes needs to be expanded. This transparency allows engineers to diagnose potential issues before they become problems, a feature that is often missing in other "black box" AI models.

The work represents a shift in how we think about solving complex physical problems with uncertainty. By treating space, time, and randomness as separate but connected ingredients, the researchers have created a framework that is both mathematically elegant and computationally efficient. It does not rely on massive amounts of pre-collected data, nor does it require the brute-force repetition of thousands of simulations. Instead, it learns the structure of the problem directly from the laws of physics and the nature of the randomness. While the method is not a magic wand that solves every possible problem instantly, it provides a robust and scalable tool for a wide range of scientific and engineering challenges. The researchers have made their code available to the public, inviting others to build upon this foundation. As the world faces increasingly complex challenges in energy, climate, and materials science, the ability to quantify uncertainty efficiently and accurately will become ever more critical, and this new approach offers a promising path forward.

Problem Statement
Uncertainty quantification (UQ) for random partial differential equations (PDEs) is critical in computational science, yet classical spectral solvers (e.g., stochastic Galerkin, collocation) suffer from the curse of dimensionality as the number of random variables increases. While Physics-Informed Neural Networks (PINNs) offer a data-free alternative, standard architectures typically treat stochastic inputs as generic coordinates, failing to explicitly exploit the separable structure across space, time, and randomness. This often leads to inefficient learning of stochastic moments and poor calibration. Existing neural operators and stochastic PINN variants either require supervised data, lack explicit stochastic basis structures, or rely on dense representations that do not scale well with stochastic dimensions.

Methodology: S2-PINN
The authors propose S2-PINN (Stochastic Separable Physics-Informed Neural Network), a framework that represents the solution $u(t, x, Z)$ of a random PDE using an explicitly separated space–time–random ansatz. The method combines three specialized components coupled via a low-rank Canonical Polyadic (CP) tensor decomposition:

Spatial Representation: A learnable Gaussian dictionary with an optional residual coordinate warp ( $W(x) = x + w_\theta(x)$ ), allowing the network to adapt to complex spatial features.
Temporal Representation: Fourier features with learnable frequencies, providing analytic time derivatives.
Stochastic Representation: A generalized Polynomial Chaos (gPC) basis (e.g., Hermite polynomials for Gaussian inputs) chosen based on the marginal distributions of the random variables $Z$ .

These three branches are coupled by a rank- $q$ CP core, reducing the parameter complexity from a dense $O(M M_t G)$ scaling to an additive $O(q(M + M_t + G))$ , where $M$ , $M_t$ , and $G$ are the number of spatial atoms, temporal features, and gPC modes, respectively.

Training and Loss Function
S2-PINN is trained using a hybrid loss function that enforces:

Strong-form PDE Residual: Pointwise minimization of the PDE residual.
gPC-Projected Residual: A term minimizing the squared expectation of the residual projected onto each gPC basis function. Theoretically, this corresponds to a stochastic Galerkin constraint.
Boundary and Initial Conditions: Standard Dirichlet/initial condition penalties.
Orthogonality Regularization: Penalties on the spatial feature Gram matrix and CP factors to ensure well-conditioned, near-orthogonal learned dictionaries.

Theoretical analysis establishes that the separable class is dense in $L^2$ , the projected residual is equivalent to a stochastic Galerkin constraint, and the mini-batch estimator for projection coefficients exhibits uniform concentration scaling logarithmically with the number of gPC modes ( $O(\log G)$ ).

Key Contributions

Architecture: Introduction of a stochastic separable PINN that couples learnable spatial atoms, temporal Fourier features, and gPC modes via a low-rank CP factorization, significantly reducing coupling complexity.
Theory: Proofs of expressivity for finite-rank separable sums, equivalence between the projected residual and stochastic Galerkin constraints, and bounds on the conditioning of the learned spatial dictionary via orthogonality penalties.
Empirical Performance: Demonstration that S2-PINN outperforms nine baselines (including standard PINNs, DeepONets, stochastic collocation, and other tensor-based methods) in mean and variance accuracy and calibration across four manufactured benchmarks (Diffusion, Allen–Cahn, Burgers, Darcy) and non-manufactured problems (Stochastic Poisson, Darcy, Navier–Stokes).
Inverse Problems: Extension of the framework to stochastic inverse problems, successfully recovering unknown PDE parameters (e.g., log-mean and standard deviation of random fields) jointly with the solution field.

Results

Accuracy: S2-PINN achieves lower relative $L^2$ errors for both mean and variance compared to baselines. For instance, on the stochastic diffusion problem, it reduces mean error by an order of magnitude compared to a standard Vanilla PINN while using roughly 30 times fewer parameters (7K vs. 200K).
Calibration: The method produces predictive intervals with empirical coverage close to nominal levels (e.g., 90% coverage at 90% nominal), outperforming MC-Dropout and Deep Ensembles in calibration metrics.
Scalability: The method remains stable up to $R=15$ random dimensions (816 gPC modes) and shows promise up to $R=50$ (23,426 modes) on diffusion problems, with memory usage scaling polynomially rather than exponentially.
Regime Dependence: The gPC projection loss is shown to be beneficial for problems with polynomial dependence on $Z$ but can introduce bias for problems with exponential dependence (e.g., log-normal Darcy), where the "no-projection" variant or higher gPC orders are required.

Significance and Claims
The paper claims that stochastic separability serves as an effective inductive bias for neural UQ of random PDEs. By explicitly decoupling spatial, temporal, and stochastic factors and enforcing consistency via both strong-form and Galerkin-projected residuals, S2-PINN circumvents the curse of dimensionality inherent in dense grids. The authors assert that this structure is particularly advantageous when solutions exhibit low-rank space–time–random structure and can be well-approximated by low-order stochastic expansions. The work supports the design principle that integrating learnable representations with explicit stochastic bases yields more accurate, calibrated, and parameter-efficient solvers for forward and inverse stochastic problems.

Limitations
The authors acknowledge that the method currently inherits the regular-domain bias of separable coordinate architectures, limiting its direct application to complex geometries without geometry-aware dictionaries or coordinate maps. Additionally, while theoretical results establish asymptotic expressivity, they do not provide finite-rank, finite-gPC end-to-end error guarantees for trained models, and the performance depends on the alignment between the solution's stochastic energy and the chosen polynomial basis.

More like this