MD-NOMAD: Mixture density nonlinear manifold decoder for emulating stochastic differential equations and uncertainty propagation

TL;DR

The paper addresses uncertainty propagation for stochastic simulators by introducing MD-NOMAD, a hybrid surrogate that combines a neural-operator backbone (NOMAD) with a mixture-density network to learn conditional probability distributions $p(G_{\\omega}(f)(x)|f)$ without Monte Carlo sampling. MD-NOMAD outputs the parameters of a $m$-component Gaussian mixture $(\\pi_i, \\\mu_i, \\\sigma_i)$ conditioned on inputs, enabling analytic computation of PDFs and moments in a one-shot fashion. The authors demonstrate the framework on a suite of stochastic ODEs and PDEs, including high-dimensional Lorenz-96, a bimodal analytical benchmark, and a 2D SPDE, showing accurate PDFs and substantial computational gains over MC and KCDE baselines. They also discuss resolution-invariance and the potential for future work, such as automatic mixture-component selection and extension to multivariate mixtures. Overall, MD-NOMAD provides a scalable, accurate, and efficient approach for surrogate modeling and uncertainty propagation in stochastic systems.

Abstract

We propose a neural operator framework, termed mixture density nonlinear manifold decoder (MD-NOMAD), for stochastic simulators. Our approach leverages an amalgamation of the pointwise operator learning neural architecture nonlinear manifold decoder (NOMAD) with mixture density-based methods to estimate conditional probability distributions for stochastic output functions. MD-NOMAD harnesses the ability of probabilistic mixture models to estimate complex probability and the high-dimensional scalability of pointwise neural operator NOMAD. We conduct empirical assessments on a wide array of stochastic ordinary and partial differential equations and present the corresponding results, which highlight the performance of the proposed framework.

Figures (16)

• Figure 1: Schematic diagram depicting the architecture of MD-NOMAD. Here, $\mathbf{x_b}$ serves as the input to the branch network, which is an FCN producing a finite-dimensional latent feature representation as its output. The sensor values or locations $\mathbf{x_d}$ (also referred to as augmenting input) are augmented to the latent vector obtained from the branch network and provided as an input to the decoder network, which is also an FCN. The output of the decoder network comprises the parameters of the $\mathscr{m}$-component mixture of a probability distribution, with the probability distribution chosen from a suitable parametric family. In this article, we use a mixture of univariate Gaussian distribution. Subsequently, the MD-NOMAD model is trained by minimizing the sum of the negative logarithm of the probability $P(G_{\omega}(\mathbf{x_b})(\mathbf{x_d}))$ over $\mathscr{n}$ training samples.
• Figure 2: Comparison between reference distribution and MD-NOMAD predicted PDF at times $t = {1.01,\, 4.24,\, 12.32,\,\text{and } 20}$ s for (a)$\lambda = 0.45$, (b)$\lambda = 0,75$ for the spatial location component of stochastic Van der Pol equation.
• Figure 3: Comparison between reference statistics and MD-NOMAD predicted statistics till $t=20$ s for (a)$\lambda = 0.45$, (b)$\lambda = 0,75$ for the spatial location component of stochastic Van der Pol equation.
• Figure 4: Comparison between reference distribution and MD-NOMAD predicted PDF at times $t = {0.22,\, 1.76,\, 2.96,\,\text{and } 4.0}$ s corresponding to $\lambda=0.22$ for the (a) second, (b) fifth, and (c) tenth solution component of the $10$-component stochastic Lorenz equations.
• Figure 5: Comparison between reference distribution and MD-NOMAD predicted PDF at times $t = {0.22,\, 1.76,\, 2.96,\,\text{and } 4.0}$ s corresponding to $\lambda=0.22$ for the (a) second, (b) twenty fifth, and (c) fiftieth solution component of the $50$-component stochastic Lorenz equations.