Paired Wasserstein Autoencoders for Conditional Sampling

Generative autoencoders learn compact latent representations of data distributions through jointly optimized encoder–decoder pairs. In particular, Wasserstein autoencoders (WAEs) minimize a relaxed optimal transport (OT) objective, where similarity between distributions is measured through a cost-minimizing joint distribution (OT coupling). Beyond distribution matching, neural OT methods aim to learn mappings between two data distributions induced by an OT coupling. Building on the formulation of the WAE loss, we derive a novel loss that enables sampling from OT-type couplings via two paired WAEs with shared latent space. The resulting fully parametrized joint distribution yields (i) learned cost-optimal transport maps between the two data distributions via deterministic encoders. Under cost-consistency constraints, it further enables (ii) conditional sampling from an OT-type coupling through stochastic decoders. As a proof of concept, we use synthetic data with known and visualizable marginal and conditional distributions.

💡 Research Summary

The paper introduces a novel framework called Paired Wasserstein Autoencoders (Paired WAEs) that extends the traditional Wasserstein Autoencoder (WAE) paradigm from matching a single data distribution to a latent prior, to jointly modeling two distinct data distributions through a shared latent space. By coupling two WAEs—each consisting of an encoder‑decoder pair (Eₓ, Dₓ) and (E_y, D_y)—the authors parametrize a joint distribution (Dₓ, D_y)#ξ, where ξ is a fixed prior on the latent space Z. This joint distribution serves as a parametric approximation of an optimal transport (OT) plan π ∈ Π(µ, ν) between the two target distributions µ and ν.

The core loss (Equation 4) combines three terms: (i) the expected transport cost under the latent prior, 𝔼_{z∼ξ}