Paired Wasserstein Autoencoders for Conditional Sampling

Moritz Piening, Matthias Chung

Keywords: Generative Autoencoder, Optimal Transport, Conditional Generative Modelling

TL;DR: We pair two WAEs with a shared latent space to learn OT-optimal maps and enable conditional sampling from OT-type couplings.

Abstract:

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.

Submission Number: 146