Consistency Regularization for Variational Auto-Encoders

1 Introduction . 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 . 2 Method . 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 . 3 Related Work . 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 .

Summary: 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.