Classification with Rejection: Scaling Generative Classifiers with Supervised Deep Infomax

1 Introduction

Non-robustness of neural network models emerges as a pressing concern since they are observed to be vulnerable to adversarial examples [Szegedy et al., 2013]. Many attack methods have been developed to find imperceptible perturbations to fool the target classifiers [Carlini and Wagner, 2017; Brendel et al., 2017]. Meanwhile, many defense schemes have also been proposed to improve the robustnesses of the target models [Goodfellow et al., 2014; Madry et al., 2017].

An important fact about these works is that they focus on discriminative classifiers, which directly model the conditional probabilities of labels given samples. Another promising direction, which is almost neglected so far, is to explore robustness of generative classifiers [Ng and Jordan, 2002]. A generative classifier explicitly model conditional distributions of inputs given the class labels. During inference, it evaluates all the class conditional likelihoods of the test input, and outputs the class label corresponding to the maximum. Conditional generative models are powerful and natural choices to model the class conditional distributions, but they suffer from two big problems: (1) it is hard to scale generative classifiers on high-dimensional tasks, like natural images classification, with comparable performance to the discriminative counterparts. Though generative classifiers have shown promising results of adversarial robustness, they hardly achieve acceptable classification performance even on CIFAR10 [Li et al., 2018; Schott and Rauber, 2018; Fetaya et al., 2019]. (2) The behaviors of likelihood-based generative models can be counter-intuitive. They surprisingly assign higher likelihoods to out-of-distribution (OOD) samples [Nalisnick et al., 2018; Choi and Jang, 2018]. [Fetaya et al., 2019] discuss the issues of likelihood as a metric for density modeling, which may be the reason of non-robust classification, e.g. OOD samples detection.

∗Contact Author

In this paper, we propose supervised deep infomax (SDIM) by introducing supervised statistical constraints into deep infomax (DIM, [Hjelm et al., 2018]), an unsupervised learning framework by maximizing the mutual information between representations and data. SDIM is trained by optimizing two objectives: (1) maximizing the mutual information (MI) between the inputs and the high-level data representations from encoder; (2) ensuring that the representations satisfy the supervised statistical constraints. The supervised statistical constraints can be interpreted as a generative classifier on highlevel data representations giving up the full generative process. Unlike full generative models making implicit manifold assumptions, the supervised statistical constraints of SDIM serve as explicit enforcement of manifold assumption: data representations (low-dimensional) are trained to form clusters corresponding to their class labels. With SDIM, we could perform classification with rejection. SDIMs reject illegal inputs based on off-manifold conjecture [Samangouei et al., 2018; Gu and Rigazio, 2014], where illegal inputs, e.g. adversarial examples, lie far away from the data manifold. Samples whose class conditionals are smaller than the pre-chosen thresholds will be deemed as off-manifold, and prediction requests on them will be rejected. The contributions of this paper are :

• We propose Supervised Deep Infomax (SDIM), an end-

to-end framework whose probabilistic constraints are equivalent to a generative classifier. SDIMs can achieve comparable classification performance with similar discrinimative counterparts at the cost of small overparameterization.

• We propose a simple but novel rejection policy based on off-manifold conjecture: SDIM outputs a class label only if the test sample's largest class conditional surpasses the pre-chosen class threshold, otherwise outputs rejection. The choice of thresholds relies only on training set, and takes no additional computations.
• Experiments show that SDIM with rejection policy can effectively reject illegal inputs, including OOD samples and adversarial examples generated by a comprehensive group of adversarial attacks.

2 Background: Deep InfoMax

Deep InfoMax (DIM, [Hjelm et al., 2018]) is an unsupervised representation learning framework by maximizing the mutual information (MI) of the inputs and outputs of an encoder. The computation of MI takes only input-output pairs with the deep neural networks based estimator MINE [Belghazi et al., 2018].

Let Eφ be an encoder parameterized by φ, working on the training set X = {xi}Ni=1 , and generating output set Y = {E(xi)}Ni=1 . DIM is trained to find the set of parameters φ such that: (1) the mutual information I(X, Y ) is maximized over sample sets X and Y. (2) the representations, depending on the potential downstream tasks, match some prior distribution. Denote J and M the joint and product of marginals of random variables X, Y respectively. MINE estimates a lower-bound of MI with Donsker-Varadhan [Donsker and Varadhan, 1983] representation of KL-divergence:

