Wasserstein Generative Learning of
Conditional Distribution

Consider a pair of random vectors $(X,Y)\in\mathcal{X}\times\mathcal{Y}$, where $\mathcal{X}\subseteq\mathbb{R}^{d}$ and $\mathcal{Y}\subseteq\mathbb{R}^{q}.$ Suppose $(X,Y)\sim P_{X,Y}$ with marginal distributions $X\sim P_{X}$ and $Y\sim P_{Y}$. Denote the conditional distribution of $Y$ given $X$ by $P_{Y\mid X}$. For a given value $x$ of $X$, we also write the conditional distribution as $P_{Y\mid X=x}$. For regression problems, $\mathcal{Y}\subseteq\mathbb{R}^{q}$ with $q\geq 1$; for classification problems, $\mathcal{Y}$ is a set of finitely many labels. We assume $\mathcal{X}\subseteq\mathbb{R}^{d}$ with $d\geq 1$. The random vectors $X$ and $Y$ can contain both continuous and categorical components. Let $\eta\in\mathbb{R}^{m}$ be a random vector independent of $(X,Y)$ with a known distribution $P_{\eta}$ that is easy to sample from, such as a normal distribution.

We are interested in finding a function $G:\mathbb{R}^{m}\times\mathcal{X}\mapsto\mathcal{Y}$ such that the conditional distribution of $G(\eta,X)$ given $X=x$ equals the conditional distribution of $Y$ given $X=x$. Since $\eta$ is independent of $X$, this is equivalent to finding a $G$ such that

Therefore, to sample from the conditional distribution $P_{Y\mid X=x}$, we can first sample an $\eta\sim P_{\eta}$ and calculate $G(\eta,x)$, then the resulting $G(\eta,x)$ is a sample from $P_{Y\mid X=x}.$ Because of this property, we shall refer to $G$ as a conditional generator. This approach has also been used in Zhou et al., (2021). The existence of such a $G$ is guaranteed by the noise-outsourcing lemma (Theorem 5.10 in Kallenberg, (2002)). For ease of reference, we state it here with a slight modification.

In Lemma 2.1, the noise distribution $P_{\eta}$ is taken to be $N(\mathbf{0},\mathbf{I}_{m})$ with $m\geq 1$, since it is convenient to generate random numbers from a standard normal distribution. It is a simple consequence of the original noise-outsourcing lemma which uses the uniform distribution on $[0,1]$ for the noise distribution. The dimension $m$ of the noise vector does not need to be the same as $q$, the dimension of $Y.$

Because $\eta$ and $X$ are independent, a $G$ satisfies (2.1) if and only if it also satisfies (2.2). Therefore, to construct the conditional generator, we can find a $G$ such that the joint distribution of $(X,G(\eta,X))$ matches the joint distribution of $(X,Y)$. This is the basis of the proposed generative approach described below.

Let $\Omega$ be a subset of $\mathbb{R}^{k},k\geq 1,$ on which the measures we consider are defined. Let $\mathcal{B}_{1}(\Omega)$ be the set of Borel probability measures on $\Omega$ with finite first moment, that is, the set of probability measures $\mu$ on $\mathbb{R}^{k}$ such that $\int_{\Omega}\|x\|_{1}d\mu(x)<\infty$, where $\|\cdot\|_{1}$ denotes the 1-norm in $\mathbb{R}^{k}$. The $1$-Wasserstein metric is defined as

where $\Gamma(\mu,\nu)$ is the set of joint probability distributions with marginals $\mu$ and $\nu.$ The $1$-Wasserstein metric is also known as the earth mover distance. A computationally more convenient form of the 1-Wasserstein metric is the Monge-Rubinstein dual (Villani,, 2008),

where $\mathcal{F}^{1}_{\text{Lip}}$ is the $1$-Lipschitz class,

When only random samples from $\mu$ and $\nu$ are available in practice, it is easy to obtain the empirical version of (2.4).

