Title: ADVERSARIALLY ROBUST REPRESENTATIONS WITH SMOOTH ENCODERS

Abstract: 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-efficient learning and transfer. To address this problem, we propose to augment the data with specifications 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.

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.

Figure 1 : An illustration of the intrinsic fragility of VAE representations. Outputs from a Variational Autoencoder with encoder f and decoder g parametrized by η and θ, respectively, trained on CelebA. Conditioned on the encoder input X a = x a the decoder output X = g(f (x a )) = (g • f )(x a ) is shown on the top row. When the original example is perturbed with a carefully selected vector d such that X b = X a + d with d ≤ , the output X turns out to be perceptually very different. 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 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; θ), 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),

log p(X = x|θ) ≥ E {log p(X = x|Z, θ)} q(Z|X=x,η) -D KL (q(Z|X = x, η)||p(Z)) ≡ B x (η, θ)

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 B x (n) (η, θ) = -D KL (π(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, . . . , N x } and Z ∈ {1, . . . , N z } and define the following system of conditional distributions as the decoder and encoder models as:

p(X = i|Z = j) ∝ exp 1 v ω (m j -i/N x ) q(Z = j|X = i) ∝ exp 1 σ i ω (µ i -j/N z )

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. (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, m j , v j ) where each column j is a probability distribution with mode at m j 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, m j , 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 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.

Q P Data True Model Q P Data True Model (a) (b)

Figure 3 : (a) Result by optimizing the ELBO for a VAE that illustrates the fragility of the encoder. Subfigure with the title 'Data' (π(X)) is a random sample from the true target 'Target' (π(X)) on the right. The resulting encoder q(Z|X) and decoder p(X|Z) are shown as 'Q' and 'P', respectively. The vertical and horizontal axes correspond to latents Z and observations X respectively. Integrating over the decoder distribution using a uniform prior p(Z) over the latents, we obtain the model distribution 'Model' p(X) = Z p(X|Z)p(Z). (b) The results obtained by a smooth encoder. Both the decoder and the representation (encoder) are more smooth while essentially having a similar fitting quality.

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 m j and m j 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 s i can be written as

ELBO = - i j s i q(Z = j|X a = i) log s i q(Z = j|X a = i) p(X = i|Z = j)p(Z = j) (2)

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

3. ROBUST REPRESENTATIONS WITH SMOOTH ENCODERS:

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 fictive data point x 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 fictitious point z = f (x ; η) should be close to each other in some sense. Assuming the existence of such an external selection mechanism, we first define the following augmented distribution

EQUATION

where

ψ(Z a , Z b ) = exp(-γ 2 c(Z a , Z b ))p(Z a )p(Z b )

. This is a pairwise conditional Markov random field (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 Z a -Z b 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 fixed 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 final two terms in (4). 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 first 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 defined as the solution to a variational problem (Solomon, 2018) :

EQUATION

The Wasserstein divergence1 , denoted by WD is defined as the solution of the optimization problem with respect to pairwise distribution Q

WD[c](Q a , Q b ) = inf Q∈Γ c(Z a , Z b )Q(Z a , Z b )dZ a dZ b (6)

where c(Z a , Z b ) is a function that specifies 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 , effectively ignoring the likelihood term for the fictive observations. Our choice does not modify the original objective of the decoder due to the fact that the marginals are fixed given η. To see this, take the exponent of ( 4) and integrate over the unobserved X

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 is an artificially generated sample, we should use only terms that depend on x and not on x . 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 modification to the standard VAE and is summarized below:

Initialize η (0) , θ (0) for τ = 1, 2, . . . do x a = GetData(), x b = Select(x a ; L, ) (see Section 3.1) µ a , Σ a = f (x a ; η), µ b , Σ b = f (x b ; η) (Compute Representation) WD 2,γ (η) = EntropyRegularizedWassersteinDivergence(µ a , Σ a , µ b , Σ b , γ) (see Apdx. B.2) u ∼ N (0, I) (Reparametrization Trick) E 1 (η, θ) = -1 2v x a -g(µ a + Σ 1/2 a u; θ) 2 (Data Fidelity) E 2 (η) = -1 2 µ a 2 + µ b 2 + Tr{Σ a + Σ b } (Prior Fidelity) E(η, θ) = E 1 (η, θ) + E 2 (η) -WD 2,γ (η) η (τ ) , θ (τ ) = Optimization Step(E, η (τ -1) , θ (τ -1) ) end 3.1 SELECTION 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 finding small perturbations to an input example that would maximally change the output, e.g., flip a classification decision, change significantly 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 finding 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

x = x + arg max d: d p ≤ WD(q(Z|X = x, η), q(Z|X = x + d, η))

This attack is assigned a certain iteration budget L for a given radius , 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.

4. EXPERIMENTS:

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 first train encoder-decoder pairs agnostic to any downstream task. Then we fix 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 classifier based on the fixed 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 specific adversarial training procedure. To evaluate the robustness of the resulting classifier, for each data point in the test set, we search for an adversarial example using an untargeted attack that tries to change the classification decision. The adversarial accuracy is reported in terms of percentage of examples where the attack is not able to find 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 δ ∞ < . 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 artificially 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 classification task, we are able to obtain close to perfect adversarial test accuracy using representations learned by the VAE and SE. However, when the classifiers 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 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 Fig- ure 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). 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 classification 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 specific 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 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 The SE representation is attacked by a perturbation that can clearly be identified as drawing a beard on the image. In this case, the attack is able to fool the classifier and the generated image from the representation is that of a person with beard and mustache. In the second example (c), the VAE representation seems to be attacked by exploiting the non-smooth nature of the encoder by mapping the latent representation to the one in the vicinity of a clearly different person with the desired features, as can be seen from the corresponding reconstruction. In contrast, in the SE case (d), an unsuccessful attack that adds a much more structured perturbation. From the reconstruction it is evident that a latent feature is attacked that seems to control the receding hairline.

EQUATION

are inferring non-robust features. Our approach seems to be a step towards obtaining more robust representations. 

VAE:



5. RELATED WORK:

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 fidelity with that of prior fidelity. 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 fitting the decoder. Following the general framework sketched in Bousquet et al. (2017) , the terms of the variational decomposition of the marginal likelihood can be modified in order to change the observation model or the regulariser. For example, Wasserstein AutoEncoders (WAE) Tolstikhin et al. (2017) , Zhang et al. (2019) or sliced Wasserstein Autoencoders Kolouri et al. (2018) propose to replace data fidelity 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 fitting. 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. 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 artificially 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) .

