Consistency Regularization for Variational Auto-Encoders
Year: 2022

Abstract:
Variational auto-encoders (vaes) are a powerful approach to unsupervised learning. They enable scalable approximate posterior inference in latent-variable models using variational inference (vi). A vae posits a variational family parameterized by a deep neural network—called an encoder—that takes data as input. This encoder is shared across all the observations, which amortizes the cost of inference. However the encoder of a vae has the undesirable property that it maps a given observation and a semantics-preserving transformation of it to different latent representations. This “inconsistency" of the encoder lowers the quality of the learned representations, especially for downstream tasks, and also negatively affects generalization. In this paper, we propose a regularization method to enforce consistency in vaes. The idea is to minimize the Kullback-Leibler (kl) divergence between the variational distribution when conditioning on the observation and the variational distribution when conditioning on a random semantic-preserving transformation of this observation. This regularization is applicable to any vae. In our experiments we apply it to four different vae variants on several benchmark datasets and found it always improves the quality of the learned representations but also leads to better generalization. In particular, when applied to the nouveau variational auto-encoder (nvae), our regularization method yields state-of-the-art performance on mnist, cifar-10, and celeba. We also applied our method to 3D data and found it learns representations of superior quality as measured by accuracy on a downstream classification task. Finally, we show our method can even outperform the triplet loss, an advanced and popular contrastive learning-based method for representation learning. 1

heading: 1 Introduction
text: Variational auto-encoders (vaes) have significantly impacted research on unsupervised learning. They have been used in several areas, including density estimation (Kingma & Welling, 2013; Rezende et al., 2014), image generation (Gregor et al., 2015), text generation (Bowman et al., 2015; Fang et al., 2019), music generation (Roberts et al., 2018), topic modeling (Miao et al., 2016; Dieng et al., 2019), and recommendation systems (Liang et al., 2018). Vaes have also been used for different representation learning problems such as semi-supervised learning (Kingma et al., 2014), anomaly detection (An & Cho, 2015; Zimmerer et al., 2018), language modeling Bowman et al. (2015), active learning (Sinha et al., 2019), continual learning (Achille et al., 2018), and motion prediction of agents (Walker et al., 2016). This widespread application of vae representations makes it critical that we focus on improving them.
vaes extend deterministic auto-encoders to probabilistic generative modeling. The encoder of a vae parameterizes an approximate posterior distribution over latent variables of a generative model. The encoder is shared between all observations, which amortizes the cost of posterior inference. Once fitted, the encoder of a vae can be used to obtain low-dimensional representations of data, (e.g. for
1Code for this work can be found at https://github.com/sinhasam/CRVAE
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
downstream tasks.) The quality of these representations is therefore very important to a successful application of vaes. Researchers have looked at ways to improve the quality of the latent representations of vaes, often tackling the so-called latent variable collapse problem—in which the approximate posterior distribution induced by the encoder collapses to the prior over the latent variables (Bowman et al., 2015; Kim et al., 2018; Dieng et al., 2018; He et al., 2019; Fu et al., 2019).
In this paper, we focus on a different problem pertaining to the latent representations of vaes for image data. Indeed, the encoder of a fitted vae tends to map an image and a semantics-preserving transformation of that image to different parts in the latent space. This “inconsistency" of the encoder affects the quality of the learned representations and generalization. We propose a method to enforce consistency in vaes. The idea is simple and consists in maximizing the likelihood of the images while minimizing the Kullback-Leibler (kl) divergence between the approximate posterior distribution induced by the encoder when conditioning on the image, on one hand, and its transformation, on the other hand. This regularization technique can be applied to any vae variant to improve the quality of the learned representations and boost generalization performance. We call a vae with this form of regularization, a consistency-regularized variational auto-encoder (cr-vae). Figure 1 illustrates the inconsistency problem of vaes and how cr-vaes address this problem on mnist. The red dots are representations of a few images and the blue dots are the representations of their transformations. We applied semantics-preserving transformations: rotation, translation, and scaling. The vae maps each image and its transformation to different parts in the latent space as evidenced by the long arrows connecting each pair (a). Even when we include the transformed images to the data and fit the vae the inconsistency problem still occurs (b). The cr-vae does not suffer from the inconsistency problem; it maps each image and its transformation to nearby areas in the latent space, as evidenced by the short arrows connecting each pair (c).
In our experiments (see Section 4), we apply the proposed technique to four vae variants, the original vae (Kingma & Welling, 2013), the importance-weighted auto-encoder (iwae) (Burda et al., 2015), the β-vae (Higgins et al., 2017), and the nouveau variational auto-encoder (nvae) (Vahdat & Kautz, 2020). We found, on four different benchmark datasets, that cr-vaes always yield better representations and generalize better than their base vaes. In particular, consistency-regularized nouveau variational auto-encoders (cr-nvaes) yield state-of-the-art performance on mnist and cifar-10. We also applied cr-vaes to 3D data where these conclusions still hold.