We apply the $1$-Wasserstein metric (2.4) to the present conditional generative learning problem, that is, we seek a conditional generator function $G:\mathbb{R}^{m}\times\mathcal{X}\to\mathbb{R}^{q}$ satisfying (2.1). The basic idea is to formulate this problem as a minimisation problem based on the $1$-Wasserstein metric. By Lemma 2.1, we can find a conditional generator $G$ by matching the joint distributions $P_{X,G(\eta,X)}$ and $P_{X,Y}.$ By (2.4), the 1-Wasserstein distance between $P_{X,G(\eta,X)}$ and $P_{X,Y}$ is

We have $W_{1}(P_{X,G(\eta,X)},P_{X,Y})\geq 0$ for every measurable $G$ and $W_{1}(P_{X,G(\eta,X)},P_{X,Y})=0$ if and only if $P_{X,G(\eta,X)}=P_{X,Y}.$ Therefore, a sufficient and necessary condition for

is $P_{X,G^{*}(\eta,X)}=P_{X,Y}$, which implies $G^{*}(\eta,x)\sim P_{Y\mid X=x},x\in\mathcal{X}.$ It follows that the problem of finding the conditional generator can be formulated as the minimax problem:

where

Let $\{(X_{i},Y_{i}),i=1,\ldots,n\}$ be a random sample from $P_{X,Y}$ and let $\{\eta_{i},i=1,\ldots,n\}$ be independently generated from $P_{\eta}$. The empirical version of $\mathcal{L}(G,D)$ based on $(X_{i},Y_{i},\eta_{i}),i=1,\ldots,n,$ is

We use a feedforward neural network $G_{{\bm{\theta}}}$ with parameter ${\bm{\theta}}$ for estimating the conditional generator G and a second network $D_{{\bm{\phi}}}$ with parameter ${\bm{\phi}}$ for estimating the discriminator $D$. We refer to Goodfellow et al., (2016) for a detailed description of neural network functions. We estimate ${\bm{\theta}}$ and ${\bm{\phi}}$ by solving the minimax problem:

The estimated conditional generator is $\hat{G}=G_{\hat{{\bm{\theta}}}}$ and the estimated discriminator is $\hat{D}=D_{\hat{{\bm{\phi}}}}$.

We show below that for $\eta\sim P_{\eta}$, $\hat{G}(\eta,x)$ converges in distribution to the conditional distribution $P_{Y\mid X=x},x\in\mathcal{X}$. Therefore, we can use $\hat{G}$ to learn this conditional distribution. Specifically, for an integer $J>1$, we generate a random sample $\{\eta_{j},j=1,\ldots,\eta_{J}\}$ from the reference distribution $P_{\eta}$ and calculate $\{\hat{G}(\eta_{j},x),j=1,\ldots,J\}$, which are approximately distributed as $P_{Y\mid X=x}.$ Since generating random samples from $P_{\eta}$ and calculating $\hat{G}(\eta_{j},x)$ are inexpensive, the computational cost is low once $\hat{G}$ is obtained. For $Y\in\mathbb{R}$, this immediately leads to the estimated conditional quantiles of $P_{Y\mid X=x}.$ Moreover, for any function $g:\mathcal{Y}\to\mathbb{R}^{k}$, by the noise outsourcing lemma, we have $\mathbb{E}(g(Y)\mid X=x)=\mathbb{E}g(G(\eta,x)).$ We can estimate this conditional expectation by $n^{-1}\sum_{j=1}^{J}g(\hat{G}(\eta_{j};x)).$ In particular, the conditional mean and conditional variance of $P_{Y\mid X=x}$ can be estimated in a straightforward way. In Section 5 below, we illustrate this approach with a range of examples.

Let $(L_{1},W_{1})$ be the depth and width of the feedforward ReLU discriminator network $D_{{\bm{\phi}}}$ and let $(L_{2}.W_{2})$ be the depth and width of the feedforward ReLU generator network $G_{{\bm{\theta}}}$. Denote the joint distribution of $(X,\hat{G}(\eta,X))$ by $P_{X,\hat{G}}$.

When $P_{X,Y}$ has a bounded support, we can drop the logarithm factor in the error bound.

The error bound for the conditional distribution follows as a corollary.