6. DISCUSSION AND CONCLUSIONS:

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 classifier, 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 significant improvements in adversarial robustness can be achieved by our approach that forces smooth representations.

A APPENDIX:

A.1 KL DIVERGENCE

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 first 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 OPTIMAL TRANSPORT AND WASSERSTEIN DISTANCE

EQUATION

The Wasserstein divergence WD is defined as the solution of the optimization problem with respect to pairwise distribution Q

EQUATION

where c(z a , z b ) is a function that specifies the 'cost' of transferring a unit of probability mass from z a to z b .

B.1 2 -WASSERSTEIN DISTANCE W:

The 2 -Wasserstein distance W 2 2 for two Gaussians has an interesting form. The optimum transport plan, where the minimum of ( 10) 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 :

z a (z b ) = µ a + Σ a Σ 1/2 b (Σ 1/2 b Σ a Σ 1/2 b ) -1/2 Σ -1/2 b (z b -µ b )

where A 1/2 denotes the matrix square root, a symmetric matrix such that (A 1/2 ) 2 = A for a symmetric positive semidefinite matrix A. The 2 -Wasserstein distance is the value attained by the optimum transport plan

W 2 2 (P a , P b ) = µ a -µ b 2 2 + Tr Σ a + Σ b -2 Σ 1/2 b Σ a Σ 1/2 b 1/2 (12) B.2 ENTROPY REGULARIZED 2 -WASSERSTEIN DISTANCE

Entropy Regularized 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 (11) subject to the semidefinite constraint Σ a -ΨΣ -1 b Ψ 0

Tr (z a -z b ) (z a -z b ) = -2 Tr(Ψ) + const (14)

The entropy of a Gaussian Q(z a , z b ) is given by the Schur formula

H[Q(z a , z b )] = D 2 log(2πe) + 1 2 log |Σ b ||Σ a -ΨΣ -1 b Ψ | (15)

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 specific entropy regularized Wasserstein distance