heading: 2 Method
text: We consider a latent-variable model pθ(x, z) = pθ(x|z) · p(z), where x denotes an observation and z is its associated latent variable. The marginal p(z) is a prior over the latent variable and pθ(x|z)
is an exponential family distribution whose natural parameter is a function of z parameterized by θ, e.g. through a neural network. Our goal is to learn the parameters θ and a posterior distribution over the latent variables. The approach of vaes is to maximize the evidence lower bound (elbo), a lower bound on the log marginal likelihood of the data,
Lvae = elbo = Eqφ(z|x) [ log ( pθ(x, z)
qφ(z|x)
)] (1)
where qφ(z|x) is an approximate posterior distribution over the latent variables. The idea of a vae is to let the parameters of the distribution qφ(z|x) be given by the output of a neural network, with parameters φ, that takes x as input. The parameters θ and φ are then jointly optimized by maximizing a Monte Carlo approximation of the elbo using the reparameterization trick (Kingma & Welling, 2013).
Consider a semantics-preserving transformation t(x̃|x) of data x (e.g. rotation or translation for images.) A good representation learning algorithm should provide similar latent representations for x and x̃. This is not the case for the vae that maximizes Equation 1 and its variants. Once fit to data, the encoder of a vae is unable to yield similar latent representations for a data x and its tranformation x̃ (see Figure 1). This is because there is nothing in Equation 1 that forces this desideratum.
We now propose a regularization method that ensures consistency of the encoder of a vae. We call a vae with such a regularization a cr-vae. The regularization proposed is applicable to many variants of the vae such as the iwae (Burda et al., 2015), the β-vae (Higgins et al., 2017), and the nvae (Vahdat & Kautz, 2020). In what follows, we use the standard vae, the one that maximizes Equation 1, as the base vae to regularize to illustrate the method. Consider an image x. Denote by t(x̃|x) the random process by which we generate x̃, a semanticspreserving transformation of x. We draw x̃ from t(x̃|x) as follows:
x̃ ∼ t(x̃|x) ⇐⇒ ∼ p( ) and x̃ = g(x, ). (2)
Here g(x, ) is a semantics-preserving transformation of the image x, e.g. translation with random length drawn from p( ) = U [−δ, δ] for some threshold δ. A cr-vae then maximizes
Lcr-vae(x) = Lvae(x) + Et(x̃|x) [Lvae(x̃)]− λ · R(x, φ) (3)
where the regularization termR(x, φ) is
R(x, φ) = Et(x̃|x) [kl (qφ(z|x̃)||qφ(z|x))] . (4)
Maximizing the objective in Equation 3 maximizes the likelihood of the data and their augmentations while enforcing consistency throughR(x, φ). MinimizingR(x, φ), which only affects the encoder (with parameters φ), forces each observation and the corresponding augmentations to lie close to each other in the latent space. The hyperparameter λ ≥ 0 controls the strength of this constraint. The objective in Equation 3 is intractable but we can easily approximate it using Monte Carlo with the reparameterization trick. In particular, we approximate the regularization term with one sample from t(x̃|x) and make the dependence to this sample explicit using the notationR(x, x̃, φ). Algorithm 1 illustrates this in greater detail. Although we show the application of consistency regularization using the vae that maximizes the elbo, Lvae(·) in Equation 3 can be replaced with any vae objective.