The proofs of Theorems 3.1 and 3.2 and Corollary 3.3 are given in the supplementary material. Assumption 1 only concerns the first moment tail of the joint distribution of $(X,Y)$. Assumption 2 requires the noise distribution to be absolutely continuous with respect to the Lebesgue measure. These assumptions are easily satisfied in practice. Moreover, we have made clear how the prefactor in the error bound depends on the dimension $d+q$ of $(X,Y)$. This is useful in the high-dimensional settings since how the prefactor depends on the dimension plays an important role in determining the quality of the bound. Here the prefactors in the error bounds depend on the dimensions $d$ and $q$ through $(d+q)^{1/2}$. The convergence rate is $n^{-1/(d+q)}$.

Unfortunately, these results suffer from the curse of dimensionality in the sense that the quality of the error bound deteriorates quickly when $d+q$ becomes large. It is generally not possible to avoid the curse of dimensionality without any conditions on the distribution of $(X,Y)$. Detailed discussions on this problem in the context of nonparametric regression using neural networks can be found in Bauer and Kohler, (2019); Schmidt-Hieber, (2020) and Jiao et al., (2021) and the references therein. In particular, Jiao et al., (2021) also provided an analysis of how the prefactor depends on the dimension of $X$. To mitigate the curse of dimensionality, certain assumptions on the distribution of $(X,Y)$ is needed. We consider one such assumption below.

If the joint distribution of $(X,Y)$ is supported on a low dimensional set, we can improve the convergence rate substantially. Low dimensional structure of complex data have been frequently observed by researchers in image analysis, computer vision and natural language processing, therefore, it is often a reasonable assumption. We use the Minkowski dimension as a measure of the dimensionality of a set (Bishop and Peres,, 2016).

Let $A$ be a subset of $\mathbb{R}^{d}$. For any $\varepsilon>0$, the covering number $\mathcal{N}(\epsilon,A,\|\cdot\|_{2})$ is the minimum number of balls with radius $\varepsilon$ needed to cover $A$ (Van der Vaart and Wellner,, 1996). The upper and the lower Minkowski dimensions of $A\subseteq\mathbb{R}^{d}$ are defined respectively as

If $\dim^{u}_{M}(A)=\dim^{l}_{M}(A)=\dim_{M}(A)$, then $\dim_{M}(A)$ is called the Minkowski dimension or the metric dimension of $A.$

Minkowski dimension measures how the covering number of the set decays when radius of the covering balls goes to 0. When $A$ is a smooth manifold, its Minkowski dimension equals its own topological dimension characterized by the local homeomorphisms. Minkowski dimension can also be used to measure the dimensionality of highly non-regular set, such as fractals (Falconer,, 2004). Nakada and Imaizumi, (2020) and Jiao et al., (2021) have shown that deep neural networks can adapt to the low dimensional structure of the data and mitigate the curse of dimensionality in nonparametric regression. We show similar results can be obtained for the proposed method if the data distribution is supported on a set with low Minkowski dimension.

In comparison with Theorem 3.2, where the rate of convergence depends on $d+q$, the convergence rate in Theorem 3.4 is determined by the Minkowski dimension $d_{A}$. Therefore, the assumption of a low Minkowski dimension on the support of data distribution alleviates the curse of dimensionality. However, unless we have better approximation error bounds for Lipschitz functions defined on low Minkowski dimensional set using neural network functions, the prefactor in Theorem 3.4 is still $(d+q)^{1/2}$, that is, even the convergence rate only depends on the Minkowski dimension $d_{A}$, the prefactor still depends on the ambient dimension $d+q.$

We use the two-moon data example to illustrate the use of the proposed method for generating conditional samples. Let $X\in\{1,2\}$ be the class label and let $Y\in\mathbb{R}^{2}$. The two-moon model is

where $\alpha\sim\text{Uniform}[0,\pi]$, $\epsilon_{1},\ldots,\epsilon_{4}$ are independent and identically distributed as $N(0,\sigma^{2}).$ We generate three sets of random samples of size $n=5,000$ with $2,500$ for each class and $\sigma=0.1,0.2$ and 0.3 from this model. The generated datasets are shown in the first row of Figure 1.