W 2 2,γ (N (m a , Σ a ), N (m b , Σ b )) = m a -m b 2 2 + Tr{Σ a + Σ b -2Ψ} (19) WD 2,γ (N (m a , Σ a ), N (m b , Σ b )) ≡ γ 2 W 2 2,γ (N (m a , Σ a ), N (m b , Σ b )) (20) - D 2 log(2πe) - 1 2 log |Σ b ||Σ a -ΨΣ -1 b Ψ | (21)

For the case of two univariate Gaussians, i.e., when the joint distribution has the form

Q(Z a , Z b ) = N m a m b , Σ a ψ ψ Σ b

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 satisfies

EQUATION

where we have defined

u γ (Σ a , Σ b ) = 1 + 4γ 2 Σ b Σ a 1/2

It can be easily checked that the other root is infeasible. For the scalar ψ case we obtain

WD 2,γ (N (m a , Σ a ), N (m b , Σ b )) = γ 2 m a -m b 2 2 + Σ a + Σ b - 1 2 (u γ (Σ a , Σ b ) -1) + 1 2 log(u γ (Σ a , Σ b ) + 1) - 1 2 log(2Σ b Σ a ) -log(2π) -1 C SUMMARY OF THE SMOOTH ENCODER ALGORITHM WITH FACTORIZED GAUSSIAN Assume a factorized encoder distribution of form q(Z a |x, η) = Dz k=1 N (Z k a ; µ k a , Σ k a ) and q(Z b |x , η) = Dz 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 θ:

B SE (η, θ) = E {log p(X = x|Z a , θ)} q(Za|xa,η) + E {log p(Z a )} q(Za|xa,η) + E {log p(Z b )} q(Z b |x b ,η) - γ 2 E {c(Z a , Z b )} q(Za,Z b |xa,x b ,η) + H(q(Z a , Z b |x a , x b , η))

Given x, we first select a fictive sample x via a selection mechanism, in this case as an adversarial attack as explained in section 3.1.

Sample a latent representation and calculate the associated prediction

z a ∼ q(Z a |X a = x, η) = N (Z a ; µ a , Σ a ) x = g(z a ; η)

The terms of the SE-ELBO can be calculated as

E {log p(x|Z a , θ)} q(Za|Xa=x,η) ≈ - D x 2 log 2πv - 1 2v Dx i=1 (x i -xi ) 2 E {log p(Z a )} q(Za|Xa=x,η) = - 1 2 Dz k=1 ((µ k a ) 2 + Σ k a ) - D z 2 log 2π E {log p(Z b )} q(Z b |X b =x ,η) = - 1 2 Dz k=1 ((µ k b ) 2 + Σ k b ) - D z 2 log 2π WD 2,γ = γ 2 E Z a -Z b 2 q(Za,Z b |xa,x b ,η) -H(q(Z a , Z b |x a , x b , η))

where

u γ (Σ a , Σ b ) = 1 + 4γ 2 Σ b Σ a WD 2,γ = γ 2 Dz k=1 µ k a -µ k a 2 2 + Σ k a + Σ k b - 1 2 Dz k=1 (u γ (Σ k a , Σ k b ) -1) -log(u γ (Σ k a , Σ k b ) + 1) + log(2Σ k b Σ k a ) -D z log(2π) -D z

C.1 EXPERIMENTAL DETAILS AND FURTHER RESULTS:

We always train decoder-encoder pairs with identical architectures using both the standard VAE ELBO and the SE ELBO with a fixed γ. Then, in each case by fixing the encoder (that is essentially using the same representation) and by only using the mean activations of the encoders, we train linear classifiers 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 300K iterations.

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 first review GANs.

To fit 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

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. (2017) 

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

with the constraint that |l(x; w) -l(x; w)| ≤ c(x, x) , i.e. l is a Lipschitz function for some L where c(x, x) ≤ L x -x . Here, D fake (x; θ) is the fitted 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 profit 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 profit of the company. This can be achieved by changing the desired delivery distribution D fake by adjusting θ, so that the transfer from the fixed 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 coefficient γ is chosen. Moreover, the variational marginals have specific 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 ) = q(X b |Z b , η)p(Z b )dZ b .