heading: 3 Related Work
text: Applying consistency regularization to vaes, as we do in this paper, has not been previously explored. Consistency regularization is a widely used technique for semi-supervised learning (Bachman et al., 2014; Sajjadi et al., 2016; Laine & Aila, 2016; Miyato et al., 2018; Xie et al., 2019). The core idea behind consistency regularization for semi-supervised learning is to force classifiers to learn representations that are insensitive to semantics-preserving changes to images, so as to improve classification of unlabeled images. Examples of semantics-preserving changes used in the literature include rotation, zoom, translation, crop, or adversarial attacks. Consistency is often enforced by minimizing the L2 distance between a classifier’s logit output for an image and the logit output for its semantics-preserving transformation (Sajjadi et al., 2016; Laine & Aila, 2016), or by minimizing
Algorithm 1: Consistency Regularization for Variational Autoencoders
input :Data x, consistency regularization strength λ, latent space dimensionality K Initialize parameters θ, φ for iteration t = 1, 2, . . . do
Draw minibatch of observations {xn}Bn=1 for n = 1, . . . , B do
Transform the data: n ∼ p( n) and x̃n = T (xn, n) Get variational mean and variance for the data:
µn = W >NN(xn;φ) + a and σn = softplus(Q>NN(xn;φ) + b)
Get S samples from the variational distribution when conditioning on xn: η(s) ∼ N (0, I) and z(s)n = µn + η(s) · σn for s = 1, . . . , S Get variational mean and variance for the transformed data: µ̃n = W
>NN(x̃n;φ) + a and σ̃n = softplus(Q>NN(x̃n;φ) + b) Get S samples from the variational distribution when conditioning on x̃n:
η(s) ∼ N (0, I) and z̃(s)n = µ̃n + η(s) · σ̃n for s = 1, . . . , S end Compute Lvae(x): Lvae(x) ≈ 1B ∑B n=1 1 S ∑S s=1 [ log pθ(xn, z (s) n )− log qφ(z(s)n |xn)
] Compute Lvae(x̃): Lvae(x̃) ≈ 1B ∑B n=1 1 S ∑S s=1 [ log pθ(x̃n, z̃ (s) n )− log qφ(z̃(s)n |x̃n)
] Compute KL consistency regularizer:
R(x, x̃, φ) = 12 ∑K k=1 ( σ̃2nk+(µ̃nk−µnk) 2 σ2nk − 1 + 2 · log σnkσ̃nk ) Compute final loss:
Lcr-vae(x) = Lvae(x) + Lvae(x̃)− λ · R(x, x̃, φ) Backpropagate through L(x, θ, φ) = −Lcr-vae(x) and take a gradient step for θ and φ
end
the kl divergence between the classifier’s label distribution induced by the image and that of its tranformation (Miyato et al., 2018; Xie et al., 2019).
More recently, consistency regularization has been applied to generative adversarial networks (gans) (Goodfellow et al., 2014). Indeed Wei et al. (2018) and Zhang et al. (2020) show that applying consistency regularization on the discriminator of a gan—also a classifier—can substantially improve its performance.
The idea we develop in this paper differs from the works above in two ways. First, it applies consistency regularization to vaes for image data. Second, it leverages consistency regularization, not in the label or logit space, as done in the works mentioned above, but in the latent space.
Although different, consistency regularization for vaes relates to works that study ways to constrain the sensitivity of encoders to various perturbations. For example, denoising auto-encoders (daes) and their variants (Vincent et al., 2008, 2010) corrupt an image x into x′, typically using Gaussian noise, and then minimize the distance between the reconstruction of x′ and the un-corrupted image x. The motivation is to learn representations that are insensitive to the added noise. Our work differs in that we do not constrain the decoder to recover the original image from the corrupted image but, rather, to constrain the encoder to recover the latent representation of the original image from the corrupted image via a kl divergence minimization constraint. Contractive auto-encoders (caes) (Rifai et al., 2011) share a similar goal with cr-vaes. A cae is an auto-encoder whose encoder is constrained by minimizing the norm of the Jacobian of the output of the encoder with respect to the input image. This norm constraint on the Jacobian forces the representations learned by the encoder to be insensitive to changes in the input. Our work differs in several main ways. First, cr-vaes are not deterministic auto-encoders, contrary to caes. We can easily sample from a cr-vae, as for any vae, which is not the case for a cae. Second, a cae does
not apply transformations to the input image, which limits the sensitivities it can learn to limit to those exhibited in the training set. Finally, caes use the Jacobian to impose a consistency constraint, which are not as easy to compute as the kl divergence we use on the variational distribution induced by the encoder.

