Adversarially Robust Representations with Smooth Encoders

1 INTRODUCTION . Representation learning is a fundamental problem in Machine learning and holds the promise to enable data-efficient learning and transfer to new tasks . Researchers working in domains like Computer Vision ( Krizhevsky et al. , 2012 ) and Natural Language Processing ( Devlin et al. , 2018 ) have already demonstrated the effectiveness of representations and features computed by deep architectures for the solution of other tasks . A case in point is the example of the FC7 features from the AlexNet image classification architecture that have been used for many other vision problems ( Krizhevsky et al. , 2012 ) . The effectiveness of learned representations has given new impetus to research in representation learning , leading to a lot of work being done on the development of techniques for inducing representations from data having desirable properties like disentanglement and compactness ( Burgess et al. , 2018 ; Achille & Soatto , 2017 ; Bengio , 2013 ; Locatello et al. , 2019 ) . Many popular techniques for generating representation are based on the Variational AutoEncoders ( VAE ) model ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . The use of deep networks as universal function approximators has facilitated very rapid advancements which samples generated from these models often being indistinguishable from natural data . While the quality of generated examples can provide significant convincing evidence that a generative model is flexible enough to capture the variability in the data distribution , it is far from a formal guarantee that the representation is fit for other purposes . In fact , if the actual goal is learning good latent representations , evaluating generative models only based on reconstruction fidelity and subjective quality of typical samples is neither sufficient nor entirely necessary , and can be even misleading . In this paper , we uncover the problematic failure mode where representations learned by VAEs exhibit over-sensitivity to semantically-irrelevant changes in data . One example of such problematic behaviour can be seen in Figure 1 . We identify a cause for this shortcoming in the classical Variational Auto-encoder ( VAE ) objective , the evidence lower bound ( ELBO ) , that fails to control the behaviour of the encoder out of the support of the empirical data distribution . We show this behaviour of the VAE can lead to extreme errors in the recovered representation by the encoder and is a key hurdle in the effective use of representations for data-efficient learning and transfer . To address this problem , we propose to augment the data with properties that enforce insensitivity of the representation with respect to families of transformations . To incorporate these specifications , we propose a regularization method that is based on a selection mechanism that creates a fictive data point by explicitly perturbing an observed true data point . For certain choices of parameters , our formulation naturally leads to the minimization of the entropy regularized Wasserstein distance between representations . We illustrate our approach on standard datasets and experimentally show that significant improvements in the downstream adversarial accuracy can be achieved by learning robust representations completely in an unsupervised manner , without a reference to a particular downstream task and without a costly supervised adversarial training procedure . It is clear that if learned representations are overly sensitive to irrelevant changes in the input ( for example , small changes in the pixels of an image or video , or inaudible frequencies added to an audio signal ) , models that rely on these representations are naturally susceptible to make incorrect predictions when inputs are changed . We argue that such specifications about the robustness properties of learned representations can be one of the tractable guiding features in the search for good representations . Based on these observations , we make the following contributions : 1 . We introduce a method for learning robust latent representations by explicitly targeting a structured model that admits the original VAE model as a marginal . We also show that in the case the target is chosen a pairwise conditional random field with attractive potentials , this choice leads naturally to the Wasserstein divergence between posterior distributions over the latent space . This insight provides us a flexible class of robustness metrics for controlling representations learned by VAEs . 2 . We develop a modification to training algorithms for VAEs to improve robustness of learned representations , using an external selection mechanism for obtaining transformed examples and by enforcing the corresponding representations to be close . As a particular selection mechanism , we adopt attacks in adversarial supervised learning ( Madry et al. , 2017 ) to attacks to the latent representation . Using this novel unsupervised training procedure we learn encoders with adjustable robustness properties and show that these are effective at learning representations that perform well across a variety of downstream tasks . 3 . We show that alternative models proposed in the literature , in particular β-VAE model used for explicitly controlling the learned representations , or Wasserstein Generative Adversarial Networks ( GANs ) can also be interpreted in our framework as variational lower bound maximization . 4 . We show empirically using simulation studies on MNIST , color MNIST and CelebA datasets , that models trained using our method learn representations that provide a higher degree of adversarial robustness even without supervised adversarial training . 2 GENERATIVE MODELS . Modern generative models are samplers p ( X|θ ) for generating realizations from an ideal target distribution π ( X ) , also known as the data distribution . In practice π ( X ) is unknown in the sense that it is hard to formally specify . Instead , we have a representative data set X , samples that are assumed to be conditionally independently drawn from the data distribution π ( X ) of interest . We will refer to the empirical distribution as π̂ ( X ) = 1|X | ∑ ξ∈X δ ( x − ξ ) . The goal is learning a parameter θ∗ such that p ( X|θ = θ∗ ) = ∫ dZp ( X|Z , θ = θ∗ ) p ( Z ) ≈ π̂ ( X ) , thereby also learning a generator . 2.1 FROM VAE TO SMOOTH ENCODERS . The VAE corresponds to the latent variable model p ( X|Z , θ ) p ( Z ) with latent variable Z and observation X . The forward model p ( X|Z = z , θ ) ( the decoder ) is represented using a neural network g with parameters θ , usually the mean of a Gaussian N ( X ; g ( z ; θ ) , vIx ) where v is a scalar observation noise variance and Ix is an identity matrix . The prior is usually a standard Gaussian p ( Z = z ) = N ( z ; 0 , Iz ) . The exact posterior over latent variables p ( Z|X = x , θ ) is approximated by a probability model q ( Z|X = x , η ) with parameters η . A popular choice here is a multivariate Gaussian N ( Z ; µ ( x ; η ) , Σ ( x ; η ) ) , where the mapping f such that ( µ , Σ ) = f ( x , η ) is chosen to be a neural network ( with parameters η to be learned from data ) . We will refer to the pair f , g as an encoder-decoder pair . Under the above assumptions , VAE ’ s are trained by maximizing the following form of the ELBO using stochastic gradient descent ( SGD ) , log p ( X = x|θ ) ≥ E { log p ( X = x|Z , θ ) } q ( Z|X=x , η ) −DKL ( q ( Z|X = x , η ) ||p ( Z ) ) ≡ Bx ( η , θ ) The gradient of the Kullback-Leibler ( KL ) divergence term above ( see A.1 ) is available in closed form . An unbiased estimate of the gradient of the first term can be obtained via sampling z from q using the reparametrization trick Kingma & Welling ( 2013 ) , aided by automatic differentiation . 2.2 A PROBLEM WITH THE VAE OBJECTIVE . Under the i.i.d . assumption , where each data point x ( n ) , for n = 1 . . . N is independently drawn from the model an equivalent batch ELBO objective can be defined ( See also E.1 ) as B ( η , θ ) ≡ 1 N N∑ n=1 Bx ( n ) ( η , θ ) = −DKL ( π̂ ( X ) q ( Z|X , η ) ||p ( X|Z , θ ) p ( Z ) ) + const ( 1 ) where the empirical distribution of observed data is denoted as π̂ . This form makes it more clear that the variational lower bound is only calculating the distance between the encoder and decoder under the support of the empirical distribution . To see how this locality leads to a fragile representation , we construct a VAE with discrete latents and observations . We let X ∈ { 1 , . . . , Nx } and Z ∈ { 1 , . . . , Nz } and define the following system of conditional distributions as the decoder and encoder models as : p ( X = i|Z = j ) ∝ exp ( 1 v ω ( mj − i/Nx ) ) q ( Z = j|X = i ) ∝ exp ( 1 σi ω ( µi − j/Nz ) ) where ω ( u ) = cos ( 2πu ) . These distributions can be visualized by heatmaps of probability tables where i and j are row and column indicies , respectively Figure 2 . This particular von-Mises like parametrization is chosen for avoiding boundary effects due to a finite latent and observable spaces . Latent O b se rv a ti o n Encoder Latent Decoder Probability Figure 2 : Example VAE model . ( left ) Heatmap of the encoder distribution ( darker colors referring to higher probability ) q ( Z = j|X = i ; µi , σi ) where each row i is a probability distribution over latents with a mode around µi and spread σi ( middle ) Heatmap of the decoder distribution p ( X = i|Z = j , mj , vj ) where each column j is a probability distribution with mode at mj and spread v. The prior p ( Z ) is chosen to be uniform and is not shown here . ( right ) The marginal model p ( X = i|m , v ) = ∑Nz j=1 p ( Z = j ) p ( X = i|Z = j , mj , v ) depicted as an histogram . The prior p ( Z ) is taken as uniform , and is not shown . Note that this parametrization emulates a high capacity network that can model any functional relationship between latent states and observations , while being qualitatively similar to a standard VAE model with conditionally Gaussian decoder and encoder functions . In reality , the true target density is not available but we would have a representative sample . To simulate this scenario , we sample a ’ dataset ’ from a discrete target distribution π ( X ) : this is merely a draw from a multinomial distribution , yielding a multinomial vector s with entries si that gives the count how many times we observe x = i . The results of such an experiment are depicted in Figure 3 ( a ) ( see caption for details ) . This picture reveals several important properties of a VAE approximation . 1 . After training , we observe that when j and j′ are close , the corresponding conditionals p ( X|Z = j ) and p ( X|Z = j′ ) are close ( hence corresponding decoder mean parameters mj and mj′ are close , hence ( see middle panel of Fig.3 ( a ) with the title P showing the decoder ) . This smoothness is perhaps surprising at a first sight : in this example , we could arbitrarily permute columns of the decoder and still get the same marginal distribution . Technically speaking , given a uniform prior p ( Z ) , the marginal likelihood p ( X|θ ) is entirely invariant with respect to permutations of the latent state . In fact if the encoder distribution wouldn ’ t be constrained we could also permute the columns of the encoder to keep the ELBO invariant . In the appendix E.2 , we provide an argument why the choice of an unimodal encoder model and optimization of the variational objective leads naturally to smooth decoder functions . 2 . The encoders found by the VAE on the other hand are not smooth at all , despite the fact that the model shows a relatively good fit . This behaviour alerts us about judging generative models only by the quality of the samples , by traversing the latent space and generating conditional samples from the decoder . The quality of the decoder seems to be not a proxy for the robustness of the representation . The fragility of representations is inherent from the ELBO objective . For the entire dataset , a batch ELBO that involves the counts si can be written as ELBO = − ∑ i ∑ j siq ( Z = j|Xa = i ) log siq ( Z = j|Xa = i ) p ( X = i|Z = j ) p ( Z = j ) ( 2 ) The last expression is proportional to the negative KL divergence between two tabular distributions : siq ( Z = j|Xa = i ) /L and p ( X = i|Z = j ) p ( Z = j ) . As such , whenever si is zero , the contribution of row i of the encoder distribution vanishes and the corresponding parameters µi and σi are not effecting the lower bound . In a sense , the objective does not enforce any structure on the encoder outside of the position of the data points in the training set . This figure shows that the outof-sample behaviour ( i.e. , for i where π̂ ( X ) = 0 ) the encoder is entirely initialization dependent , hence no learning takes place . We would also expect that the resulting representations would be fragile , in the sense that a small perturbation of an observation can result in a large change in the encoder output .

Summary: This paper studies the vulnerability of representations learned by variational auto-encoders (VAE). It first show that the learned representation of VAE is susceptible to small changes, similar to the adversarial examples in supervised learning setting. Then propose a regularization method, called smooth encoder, to improve the robustness of the representation. Experiments are conducted on several benchmark datasets to show the effectiveness of the method.