The artificial extended target becomes

EQUATION

It can be seen that the ELBO in this case becomes

EQUATION

Now, by taking the coupling γ sufficiently 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 artificial 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. (2018) . 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.2 DISENTANGLED REPRESENTATIONS 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

log p(X = x|θ) ≥ E {log p(X = x|Z, θ)} q(Z|Xa=x,η) -βD KL (q(Z|X a = x, η)||p(Z))(30)

The justification 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)) < and β appears in a Lagrangian formulation. Hoffman & Johnson (2016) have also provided an alternative justification.

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

EQUATION

∝ dX 1:β-1 dZdZ 1:β-1 p(X|Z, θ)p(Z)

EQUATION

Now, instead of integrating out Z 1:β-1 , we choose a variational distribution with identical marginals of form

EQUATION

The variational lower bound becomes identical to the β-VAE objective as

EQUATION

+ β-1 i=1 E {log p(Z i )} q(Z i |X b,i =x,η) + H(Q(Z, Z 1:β-1 )) (36) = E {log p(X = x|Z, θ)} q(Z|Xa=x,η) -βD KL (q(Z|X a = x, η)||p(Z)) (37)

where the last step follows due to the functional form of the variational distribution.

E TECHNICAL RESULTS:



E.1 BATCH ELBO:

In section 2.2, we have defined 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 first define the empirical data distribution π

(X) = 1 N N i=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)) ≡ B x (η, θ) 1 N N i=1 log p(X = x i |θ) = 1 N N i=1 E {log p(X = x i |Z, θ)} q(Z|X=xi,η) - 1 N N i=1

E {log q(Z|X = x i , η)} q(Z|X=xi,η)

+ 1 N N i=1

E {log p(Z)} q(Z|X=xi,η)

= E {log p(X|Z, θ)} q(Z|X,η)π(X) -E {log q(Z|X, η)} q(Z|X,η)π(X)

+E {log p(Z)} q(Z|X,η)π(X)

-E {log π(X)} q(Z|X,η)π(X) + E {log π(X)} π(X)

= -D KL (q(Z|X, η)π(X)||p(X|Z, θ)p(Z)) + const

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.2 WHY 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 sufficiently 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 fidelity term x -g(z; θ) 2 /v.

For a moment, assume that the encoder is fixed 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 x -g(z i ; θ) ≤ 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 z a -f µ (x; η) K ≈ D z with high probability where K = f Σ (x; η) -1 and x K ≡ √ x Kx.

For a point z b independently drawn from q(Z|X = x; η), by the triangle inequality we have g(z a ; θ) -g(z b ; θ) ≤ 2E (38)

where the Mahalanobis distance

2 D z ≈ z a -z b K ≤ 1 √ λ min z a -z b

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 z a -z b ≈ 2 √ D z λ min so we expect the ratio to be bounded g(z a ; θ) -g(z b ; θ) / z a -z b ≤ E/ D z λ min (39) 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 fidelity 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.3 SMOOTHNESS 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:

x * = arg max

x : x -x p ≤ WD(q(Z|X = x, η), q(Z|X = x , η))

Assuming that PGD finds the global maximum at the boundary of the -ball where x -x * p = , under constant variance assumption for the encoder we can see that the Wasserstein divergence simply becomes the square distance between mean mappings WD(q(Z|X = x, η), q(Z|X = x * , η)) = f µ (x; η) -f µ (x * ; η) 2 2

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 L(x) = f µ (x; η) -f µ (x * ; η)

x -x * p ≤ E and promotes smoothness where E is an upper bound on the change in the representation f µ (x; η)f µ (x * ; η) ≤ E.

:

We use the term divergence to distinguish the optimal transport cost from the corresponding metric. This distinction is reminiscent to the distinction between Euclidian divergence •

and the Euclidian distance •

Figures and Tables:

Figure2: 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, m j , v j ) where each column j is a probability distribution with mode at m j 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) =

Figure 5: Simulation Results on MNIST. The goal is illustrating the effect of the architecture (MLP and ConvNet). In all the examples, the SE is trained by a selection radius a budget of 50 PGD iterations. The linear classifier is always trained without any adversarial training, by fixing the encoder parameters. The blue dot, green triangle and red squares correspond to the standard VAE and SE trained with a selection radius of 0.1 and 0.2 respectively. The dashed and dotted lines show the nominal accuracy of the VAE and SE when there are no attacks (The SE nominal accuracy is virtually identical for different selection radii, hence only a single level is shown.)