We use the simulated data to estimate the conditional generator, which is parameterized as a two-layer fully-connected network with 30 and 20 nodes. The discriminator is also a two-layer fully-connected network with 40 and 20 nodes. The noise $\eta\sim N(\mathbf{0},\mathbf{I}_{2}).$ The activation function for the hidden layer of the generator and the discriminator is ReLU. For the output layer of the generator, the activation function is the hyperbolic tangent function. The estimated samples $\{\hat{G}(\eta_{j},x),j=1.\ldots,5,000\}$, $x\in\{1,2\},$ are shown in the second row of Figure 1. It can be seen that the scatter plots of the estimated samples are similar to those of the simulated data. This provides a visual validation of the estimated samples in this toy example.

We consider the problem of estimating conditional mean and conditional standard deviation in nonparametric conditional density models. We also compare the proposed Wasserstein generative conditional sampling method, referred to as WGCS in Table 1, with three existing conditional density estimation methods, including the nearest neighbor kernel conditional density estimator (NNKCDE) (Dalmasso et al.,, 2020), the conditional kernel density estimator (CKDE, implemented in R package np (Hall et al.,, 2004), and a basis expansion method FlexCode (Izbicki et al.,, 2017). We simulated data from the following three models:

• •

  Model 1 (M1). A nonlinear model with an additive error term:

  $\displaystyle Y=X_{1}^{2}+\exp(X_{2}+X_{3}/3)+\sin(X_{4}+X_{5})+\varepsilon,\ \varepsilon\sim N(0,1).$

• •

  Model 2 (M2). A model with a multiplicative non-Gassisan error term:

  $\displaystyle Y=(5+X_{1}^{2}/3+X_{2}^{2}+X_{3}^{2}+X_{4}+X_{5})*\exp(0.5\times\varepsilon),$

  where $\varepsilon\sim 0.5N(-2,1)+0.5N(2,1).$

• •

  Model 3 (M3). A mixture of two normal distributions:

  $\displaystyle Y=\mathbb{I}_{\{U\leq{1}/{3}\}}N(-1-X_{1}-0.5X_{2},0.5^{2})+\mathbb{I}_{\{U>{1}/{3}\}}N(1+X_{1}+0.5X_{2},1),$

  where $U\sim\text{Unif}(0,1)$ and is independent of $X$.

In each model, the covariate vector $X$ is generated from $N(\mathbf{0},\mathbf{I}_{100})$. So the ambient dimension of $X$ is 100, but (M1) and (M2) only depend on the first 5 components of $X$ and (M3) only depends on the first 2 components of $X$. The sample size $n=5,000$.

For the proposed method, the conditional generator $G$ is parameterized using a one-layer neural network in models (M1) and (M2); it is parameterized by a two-layer fully connected neural network in (M3). The discriminator $D$ is parameterized using a two-layer fully connected neural network. The noise vector $\eta\sim N(0,1).$

For the conditional density estimation method NNKCDE, the tuning parameters are chosen using cross-validation. The bandwidth of the conditional kernel density estimator CKDE is determined based on the standard formula $h_{j}=1.06\sigma_{j}n^{-1/(2k+d)},$ where $\sigma_{j}$ is a measure of spread of the $j$th variable, $k$ the order of the kernel, and $d$ the dimension of $X$. The basis expansion based method FlexCode uses Fourier basis. The maximum number of bases is set to 40 and the actual number of bases is selected using cross-validation.

We consider the mean squared error (MSE) of the estimated conditional mean $\mathbb{E}(Y|X)$ and the estimated conditional standard deviation $\text{SD}(Y|X)$. We use a test data set $\{x_{1},\dots,x_{K}\}$ of size $K=2,000$ . For the proposed method, we first generate $J=10,000$ samples $\{\eta_{j}:j=1,\ldots,J\}$ from the reference distribution $P_{\eta}$ and calculate conditional samples $\{\hat{G}(\eta_{j},x_{k}),j=1,\dots,J,k=1,\ldots,K\}$ The estimated conditional standard deviation is calculated as the sample standard deviation of the conditional samples. The MSE of the estimated conditional mean is $\text{MSE(mean)}=(1/K)\sum_{k=1}^{K}\{\hat{\mathbb{E}}(Y|X=x_{k})-\mathbb{E}(Y|X=x_{k})\}^{2}.$ For the proposed method, the estimate of the conditional mean is obtained using Monte Carlo. For other methods, the estimate is calculated by numerical integration. Similarly, the MSE of the estimated conditional standard deviation is $\text{MSE(sd)}=(1/K)\sum_{k=1}^{K}\{\hat{\text{SD}}(Y|X=x_{k})-\text{SD}(Y|X=x_{k})\}^{2}.$ The estimated conditional standard deviation of other methods are computed by numerical integration.

We repeat the simulations 10 times. The average MSEs and simulation standard errors are summarized in Table 1. Comparing with CKDE and FlexCode and NNKCED. WGCS has the smallest MSEs for estimating conditional mean and conditional SD in most cases.

We use the wine quality dataset (Cortez et al.,, 2009) to illustrate the application of the proposed method to prediction interval construction. This dataset is available at UCI machine learning repository (https://archive.ics.uci.edu/ml/datasets/Wine+Quality). There are eleven physicochemical quantitative variables including fixed acidity, volatile acidity, citric acid, residual sugar, chlorides, free sulfur dioxide, total sulfur dioxide, density, pH, sulphates, alcohol as predictors and a sensory quantitative variable quality that measures wine quality, which is a score between 0 and 10. The goal is to build a prediction model for the wine quality score based on the eleven physicochemical variables. Such a model can be used to help the oenologist wine tasting evaluations and improve wine production. We use the proposed method to learn the conditional distribution of quality given the eleven variables. An advantage of the proposed method over the standard nonparametric regression is that it provides a prediction interval for the quality score, not just a point prediction.

The sample size of this dataset is 4,898. We use 90% of the samples as the training set and 10% as the test set. All variables are centered and standardized first before training the model. In our proposed method, the conditional generator is a two-layer fully-connected feedforward network with 50 and 20 nodes, the discriminator is a two-layer network with 50 and 25 nodes. The ReLU activation function is used in both networks. The noise $\eta\sim N(\mathbf{0},\mathbf{I}_{3}).$

To examine the prediction performance of the proposed method, we construct the 90% prediction interval for the wine quality score in the test set. The prediction intervals are shown in Figure 2. The actual quality scores are plotted as solid dots. The actual coverage for all 490 in the test set is 90.80%, close to the nominal level of 90%.

We use the California housing dataset to demonstrate that the generated conditional samples from the proposed method can be used for visualizing conditional distributions of bivariate response data. This dataset is available at StatLib repository (http://lib.stat.cmu.edu/datasets/). It contains 20,640 observations on housing prices with 9 covariates including median house value, median income, median house age, total rooms, total bedrooms, population, households, latitude, and longitude. Each row of the data matrix represents a neighborhood block in California. We use the logarithmic transformation of the variables except longitude, latitude and median house age. All columns of the data matrix are centered and standardized to have zero mean and unit variance. The generator is a two-layer fully-connected neural network with 60 and 25 nodes and the discriminator is a two-layer network with 65 and 30 nodes. ReLU activation is used for the hidden layers. The random noise vector $\eta\sim N(\mathbf{0},\mathbf{I}_{4})$.

We train the WGCS model with (median income, median house value) $\in\mathbb{R}^{2}$ as the response vector and other variables as the predictors in the model. We use 18,675 observations (about the 90% of the dataset) as training set and the remaining 2,064 (about 10% of the dataset) observations as the test set.

Figure 3 shows the contour plot of the conditional distributions of median income and median house value (in logarithmic scale) for 8 single blocks in the test set. The colored density function represents the kernel density estimates based on the samples generated using the proposed method. We see that the blue cross is in the main body of the plot, which shows that the proposed method provides reasonable prediction of the median income and house values. We also see that big cities like San Francisco and Los Angles have higher house values with larger variations, smaller cities such as Davis and Concord tend to have lower house values.

We now illustrate the application of the proposed method to high-dimensional data problems and demonstrate that it can easily handle the models when either of both of $X$ and $Y$ are high-dimensional. We use the MNIST handwritten digits dataset (LeCun and Cortes,, 2010), which contains 60,000 images for training and 10,000 images for testing. The images are stored in $28\times 28$ matrices with gray color intensity from 0 to 1. Each image is paired with a label in $\{0,1\dots,9\}$. We use WGCS to perform two tasks: generating images from labels and reconstructing the missing part of an image.

We illustrate using the proposed method for image reconstruction when part of the image is missing with the MNIST dataset. Suppose we only observe ${1}/{4},{1}/{2}$ or ${3}/{4}$ of an image and would like to reconstruct the missing part of the image. For this problem, let $X$ be the observed part of the image and let $Y$ be the missing part of the image. Our goal is to reconstruct $Y$ based on $X$. For discriminator, we use two convolutional networks to process $X$ and $Y$ separately. The filters are then concatenated together and fed into another convolution layer and fully-connected layer before output. For the generator, $X$ is processed by a fully-connected layer followed by 3 deconvolution layers. The random noise vector $\eta\sim N(\mathbf{0},\mathbf{I}_{100}).$

In Figure 4, three panels from left to right correspond to the situations when ${1}/{4},{1}/{2}$ or ${3}/{4}$ of an image are given, the problem is to reconstruct the missing part of the image. In each panel, the first column contains the true images in the testing set; the second column shows the observed part of the image; the third to the seventh columns show the reconstructed images. The digits “0”, “1” and “7” are easy to reconstruct, even when only 1/4 of their images are given. For the other digits, if only 1/4 of their images are given, it is difficult to reconstruct them. However, as the given part increases from 1/4 to 1/2 and then 3/4 of the images, the proposed method is able to reconstruct the images correctly, and the reconstructed images become less variable and more similar to the true image. For example, for the digit “2”, if only the left lower 1/4 of the image is given, the reconstructed images tend to be incorrect; the reconstruction is only successful when 3/4 of the image is given.

We illustrate the application of the proposed method to the problem of image generation given label information with the CelebFaces Attributes (CelebA) dataset (Liu et al.,, 2015). This dataset contains more then 200K colored celebrity images with 40 attributes annotation. Here we use 10 binary features, including: Gender, Young, Bald, Bangs, Blackhair, Blondhair, Eyeglasses, Mustache, Smiling, WearingNecktie. We take these features as $X$. So $X$ is a vector consisting of 10 binary components. We consider six types of images; the attributes of these six types are shown in Table A.1. We used the aligned and cropped images and further resize them to $96\times 96$ as our training set. We take these colored images as the values of $Y$. Therefore, the dimension of $X$ is 10 and the dimension of $Y$ is $96\times 96\times 3=27,648$.

The architecture of the discriminator is specified as follows: the 10 dimensional one-hot vector is first expanded to $96\times 96\times 10$ by replicating each number to a $96\times 96$matrix. Then it is concatenated with image $Y$ on the last axis. Thus the processed data has dimension $96\times 96\times 13$. This processed data is then fed into 5 consecutive convolution layers initialized by truncated normal distribution with $SD=0.01$ with 64,128,256,512,1024 filters respectively. The activation for each convolution layer is a leaky ReLU with $\alpha=0.2$. The strides is 2 and the kernel size is 5. After convolution, it is flattened and connected to a dense layer with one unit as output.

The architecture of the conditional generator is as follows: the noise vector and the feature vector $X$ are concatenated and fed to a dense layer with $3\times 3\times 1024$ units with ReLU activation and then reshaped to $3\times 3\times 1024$. Then it goes through 5 layers of deconvolution initialized by truncated normal distribution with $SD=0.01$ with 512,256,128,64,3 filters respectively. The activation for each intermediate convolution layer is ReLU and hyperbolic tangent function for the last layer. The strides is 2 and kernel size is 5. The optimizer is Adam with $\beta_{1}=0.5$ and learning rate decrease from $0.0001$ to $0.000005$. The noise dimension is 100.

In Figure 5, the left panel shows real images and the right panel shows generated images. Each row corresponds to a specific type of face. The attributes of the six types of face images are given in Table A.1. We can see that the generated images have similar qualities as the original ones.