Enabling stratified sampling in high dimensions via nonlinear dimensionality reduction

We consider the problem of propagating the uncertainty from a possibly large number of random inputs through a computationally expensive model. Stratified sampling is a well-known variance reduction strategy, but its application, thus far, has focused on models with a limited number of inputs due to the challenges of creating uniform partitions in high dimensions. To overcome these challenges, we propose a simple methodology for constructing an effective stratification of the input domain that is adapted to the model response. Our approach leverages neural active manifolds, a recently introduced nonlinear dimensionality reduction technique based on neural networks that identifies a one-dimensional manifold capturing most of the model variability. The resulting one-dimensional latent space is mapped to the unit interval, where stratification is performed with respect to the uniform distribution. The corresponding strata in the original input space are then recovered through the neural active manifold, generating partitions that tend to follow the level sets of the model. We show that our approach is effective in high dimensions and can be used to further reduce the variance of multifidelity Monte Carlo estimators.