Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Conditional Sampling via Wasserstein Autoencoders and Triangular Transport

Mohammad Al-Jarrah, Michele Martino, Marcus Yim, Bamdad Hosseini, Amirhossein Taghvaei

Abstract

We present Conditional Wasserstein Autoencoders (CWAEs), a framework for conditional simulation that exploits low-dimensional structure in both the conditioned and the conditioning variables. The key idea is to modify a Wasserstein autoencoder to use a (block-) triangular decoder and impose an appropriate independence assumption on the latent variables. We show that the resulting model gives an autoencoder that can exploit low-dimensional structure while simultaneously the decoder can be used for conditional simulation. We explore various theoretical properties of CWAEs, including their connections to conditional optimal transport (OT) problems. We also present alternative formulations that lead to three architectural variants forming the foundation of our algorithms. We present a series of numerical experiments that demonstrate that our different CWAE variants achieve substantial reductions in approximation error relative to the low-rank ensemble Kalman filter (LREnKF), particularly in problems where the support of the conditional measures is truly low-dimensional.

Conditional Sampling via Wasserstein Autoencoders and Triangular Transport

Abstract

We present Conditional Wasserstein Autoencoders (CWAEs), a framework for conditional simulation that exploits low-dimensional structure in both the conditioned and the conditioning variables. The key idea is to modify a Wasserstein autoencoder to use a (block-) triangular decoder and impose an appropriate independence assumption on the latent variables. We show that the resulting model gives an autoencoder that can exploit low-dimensional structure while simultaneously the decoder can be used for conditional simulation. We explore various theoretical properties of CWAEs, including their connections to conditional optimal transport (OT) problems. We also present alternative formulations that lead to three architectural variants forming the foundation of our algorithms. We present a series of numerical experiments that demonstrate that our different CWAE variants achieve substantial reductions in approximation error relative to the low-rank ensemble Kalman filter (LREnKF), particularly in problems where the support of the conditional measures is truly low-dimensional.

Paper Structure

This paper contains 15 sections, 2 theorems, 48 equations, 5 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Problem Formulation & Background
  3. Problem Setup
  4. Background on Wasserstein AutoEncoders
  5. Conditioning via Wasserstein AutoEncoders
  6. Joint vs. Conditional OT Costs
  7. The WAE for Conditioning
  8. Conditional Wasserstein Autoencoders
  9. Numerical experiments
  10. Numerical procedure
  11. Experiment Setup
  12. Nonlinear Manifold Embedding With Cubic Observations
  13. Spherical Posterior Example
  14. Incompressible Flow Field Reconstruction
  15. Discussion

Key Result

Proposition 1

For any encoder approximation $\Phi_\mathcal{Y}:\mathcal{Y}\to\mathcal{Z}$ encoding $Z=\Phi_\mathcal{Y}(Y)$ and transport map $\overline{G}_\mathcal{X}:\mathcal{Z}\times\mathcal{U}\to\mathcal{X}$, it holds that: Moreover, if $\Phi_\mathcal{Y}$ is a valid solution of Problem 2, it follows that $\mathcal{E}(\Phi_\mathcal{Y}) = 0$ and $\mathcal{R}_Y(\overline{G}_\mathcal{X};\Phi_\mathcal{Y}) = \math

Figures (5)

  • Figure D1: Numerical results for the synthetic example in Sec. \ref{['sec: Low-Dimensional Latent Structure Embedded in a High-Dimensional State']}. The figure shows the sample distributions for the last three states left to right for $y = \mathbf{1}_{d_\mathcal{Y}}$.
  • Figure D2: Numerical results for the spherical posterior example in Sec. \ref{['sec:cone_example']}. Both panels show the sample distributions for each method alongside the true distribution for two distinct values of the observation $y$. The left panel corresponds to $y = 1$ and the right panel corresponds to $y = 2$.
  • Figure D3: Simulation results for the flow field example in Sec. \ref{['sec: flow_example']}. (upper) $u_x$ component over time showing a sinusoidal stream and deceleration. (lower) $u_y$ component over time showing alternating sign and symmetric vortex formation. Taken with $\text{Re}=281$ and a cylindrical obstruction.
  • Figure D4: Numerical results for the flow field example in Sec. \ref{['sec: flow_example']}. (a) and (b) are the original field with the $d_{\mathcal{Y}}=9 \times 9$ observed pixels boxed in green. The bottom panels are reconstructed from the noisy observation and averaged per-component over 1000 samples.
  • Figure D5: Numerical results for the flow field example in Sec. \ref{['sec: flow_example']} showing 16 out of 1000 samples. The left panel corresponds to the $u_x$ component and the right panel corresponds to the $u_y$ component.

Theorems & Definitions (6)

  • Definition 1: Block-Triangular Map
  • Proposition 1
  • proof
  • Theorem 1
  • proof
  • Remark 1