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