heading: 4 Empirical Study
text: In this section we show that a cr-vae improves the learned representations of its base vae and positively affects generalization performance We also show that the proposed regularization method is amenable to different vae variants by applying it not only to the original vae but also to the iwae, the β-vae, and the nvae. We showcase the importance of the KL regularization term by conducting an ablation study. We found that only regularizing with data augmentation improves performance but that accounting for the kl term (λ > 0) further improves the quality of the learned representations and generalization.
We will conduct three sets of experiments. In the first experiment, we will apply the regularization method proposed in this paper to standard vaes such as the original vae, the iwae, and the β-vae. We use mnist, omniglot, and celeba as datasets for this experiment. For celeba, we choose the 32x32 resolution for this experiment. Our results show that adding consistency regularization always improves upon the base vae, both in terms of the quality of the learned representations and generalization. We conduct an ablation study and also report performance of the different vae variants above when they are fitted with the original data and their augmentations. The results from this ablation highlight the importance of setting λ > 0.
In the second set of experiments we apply our method to a large-scale vae, the latest nvae (Vahdat & Kautz, 2020). We use mnist, cifar-10, and celeba as datasets for this experiment. We increased the resolution for the celeba dataset for this experiment to 64x64. We reach the same conclusions as for the first sets of experiments; cr-vaes improve the learned representations and generalization of their base vaes. In this particular setting, the cr-nvae achieves state-of-the-art generalization performance on both mnist and cifar-10. This state-of-the-art performance couldn’t be reach simply by training the nvae with augmentations, as our results show. Finally, in a third set of experiments, we apply our regularization technique to a 3D point-cloud dataset called ShapeNet (Chang et al., 2015). We adapt a high-performing auto-encoding method called FoldingNet (Yang et al., 2018) to its vae counterpart and apply the method we described in this paper to that vae variant on the ShapeNet dataset. We found that adding consistency regularization yields better learned representations.
We next describe in great detail the set up for each of these experiments and the results showcasing the usefulness of the regularization method we propose in this paper.
4.1 Application to standard vaes on benchmark datasets We apply consistency regularization, as described in this paper, to the original vae, the iwae, and the β-vae. We now describe the set up and results for this experiment. Datasets. We study three benchmark datasets that we briefly describe below. We first consider mnist. mnist is a handwritten digit recognition dataset with 60, 000 images in the training set and 10, 000 images in the test set (LeCun, 1998). We form a validation set of 10, 000 images randomly sampled from the training set.
We also consider omniglot, a handwritten alphabet recognition dataset (Lake et al., 2011). This dataset is composed of 19, 280 images. We use 16, 280 randomly sampled images for training and 1, 000 for validation and the remaining 2, 000 samples for testing.
Finally we consider celeba. It is a dataset of faces, consisting of 162, 770 images for training, 19, 867 images for validation, and 19, 962 images for testing (Liu et al., 2018). We set the resolution to 32x32 for this experiment.
Transformations t(x̃|x). We consider three transformations variants for image data t(x̃|x). The first randomly translates an image [−2, 2] pixels in any direction. The second transformation randomly
rotates an image uniformly in [−15, 15] degrees clockwise. Finally the third transformation randomly scales an image by a factor uniformly sampled from [0.9, 1.1].
Evaluation metrics. The regularizationmethodwe propose in this paper is mainly aimed at improving the learned representations of vaes. To assess these representations we use three metrics: mutual information, number of active latent units, and accuracy on a downstream classification task. We also evaluate the effect of the proposed method on generalization to unseen data. For that we also report negative log-likelihood. We define each of these metrics next.
Mutual information (MI). The first quality metric is the mutual information I(z;x) between the observations and the latents under the joint distribution induced by the encoder,
I(z;x) = Epd(x) [KL(qφ(z|x)||p(z))− kl(qφ(z)||p(z))] (5)
where pd(x) is the empirical data distribution and qφ(z) is the aggregated posterior, the marginal over z induced by the joint distribution defined by pd(x) and qφ(z|x). The mutual information is intractable but we can approximate it with Monte Carlo. Higher mutual information corresponds to more interpretable latent variables.
Number of active latent units (AU). The second quality metrics we consider is the number of active latent units (AU). It is defined in Burda et al. (2015) and measures the “activity" of a dimension of the latent variables z. A latent dimension is “active" if
Covx(Eu∼qφ(u|x)) > δ (6)
where δ is a threshold defined by the user. For our experiments we set δ = 0.01. The higher the number of latent active units, the better the learned representations.
Accuracy on downstream classification. This metric is calculated by fitting a given vae, taking the learned representations for each data in the test set and computing the accuracy from the prediction of the labels of the images in that same test set by a classifier fitted on the training set. This metric is only applicable to labelled datasets.
Negative log-likelihood. We use negative held-out log-likelihood to assess generalization. Consider an unseen data x∗, its negative held-out log-likelihood under the fitted model is
log pθ(x ∗) = − log ( Eqφ(z|x∗) [ pθ(x ∗, z)
qφ(z|x∗)
]) . (7)
This is intractable and we approximate it using Monte Carlo,
log pθ(x ∗) ≈ − log 1
S S∑ s=1 pθ(x ∗, z(s)) qφ(z(s)|x∗) (8)
where z(1), . . . , z(S) ∼ qφ(z|x∗). Settings. The vaes are built on the same architecture as Tolstikhin et al. (2017). The networks are trained with the Adam optimizer with a learning rate of 10−4 (Kingma & Ba, 2014) and trained for 100 epochs with a batch size of 64. We set the dimensionality of the latent variables to 50, therefore the maximum number of active latent units in the latent space is 50. We found λ = 0.1 to be best according to cross-validation using held-out log-likelihood and exploring the range [1e−4, 1.0] datasets. In an ablation study we explore λ = 0. For the β-vae we set λ = 0.1 · β and study both β = 0.1 and β = 10, two regimes under which the β-vae performs qualitatively very differently (Higgins et al., 2017). All experiments were done on a GPU cluster consisting of Nvidia P100 and RTX. The training took approximately 1 day for most experiments.
Results. Table 1 shows that on all the three benchmark datasets all the different vae variants we studied, consistency regularization as developed in this paper always improves the quality of the learned representations as measured by mutual information and the number of active latent units. These results are confirmed by the numbers shown in Table 2 where cr-vaes always lead to better accuracy on downstream classification.
We proposed consistency regularization as a way to improve the quality of the learned representations. Incidentally, Table 3 also shows that it can improve generalization as measured by negative loglikelihood.
Ablation Study. We now look at the impact of each factor that goes into the regularization method we introduced in this paper using mnist. We test the impact of the regularization term λ and the impact of the choice of augmentation on all metrics. Table 4 and Table 5 show the results.
Table 4 shows that even small consistency regularization (a small λ value) results in improvement over the base vae but that a large enough λ value can hurt performance. Table 5 shows that rotations and translations are more important than scaling, but the combination of all three augmentations works best for cr-vaes. Comparison to Contrastive Learning. We look at how cr-vaes compare against a popular and advanced contrastive-learning-based technique, the triplet loss (Schroff et al., 2015) using mnist. Table 6 shows that the cr-vae outperforms the triplet loss on both generalization performance and quality of learned representations. Table 6 also confirms existing literature showing simply applying augmentations can outperform complex contrastive learning-based methods such as the triplet loss (Kostrikov et al., 2020; Sinha & Garg, 2021).
4.2 Application to the large-scale nvae on benchmark datasets Along with standard VAE variants, we also experiment with a large scale state-of-the-art vae, the nvae(Vahdat & Kautz, 2020). Similar to before, we simply add consistency regularization using the image-based augmentations techniques to the NVAE model and experiment on benchmark datasets: mnist (LeCun, 1998), cifar-10 (Krizhevsky et al., 2009) and celeba (Liu et al., 2018). The results for large scale generative modeling are tabulated in Table 8 and Table 7, where we see that using cr-nvae we are able to learn representations that yield better accuracy on downstream
classification and set new state-of-the-art values on each of the datasets, improving upon the baseline log-likelihood values. This shows the ability of consistency regularization to work at scale on challenging generative modeling tasks.