where Tω(x, y) ∈ R is a family of functions with parameters ω represented by a neural network. Since in representation learning we are more interested in maximizing MI, than its exact value, non-KL divergences are also favorable candidates. We can get a family of variational lower-bounds using f-divergence representations [Nguyen et al., 2010]:

where f ∗ is the Fenchel conjugate of a specific divergence f. For KL-divergence, f ∗ (t) = e(t-1) . A full f ∗ list is provided in Tab. 6 of [Nowozin et al., 2016].

3 Supervised Deep InfoMax

All the components of SDIM framework are summurized in Fig. 1. The focus of Supervised Deep InfoMax (SDIM) is on introducing supervision to probabilistic constraints of DIM for (generative) classification. We choose to maximize the local MI, which has shown to be more effective in classification tasks than maximizing global MI [Hjelm et al., 2018].

Equivalently, we minimize JMI:

where Lφ(x) is a local M ×M feature map of x extracted from some intermediate layer of encoder E, and ˜ I can be any possible MI lower-bounds.

3.1 Explicit Enforcement of Manifold Assumption

By adopting a generative approach p(x, y) = p(y)p(x|y), we assume that the data follows the manifold assumption: the (high-dimensional) data lies on low-dimensional manifolds corresponding to their class labels. Denote ˜x the compact representation generated with encoder Eφ(x). In order to explicitly enforce the manifold assumption , we admit the existence of data manifold in the representation space. Assume that y is a discrete random variable representing class labels, and p(˜x|y) is the real class conditional distribution of the data manifold given y. Let pθ(˜x|y) be the class conditionals we model parameterized with θ. We approximate p( ˜x|y) by minimizing the KL-divergence between p(˜x|y) and our model pθ(˜x|y), which is given by:

where the first item on RHS is a constant independent of the model parameters θ. Eq. 4 equals to maximize the expectation E˜x,y∼p( ˜x,y)[log pθ ( ˜x|y)].

In practice, we minimize the following loss JNLL, equivalent to empirically maximize the above expectation over {˜xi = Eφ(xi), yi}Ni=1 :

Besides the introduction of supervision, SDIM differs from DIM in its way of enforcing the statistical constraints: DIM use adversarial learningto push the representations to the desired priors, while SDIM directly maximizes the parameterized class conditional probability.

Maximize Likelihood Margins Since a generative classifier, at inference, decides which class a test input x belongs to according to its class conditional probability. On one hand, we maximize samples' true class conditional probabilities (classes they belong to) using JNLL; On the other hand, we also hope that samples' false class conditional probabilities (classes they do not belong to) can be minimized. This is assured by the following likelihood margin loss JLM:

where K is a positive constant to control the margin. For each encoder output ˜xi, the C-1 true-false class conditional gaps are squared 1 , which quadratically increases the penalties when the gap becomes large, then are averaged.

Putting all these together, the complete loss function we

where α, β, γ are scaling factors.

Parameterization of Class Conditional Probability Each of the class conditional distribution is represented as an isotropic Gaussian. So the generative classifier is simply a embedding layer with C entries, and each entry contains the trainable mean and variance of a Gaussian. This minimized parameterization encourages the encoder to learn simple and stable low-dimensional representations that can be easily explained by even unimodal distributions. Considering that we maximize the true class conditional probability, and minimize the false class conditional probability at the same time, we do not choose conditional normalizing flows, since the parameters are shared across class labels, and the training can be very difficult. In [Schott and Rauber, 2018], each class conditional probability is represented with a VAE, thus scaling to complex datasets with huge number of classes, e.g. ImageNet, is almost impossible.

3.2 Decision Function with Rejection

A generative approach models the class-conditional distributions p(x|y), as well as the class priors p(y). For classification, we compute the posterior probabilities p(y|x) through Bayes' rule:

1 Using squared margin, we achieve slightly better results in our experiments than simple margin.

The prior p(y) can be computed from the training set, or we simply use uniform class prior for all class labels by default. Then the prediction of test sample x∗ from posteriors is:

