Neural Chaos: A Spectral Stochastic Neural Operator

TL;DR

Uncertainty quantification for stochastic operators often relies on Polynomial Chaos Expansion (PCE), which faces scalability and dependence challenges. The authors introduce Neural Chaos, a spectral stochastic neural operator that learns neural-network-based, orthogonal basis functions directly from data to operate on the joint distribution of random inputs without assuming independence or marginal forms. They present two learning schemes—continuous spectral dictionary learning and discrete-continuous variants—that enforce near-orthogonality and enable compact, accurate operator learning, with connections to DeepONet. Through five numerical examples, including high-dimensional and dependent-random-variable cases, Neural Chaos achieves comparable or superior accuracy with far fewer basis functions than classical PCE, highlighting strong potential for non-intrusive, scalable stochastic surrogates.

Abstract

Building surrogate models with uncertainty quantification capabilities is essential for many engineering applications where randomness, such as variability in material properties, is unavoidable. Polynomial Chaos Expansion (PCE) is widely recognized as a to-go method for constructing stochastic solutions in both intrusive and non-intrusive ways. Its application becomes challenging, however, with complex or high-dimensional processes, as achieving accuracy requires higher-order polynomials, which can increase computational demands and or the risk of overfitting. Furthermore, PCE requires specialized treatments to manage random variables that are not independent, and these treatments may be problem-dependent or may fail with increasing complexity. In this work, we adopt the spectral expansion formalism used in PCE; however, we replace the classical polynomial basis functions with neural network (NN) basis functions to leverage their expressivity. To achieve this, we propose an algorithm that identifies NN-parameterized basis functions in a purely data-driven manner, without any prior assumptions about the joint distribution of the random variables involved, whether independent or dependent. The proposed algorithm identifies each NN-parameterized basis function sequentially, ensuring they are orthogonal with respect to the data distribution. The basis functions are constructed directly on the joint stochastic variables without requiring a tensor product structure. This approach may offer greater flexibility for complex stochastic models, while simplifying implementation compared to the tensor product structures typically used in PCE to handle random vectors. We demonstrate the effectiveness of the proposed scheme through several numerical examples of varying complexity and provide comparisons with classical PCE.

Figures (32)

• Figure 1: Neural Chaos architecture, which mirrors the structure of a spectral expansion. Spectral terms are added one by one such that each new term is orthogonal to the previous ones. Each spectral term has a multiplicative structure with $h_i(\bm{x},\boldsymbol{\xi}) = \Phi_i(\bm{x})\Psi_i(\boldsymbol{\xi})$ comprising deterministic and stochastic basis functions, each parameterized by a neural network.
• Figure 2: Example 1 -- Mean squared error between the data and the learned spectral expansion via Algorithm \ref{['algo:SSDL-disc-cont']} for increasing number of terms used in the expansion and different input distributions: (a) Uniform, (b) Normal, (c) Gamma, (d) Poisson. The slope indicates the rate of exponential error decay.
• Figure 3: Example 1 -- Distribution of error between the data and the model built via Algorithm \ref{['algo:SSDL-disc-cont']} using neural network basis functions for different input distributions: (a) Uniform, (b) Normal, (c) Gamma, (d) Poisson. The training data includes 700 realizations, while the test data includes 300 realizations.
• Figure 4: Example 1 -- Learned stochastic basis functions for data sampled from different input distributions: (a) Uniform, (b) Normal, (c) Gamma, (d) Poisson. The dots represent the basis vectors learned from Algorithm \ref{['algo:SSDL-disc-cont']}, while the solid lines correspond to their respective neural network basis functions.
• Figure 5: Example 1 -- Learned deterministic basis functions for data sampled from different input distributions: (a) Uniform, (b) Normal, (c) Gamma, (d) Poisson. The dots represent the basis vectors learned from Algorithm \ref{['algo:SSDL-disc-cont']}, while the solid lines correspond to their respective neural network basis functions.

Theorems & Definitions (6)

• Lemma 2.1
• proof
• Lemma 2.2
• proof
• Lemma 2.3
• Lemma 2.4