Learning Consistent Deep Generative Models from Sparse Data via Prediction Constraints

The variational autoencoder (Kingma & Welling, 2014) is an unsupervised model with two components: a generative model and an inference model. The generative model defines for each example a joint distribution $p_{\theta}(x,z)$ over “features” (observed vector $x\in\mathbb{R}^{D}$) and “encodings” (hidden vector $z\in\mathbb{R}^{C}$). The “inference model” of the VAE defines an approximate posterior $q_{\phi}(z\mid x)$, which is trained to be close to the true posterior ($q_{\phi}(z\mid x)\approx p_{\theta}(z\mid x)$) but much easier to evaluate. As in Kingma & Welling (2014), we assume the following conditional independence structure:

$\displaystyle p_{\theta}(x,z)=\mathcal{N}(z\mid 0,I_{C})\cdot\mathcal{F}(x\mid\mu_{\theta}(z),\sigma_{\theta}(z)),\quad q_{\phi}(z\mid x)=\mathcal{N}(z\mid\mu_{\phi}(x),\sigma_{\phi}(x)).$

(1)

The likelihood $\mathcal{F}$ is often multivariate normal, but other distributions may give robustness to outliers. The (deterministic) functions $\mu_{\theta}$ and $\sigma_{\theta}$, with trainable parameters $\theta$, define the mean and covariance of the likelihood. Given any observation $x$, the posterior of $z$ is approximated as normal with mean $\mu_{\phi}$ and (diagonal) covariance $\sigma_{\phi}$ parameterized by $\phi$. These functions can be represented as multi-layer perceptrons (MLPs), convolutional neural networks (CNNs), or other (deep) neural networks.

We would ideally learn generative parameters $\theta$ by maximizing the marginal likelihood of features $x$, integrating latent variable $z$. Since this is intractable, we instead maximize a variational lower bound:

$\displaystyle\max_{\theta,\phi}~{}~{}\textstyle\sum_{x\in\mathcal{D}}\mathcal{L}^{\text{VAE}}(x;\theta,\phi),\qquad\mathcal{L}^{\text{VAE}}(x;\theta,\phi)=\mathbb{E}_{q_{\phi}(z|x)}\left[\log\frac{p_{\theta}(x,z)}{q_{\phi}(z|x)}\right]\leq\log p_{\theta}(x).$

(2)

This expectation can be evaluated via Monte Carlo samples from the inference model $q_{\phi}(z|x)$. Gradients with respect to $\theta,\phi$ can be similarly estimated by the reparameterization “trick” of representing $q_{\phi}(z\mid x)$ as a linear transformation of standard normal variables (Kingma & Welling, 2014).

Throughout this paper, we denote variational parameters by $\phi$. Because the factorization of $q$ changes for more complex models, we will write $\phi^{z|x}$ to denote the parameters specific to factor $q(z|x)$.

One way to employ the VAE for a semi-supervised task is a two-stage “VAE-then-MLP”. First, train a VAE to maximize the unsupervised likelihood (2) of all observed features $x$ (both labeled $\mathcal{D}^{S}$ and unlabeled $\mathcal{D}^{U}$). Second, we define a label-from-code predictor $\hat{y}_{w}(z)$ that maps each learned code representation $z$ to a predicted label $y\in\mathcal{Y}$. We use an MLP with weights $w$, though any predictor could do. Let $\ell_{S}(y,\hat{y})$ be a loss function, such as cross-entropy, appropriate for the prediction task. We train the predictor to minimize the loss: $\min_{w}\sum_{x,y\in\mathcal{D}^{S}}\mathbb{E}_{q_{\phi}(z|x)}\left[\ell_{S}(y,\hat{y}_{w}(z))\right]$. Importantly, this second stage uses only the small labeled dataset and relies on fixed parameters $\phi$ from stage one.

While “VAE-then-MLP” is a simple common baseline (Kingma et al., 2014), it has a key disadvantage: Labels are only used in the second stage, and thus a misspecified generative model in stage one will likely produce inferior predictions. Fig. 1 illustrates this weakness.

To overcome the weakness of the two-stage approach, previous work by Kingma et al. (2014) presented a VAE-inspired model called “M2” focused on the joint generative modeling of labels $y$ and data $x$. M2 has two components: a generative model $p_{\theta}(x,y,z)$ and an inference model $q_{\phi}(y,z\mid x)$. Their generative model is factorized to sample labels (with frequencies $\pi$) first, and then features $x$:

$\displaystyle p_{\theta}(x,y,z)=\mathcal{N}(z\mid 0,I_{C})\cdot\text{Cat}(y\mid\pi)\cdot\mathcal{F}(x\mid\mu_{\theta}(y,z),\sigma_{\theta}(y,z)).$

(3)

The M2 inference model sets $q_{\phi}(y,z\mid x)=q_{\phi^{y|x}}(y\mid x)q_{\phi^{z|x,y}}(z\mid x,y)$, where $\phi=(\phi^{y|x},\phi^{z|x,y})$.

To train M2, Kingma et al. (2014) maximize the likelihood of all observations (labels and features):

$\displaystyle\max_{\theta,\phi^{y|x},\phi^{z|x,y}}\quad\textstyle\sum_{x,y\in\mathcal{D}^{S}}\mathcal{L}^{S}(x,y;\theta,\phi^{z|x,y})+\textstyle\sum_{x\in\mathcal{D}^{U}}\mathcal{L}^{U}(x;\theta,\phi^{y|x},\phi^{z|x,y}).$

(4)

The first, “supervised” term in Eq. (4) is a variational bound for the feature-and-label joint likelihood:

$\displaystyle\mathcal{L}^{S}(x,y;\theta,\phi^{z|x,y})$

$\displaystyle=\textstyle\mathbb{E}_{q_{\phi^{z|x,y}}(z|x,y)}\left[\log\frac{p_{\theta}(x,y,z)}{q_{\phi^{z|x,y}}(z|x,y)}\right]\leq\log p_{\theta}(x,y).$

(5)

The second, “unsupervised” term is a variational lower bound for the features-only likelihood $\log p_{\theta}(x)\geq\mathcal{L}^{U}$, where $\mathcal{L}^{U}=\mathbb{E}_{q_{\phi}(y,z|x)}\left[\log\frac{p_{\theta}(x,y,z)}{q_{\phi}(y,z|x)}\right]$ can be simply expressed in terms of $\mathcal{L}^{S}$:

$\displaystyle\mathcal{L}^{U}(x;\theta,\phi^{y|x},\phi^{z|x,y})$

$\displaystyle=\textstyle\sum_{y\in\mathcal{Y}}q_{\phi^{y|x}}(y\mid x)\left(\mathcal{L}^{S}(x,y;\theta,\phi^{z|x,y})-\log q_{\phi^{y|x}}(y\mid x)\right).$

(6)

As with the unsupervised VAE, both terms in the objective can be computed via Monte Carlo sampling from the variational posterior, and gradients can be estimated via the reparameterization trick.

We develop a framework for jointly learning a strong generative model of features $x$, and making label-given-feature predictions $\hat{y}(x)$ of uncompromised quality, by requiring predictions to meet a user-specified quality threshold. Our prediction constrained training objective enables end-to-end estimation of all parameters while incorporating the same task-specific prediction rules and loss functions that will be used in heldout evaluation (“test”) scenarios. Our goals are similar to previous work on end-to-end approximate inference for task-specific losses with simpler probabilistic models (Lacoste-Julien et al., 2011, Stoyanov et al., 2011), but our approach yields simpler algorithms.

Generative model. Our generative model does not include labels $y$, only features $x$ and encodings $z$. Their joint distribution $p_{\theta}(x,z)$ factorizes as the unsupervised VAE of Eq. (1), and we also use the inference model $q_{\phi}(z\mid x)$ defined in Eq. (1). While M2 included the labels $y$ in its generative model (Kingma et al., 2014), our goals are different: we wish to make label-given-feature predictions, but we are not interested in label marginals or other distributions over $y$ that do not condition on $x$.

Label-from-feature prediction. To predict labels $y$ from features $x$, we use a predictor similar to the two-stage method of Sec. 2.2. We first sample an encoding $z\sim q_{\phi}(z|x)$ from the learned inference model, and then transform this encoding $z$ to a label via the predictor function $\hat{y}_{w}(z)$ with parameter $w$. By sharing random variable $z$, the generative model is involved in label-from-feature predictions.

Constrained PC objective. Unlike the two-stage model, our approach does not do post-hoc prediction with a previously learned generative model. Instead, we train the predictor simultaneously with the generative model via a new, prediction-constrained (PC) objective:

$$\max_{\theta,\phi^{z|x},w}\quad\sum_{x\in\mathcal{D}^{U}\cup\mathcal{D}^{S}}\mathcal{L}^{\text{VAE}}(x;\theta,\phi^{z|x}),~{}~{}\textnormal{subj. to:}~{}\;\frac{1}{M}\!\!\!\sum_{x,y\in\mathcal{D}^{S}}\underbrace{\mathbb{E}_{q_{\phi}(z|x)}[\ell_{S}(y,\hat{y}_{w}(z))]}_{\mathcal{P}(x,y;\phi^{z|x},w)}\leq\epsilon.$$

(8)

The constraint requires that any feasible solution achieve average prediction loss less than $\epsilon$ on the labeled training set. Both the loss function and scalar threshold $\epsilon>0$ can be set to reflect task-specific needs (e.g., classification must have a certain false positive rate or overall accuracy). The loss function may be any differentiable function, and need not equal the log-likelihood of discrete labels as assumed by previous work specialized to supervision of topic models (Hughes et al., 2018).

Unconstrained PC objective. Using the KKT conditions, we define an equivalent unconstrained objective that maximizes the unsupervised likelihood but penalizes inaccurate label predictions:

$\displaystyle\max_{\theta,\phi^{z|x},w}\quad\textstyle\sum_{x\in\mathcal{D}^{U}\cup\mathcal{D}^{S}}\mathcal{L}^{\text{VAE}}(x;\theta,\phi^{z|x})-\lambda\textstyle\sum_{x,y~{}\in\mathcal{D}^{S}}\mathcal{P}(x,y;\phi^{z|x},w).$

(9)

Here $\lambda$ is a positive Lagrange multiplier chosen to ensure that the target prediction constraint is achieved; smaller loss tolerances $\epsilon$ require larger penalty multipliers $\lambda$. This PC objective, and gradients for parameters $\theta,\phi,w$, can be estimated via Monte Carlo samples from $q_{\phi}(z\mid x)$.

Justification. While the PC objective of Eq. (9) may look superficially similar to Eq. (7), we emphasize two key differences. First, our objective couples a generative likelihood and a prediction loss via the shared variational parameters $\phi^{z|x}$. This makes both generative and discriminative performance depend on the same learned encoding $z$. (Later we show how to partition $z$ so some entries are discriminative, while others affect generative “style” only.) In contrast, the M2 objective uses a label-given-features conditional to make predictions that does not share any of its parameters $\phi^{y|x}$ with the supervised likelihood $\mathcal{L}^{S}$. Second, our objective is more affordable: no term requires an expensive marginalization over labels. This is key to scaling to big unlabeled datasets, and also enables tractable learning from datasets whose labels are continuous or multi-dimensional.

Hyperparameters. The major hyperparameter influencing PC training is the constraint multiplier $\lambda\geq 0$. Setting $\lambda=0$ leads to unsupervised maximum likelihood training (or MAP training, given priors on $\theta$) of a classic VAE. Setting $\lambda=1$ and choosing a probabilistic loss $-\log p(y\mid z)$ produces a “supervised VAE” that maximizes the joint likelihood $p_{\theta}(x,y)$. But as illustrated in Fig. 2, because features $x$ have much higher dimension than labels $y$, the resulting model may have weak predictive performance. Satisfying the strong prediction constraint of Eq. (8) typically requires $\lambda\gg 1$, and in practice we use validation data to select the best of several candidate $\lambda$ values. If a task motivates a concrete tolerance $\epsilon$, we can test an increasing sequence of $\lambda$ values until the constraint is satisfied.

We emphasize that although Eq. (9) is easier to optimize, we prefer to think of the constrained problem in Eq. (8) as the “primary” objective, because our applied goals are to satisfy discriminative quality first; a generative model that predicts poorly is not plausible. Furthermore, the constrained objective is far more natural for semi-supervised learning. The choice of $\epsilon$ need not be concerned by the relative sizes of the labeled and unlabeled datasets. In contrast, if either $|\mathcal{D}^{U}|$ or $|\mathcal{D}^{S}|$ changes, the value of $\lambda$ may need to change dramatically to reach the same prediction quality.

While the PC objective is effective given sufficient labeled data, it may generalize poorly when labels are very sparse (see Fig. 1). This fundamental problem arises because in the PC objective of Eq. (8), the parameters $w$ of the predictor $\hat{y}_{w}(z)$ are only directly informed by the labeled training data.