The drawback of the above decision function is that it always gives a prediction even for illegal inputs. Instead of simply outputting the class label that maximizes class conditional probability of x∗ , we set a threshold for each class conditional probability, and define our decision function with rejection to be:

The model gives a rejection when log p(x∗ |y∗ ) is smaller than the threshold δy∗. Note that here we can use p(x∗ |y∗ ) and p(˜x∗ |y∗ ) interchangeably.

Classification with rejection is not novel [Geifman and ElYaniv, 2017], and previous works closely related to ours are the generative models based ones. The most recent one [Nalisnick et al., 2019], which propose a hybrid model modeling distribution of features p(features) and predictive distribution p(targets|features) at the same time. Normalizing flow is used to learn invertible features as inputs of discriminative model, i.e. predictive distribution, and provides evaluation of features x∗ . Inputs out of the training data distribution are rejected by setting a threshold for p(x∗ ). For SDIM, illegal inputs are rejected by setting thresholds for each of the class conditional. The class conditionals are modeled on the data representations, rather than raw inputs. In terms of robustness, the hybrid model in [Nalisnick et al., 2019] can successfully detect OOD samples, and applicability as well as performance on adversarial examples is not clear. While SDIM reject illegal inputs including OOD dataset samples and adversarial examples with more fine-grained class conditionals.

4 Experiments

Datasets We evaluate the effectiveness of the rejection policy of SDIM on four image datasets: MNIST, FashionMNIST (both resized to 32 × 32 from 28 × 28); and CIFAR10, SVHN. For out-of-distribution samples detection, we use the dataset pairs on which likelihood-based generative models fail [Nalisnick et al., 2018; Choi and Jang, 2018]: FashionMNIST (in)-MNIST (out) and CIFAR10 (in)-SVHN (out). Adversarial examples detection are evaluated on MNIST and CIFAR10. Throughout our experiments, we use α = β = γ = 1 in the loss function.

Choice of thresholds It is natural that choosing thresholds based on what the model knows, i.e. training set, and can reject what the model does not know, i.e. possible illegal inputs. We set one threshold for each class conditional. For each class conditional probability, we choose to evaluate on two different thresholds: 1st and 2nd percentiles of class conditional log-likelihoods of the correctly classified training samples. Compared to the detection methods proposed in [Li et al., 2018], our choice of thresholds is much simpler, and takes no additional computations.

Models A typical SDIM instance consists of three networks: an encoder, parameterized by φ, which outputs a d-dimensional representation; mutual information evaluation networks, i.e. Tω in Eqn. (1) and Eqn. (2); and C-way class conditional embedding layer, parameterized by θ, with each entry a 2d-dimensional vector. We set d = 64 in all our experiments.

For encoder of SDIM, we use ResNet [He et al., 2016] on 32 × 32 with a stack of 8n + 2 layers, and 4 filter sizes {32, 64, 128, 256}. The architecture is summarized as:

The last layer of encoder is a d-way fully-connected layer. To construct a discriminative counterpart, we simply set the output size of the encoder's last layer to C for classification. We use ResNet10 (n = 1) on MNIST, FashionMNIST, and ResNet26 (n = 3) on CIFAR10, SVHN.

4.1 Evaluation on Clean Data

We report the classification accuracy (see Tab. 2 and Tab. 3) of SDIMs and the discriminative counterparts on clean test sets . Results show that SDIMs achieve the same level of accuracy as the discriminative counterparts with slightly increased number of parameters (17% increase for ResNet10, and 5% increase for ResNet26). We are aware of the existence of better results reported on these datasets , but pushing the state-of-the-art is not the focus of this paper.

Scalability