heading: 4.3 Application to the FoldingNet on 3D point-cloud data
text: Along with working with image data, we additionally experiment with 3D point cloud data using a FoldingNet Yang et al. (2018) and the ShapeNet dataset Chang et al. (2015) which consists of 55 distinct object classes. FoldingNet learns a deep AutoEncoder to learn unsupervised representations from the point cloud data. To add consistency regularization, we first substitute the AutoEncoder to a
vae by adding the KL term from the ELBO to the baseline FoldingNet. We then add the additional consistency regularization KL term to the latent space of FoldingNet.
For the ShapeNet point cloud data, we perform data augmentation using a similar scheme to what we did for the previous experiments, we randomly translate, rotate and add jitter to the (x, y, z) coordinates of the point cloud data. We follow the same scheme detailed in FoldingNet (Yang et al., 2018).
We train both the FoldingNet turned in a vae and the CR-FoldingNet with these augmentations. To train CR-FoldingNet, we additionally apply the consistency regularization term as proposed in Equation 3. The results on the validation set for reconstruction (as measured by Chamfer distance) and accuracy are shown in Table 9.
We also visualize the point clouds reconstructions and interpolations between 3 different object classes using a CR-FoldingNet in Figure 2. We perform 4 interpolation steps for each of the objects, to highlight the interpretable learned latent space. Additionally, we perform the same interpolation on the baseline FoldingNet model. We show these interpolations in the appendix.