Revisiting the generative model, let $x\sim p_{\theta}(\cdot\mid z^{\prime})$ and $\bar{x}\sim p_{\theta}(\cdot\mid z^{\prime})$ be two observations sampled from the same latent code $z^{\prime}$. Even if the true label $y$ of $x$ is uncertain, we know that for this model to be useful for predictive tasks, $\bar{x}$ must have the same label as $x$. We formalize this relationship via a consistency constraint requiring label predictions for common-code data pairs $(x,\bar{x})$ to approximately match (see Fig. 3). As we show, this regularization may dramatically boost performance.

Given features $x$, our method predicts labels by sampling $z\sim q_{\phi}(z\mid x)$ from the approximate posterior, and then applying our predictor $\hat{y}_{w}(z)$. Alternatively, given $x$ we can first simulate alternative features $\bar{x}$ with matching code $z$ by sampling from the inference and generative models, and then predict the label associated with $\bar{x}$. We constrain the label predictions $y$ for $x$, and $\bar{y}$ for $\bar{x}$, to be similar via a consistency penalty function $\ell_{C}(y,\bar{y})$. For the classification tasks considered below, we use a cross-entropy consistency penalty $\ell_{C}$. Given this penalty, we constrain the maximum values of the following consistency costs on unlabeled and labeled examples, respectively:

	$\displaystyle\mathcal{C}^{U}(x;\theta,\phi,w)$	$\displaystyle\triangleq\mathbb{E}_{q_{\phi}(z|x)}\left[\mathbb{E}_{p_{\theta}(\bar{x}|z)}\left[\mathbb{E}_{q_{\phi}(\bar{z}|\bar{x})}\left[\ell_{C}(\hat{y}_{w}(z),\hat{y}_{w}(\bar{z}))\right]\right]\right],$		(10)
	$\displaystyle\mathcal{C}^{S}(x,y;\theta,\phi,w)$	$\displaystyle\triangleq\mathbb{E}_{q_{\phi}(z|x)}\left[\mathbb{E}_{p_{\theta}(\bar{x}|z)}\left[\mathbb{E}_{q_{\phi}(\bar{z}|\bar{x})}\left[\ell_{C}(y,\hat{y}_{w}(\bar{z}))\right]\right]\right].$		(11)

Consistent PC: Unconstrained objective. To train parameters, we apply our consistency costs to unlabeled and labeled feature vectors, respectively. The overall objective becomes:

$\displaystyle\max_{\theta,\phi,w}$

$\displaystyle\sum_{x\in\mathcal{D}^{U}\cup\mathcal{D}^{S}}\mathcal{L}^{\text{VAE}}(x;\theta,\phi)-\sum_{x\in\mathcal{D}^{U}}\gamma\mathcal{C}^{U}(x;\theta,\phi,w)+\!\sum_{x,y\in\mathcal{D}^{S}}-\lambda\mathcal{P}(x,y;\phi,w)-\gamma\mathcal{C}^{S}(x,y;\theta,\phi,w),$

where $\mathcal{L}$ is the unsupervised likelihood, $\mathcal{P}$ is the predictor loss, and $\mathcal{C}$ are the consistency constraints. Here, $\gamma>0$ is a scalar Lagrange multiplier for the consistency terms, with similar interpretation as $\lambda$.

Aggregate Label Consistency. For SSL applications, we find it is also useful to regularize our model with an aggregate label consistency constraint, which forces the distribution of label predictions for unlabeled data to be aligned with a known target distribution $\pi$. This discourages predictions on ambiguous unlabeled examples from collapsing to a single value. We define the aggregate consistency loss as: $\ell_{A}(\pi,\mathbb{E}_{x\sim\mathcal{D}^{U},z\sim q(z|x)}[\hat{y}_{w}(z)])$, and again use a cross-entropy penalty. If the target distribution of labels $\pi$ is unknown, we set it to the empirical distribution of the labeled data.

Related work on consistency. Recently popular SSL image classifiers focused on discriminative goals will train the weights of a CNN to minimize a modified objective that penalizes both label accuracy and a notions of consistency or smoothness on unlabeled data. Examples include consistency under adversarial perturbations (Miyato et al., 2019), label-invariant transformations (Laine & Aila, 2017), and when interpolating between training features (Berthelot et al., 2019). This regularization can deliver competitive discriminative performance, but does not meet our goal of generative modeling. Recently, Unsupervised Data Augmentation (UDA, Xie et al. (2020)), achieved state-of-the-art vision and text SSL classification by enforcing label consistency on augmented samples of unlabeled features. UDA relies on the availability of well-engineered augmentation routines for specific domains (e.g. image processing library transforms for vision or back-translation for text). In contrast, we learn a generative model that produces feature vectors for which predictions need to be consistent. Our approach is more applicable to new domains where advanced augmentation routines are not available.

