A DVERSARIALLY R OBUST R EPRESENTATIONS WITH S
MOOTH E NCODERS
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper studies the undesired phenomena of over-sensitivity of representations learned by deep networks to semantically-irrelevant changes in data.
We identify a cause for this shortcoming in the classical Variational Auto-encoder (VAE) objective, the evidence lower bound (ELBO).
We show that the ELBO fails to control the behaviour of the encoder out of the support of the empirical data distribution and this behaviour of the VAE can lead to extreme errors in the learned representation.
This is a key hurdle in the effective use of representations for data-efﬁcient learning and transfer.
To address this problem, we propose to augment the data with speciﬁcations that enforce insensitivity of the representation with respect to families of transformations.
To incorporate these speciﬁcations, we propose a regularization method that is based on a selection mechanism that creates a ﬁctive 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 signiﬁcant 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.
I NTRODUCTION
Representation learning is a fundamental problem in Machine learning and holds the promise to enable data-efﬁcient 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 classiﬁcation 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 signiﬁcant convincing evidence that a generative model is ﬂexible enough to capture the variability in the data distribution, it is far from a formal guarantee that the representation is ﬁt for other purposes.
In fact, if the actual goal is learning good latent representations, evaluating generative models only based on reconstruction ﬁdelity and subjective quality of typical samples is neither sufﬁcient 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 Vari- ational 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-efﬁcient 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 speciﬁcations, we propose a regularization method that is based on a selection mechanism that creates a ﬁctive 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 signiﬁcant 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.
Such examples suggest that either the representations Z a and Z b are very different (the encoder is not smooth), or the decoder is very sensitive to small changes in the representation (the decoder is not smooth), or both.
We identify the source of the problem primarily as the encoder and propose a practical solution.
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 speciﬁcations 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:
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 ﬁeld with attractive potentials, this choice leads naturally to the Wasserstein divergence between posterior distributions over the latent space.
This insight provides us a ﬂexible class of robustness metrics for controlling representations learned by VAEs.
We develop a modiﬁcation 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.
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.
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.
G ENERATIVE M ODELS
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| (cid:80) ξ ∈X δ ( x − ξ ) .
The goal is learning a parameter θ ∗ such that p ( X
| θ = θ ∗ )
= (cid:82) dZp ( X | Z, θ = θ ∗ ) p ( Z )
≈ ˆ π ( X ) , thereby also learning a generator.
2. F ROM VAE TO S MOOTH E NCODERS
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 ; θ ) , vI x ) where v is a scalar observation noise variance and I x is an identity matrix.
The prior is usually a standard Gaussian p ( Z = z ) =
N ( z ; 0 , I z ) .
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),
[Equation]
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 ﬁrst term can be obtained via sampling z from q using the reparametrization trick Kingma & Welling (2013), aided by automatic differentiation.
2. A P ROBLEM 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 deﬁned as
[Equation]
where the empirical distribution of observed data is denoted as ˆ π (See E.1 for a derivation).
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 , . . .
, N x } and Z ∈ { 1 , . . .
, N z } and deﬁne the following system of conditional distributions as the decoder and encoder models as:
[Equation]
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 ﬁnite latent and observable spaces.
[Figure]
O b s e r v a t i o n
[Figure]
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 s i 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.
[Figure]
After training, we observe that when j and j (cid:48) are close, the corresponding conditionals p ( X | Z = j ) and p ( X | Z = j (cid:48) ) are close (hence corresponding decoder mean parameters m j and m j (cid:48) are close, hence (see middle panel of Fig.3(a) with the title P showing the decoder).
This smoothness is perhaps surprising at a ﬁrst 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.
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 ﬁt.
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 s i can be written as
[Equation]
The last expression is proportional to the negative KL divergence between two tabular distributions: s i q ( Z = j | X a = i ) /L
and p ( X = i | Z = j ) p ( Z = j ) .
As such, whenever s i 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 ﬁgure 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.
R OBUST R EPRESENTATIONS WITH S MOOTH E NCODERS
In this section, we will adopt a strategy for training the encoder that is guaranteed not to change the original objective of the decoder when maximizing the lower bound while obtaining a smoother representation.
The key idea of our approach is that we assume an external selection mechanism that is able to provide new ﬁctive data point x (cid:48) in the vicinity of each observation in our data set x .
Here, “in the vicinity” means that we desire that the corresponding latent state of the original datapoint z
= f ( x ; η ) and the latent state of the ﬁctitious point z (cid:48) =
f ( x (cid:48) ; η ) should be close to each other in some sense.
Assuming the existence of such an external selection mechanism, we ﬁrst deﬁne the following augmented distribution
[Equation]
where
[Equation]
This is a pairwise conditional Markov random ﬁeld (CRF) model (Lafferty et al., 2001; Sutton &
McCallum, 2012), where we take c ( Z a , Z b ) as a pairwise cost function.
A natural choice here would be, for example, the Euclidean square distance (cid:107) Z a − Z b (cid:107) 2 .
Moreover, we choose a nonnegative coupling parameter γ ≥ 0 .
For any pairwise Q ( Z a , Z b ) distribution, the ELBO has the following form
[Equation]
It may appear that the SE has to maintain a pairwise approximation distribution Q ( Z a , Z b ) .
However, this turns out to be not necessary.
Given the encoder, the marginals of Q ( Z a , Z b ) are ﬁxed as Q a ( Z a ) =
q ( Z | X a = x, η ) and Q b ( Z b )
=
q ( Z | X b = x, η ) , so the only remaining terms that depend on the pair distribution are the ﬁnal two terms in (5) .
We note that this two terms are just the objective function of the entropy regularized optimal transport problem (Cuturi, 2013; Amari et al., 2017).
If we view Q ( Z a , Z b ) as a transport plan, the ﬁrst term is maximal when the expected cost is minimal while the second term is maximal when the variational distribution is factorized as
Q ( Z a , Z b ) =
Q a ( Z a ) Q b ( Z b ) .
In retrospection, this link is perhaps not that surprising as the Wasserstein distance, the solution of the optimal transport problem, is itself deﬁned as the solution to a variational problem (Solomon, 2018): Consider a set Γ of joint densities Q ( Z a , Z b ) with the property that Q has ﬁxed marginals
Q a ( Z a ) and Q b ( Z b ) , i.e.,
[Equation]
The Wasserstein divergence 1 , denoted by WD is deﬁned as the solution of the optimization problem with respect to pairwise distribution Q
[Equation]
where c ( Z a , Z b ) is a function that speciﬁes the ‘cost’ of transferring a unit of probability mass from Z a to Z b .
It is important to note that with our choice of the particular form of the variational distribution
Q ( Z a , Z b ) we can ensure that we are still optimizing a lower bound of the original problem.
We can achieve this by simply integrating out the X (cid:48) , effectively ignoring the likelihood term for the ﬁctive observations.
Our choice does not modify the original objective of the decoder due to the fact that the marginals are ﬁxed given η .
To see this, take the exponent of (5) and integrate over the unobserved X (cid:48) [Equation]
[Equation]
we name this lower bound B SE as the Smooth Encoder ELBO (SE-ELBO). The gradient of B SE with respect to the decoder parameters θ is identical to the gradient of the original VAE objective B .
This is intuitive as x (cid:48) is an artiﬁcially generated sample, we should use only terms that depend on x and not on x (cid:48) .
Another advantage of this choice is that it is possible to optimize the decoder and encoder concurrently as in the standard VAE. Only an additional term enters for the regularization of the encoder where
the marginals obtained via amortized inference q
( Z a | x a , η ) and q ( Z b | x b , η ) are forced to be close in a regularized Wasserstein distance sense, with the coupling strength γ .
Effectively, we are doing data augmentation for smoothing the representations obtained by the encoder without changing the actual data distribution.
In the appendix E.3, we also provide an argument about the smoothness of the corresponding encoder mapping, justifying the name.
The resulting algorithm is actually a simple modiﬁcation to the standard VAE and is summarized below:
Initialize η (0) , θ (0) for τ = 1 , 2 , . . .
do x a = GetData () x
b = Select ( x a ; L, (cid:15) )
(see Section 3.1) µ a , Σ a = f
( x a ; η ) , µ b , Σ b = f ( x b ; η ) (Compute Representation)
WD 2 ,γ ( η ) =
EntropyRegularizedWassersteinDivergence ( µ a , Σ a , µ b , Σ b , γ ) (see Appendix B.2) u ∼ N (0 , I ) (Reparametrization Trick)
E 1 ( η, θ )
= − 1 2 v (cid:107) x a − g ( µ a + Σ 1 / 2 a u ; θ ) (cid:107)
(Data Fidelity)
E 2 ( η ) = − 12 (cid:0) (cid:107) µ a (cid:107) 2 + (cid:107) µ b (cid:107)
+ Tr { Σ a + Σ b } (cid:1) (Prior Fidelity)
E ( η, θ ) =
E 1 ( η, θ ) + E 2 ( η ) − WD 2 ,γ ( η ) η ( τ ) , θ ( τ )
=
Optimization Step ( E, η
( τ − 1) , θ ( τ − 1) ) end
3. S ELECTION MECHANISM VIA ADVERSARIAL ATTACKS
Adversarial attacks are one of the most popular approaches for probing trained models in supervised tasks, where the goal of an adversarial attack is ﬁnding small perturbations to an input example that would maximally change the output, e.g., ﬂip a classiﬁcation decision, change signiﬁcantly a prediction (Szegedy et al., 2013).
The perturbed input is named as an adversarial example and these extra examples are used, along with the original data points, for training adversarially robust models (Madry et al., 2017; Kurakin et al., 2016).
As extra samples are also included, such a training procedure is referred as data augmentation .
However, in unsupervised learning and density estimation, data augmentation is not a valid approach as the underlying empirical distribution would be altered by the introducing new points.
However, as we let the encoder to target a different distribution than the actual decoder, we can actually use the extra, self generated samples to improve desirable properties of a given model.
Hence this approach could also be interpreted as a ’self-supervised’ learning approach where we bias our search for a ’good encoder’ and the data selection mechanism acts like a critique, carefully providing examples that should lead to similar representations.
In this paper we will restrict ourselves to Projected Gradient Descent (PGD) attacks popular in adversarial training Carlini & Wagner (2016) as a selection mechanism, where the goal of the attacker is ﬁnding a point that would introduce the maximum difference in the Wasserstein distance of the latent representation.
In other words, we implement our selection mechanism where the extra data point is found by approximately solving the following constrained optimization problem
[Equation]
This attack is assigned a certain iteration budget L for a given radius (cid:15) , that we refer as selection iteration budget and the selection radius , respectively.
We note a similar attack mechanism is proposed for generative models as described in (Kos et al., 2017), where one of the proposed attacks is directly optimizing against differences in source and target latent representations.
Note that our method is not restricted to a particular selection mechanism; indeed two inputs that should give a similar latent representation could be used as candidates.
E XPERIMENTS
Goal and Protocol
In our experiments, we have tested and compared the adversarial accuracy of representations learned using a VAE and our smooth encoder approach.
We adopt a two step experimental protocol, where we ﬁrst train encoder-decoder pairs agnostic to any downstream task. Then we ﬁx the representation, that is we freeze the encoder parameters and only use the mean of the encoder as the representation, then train a simple linear classiﬁer based on the ﬁxed representation using standard techniques. In this supervised stage, no adversarial training technique is employed.
Ideally, we hope that such an approach will provide a degree of adversarial robustness, without the need for a costly, task speciﬁc adversarial training procedure.
To evaluate the robustness of the resulting classiﬁer, for each data point in the test set, we search for an adversarial example using an untargeted attack that tries to change the classiﬁcation decision.
The adversarial accuracy is reported in terms of percentage of examples where the attack is not able to ﬁnd an adversarial example.
The VAE and SE decoder and encoder are implemented using standard MLP and ConvNet architectures.
The selection procedure for SE training is implemented as a projected gradient descent optimization (a PGD attack) with selection iteration budget of L iterations to maximize the Wasserstein distance between q ( Z | X = x ) and q ( Z | X = x + δ ) with respect to the perturbation δ where (cid:107) δ (cid:107) ∞ < (cid:15) .
Further details about the experiment can be found in the appendix C.1.
Results: We run simulations on ColorMNIST, MNIST and CelebA datasets.
The ColorMNIST is constructed from the MNIST dataset by coloring each digit artiﬁcially with all of the colors corresponding to the seven of the eight corners of the RGB cube (excluding black).
We present the results with the strongest attack we have experimented: a PGD attack with 100 iterations and 10 restarts.
We observe that for weaker attacks (such as 50 iterations with no restarts), the adversarial accuracy is typically much higher.
For the ColorMNIST dataset, the results are shown in Figure 4 where we test the adversarial accuracy of representations learned by our method and compare it to a VAE.
We observe that the adversarial accuracy of a VAE representation quickly drops towards zero while SE can maintain adversarial accuracy in both tasks.
In particular, we observe that for the arguably simpler color classiﬁcation task, we are able to obtain close to perfect adversarial test accuracy using representations learned by the VAE and SE.
However, when the classiﬁers are attacked using PGD, the adversarial accuracy quickly drops with increasing radius size, while the accuracy degrades more gracefully in the SE case.
In Figure 5, we show the robustness behaviour of the method for different architectures.
A ConvNet seems to perform relatively better than an MLP but these results show that the VAE representation [Figure] is not robust, irrespective of the architecture.
We have also carried out controlled experiments with random selection instead of the more costly untargetted adversarial attacks (See appendix C.1 Figure 7(a) for further results).
We observe some limited improvements with SE using random selection in adversarial accuracy compared to VAE but training a SE with adversarial selection seems to be much more effective.
We note that the selection iteration budget was lower ( L = 20 with no restarts) than the attack iteration budget ( 100 with 10 restarts) during evaluation.
It was not practical to train the encoder with more powerful selection attacks, thus it remains to be seen if the tendency of increased accuracy with increased iteration budgets would continue.
We also observe that essentially the same level of adversarial accuracy can be obtained with a small fraction of the available labels (See appendix C.1 Figure 8 for further results).
[Figure]
We have also repeated our experiments on the CelebA dataset, a large collection of high resolution face images labeled with 40 attribute labels per example.
We have used 17 of the attribute labels as the targets of 17 different downstream classiﬁcation tasks.
The results are shown in Table.2.
The results clearly illustrate that we can achieve much more robust representations than a VAE.
It is also informative to investigate speciﬁc adversarial examples to understand the failure modes.
In Figure 6 we show two illustrative examples from the CelebA.
Here we observe that attacks to the SE representations are much more structured and semantically interpretable.
In our exploratory investigations, we qualitatively observe that the reconstructions corresponding to the adversarial [Figure] examples are almost always legitimate face images with clearly recognizable features.
This also seems to support our earlier observation that VAE decoders are typically smooth while the encoders are inferring non-robust features.
Our approach seems to be a step towards obtaining more robust representations.
[Table]
Mustache
[Table]
Necklace
[Table]
Straight
[Table]
R ELATED W ORK
The literature on deep generative models and representation learning is quite extensive and is rapidly expanding.
There is a plethora of models, but some approaches have been quite popular in recent years:
Generative Adversarial Networks (GANs) and VAEs.
While the connection of our approach to VAE’s is evident, there is also a connection to GANs.
In the appendix, we provide the details where we show that a GAN decoder can be viewed as an instance of a particular smooth encoder.
Our method is closely related to the β -VAE (Higgins et al., 2017), used for controlling representations replaces the original variational objective (1) with another one for explicitly trading the data ﬁdelity with that of prior ﬁdelity.
In the appendix, we show that the method can be viewed as an instance of the smooth encoders.
Wasserstein distance minimization has been applied in generative models as an alternative objective for ﬁtting the decoder.
Following the general framework sketched in Bousquet et al.
(2017), the terms of the variational decomposition of the marginal likelihood can be modiﬁed in order to change the observation model or the regulariser.
For example, Wasserstein AutoEncoders (WAE) Tolstikhin et al.
(2017), Zhang et al.
or
sliced Wasserstein Autoencoders Kolouri et al.
(2018) propose to replace data ﬁdelity and/or the KL terms with a Wasserstein distance.
Our approach is different from these approaches as we do not propose to replace the likelihood as a fundamental principle for data ﬁtting.
In contrast, the Wasserstein distance formulation naturally emerges from the particular model choice and the corresponding variational approximation.
Our approach involves an adversarial selection step.
The word ’Adversarial’ is an overloaded term in generative modelling so it is important to mention differences between our approach.
Adversarial
Variational Bayes is a well known technique in the literature that aims to combine the empirical success of GANs with the probabilistic formulation of VAEs, where the limiting functional form of the variational distribution can be replaced by blackbox inference (Mescheder et al., 2017).
This approach also does not modify the original VAE objective, however, the motivation here is different as the aim is developing a more richer family.
In our view, for learning useful representations, when the decoder is unknown, the advantage of having a more powerful approximating family is not clear yet; this can even make the task of learning a good representation harder.
Adversarial Autoencoders (Makhzani et al., 2015), Adversarially Learned Inference (ALI) (Dumoulin et al., 2016) and BiGANs (Bidirectional GANs) (Donahue et al., 2016) are also techniques that combine ideas from GANs and VAEs for learning generative models.
The key idea is matching an encoder process q ( z | x ) p ( x ) and to the decoder process p ( z ) p ( x | z ) using an alternative objective, rather than by minimizing the KL divergence as done by the variational inference (see ( ?? )).
In this formulation, p ( x ) is approximated by the empirical data distribution, and p ( z ) is the prior model of a VAE.
The encoder q ( z | x ) and decoder p ( x | z ) are modelled using deep networks.
This approach is similar to Wasserstein autoencoders that propose to replace the likelihood principle.
The idea of improving VAEs by capturing the correlation structure between data points using MRFs and graphical models has been also been recently proposed (Tang et al., 2019) under the name
Correlated Variational Auto-Encoders (CVAEs).
Our approach is similar, however we introduce the correlation structure not between individual data points but only between true data points and artiﬁcially selected data points.
We believe that correctly selecting such a correlation structure of the individual data points can be quite hard in practice, but if such prior knowledge is available, CVAE can be indeed a much more powerful model than a VAE.
We note that a proposal for automatically learning such a correlation structure is also recently proposed by (Louizos et al., 2019).
D ISCUSSION AND C ONCLUSIONS
In this paper, we have introduced a method for improving robustness of latent representations learned by a VAE.
It must be stressed that our goal is not building the most powerful adversarially robust supervised classiﬁer, but obtaining a method for learning generic representations that can be used for several tasks; the tasks can be even unknown at the time of learning the representations.
While the nominal accuracy of an unsupervised approach is expected to be inferior to a supervised training method that is informed by extra label information, we observe that signiﬁcant improvements in adversarial robustness can be achieved by our approach that forces smooth representations.
A A PPENDIX A. KL D IVERGENCE
The KL divergence between two Gaussian distributions translates to a well known divergence in the parameters (in the general case this is a Bregman divergence)
[Equation]
where P a = N ( µ a , Σ a ) and P b = N ( µ b , Σ b ) are Gaussian densities with mean µ · and covariance matrix Σ · , and | · | denotes the determinant for a matrix argument, and Tr denotes the trace.
The
KL divergence consists of two terms, the ﬁrst term is the scale invariant divergence between two covariance matrices also known as a Itakuro-Saito divergence and the second term is a Mahalonobis distance between the means.
The KL divergence is invariant to the choice of parametrization or the choice of the coordinate system.
B O PTIMAL T RANSPORT AND W ASSERSTEIN DISTANCE
Consider a set Γ of joint densities Q ( Z a , Z b ) with the property that Q has ﬁxed marginals Q a ( Z a ) and Q b ( Z b ) , i.e.,
[Equation]
The Wasserstein divergence WD is deﬁned as the solution of the optimization problem with respect to pairwise distribution Q
[Equation]
11) where c ( z a , z b ) is a function that speciﬁes the ‘cost’ of transferring a unit of probability mass from z a to z b .
B.
[Equation] -W ASSERSTEIN DISTANCE W
The (cid:96) 2 -Wasserstein distance W 22 for two Gaussians has an interesting form.
The optimum transport plan, where the minimum of (11) is attained, is given
[Equation]
where Ψ = Σ a Σ 1 / 2 b (Σ 1 / 2 b Σ a Σ 1 / 2 b ) − 1 / 2 Σ 1 / 2 b .
It can be checked that this optimal Guassian density is degenerate in the sense that there exists a linear mapping between z a and z b :
[Equation]
where A 1 / 2 denotes the matrix square root, a symmetric matrix such that ( A 1 / 2 )
= A for a symmetric positive semideﬁnite matrix A .
The (cid:96) 2 -Wasserstein distance is the value attained by the optimum transport plan
[Equation]
B. E NTROPY R EGULARIZED (cid:96) 2 -W ASSERSTEIN DISTANCE
Entropy Regularized (cid:96) 2 -Wasserstein is the value attained by the minimizer of the following functional
[Equation]
where H is the entropy of the joint distribution Q .
Using the form in (12) subject to the semideﬁnite constraint Σ a − ΨΣ − 1 b Ψ [Equation]
[Equation]
The entropy of a Gaussian Q ( z a , z b ) is given by the Schur formula [Equation]
[Equation]
Here, D is the dimension of the vector ( z a , z b ) .
The entropy regularized problem has a solution where we need to minimize
[Equation]
Taking the derivative and setting to zero
[Equation]
we obtain a particular Matrix Ricatti equation [Equation] that gives us a closed form formula for the speciﬁc entropy regularized Wasserstein distance
[Equation]
For the case of two univariate Gaussians
[Equation]
i
[Equation] e., when the joint distribution has the form [Equation] the solution is given by the solution of the scalar quadratic equation.
[Equation]
We take the root that gives a feasible solution as the minimizer.
In the scalar case, this is the solution that satisﬁes Σ a − ψ 2 / Σ b ≥ 0 , or equivalently Σ a Σ b ≥ ψ
[Equation]
25) where we have deﬁned
[Equation]
It can be easily checked that the other root is infeasible.
For the scalar ψ case we obtain
[Equation]
C S UMMARY OF THE S MOOTH E NCODER ALGORITHM WITH FACTORIZED G AUSSIAN
Assume a factorized encoder distribution of form q ( Z a | x, η ) =
(cid:81)
D z k =1 N ( Z ka ; µ ka , Σ ka ) and q ( Z b | x (cid:48) , η )
= (cid:81)
D z k =1 N ( Z k b ; µ k b , Σ k b ) where D z is the dimension of the latent representation, and µ k and Σ k are the k ’th component of the output of a neural network with parameters η .
Similarly, x i denotes the i ’th component of the observation vector x of size D x .
For optimization, we need an unbiased estimate of the gradient of the SE-ELBO with respect to encoder parameters η and decoder parameters θ :
[Equation]
Given x , we ﬁrst select a ﬁctive sample x (cid:48) via a selection mechanism, in this case as an adversarial attack as explained in section 3.1.
representation and calculate the associated prediction
[Equation]
The terms of the SE-ELBO can be calculated as
E { log p ( x | Z a , θ ) } q ( Z a | X a = x,η )
[Equation]
where
[Equation]
C. E XPERIMENTAL D ETAILS AND F URTHER R ESULTS
We always train decoder-encoder pairs with identical architectures using both the standard VAE
ELBO and the SE ELBO with a ﬁxed γ .
Then, in each case by ﬁxing the encoder (that is essentially using the same representation) and by only using the mean activations of the encoders, we train linear classiﬁers using standard training for solving several downstream tasks.
For both encoder and decoder networks we use a 4 layer multi layer perceptron (MLP) and a convolutional network (ConvNET) architectures with 200 units of ReLU activations at each layer.
We carried out experiments with latent space dimensions of 32 , 64 and 128 , corresponding to an output sizes of an encoder with 64 , 128 and 256 units, with two units per dimensions to encode the mean and the log-variance parameters of a fully factorized Gaussian condition distribution.
The training is done using the Adam optimizer.
Each network (both the encoder and decoder) are randomly initialized and trained for 300 K iterations. Tasks MNIST Digit
Color-MNIST Color, Digit CelebA see table Training
Representation Dimension, dim( Z ) 32, 64, 128, VAE or SE Observation noise variance, v 0.25, 0.5, 1., 2. Architecture
MLP, ConvNET Training Coupling Strength γ 0.01, 0.1, 1, 5, 10, SE
Only
Selection PGD Radius (cid:15) 0.01, 0.1, 0.2, 0.
Selection PGD Iteration Budget L 1, 5, 10, 20,
Evaluation
[Figure]
[Figure]
100 (10 random restarts)
[Figure] [Figure]
v e r s a r i a l A
[Figure] c y [Figure] D R ELATED W ORK D. G ENERATIVE A DVERSARIAL N ETWORKS (GAN S )
GANs are presented as neural sampling models of observations x of form x = f ( ζ ; η ) where f is typically a deep neural network with parameters η [Figure] and ζ is a realization from some simple distribution p ( Z ) .
In the context of GANs, the function f is called a generator.
When the dimension of x is bigger than the dimension of ζ , the density p ( x ) induced by the transformation f is inevitably a degenerate distribution.
Since f is continuous, and it is concentrated on a subset of the data space
X f ≡ { x : ∃ ζ, x = f ( ζ ; η ) } .
Our use of letter f and parameters η is deliberate and we will illustrate in the sequel that the generator network of a GANs is actually analogous to a smooth encoder, where the roles of the latent variables and observations are switched, but we will ﬁrst review GANs.
To ﬁt a degenerate distribution to a dataset, the GAN approach adopts a strategy where the generator is co-trained with a second neural network d ( x ; w ) with parameters w with the following objective
[Equation]
[Equation] where D real ( x ) is the empirical data distribution.
This objective is (unfortunately) referred as an adversarial objective in the literature, not to be confused with adversarial attack mechanism in the context of supervised learning Madry et al.
The function d is called a discriminator .
After replacing expectations with ﬁnite sample averages, this objective enforces that in a dataset that contains both synthetically generated (fake) and real examples, the classiﬁcation function d should increase the correct classiﬁcation rate by discriminating fakes from real examples while the generator f should decrease the detection rate of fake examples.
When 0 ≤ d ( · ) ≤ 1 , which is the case for a classiﬁer, one can also write the objective as
[Equation]
where l ( x ; w ) =
log d ( x ; w ) .
This form also highlights an interesting alternative formulation and an interpretation in terms of optimal transport.
In fact, not long after the seminal work of Goodfellow et al.
(2014), the mechanism beyond the GAN objective and its direct connection to the theory of optimal transport has been recognized by the seminal paper Arjovsky et al.
(2017) where the problem is further framed as
[Equation]
[Equation] with the constraint that | l ( x ; w ) − l (¯ x ; w ) | ≤ (cid:107) c ( x, ¯ x ) (cid:107) , i.e. l is a Lipschitz function for some
L where (cid:107) c
( x, ¯ x ) (cid:107) ≤ L (cid:107) x − ¯ x (cid:107) .
Here, D fake (¯ x ; θ ) is the ﬁtted density of ¯ x = f ( ζ ; η ) .
This is the dual formulation of the optimal transport problem, that can be understood as an economic transaction between a customer and a shipment company.
Here, the difference l ( x ; w ) − l (¯ x ; w ) can be interpreted as the proﬁt made by the company for the shipment of one unit of mass from x and to ¯ x , and the Lipschitz condition ensures that it makes still sense for the customer to make use of the services of the company rather than simply doing the transport of her own Solomon (2018).
The customer wants to pay less, so she should minimize the proﬁt of the company.
This can be achieved by changing the desired delivery distribution D fake by adjusting θ , so that the transfer from the ﬁxed source distribution D real is minimized.
Ideally, when D fake = D real , there is nothing to transfer and no cost is incurred.
This objective also minimizes the Wasserstein distance between the actual data distribution D real and the fake data distribution D fake as given by the generator.
Once the GAN objective can be viewed as minimizing a particular Wasserstein distance, it is rather straightforward to view it as a maximizer of a particular ELBO corresponding to a particular smooth encoder, albeit in one where the positions of the observations and the latents are exchanged and a very large coupling coefﬁcient γ is chosen.
Moreover, the variational marginals have speciﬁc forms:
One marginal Q a ( X ) is chosen as the empirical data distribution and the other marginal is chosen as having the form of a neural sampler Q b ( X b )
= (cid:82) q ( X b | Z b , η ) p ( Z b )
dZ b .
The artiﬁcial extended target becomes
[Equation]
It can be seen that the ELBO in this case becomes
[Equation]
[Equation] Now, by taking the coupling γ sufﬁciently large, the coupling term dominates the lower bound and we obtain the Wasserstein minimization objective.
The random draws from p ( Z ) become the selection mechanism.
Moreover, the terms that depend on the artiﬁcial target p
( Z | X, θ ) become also irrelevant so in this regime the problem becomes just solving the optimal transport problem between
Q a and Q b .
A link between entropic GANs and VAEs is also pointed at in the literature, albeit for calculating a likelihood for GANs Balaji et al.
However, our motivations as well as the interpretation of the connection is quite different and we view the GAN decoder as an instance of the smooth encoder.
D. D ISENTANGLED R EPRESENTATIONS AND β -VAE
Targeting the encoder to an augmented distribution different than the decoder us the freedom to express some extensions of VAE in the same framework.
One of such extensions is the β -VAE, quite often used for controlling representations replaces the original variational objective (1) with the following objective
[Equation]
The justiﬁcation in the original paper Higgins et al.
(2017) is obtained from an implicit robustness criteria where D KL ( q ( Z | X a = x, η ) ||
p
( Z ))
< (cid:15) and β appears in a Lagrangian formulation.
Hoffman & Johnson (2016) have also provided an alternative justiﬁcation.
In our formulation, β can be simply interpreted as a dispersion term that is related to the number of points selected by the selection mechanism.
To see this, suppose the selection mechanism chooses β −
1 points x b,i where i = 1 . . .
β − 1 that are identical to the true observation x = x b,i =
x (cid:48)
i for i = 1 . . .
β − 1 .
[Equation]
Now
[Equation]
instead of integrating out Z (cid:48) 1: β − [Equation] we choose a variational distribution with identical marginals of form [Equation]
[Equation]
The variational lower bound becomes identical to the β -VAE objective as
[Equation]
where the last step follows due to the functional form of the variational distribution.
E T ECHNICAL R ESULTS E. B ATCH ELBO
In section 2.2, we have deﬁned a batch ELBO (2).
To see the connection to VAE ELBO (1) log p ( X = x | θ ) ≥ E { log p ( X = x
| Z, θ ) } q ( Z | X = x,η ) − D KL ( q ( Z | X = x, η ) || p
( Z ))
≡ B x ( η, θ ) we ﬁrst deﬁne the empirical data distribution π ( X ) = 1 N (cid:80) Ni =1 δ ( X − x i ) .
We can now write log p ( X = x | θ ) ≥ E { log p ( X = x
| Z, θ ) } q ( Z | X = x,η ) − D KL ( q ( Z | X = x, η ) || p
( Z ))
[Equation]
x
[Equation]
η
[Equation]
θ )
N
N
[Equation]
This result shows that the ELBO is minimizing the KL distance between one exact and one approximate factorization of the joint distribution p ( X, Z )
=
p ( X | Z, θ ) p ( Z )
≈ q ( Z | X, η ) π
( X ) .
E. W HY IS THE DECODER TYPICALLY SMOOTH AFTER THE VAE TRAINING ?
In the context of a VAE, the smoothness of the decoder is implicitly enforced by the highly constrained encoder distribution and the dynamics of an SGD based training.
In the sequel, we will illustrate that, if two latent coordinates are sufﬁciently close, the decoder mean mapping is forced to be bounded.
In a standard VAE, the encoder output for each data point is conditionally Gaussian as q ( Z | X = x ; η ) =
N ( f µ ( x ; η ) , f Σ ( x ; η )) .
The decoder is chosen as p ( X | Z = z ; η ) =
N ( g ( z ; θ ) , vI ) .
The decoder parameters θ under the ELBO depend only on the data ﬁdelity term (cid:107) x − g ( z ; θ ) (cid:107) 2 /v .
For a moment, assume that the encoder is ﬁxed and focus on a single data point x .
During training, a set of latent state vectors z i for i = 1 . . .
T are sampled from the conditionally Gaussian encoder distribution.
When the dimension of the latent space D z is large, these samples z i will be with high probability on the typical set.
The typical set of a nondegenerate Gaussian distribution is approximately the surface of a Mahalanobis ball, a compact hyper-ellipsoid M ( x ) centered at f µ ( x ; η ) with scaling matrix f Σ ( x ; η ) 1 / 2 .
If we assume that the training procedure is able to reduce the error in the sense that (cid:107) x − g
( z i ; θ ) (cid:107) ≤
E for all z i where E is a bound on the error magnitude for z i sampled from the encoder, the decoder is forced to give the same output for each point approximately on M ( x ) .
For a point z a drawn from q ( Z | X = x ; η ) we have
[Equation]
with high probability where K = f Σ ( x ; η ) − 1 and (cid:107) x
(cid:107) K ≡ √ x (cid:62) Kx .
For a point z b independently drawn from q ( Z | X = x ; η ) , by the triangle inequality we have
[Equation]
where the Mahalanobis distance
[Equation]
where λ min is the smallest eigenvalue of the covariance matrix.
Hence the distance is also bounded when the variance is not degenerate and minimum distance will be on the order of (cid:107) z a − z b (cid:107) ≈ 2 √ D z λ min
so we expect the ratio to be bounded
[Equation]
[Equation]
We see that the ELBO objective enforces the decoder to be invariant on the typical set of q ( Z | X = x ; η ) , where most of the probability mass is concentrated.
Now, for each data point x , the corresponding latent space hyper-ellipsoid M ( x ) are forced to be large in the sense of having a large determinant by the entropy term of the encoder that promotes large log-determinant.
The size of M ( x ) is also controlled by the prior ﬁdelity term, avoiding blowing up.
Hence the union ∪ x ∈X M ( x ) , where X is the dataset, will approximately cover the latent space when the encoder has converged and on each hyper-ellipsoid M ( x ) the decoder will be enforced to be smooth.
E. S MOOTHNESS OF THE SMOOTH ENCODER
In this section we show that the smooth encoder training forces a small Lipschitz constant for the encoder mean mapping.
To simplify the argument, we will assume that the variance mapping of the encoder would be a constant function that does not vary with x , i.e., f Σ ( x ; η )
= Σ( η ) .
The latter assumption could be removed by considering a metric on the joint space of the means and covariance.
Using the adversarial selection mechanism, during training we solve the following problem using
PGD:
[Equation]
Assuming that PGD ﬁnds the global maximum at the boundary of the (cid:15) -ball where (cid:107) x − x ∗ (cid:107) p
= (cid:15) , under constant variance assumption for the encoder we can see that the Wasserstein divergence simply becomes the square distance between mean mappings
[Equation]
We know that the SE ELBO objective has to minimize this distance for any coupling term γ so the procedure actually tries to reduce the local Lipschitz constant L ( x ) around data point x
[Equation]
and promotes smoothness where E is an upper bound on the change in the representation (cid:107) f µ ( x ; η ) − f µ ( x ∗ ; η ) (cid:107) ≤ E .