heading: 5 Conclusion
text: We proposed a simple regularization technique to constrain encoders of vaes to learn similar latent representations for an image and a semantics-preserving transformation of the image. The idea consists in maximizing the likelihood of the pair of images while minimizing the kl divergence between the variational distribution induced by the encoder when conditioning on the image on one hand, and its transformation, on the other hand. We applied this technique to several vae variants on several datasets, including a 3D dataset. We found it always leads to better learned representations and also better generalization to unseen data. In particular, when applied to the nvae, the regularization technique we developed in this paper yields state-of-the-art results on mnist and cifar-10.
Broader Impact
In this paper, we propose a simple method that performs a KL-based consistency regularization scheme using data augmentation for vaes. The broader impact of the study includes practical applications such as graphics and computer vision applications. The method we propose improves the learned representations of vaes, and as an artifact, also improves their generalization to unseen data. In this regard, any implications of vaes also apply to this work. For example, the generative model fit by a vae may be used to generate artificial data such as images, text, and 3D objects. Biases may arise as a result of poor data selection. Furthermore, text generated from generative systems may amplify harmful speech contained in the data. However, the method we propose can also improve the performance of vaes when used in certain practical domains as we discussed in the introduction of the paper.

heading: 6 Acknowledgements
text: We thank Kevin Murphy, Ben Poole, and Augustus Odena for their comments on this work.

Review summaries:
The paper proposes a novel method of using data augmentation when training Variational Autoencoders. Given a family of semantics-preserving transformations, the idea is to enforce the invariance of latent code w.r.t. these transformations by adding a KL penalty term. This penalty improves the quality of learned latent representation and of the generative model significantly.  

The paper improves VAEs by ensuring that the learned encoder maps augmented inputs $\tilde{x}$ and unaugmented inputs $x$ to the same representations. To do so, they minimize the KL divergence between the encoder conditioned on $\tilde{x}$ and $x$. They empirically show that such consistency regularised VAE (CR-VAE) achieves very strong density modeling results.

This paper proposes a study on enforcing semantic consistency on VAEs. The authors show that the current VAEs are sensitive to semantical-preserving transformations and it will lower the quality of representations and generalization. To solve this issue, the authors propose a new algorithm, which first transforms a training sample into multiple semantically similar samples and then minimize the KL-divergence between the encoded distributions of original and new samples. Experiments on both 2D and 3D datasets show that the new algorithm is effective in increasing representation quality of VAEs.

The paper  investigates a problematic behavior of current VAE models that lies in the inconsistency of current amortized encoder networks in the sense that the latent representations of semantically similar transformations (translation, rotation etc) of image data are dissimilar. Moreover, it proposes a tractable variant of the ELBO objective that is based on data augmentation and adds a KL-regularization term that penalizes the discrepancy in the posterior of the original and the transformed image data. The VAEs trained with this objective exhibit significant benefits (in terms of the utilization of the latent space and  marginal likelihood ) when compared to naive data augmentation demonstrating sota results in image generation and classification tasks.