Deep Polynomial Chaos Expansion

Abstract

Polynomial chaos expansion (PCE) is a classical and widely used surrogate modeling technique in physical simulation and uncertainty quantification. By taking a linear combination of a set of basis polynomials - orthonormal with respect to the distribution of uncertain input parameters - PCE enables tractable inference of key statistical quantities such as (conditional) means, variances, covariances, and Sobol sensitivity indices, which are essential for understanding the modeled system and identifying influential parameters and their interactions. The applicability of PCE to high-dimensional problems is limited by poor scalability, as the number of basis functions grows exponentially with the number of parameters. In this paper, we address this challenge by combining PCE with ideas from tractable probabilistic circuits, resulting in deep polynomial chaos expansion (DeepPCE) - a deep generalization of PCE that scales effectively to high-dimensional input spaces. DeepPCE achieves predictive performance comparable to that of multilayer perceptrons (MLPs), while retaining PCE's ability to compute exact statistical inferences via simple forward passes. In contrast, such computations in MLPs require costly and often inaccurate approximations, such as Monte Carlo integration.

Figures (8)

• Figure 1: A simple smooth and structured decomposable circuit with input variables $\{X_1, X_2\}$. The input nodes encode tractable functions with inputs $X_1$ or $X_2$, followed by products over nodes with distinct scopes $\{X_1\}$ or $\{X_2\}$. The root node computes a sum over product nodes with the same scope $\{X_1, X_2\}$.
• Figure 2: A DeepPCE over inputs $\{{X_1, X_2, X_3, X_4}\}$ with two input nodes, encoding PCEs over scopes $\{{X_1, X_2}\}$ and $\{{X_3, X_4}\}$. Inputs are expanded as orthogonal polynomials $\{\psi_i\}_{i=0}^K$. A layer of outer-product nodes $\otimes$ forms the multivariate basis functions (tensor products) $\{\boldsymbol{\Phi}_{\boldsymbol{\alpha}_0}, \dots, \boldsymbol{\Phi}_{\boldsymbol{\alpha}_M} \}$. A sum layer $\oplus$ with trainable network weights composes multiple PCEs over the same input scope but with distinct weights, overparameterizing the PCE layer for better expressivity. The input layer is followed by one or several blocks of Kronecker $\otimes$ or Hadamard products $\odot$ and weighted sums $\oplus$ over these products. Repeated sum-product layers can then be stacked to deeper structures. The depth of the DeepPCE is determined by the input dimensionality and the size of the input scopes.
• Figure 3: Test performances using the Bratley function with $D=[100, 250, 500, 750, 1000, 2000]$. In contrast to classical PCE approaches, the DeepPCE manages to scale beyond 750 input dimensions and also outperforms the baseline PCEs in terms of predictive performance in lower dimensional settings.
• Figure 4: First order Sobol indices (left) and total effect Sobol indices (right) for Sobol G* function with $D = 100$ (panels 1 - 5) and the measured wall clock time for an inference pass as well as mean average errors for the Sobol indices (panel 6). The DeepPCE closely approximates to true variance contributions of each variable, outperforming all classical PCE variants. The MLP's Sobol indices are obtained by Monte Carlo, using a total of $10^8$ samples to approximate the Sobol indices for each input $X_i$. While it shows the worst performance of all models tested and fails to accurately represent the true variance contributions, the computation time to approximate the Sobol indices is longer by a factor of $10^3$ to $10^4$ compared to the DeepPCE (wall clock time).
• Figure 5: Random test samples for the Darcy flow dataset (left) and the steady-state diffusion dataset (right) from DeepPCE, MLP and UNet. The first column of panels show random input samples, the second column the corresponding true targets. The results show that the DeepPCE manages to closely approximate the PDE results, offering predictive performance that is competitive against the UNet and the MLP, while also offering tractable and exact inference.

Theorems & Definitions (2)

• Definition 1: (Probabilistic) Circuit
• Definition 2: (Structured) Decomposability