Paired Wasserstein Autoencoders for Conditional Sampling

TL;DR

This paper introduces a novel paired Wasserstein autoencoder framework that learns cost-optimal transport maps and enables conditional sampling between two data distributions using shared latent spaces.

Contribution

It develops a new loss function for paired WAEs that facilitates learning OT couplings and enables conditional sampling via stochastic decoders.

Findings

Successfully learned cost-optimal transport maps.

Enabled conditional sampling from OT-type couplings.

Validated approach on synthetic data with known distributions.

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.