Figure 6: Qualitative results on CelebA. Attacks to downstream tasks (a), (b) 'Mustache' classification, (c), (d) 'Bald' classification. (a) In the VAE case, the attack is successful but the perturbation does not have a visible structure. (b)The SE representation is attacked by a perturbation that can clearly be identified as drawing a beard on the image. In this case, the attack is able to fool the classifier and the generated image from the representation is that of a person with beard and mustache. In the second example (c), the VAE representation seems to be attacked by exploiting the non-smooth nature of the encoder by mapping the latent representation to the one in the vicinity of a clearly different person with the desired features, as can be seen from the corresponding reconstruction. In contrast, in the SE case (d), an unsuccessful attack that adds a much more structured perturbation. From the reconstruction it is evident that a latent feature is attacked that seems to control the receding hairline.

Simulation Results on ColorMNIST. The goal is the comparison of adversarial accuracy of VAE representations with SE representations trained with a selection radius of 0.1 and 0.2 and a selection budget of 20 PGD iterations. Vertical axis shows the adversarial accuracy as a function of attack radius. The dashed and dotted lines show the nominal accuracy of the VAE and SE when there are no attacks (The SE nominal accuracy is virtually identical for different selection radii, hence only a single level is shown.)

Table content: <table><tr><td>Adversarial Accuracy (%)</td><td>20 40 60 80 100</td><td colspan="4">Color Mnist (Arch: Conv, Task:Color, Sel. Iter. Budget:20) VAE SE, SelRad = 0.1 SE, SelRad = 0.2 VAE Nominal SE Nominal</td><td>Adversarial Accuracy (%)</td><td>20 40 60 80 100</td><td>Color Mnist (Arch: Conv, Task:Digit, Sel. Iter. Budget:20) VAE SE, SelRad = 0.1 SE, SelRad = 0.2 VAE Nominal SE Nominal</td></tr><tr><td/><td>0</td><td/><td/><td/><td/><td/><td>0</td></tr><tr><td/><td/><td>0.01</td><td>0.10</td><td>0.20 Attack Radius</td><td>0.30</td><td/><td/><td>0.01</td><td>0.10</td><td>0.20 Attack Radius</td><td>0.30</td></tr><tr><td/><td/><td/><td colspan="2">(a) Task: Color</td><td/><td/><td/><td>(b) Task: Digit</td></tr><tr><td>Figure 4:</td><td/><td/><td/><td/><td/><td/><td/></tr></table>

. The function d is called a discriminator. After replacing expectations with finite sample averages, this objective enforces that in a dataset that contains both synthetically generated (fake) and real examples, the classification function d should increase the correct classification 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 classifier, one can also write the objective as

Table content: <table><tr><td>min θ</td><td>max w</td><td>E {l(x; w)} Dreal -E {l(f (ζ; η); w))} p(ζ)</td></tr></table>

Peer review (iclr2020):
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. Overall I find the idea interesting and the experimental results promising. The following are my detailed comments. a About the theory The illustration of the problem in VAE is interesting. However, one missing point is to theoretically quantify the effect of the proposed regularization (in some simple cases). In particular, it is claimed that the regularization could make the encoder smoother and the experimental results clearly justifies it. What would be better is to show in which sense/measure the encoder is smoother and provide some theoretical guarantee about it. (for instance smaller Lipschitz constant?) b About the Experiment The experimental section is clear and promising. I just have one question about the evaluation on the robustness of the VAE representation. In particular, a linear classifier is concatenated right after the VAE representation and it is not clear to me where it is concatenated. Is it right after the layer of mu and Sigma or in later layers? If it is in the later layers, the VAE is outputting a distribution, then how does the accuracy measured? Minor comment: I think it is unnecessary to introduce the new term selection strategy because it is just an adversarial training with respect to a different loss. In particular, the loss is the Wasserstein distance between the latent space vectors instead of a supervised loss. For simplicity, it could be just named as latent space adversarial training. (this is just a suggestion, which will not change my decision)