In broader machine learning, “cycle-consistency” has improved generative adversarial methods for images (Zhu et al., 2017, Zhou et al., 2016) or biomedical data (McDermott et al., 2018). Others have developed cycle-consistent objectives for VAEs (Jha et al., 2018) which focus on consistency in code vectors $z$. In contrast, our work focuses on semi-supervised learning and enforces cycle consistency in labels $y$. Recently, Miller et al. (2019) developed discriminative regularization for VAEs. Their objective is not designed for SSL and uses a direct feature-to-label prediction model that must be consistent with reconstructed predictions. Our approach uses code-to-label prediction and SSL.

As our approach is applicable to any generative model, we can incorporate prior knowledge of the data domain to improve both generative and discriminative performance. We consider two examples: likelihoods that model noisy pixels, and explicit affine transformations to model image deformations.

Robust Likelihoods. Instead of a normal (or other common) likelihood we use a "Noise-Normal" likelihood to model images more robustly. We assume that pixel intensities $x$ have values in the interval $[-1,1]$ and rescale our datasets to match. Our Noise-Normal likelihood is defined as a 2-component mixture of a truncated Normal and a uniform “noise” distribution, with pixel-specific mixture weights. Define the standard normal PDF as $\phi(\cdot)$ and standard normal CDF as $\Phi(\cdot)$. We write the probability density function of the Noise-Normal distribution with parameters ($\rho$, $\mu$, $\sigma$) as:

$\displaystyle\mathcal{F}(x\mid\rho,\mu,\sigma)=\rho\left(\frac{\phi\left(\frac{x-\mu}{\sigma}\right)}{\Phi\left(\frac{1-\mu}{\sigma}\right)-\Phi\left(\frac{-1-\mu}{\sigma}\right)}\right)+\left(1-\rho\right)\left(\frac{1}{2}\right).$

(12)

Following Eq. (1), our (unsupervised) VAE now uses the revised generative and inference models:

$\displaystyle p_{\theta}(x,z)=\mathcal{N}(z\mid 0,I_{C})\cdot\mathcal{F}(x\mid\rho_{\theta}(z),\mu_{\theta}(z),\sigma_{\theta}(z)),\quad q_{\phi}(z\mid x)=\mathcal{N}(z\mid\mu_{\phi}(x),\sigma_{\phi}(x)).$

(13)

This approach allows our model to avoid sensitivity to outliers and noise in the observed images, and boosts the performance of our CPC method for SSL.

Spatial Transformer VAE. Our spatial transformer VAE retains the structure of a standard VAE, but reinterprets the latent code $z$ as two components. We denote the first 6 latent dimensions as $z_{t}$, and associate these with 6 affine transformation parameters capturing image translation, rotation, scaling, and shear. The generative model maps each value into a fixed range and creates an affine transformation matrix, $M_{t}$, by applying the transformations in a fixed order. The remainder of the latent code, $z_{*}$, is used to generate parameters for independent, per-pixel likelihoods. Assuming normal likelihoods, the output parameters for the pixel with coordinates $(i,j)$ are $\mu_{\theta}(z_{*})_{ij}$, $\sigma_{\theta}(z_{*})_{ij}$.

We re-orient our per-pixel likelihoods according to the affine transform $M_{t}$. The parameters of the likelihood for pixel $(i,j)$ will use the decoder outputs at coordinate $(i^{\prime},j^{\prime})$, where $M_{t}$ defines the affine mapping from $(i^{\prime},j^{\prime})$ to $(i,j)$. We apply this transformation in a (sub) differentiable way via a spatial transformer layer (Jaderberg et al., 2015) that takes as input $M_{t}$ and the appropriate parameter maps $\mu_{\theta}(z_{*})$, $\sigma_{\theta}(z_{*})$, and outputs a final set of parameters for the individual pixel likelihoods. As $(i^{\prime},j^{\prime})$ may not correspond to integer coordinates, we use bilinear interpolation over an appropriate representation of the likelihood parameters, and appropriately pad the size of the decoder output.

We further account for this special structure in our prediction and consistency constraints. For many applications, we have prior knowledge that small affine transforms should not affect the class of an image, and thus we can define consistency constraints that condition on $z_{*}$ but not $z_{t}$.