Recent studies have demonstrated the vulnerability of deep convolutional neural networks against adversarial examples. Inspired by the observation that the intrinsic dimension of image data is much smaller than its pixel space dimension and the vulnerability of neural networks grows with the input dimension, we propose to embed high-dimensional input images into a low-dimensional space to perform classification. However, arbitrarily projecting the input images to a low-dimensional space without regularization will not improve the robustness of deep neural networks. We propose a new framework, Embedding Regularized Classifier (ER-Classifier), which improves the adversarial robustness of the classifier through embedding regularization. Experimental on several benchmark datasets show that, our proposed framework achieves state-of-the-art performance against strong adversarial attack methods. Deep neural networks (DNNs) have been widely used for tackling numerous machine learning problems that were once believed to be challenging. With their remarkable ability of fitting training data, DNNs have achieved revolutionary successes in many fields such as computer vision, natural language progressing, and robotics. However, they were shown to be vulnerable to adversarial examples that are generated by adding carefully crafted perturbations to original images. The adversarial perturbations can arbitrarily change the network's prediction but often too small to affect human recognition . This phenomenon brings out security concerns for practical applications of deep learning. Two main types of attack settings have been considered in recent research (Goodfellow et al.; a; ;): black-box and white-box settings. In the black-box setting, the attacker can provide any inputs and receive the corresponding predictions. However, the attacker cannot get access to the gradients or model parameters under this setting; whereas in the white-box setting, the attacker is allowed to analytically compute the model's gradients, and have full access to the model architecture and weights. In this paper, we focus on defending against the white-box attack which is the harder task. Recent work presented both theoretical arguments and an empirical one-to-one relationship between input dimension and adversarial vulnerability, showing that the vulnerability of neural networks grows with the input dimension. Therefore, reducing the data dimension may help improve the robustness of neural networks. Furthermore, a consensus in the highdimensional data analysis community is that, a method working well on the high-dimensional data is because the data is not really of high-dimension . These high-dimensional data, such as images, are actually embedded in a low dimensional space. Hence, carefully reducing the input dimension may improve the robustness of the model without sacrificing performance. Inspired by the observation that the intrinsic dimension of image data is actually much smaller than its pixel space dimension and the vulnerability of a model grows with its input dimension , we propose a defense framework that embeds input images into a low-dimensional space using a deep encoder and performs classification based on the latent embedding with a classifier network. However, an arbitrary projection does not guarantee improving the robustness of the model, because there are a lot of mapping functions including non-robust ones from the raw input space to the low-dimensional space capable of minimizing the classification loss. To constrain the mapping function, we employ distribution regularization in the embedding space leveraging optimal transport theory. We call our new classification framework Embedding Regularized Classifier (ER-Classifier). To be more specific, we introduce a discriminator in the latent space which tries to separate the generated code vectors from the encoder network and the ideal code vectors sampled from a prior distribution, i.e., a standard Gaussian distribution. Employing a similar powerful competitive mechanism as demonstrated by Generative Adversarial Networks , the discriminator enforces the embedding space of the model to follow the prior distribution. In our ER-Classifier framework, the encoder and discriminator structures together project the input data to a low-dimensional space with a nice shape, then the classifier performs prediction based on the lowdimensional embedding. Based on the optimal transport theory, the proposed ER-Classifier minimizes the discrepancy between the distribution of the true label and the distribution of the framework output, thus only retaining important features for classification in the embedding space. With a small embedding dimension, the effect of the adversarial perturbation is largely diminished through the projection process. We compare ER-Classifier with other state-of-the-art defense methods on MNIST, CIFAR10, STL10 and Tiny Imagenet. Experimental demonstrate that our proposed ER-Classifier outperforms other methods by a large margin. To sum up, this paper makes the following three main contributions: • A novel unified end-to-end robust deep neural network framework against adversarial attacks is proposed, where the input image is first projected to a low-dimensional space and then classified. • An objective is induced to minimize the optimal transport cost between the true class distribution and the framework output distribution, guiding the encoder and discriminator to project the input image to a low-dimensional space without losing important features for classification. • Extensive experiments demonstrate the robustness of our proposed ER-Classifier framework under the white-box attacks, and show that ER-Classifier outperforms other state-ofthe-art approaches on several benchmark image datasets. As far as we know, our approach is the first that applies optimal transport theory, i.e., a Wasserstein distance regularization, to a bottleneck embedding layer of a deep neural network in a purely supervised learning setting without considering any reconstruction loss, although optimal transport theory or a discriminator loss has been applied to generative models in an unsupervised learning setting ; Our method is also the first that establishes the connection between a Wasserstein distance regularization and the robustness of deep neural networks for defending against adversarial examples. In this section, we summarize related work into three categories: attack methods, defense mechanisms and optimal transport theory. We first discuss different white-box attack methods, followed by a description of different defense mechanisms against, and finally optimal transport theory. Under the white-box setting, attackers have all information about the targeted neural network, including network structure and gradients. Most white-box attacks generate adversarial examples based on the gradient of loss function with respect to the input. An algorithm called fast gradient sign method (FGSM) was proposed in (Goodfellow et al.) which generates adversarial examples based on the sign of gradient. Many other white-box attack methods have been proposed recently (; ; ; b), and among them C&W and PGD attacks have been widely used to test the robustness of machine learning models. C&W attack: The adversarial attack method proposed by Carlini and Wagner (b) is one of the strongest white-box attack methods. They formulate the adversarial example generating process as an optimization problem. The proposed objective function aims at increasing the probability of the target class and minimizing the distance between the adversarial example and the original input image. Therefore, C&W attack can be viewed as a gradient-descent based adversarial attack. PGD attack: The projected gradient descent attack is proposed by , which finds adversarial examples in an -ball of the image. The PGD attack updates in the direction that decreases the probability of the original class most, then projects the back to the -ball of the input. An advantage of PGD attack over C&W attack is that it allows direct control of distortion level by changing, while for C&W attack, one can only do so indirectly via hyper-parameter tuning. Both C&W attack and PGD attack have been frequently used to benchmark the defense algorithms due to their effectiveness . In this paper, we mainly use l ∞ -PGD untargeted attack to evaluate the effectiveness of the defense method under white-box setting. Instead of crafting different adversarial perturbation for different input image, an algorithm was proposed by to construct a universal perturbation that causes natural images to be misclassified. However, since this universal perturbation is image-agnostic, it is usually larger than the image-specific perturbation generated by PGD and C&W. Many works have been done to improve the robustness of deep neural networks. To defend against adversarial examples, defenses that aim to increase model robustness fall into three main categories: i) augmenting the training data with adversarial examples to enhance the existing classifiers (; ; Goodfellow et al.); ii) leveraging model-specific strategies to enforce model properties such as smoothness ; and, iii) trying to remove adversarial perturbations from the inputs (; ;). We select three representative methods that are effective under white-box setting. Adversarial training: Augmenting the training data with adversarial examples can increase the robustness of the deep neural network. Madry et al. recently introduced a minmax formulation against adversarial attacks. The proposed model is not only trained on the original dataset but also adversarial example in the -ball of each input image. Random Self-Ensemble: Another effective defense method under white-box setting is RSE . The authors proposed a "noise layer", which fuses output of each layer with Gaussian noise. They empirically show that the noise layer can help improve the robustness of deep neural networks. The noise layer is applied in both training and testing phases, so the prediction accuracy will not be largely affected. Defense-GAN: Defense-GAN leverages the expressive capability of GANs to defend deep neural networks against adversarial examples. It is trained to project input images onto the range of the GAN's generator to remove the effect of the adversarial perturbation. Another defense method that uses the generative model to filter out noise is MagNet proposed by . However, the differences between ER-Classifier and the two methods are obvious. ER-Classifier focuses on reducing the dimension, and performing classification based on the low-dimensional embedding, while Defense-GAN and MagNet mainly apply the generative model to filter out the adversarial noise, and both Defense-GAN and MagNet perform classification on the original dimension space. showed that Defense-GAN is more robust than MagNet, so we only compare with Defense-GAN in the experiments. Other Related Methods: regularizes the latent space with Gaussian Mixture Model and applies KL-divergence to do the optimization. However, our method employs a simple but nice-shaped Gaussian prior for Wasserstein distance minimization to constrain the global shape of the latent embeddings, while permitting high freedom for the shapes of individual class distributions of latent embeddings. We want the classifier to decide the optimal class-specific distributions of latent embeddings. shares a similar idea to adversarial learning, but it employs virtual labels generated by a current classifier to identify search directions that can smooth the output label distribution of the classifier and is best suitable for semi-supervised learning. Please note that both methods in are designed for improving generalization performance but not for defending against adversarial examples. A recent paper shows that adversarial examples are purely human phenomenon and models tend to learn features that are not robust yet generalize well. We show that our Wasserstein distance regularization helps to identify robust features, which will be discussed later. Notations In this paper, we use l ∞ and l 2 distortion metrics to measure similarity. We report l ∞ distance in the normalized space, so that a distortion of 0.031 corresponds to 8/256, and l 2 distance as the total root-mean-square distortion normalized by the total number of pixels. We use calligraphic letters for sets (i.e., X), capital letters for random variables (i.e., X), and lower case letters for their values (i.e., x). The probability distributions are denoted with capital letters (i.e., P X) and corresponding densities with lower case letters (i.e., p X). We propose a novel defense framework, ER-Classifier, which aims at projecting the image data to a low-dimensional space to remove noise and stabilize the classification model by minimizing the optimal transport cost between the true label distribution P Y and the distribution of the ER-Classifier output (P C). An overview of the framework is shown in Figure 1. The encoder and discriminator structures together help diminish the effect of the adversarial perturbation by projecting input data to a space of lower dimension, then the classifier part performs classification based on the lowdimensional embedding. Mathematically, input images X ∈ X = R d are projected to a low-dimensional embedding vector Z ∈ Z = R k through the encoder Q φ. The discriminator D γ discriminates between the generated codeZ ∼ Q φ (Z|X) and the ideal code Z ∼ P Z. The classifier C τ performs classification based on the generated codeZ, producing output U ∈ U = R m, where m is the number of classes. The label of X is denoted as Y ∈ U. The Kantorovich's distance induced by the optimal transport problem is given by is the set of all joint distributions of (Y, U) with marginals P Y and P C, and c(y, u): U × U → R + is any measurable cost function. W c (P Y, P C) measures the divergence between probability distributions P Y and P C. When the probability measures are on a metric space, the p-th root of W c is called the p-Wasserstein distance. To minimize the Wasserstein distance between the distribution of the true label (P Y) and the distribution of the ER-Classifier output (P C), we can prove that it is sufficient to find a conditional distribution Q(Z|X) such that its marginal distribution Q Z is identical to a prior distribution P Z. The theorem and the proof are deferred to the Appendix. In this paper, we apply standard Gaussian as our prior distribution P Z, but other priors may be used for different cases. The final objective of ER-Classifier is: where Q can be a deterministic encoder as focused by this paper due to its simplicity or stochastic encoder as the one in a standard Variational Autoencoder, λ > 0 is a hyper-parameter and D is an arbitrary divergence between Q Z and P Z. To estimate the divergences between Q Z and P Z, we apply a GAN-based framework, fitting a discriminator to minimize the 1-Wasserstein distance between Q Z and P Z: We have also tried the Jensen-Shannon divergence, but as expected, Wasserstein distance provides more stable training and better . When training the framework, the weight clipping method proposed in Wasserstein GAN is applied to help stabilize the training of discriminator D γ. The training algorithm is summarized in Algorithm 1. At training stage, the encoder Q φ first maps the input x to a low-dimensional space, ing in generated code (z). Another ideal code (z) is sampled from the prior distribution, and the discriminator D γ discriminates between the ideal code (positive data) and the generated code (negative data). The classifier (C τ) predicts the image label based on the generated code (z). Sample {(x 1, y 1),..., (x n, y n)} from the training set 5: Sample {z 1, ..., z n} from the prior P Z 6: Update D γ by ascending the following objective by 1-step Adam: Update Q φ and C τ by descending the following objective by 1-step Adam: Update Q φ by ascending the following objective by 1-step Adam: 10: end while At inference time, only the encoder Q φ and the classifier C τ are used. The input image x is first mapped to a low-dimensional space by the encoder (z = Q φ (x)), then the latent codez is fed into the classifier to obtain the predicted label. The main goal of ER-Classifier is leveraging input space dimension reduction to remove adversarial perturbations. Therefore, other defense methods can also benefit from this property. Our framework is trained with min-max robust optimization . There are two Wasserstein distances (W-distances) in our framework. One is the W-distance between the aggregated latent embedding distribution Q(z) and the prior distribution P Z, and the other one is the W-distance between the true label distribution P Y and the distribution of the ER-Classifier output (P C). In Algorithm 1, we are minimizing the first one. The theorem in the Appendix shows that minimizing the first W-distance in combination with minimizing a standard cross-entropy loss as done in Algorithm 1 is equivalent to minimizing the second W-distance, which guarantees that the training process is not distracted from the main goal of the framework, classification. That is to say, Algorithm 1 will in a classifier with the following property: the global output distribution of the classifier will match the global ground-truth label distribution in the data no matter whether the encoder Q φ is deterministic or stochastic (the second W-distance is automatically minimized). It's hard to analyze the importance of the theorem in the Appendix if we just look at a deterministic encoder. Let's convert this deterministic encoder to a stochastic encoder that outputs a Gaussian z with a fixed variance and the mean being the same as its corresponding deterministic version. The theory tells us that, by minimizing the first W-distance over all sampled z's from this stochastic encoder and the standard cross-entropy loss, we will automatically minimize the second W-distance and preserve the global label frequency in the dataset, even though these z's are only -close to the deterministic encoding features of training data. Moreover, we find that minimizing the W-distance helps the encoder identify some robust features instead of non-robust features , because our proposed regularization constrains the -ball around each Q φ (Z|X) to contribute to preserving the global label distribution in the data, even with X integrated out. From this perspective, we can view our proposed framework as "a supervised variant" of Generative Adversarial Network or Wasserstein Autoencoder in which the Generator or Decoder is replaced by a Classifier that generates labels from low-dimensional latent embeddings preserving global label frequency in the training dataset. Replacing W-distance with KL divergence loses all these nice properties. In our framework, we use a simple but nice-shaped Gaussian prior P Z for W-distance minimization to constrain the global shape of the latent embeddings, while permitting high freedom for the shapes of individual class distributions of latent embeddings. We want the classifier to decide the optimal class-specific distributions of latent embeddings. In addition, it is interesting to explore how to set -ball to make sure the stochastic encoder to best align the latent embedding z to human-perceived robust features, which will be left as future work. In this section, we compare the performance of our proposed algorithm (ER-Classifier) with other state-of-the-art defense methods on several benchmark datasets: • MNIST : handwritten digit dataset, which consists of 60, 000 training images and 10, 000 testing images. Theses are 28 × 28 black and white images in ten different classes. classes, and each class has 500 training images, 50 testing images, making it a challenging benchmark for defense task. The resolution of the images is 64 × 64. Various defense methods have been proposed to improve the robustness of deep neural networks. Here we compare our algorithm with state-of-the-art methods that are robust in white-box setting. Madry's adversarial training (Madry's Adv) has been recognized as one of the most successful defense method in white-box setting, as shown in . Random Self-Ensemble (RSE) method introduced by adds stochastic components in the neural network, achieving similar performance to Madry's adversarial training algorithm. Another method we would like to compare with is Defense-GAN . It first trains a generative adversarial network to model the distribution of the training data. At inference time, it finds a close output to the input image and feed that output into the classifier. This process "projects" input images onto the range of GAN's generator, which helps remove the effect of adversarial perturbations. In , the author demonstrated the performance of Defense-GAN on MNIST and Fashion-MNIST, so we will compare our method with Defense-GAN on MNIST. Since the main goal of ER-Classifier is using dimension reduction to improve adversarial robustness, other defense methods can also benefit from this property. The proposed ER-Classifier is trained with min-max robust optimization . To demonstrate the dimension reduction ability of ER-Classifier, we include a variant ER-Classifier − which trains ER-Classifier without min-max robust optimization. In this section, we evaluate the defense methods against l ∞ -PGD untargeted attack, which is one of the strongest white-box attack methods. Models are evaluated under different distortion level , Based on Figure 2 and − tends to perform better than other state-of-the-art defense methods on MNIST, CIFAR10 and Tiny Imagenet. This phenomenon is obvious on CIFAR10 and it even performs better than ER-Classifier when the attack strength is strong. The reason might be that without min-max robust optimization, it is easier to regularize the embedding space. Testing Accuracy Defense-GAN 55.0 ER-Classifier 99.1 We also compare Defense-GAN with our method ERClassifier on MNIST. Although Defense-GAN was shown to be partly broken by , both ER-Classifier and Defense-GAN share the similar idea of projecting the input to a learned manifold, and comparing to Defense-GAN is important to demonstrate the advantage of our novel Wassserstein distance regularization. Please note that Defense-GAN is not our major comparison baseline in this paper. Both ER-Classifier and Defense-GAN are evaluated against the l 2 -C&W untargeted attack, one of the strongest white-box attack proposed in (b). Defense-GAN is evaluated using the method proposed in , and the code is available on github 1. ER-Classifier is evaluated against l 2 -C&W untargeted attack with the same hyper-parameter values as those used in the evaluation of Defense-GAN. The under l 2 ≤ 0.005 threshold are shown in Table 2. Based on Table 2, ER-Classifier is much more robust than Defense-GAN un- der the l 2 ≤ 0.005 threshold. Since did not evaluate Defense-GAN on CIFAR10, STL10 and Tiny Imagenet, without details of GAN structure, we can not compare with Defense-GAN on these datasets. We evaluate Madry's adversarial training, ER-Classifier, and ER-Classifier −, against a recently proposed black-box attack method called Nattack 2 on CIFAR10. Nattack is only performed on the first 100 images of CIFAR10 since the attack process takes a long time. We report the accuracy = number of correctly classified / number of attacked images (exactly 100). The accuracy of Madry's adv, ER-Classifier, and ER-Classifier − is, respectively, 38%, 43%, and 32%. ER-Classifier still outperforms Madry's adv. ER-Classifier framework consists of three parts, and the classification task is done by the encoder Q φ and classifier C τ. Without the discriminator part, the encoder can also project the input images to a low-dimensional space. However, arbitrarily projecting the images to a low-dimensional space with only the encoder part cannot improve the robustness of the model. In contrast, sometimes it even decreases the robustness of the model. To show that arbitrarily projecting the input images to a low-dimensional space can not improve the robustness, we fit a framework with only the encoder and classifier part (E-CLA), where the encoder and classifier have the same structures as in ER-Classifier, and compare E-CLA with the ER-Classifier framework. For a fair comparison, both structures are trained without min-max robust optimization. The are shown in Figure 3. Based on Figure 3, we can observe that ER-Classifier is much more robust than just the encoder and classifier structure on MNIST, CIFAR10 and Tiny Imagenet. It is also more robust on STL10 but not that much. The reason might be that there are only 5, 000 training images in STL10 and the resolution is 96 × 96. Therefore, it is harder to learn a good embedding with limited amount of images. However, even when the number of training images is limited, ER-Classifier is still much more robust than the E-CLA structure. This observation demonstrates that regularization on the embedding space helps improve the adversarial robustness. Notice that the performance of E-CLA structure is similar to the performance of model without defense method on CIFAR10, STL10 and Tiny Imagenet, and worse on MNIST, which means the robustness of ER-Classifier does not come from the structure design. Variational auto-encoder can project the images to low-dimensional space and use Kullback-Leibler divergence loss to regularize the embedding distribution, which does not need discriminator structure. Therefore, we also tried VAE-CLA, which applies Variational auto-encoder structure to do the projection and regularization. The experimental in Figure 3 show that VAE-CLA does not perform as well as ER-Classifier. Based on the observation of the Kullback-Leibler loss and classification loss during the training process, it seems difficult for VAE-CLA to balance between the two tasks. The reason might be that Kullback-Leibler distances are not sensible cost functions when learning distributions supported by low dimensional manifolds . However, the selection of prior is important as it imposes different restrictions on the embedding space. Three different prior distributions are tested on MNIST and CIFAR10 datasets. They are standard Gaussian, Uniform(−3, 3) and Cauchy, where Cauchy has the same support as standard Gaussian but is heavy tailed and 99.7% of the standard Gaussian points lies within [−3, 3]. All the models are trained without min-max robust optimization, and the experimental are shown in Figure 4. Based on the , all three priors work well, but standard Gaussian performs best on both datasets. Ding et al. prove that adversarial robustness is sensitive to the input data distribution, and if the data is uniformly distributed in the input space, no algorithm can achieve good robustness. They also empirically show that cornered/concentrated data distributions tend to achieve better robustness. This helps explain why regularizing the embedding space can help improve robustness. Though the projection process reduces the input dimension, the embedding space is still large. Prior distribution helps push the embedding space to be more concentrated, reducing the valid perturbation space. Details of hyper-parameter selection, model structure and code are included in the supplementary part. Embedding space visualization can also be found in the supplementary material. In this paper, we propose a new defense framework, ER-Classifier, which projects the input images to a low-dimensional space to remove adversarial perturbation and stabilize the model through minimizing the discrepancy between the true label distribution and the framework output distribution. We empirically show that ER-Classifier is much more robust than other state-of-the-art defense methods on several benchmark datasets. Future work will include further exploration of the low-dimensional space to improve the robustness of deep neural network. Mathematically, input images X ∈ X = R d are projected to a low-dimensional embedding vector Z ∈ Z = R k through the encoder Q φ. The discriminator D γ discriminates between the generated codeZ ∼ Q φ (Z|X) and the ideal code Z ∼ P Z. The classifier C τ performs classification based on the generated codeZ, producing output U ∈ U = R m, where m is the number of classes. The label of X is denoted as Y ∈ U. The ER-Classifier framework embeds important classification features by minimizing the discrepancy between the distribution of the true label (P Y) and the distribution of the framework output (P C). In the framework, the classifier (P C (U |Z)) maps a latent code Z sampled from a fixed distribution in a latent space Z, to the output U ∈ U = R m. The density of ER-Classifier output is defined as follow: In this paper we apply standard Gaussian as our prior distribution P Z, but other priors may be used for different cases. Assume there is an oracle f: X → U assigning the image data (X ∈ X) its true label (Y ∈ U). We want to minimize the optimal transport cost between the distribution of the true label (P Y) and the distribution of the ER-Classifier output (P C). There are various ways to define the distance or divergence between the target distribution and the model distribution. In this paper, we turn to the optimal transport theory , which provides a much weaker topology than many others. In real applications, data is usually embedded in a space of a much lower dimension, such as a non-linear manifold. Kullback-Leibler divergence, Jensen-Shannon divergence and Total Variation distance are not sensible cost functions when learning distributions supported by lower dimensional manifolds . In contrast, the optimal transport cost is more sensible in this setting. Kantorovich's distance induced by the optimal transport problem is given by where P(Y ∼ P Y, U ∼ P C) is the set of all joint distributions of (Y, U) with marginals P Y and P C, and c(y, u): U × U → R + is any measurable cost function. W c (P Y, P C) measures the divergence between probability distributions P Y and P C. When the probability measures are on a metric space, the p-th root of W c is called the p-Wasserstein distance. To minimize the optimal transport cost between the distribution of the true label (P Y) and the distribution of the ER-Classifier output (P C), it is sufficient to find a conditional distribution Q(Z|X) such that its marginal distribution Q Z is identical to the prior distribution P Z. Theorem 1 For P C as defined above with a deterministic P C (U |Z) and any function C: where Γ ∈ P(Y ∼ P Y, U ∼ P C) is the set of all joint distributions of (Y, U) with marginals P Y and P C, and (y, u): U × U → R + is any measurable cost function. Q Z is the marginal distribution of Z when X ∼ P X and Z ∼ Q(Z|X). (The proof is presented later.) Therefore, optimizing over the objective on the r.h.s is equivalent to minimizing the discrepancy between the true label distribution (P Y) and the output distribution P C, thus the important classification features are embedded in the low-dimensional space. This is the core idea of the paper, summarizing the high-dimensional data in a space of much lower dimension without losing important features for classification. To implement the r.h.s objective, the constraint on Q Z can be relaxed by adding a penalty term. The final objective of ER-Classifier is: where Q is any nonparametric set of probabilistic encoders, λ > 0 is a hyper-parameter and D is an arbitrary divergence between Q Z and P Z. To estimate the divergences between Q Z and P Z, we apply a GAN-based framework, fitting a discriminator to minimize the 1-Wasserstein distance between Q Z and P Z: We have also tried the Jensen-Shannon divergence, but as expected, Wasserstein distance provides more stable training and better . When training the framework, the weight clipping method proposed in Wasserstein GAN is applied to help stabilize the training of discriminator D γ. The proof of Theorem 1 is adapted from the proof of Theorem 1 in . Consider certain sets of joint probability distributions of three random variables (X, U, Z) ∈ X × U × Z. X can be taken as the input images, U as the output of the framework, and Z as the latent codes. P C,Z (U, Z) represents a joint distribution of a variable pair (U, Z), where Z is first sampled from P Z and then U from P C (U |Z). P C defined in is the marginal distribution of U when (U, Z) ∼ P C,Z. The joint distributions Γ(X, U) or couplings between values of X and U can be written as Γ(X, U) = Γ(U |X)P X (X) due to the marginal constraint. Γ(U |X) can be decomposed into an encoding distribution Q(Z|X) and the generating distribution P C (U |Z), and Theorem 1 mainly shows how to factor it through Z. In the first part, we will show that if P C (U |Z) are Dirac measures, we have where P(X ∼ P X, U ∼ P C) denotes the set of all joint distributions of (X, U) with marginals P X, P C, and likewise for P(X ∼ P X, Z ∼ P Z). The set of all joint distributions of (X, U, Z) such that X ∼ P X, (U, Z) ∼ P C,Z, and (U ⊥ ⊥ X)|Z are denoted by P X,U,Z. P X,U and P X,Z denote the sets of marginals on (X, U) and (X, Z) induced by P X,U,Z. From the definition, it is clear that P X,U ⊆ P(P X, P C). Therefore, we have The identity is satisfied if P C (U |Z) are Dirac measures, such as U = C(Z). This is proved by the following Lemma in . Lemma 1 P X,U ⊆ P(P X, P C) with identity if P C (U |Z = z) are Dirac for all z ∈ Z. (see details in .) In the following part, we show that Based on the definition, P(P X, P C), P X,U,Z and P X,U depend on the choice of conditional distributions P C (U |Z), but P X,Z does not. It is also easy to check that P X,Z = P(X ∼ P X, Z ∼ P Z). The tower rule of expectation, and the conditional independence property of P X,U,Z implies Finally, since Y = f (X), it is easy to get Now, and are proved and the three together prove Theorem 1. Our proposed framework readily applies to non-deterministic case. If the classifier part is nondeterministic, Lemma 1 provides only the inclusion of sets P X,U ⊆ P(P X, P U), and we can get an upper bound on the Wasserstein distance between the ground-truth and predicted label distributions: where we assume the conditional distributions P C (U |Z = z) have mean values C(z) ∈ R d and marginal variances σ 2 1,..., σ 2 d ≥ 0 for all z ∈ Z, where C: Z → X, and (y, u) = y − u 2. The above upper bound is derived by: and In equation 11, the second term of the second last row becomes 0 since the optimization will drive f (X) − C(Z) to zero. One important hyper-parameter for the ER-Classifier is the dimension of the embedding space. If the dimension is too small, important features are "collapsed" onto the same dimension, and if the dimension is too large, the projection will not extract useful information, which in too much noise and instability. The maximum likelihood estimation of intrinsic dimension proposed in 3 is used to calculate the intrinsic dimension of each image dataset, serving as a guide for selecting the embedding dimension. The sample size used in calculating the intrinsic dimension is 1, 000, and changing the sample size does not influence the much. Based on the intrinsic dimension calculated by , we test several different values around the suggested intrinsic dimension and evaluate the models against l ∞ -PGD attack. All models are trained without min-max robust optimization, and the experimental are shown in Figure 5. The final embedding dimension is chosen based on robustness, number of parameters, and testing accuracy when there is no attack. The final embedding dimensions and suggested intrinsic dimensions are shown in Table 3 Table 3: Pixel space dimension, intrinsic dimension calculated by , and final embedding dimension used. Based on Figure 5, the embedding dimension close to the calculated intrinsic dimension usually offers better except on MNIST. One explanation may be that MNIST is a simple handwritten digit dataset, so performing classification on MNIST may not require that many dimensions. Epsilon is an important hyper-parameter for adversarial training. When doing Madry's adversarial training, we test the model robustness with different and choose the best one. The experiment are shown in Figure 6. Based on Figure 6, we use = 0.3, 0.03, 0.03 in Madry's adversarial training on MNIST, CIFAR10 and STL10 respectively. For Tiny Imagenet, we use = 0.01. To make a fair comparison, we use the same when training ER-Classifier. In this section, we compare the embedding learned by Encoder+Classifier structure (E-CLA) and the embedding learned by ER-Classifier on several datasets without min-max robust optimization. We first generate embedding of testing data using the encoder (z = Q φ (x)), then project the embedding points (z) to 2-D space by tSNE . Then we generate adversarial images (x adv) against E-CLA and ER-Classifier using l ∞ -PGD attack. The adversarial embedding is generated by feeding the adversarial images into the encoder (z adv = Q φ (x adv)). Finally, we project the adversarial embedding points (z adv) to 2-D space. The are shown in Figure 7. The plots in the first and second rows are embedding visualization plots for E-CLA, and the plots in the third and last rows are the embedding visualization plots for ER-Classifier. In adversarial embedding visualization plots, the misclassified point is marked as "down triangle", which means the PGD attack successfully changed the prediction, and the correctly classified point is marked as "point", which means the attack fails. Based on Figure 7, we can see that E-CLA can learn a good embedding on legitimate images of MNIST. Embedding points for different classes are separated on the 2D space, but under adversarial attack, some embedding points of different classes are mixed together. However, ER-Classifier can generate good separated embeddings on both legitimate and adversarial images. On CIFAR10, the E-CLA can not generate good separated embeddings on either legitimate images or adversarial images, while ER-Classifier can generate good separated embeddings for both. Code for reproduction will be made available online at github later. The pseudocode for training ER-Classifier is shown in Listing 1. MNIST, STL10 and TinyImagenet classifier structures used for baseline methods are shown in Fig
A general and easy-to-use framework that improves the adversarial robustness of deep classification models through embedding regularization.