The most important fact about SDIM is it provides a endto-end framework to train generative classifiers that achieve same-level performance as the corresponding discriminative counterparts. In contrast, despite the great success of fully conditional generative models in data synthesis, they perform poorly on classification tasks. In particular, they perform quite well on MNIST, but are still far away from achieving acceptable performance even on CIFAR10. For example, methods (GFZ & GFY) in DeepBayes [Li et al., 2018] achieve < 50% accuracy, and they also mention that a conditional PixelCNN++ [Salimans et al., 2017] (with much deeper networks) achieves 72.4% clean test accuracy. The test accuracy of ABS in [Schott and Rauber, 2018] is only 54%. Glow with class conditional mixture of Gaussian [Fetaya et al., 2019] achieves 56.8%, could be improved to 80 - 85% with reweighting or split prior. Many works have demonstrated that for classification tasks, discovering discriminative features (patches) is much more important than reconstructing the all the image pixels [Brendel and Bethge, 2019; Hjelm et al., 2018]. The fact that methods in [Li et al., 2018] improve the accuracy from < 50% to 92% by feeding the features learned by powerful discriminative classifier-VGG16 to their generative classifiers, also support this argument.

One direct effect of the poor classification performance of fully conditional generative models is that adversarial robustness evaluation on them is limited to MNIST, and becomes much less convincing on CIFAR10. A model with high test error implies that even for correctly classified samples, their distances to the decision boundaries are much small, thus are easy to craft small adversarial perturbations leading to misclassification [Gilmer et al., 2018].

Furthermore, due to giving up the full generative process, SDIM is considerably smaller than fully conditional generative models. The GBZ of [Li et al., 2018] on MNIST has

1.5M parameters 2 . The ABS model in [Schott and Rauber, 2018] is 0.85M 3 . While SDIM on MNIST is only 0.36M. Note that their models don't scale on CIFAR10. So if we only care about the generative classification performance with no need to generate samples, it is unnecessary to model the full generative process, and giving up saves us a lot of computational resources.

Decision with Rejection

Given pre-chosen thresholds, there are chances that legal inputs are wrongly rejected. Thus we also investigate the implications of the proposed rejection decision function with different thresholds on clean test sets. The results in Tab. 4 show that choosing a higher percentile as threshold will reject more prediction requests. At the same time, the classification accuracy of SDIMs on the left test sets become increasingly better. This demonstrate that our rejection method help SDIMs reliably reject the low-confidence requests, and avoid wrong predictions.

4.2 Adversarial Examples Detection

We compare SDIM to GBZ [Li et al., 2018], which consistently performs the best in Deep Bayes. Note that here we investigate the inherent robustness of generative classifier, so we do not incorporate the specially designed methods outside of the classifiers for adversarial examples detection.

FGSM and PGD We evaluate SDIM on L∞ versions of them. We find that SDIMs are much more robust against FGSM and PGD than baseline (see Fig. 2), since the gradients numerically vanish as a side effect of the likelihood margin loss JLM of SDIM. Recall that the class conditionals are optimized to keep a considerable margin. Before evaluating the cross entropy loss, softmax is applied on the class conditionals log p(x|c) to generate a even sharper distribution. So for the samples that are correctly classified, their losses are numerically zeros, and the gradient on inputs ∇Jx(x, y) are also numerically zeros. This phenomena is similar to what some defenses using gradient obfuscation want to achieve. Defensive distillation [Carlini and Wagner, 2016] masks the gradients of cross-entropy by increasing the temperature of softmax. However, obfuscated gradients [Athalye et al., 2018] give a false sense of security. For CW attacks, which do not use cross-entropy, and operate on logits directly, this could be ineffective.

2 The number of parameters is estimated based on the details in Appendix D of [Li et al., 2018]. The encoder network (3 convs and 2-layer MLP) is ∼ 0.5M. p(y|z) is 37k, and p(x|z) (2-layer MLP and 3 convs) is ∼1M.

3 The number of parameters is calculated from the open sourced code: https://github.com/bethgelab/AnalysisBySynthesis.

CW Attack We evaluate SDIM and baselines on CW-L2 with loss factors c = {1, 10, 100, 1000}. On adversarial examples of MNIST, SDIM performs on par with the baseline, while on CIFAR10, the rejection/detection rate of SDIM is slightly better than the GBZ on simpler CIFAR10-binary (see Fig. 3). GBZ does not scale on CIFAR10, so the rejection rate is not available.

Black-Box Attacks We also evaluate SDIM on blackbox attacks: local search (score-based) attack [Narodytska and Kasiviswanathan, 2016], boundary(decision-based) attack [Brendel et al., 2017]. Generally, we find that blackbox adversarial examples are more easily to detect even on CIFAR10 (see Tab. 5).

Discussions on off-manifold conjecture [Gilmer et al., 2018] challenges whether the off-manifold conjecture holds in general. They experiment on synthetic dataset-two highdimensional concentric spheres with theoretical analyses, showing that even for a trained classifier with close to zero test error, there may be a constant fraction of the data manifold misclassified, which indicates the existence of adversarial examples within the manifold. But there are still several concerns to be addressed: First, as also pointed out by the authors, the manifolds in natural datasets can be quite complex than that of simple synthesized dataset. [Fetaya et al., 2019] draws similar conclusion from analyses on synthesized data with particular geometry. So the big concern is whether the conclusions in [Gilmer et al., 2018; Fetaya et al., 2019] still hold for the manifolds in natural datasets. A practical obstacle to verify this conclusion is that works modeling the full generative processes are based on manifold assumption, but provide no explicit manifolds for analytical analyses like [Gilmer et al., 2018; Fetaya et al., 2019]. While SDIM enables explicit and customized manifolds on high-level data representations via probabilistic constraints, thus enables analytical analyses. In this paper, samples of different classes are trained to form isotropic Gaussians corresponding to their classes in representation space (other choices are also possible). The relation between the adversarial robustness and the forms and dimensionalities of data manifolds is to be explored in the future. Second, in their experiments, all models evaluated are discriminative classifiers. Considering the recent promising results of generative classifiers against adversarial examples, would using generative classifiers lead to different results? One thing making us feel optimistic is that even though the existence of adversarial examples is inevitable, [Gilmer et al., 2018] suggest that adversarial robustness can be improved by minimizing the test errors, which is also supported by our experimental differences on MNIST and CIFAR10.

4.3 Out-Of-Distribution Samples Detection

Class-wise OOD detections are performed, and mean detection rates over all in-distribution classes are reported in Tab. 6. For each in-distribution class c, we evaluate the loglikelihoods of the whole OOD dataset. Samples whose loglikelihoods are lower the class threshold δc will be detected as OOD samples. Same evaluations are applied on conditional Glows with 10th percentile thresholds, but the results are not good. The results are clear and confirm that SDIMs, generative classifiers on high-level representations, are more effective on classification tasks than fully conditional generative models on raw pixels. Note that fully generative models including VAE used in [Li et al., 2018; Schott and Rauber, 2018] fail on OOD detection [Nalisnick et al., 2018; Choi and Jang, 2018]. The stark difference of SDIM from full generative models (flows or VAEs) is that SDIMmodels samples' likelihood in the high-level representation spaces, rather than directly on the raw pixels.

Summary SDIMmodels performon par with or better than strongest variant GBZ in [Li et al., 2018] on detection of various types of adversarial examples. However, the performance of SDIM could extend to complex datasets, while GBZ is limited to MNIST. Modeling likelihood on image representations, SDIM models easily detect OOD samples. While fully generative models, e.g. GBZ, who model likelihood on raw image pixels, are known to fail on this task.

5 Conclusions and Future Directions

Though some promising results of the robustness of generative classifiers have been observed, it is challenging to scale them on complex datasets. In this paper, we introduce supervised probabilistic constraints to DIM. Giving up the full generative process, SDIMs are equivalent to generative classifiers on high-level data representations. Unlike full conditional generative models which achieve poor classification performance even on CIFAR10, SDIMs attain same-level performance as the comparable discriminative counterparts on complex datasets. The training of SDIM is also computationally similar to discriminative classifiers, and does not require prohibitive computational resources. Our proposed rejection policy based on off-manifold conjecture, a built-in defense mechanism of SDIM, can effectively reject illegal inputs, including adversarial examples, and OOD samples. We demonstrate that likelihoods modeled on high-level data representations, rather than raw pixel intensities, are more robust on downstream tasks without the requirement of generating real samples.

The rejection mechanism in this paper is different but complementary to other defense mechanisms, e.g. adversarial training. Performing classification with rejection, classifiers can refuse to make low-confidence predictions, while adversarial training aims to inherently improve models' recognition robustness. It is to be explored how to combine them to build more trustworthy models.