Reject Illegal Inputs: 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 ; Goodfellow et al. , 2014 ) . Many attack methods have been developed to find imperceptible perturbations to fool the target classifiers ( Moosavi-Dezfooli et al. , 2016 ; Carlini & 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 ; Tramèr et al. , 2017 ; Madry et al. , 2017 ; Samangouei et al. , 2018 ) . 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 & 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 et al. , 2018 ; Fetaya et al. , 2019 ) . ( 2 ) The behaviors of likelihood-based generative models can be counter-intuitive and brittle . They may assign surprisingly higher likelihoods to out-of-distribution ( OoD ) samples ( Nalisnick et al. , 2018 ; Choi & 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 . 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 high-level 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 ( Nalisnick et al. , 2019 ; Geifman & El-Yaniv , 2017 ) . SDIMs reject illegal inputs based on off-manifold conjecture ( Samangouei et al. , 2018 ; Gu & 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 over-parameterization . • 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 prechosen 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 esimator 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 & Varadhan , 1983 ) representation of KL-divergence : I ( X , Y ) = DKL ( J||M ) ≥ EJ [ Tω ( x , y ) ] − logEM [ eTω ( x , y ) ] ( 1 ) 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 , nonKL divergences are also favorable candidates . We can get a family of variational lower-bounds using f -divergence representations ( Nguyen et al. , 2010 ) : If ( X , Y ) ≥ EJ [ Tω ( x , y ) ] − EM [ f∗ ( Tω ( x , y ) ) ] ( 2 ) 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 ) . Noise-Contrastive Estimation ( Gutmann & Hyvärinen , 2010 ) can also be used as lower-bound of MI in “ infoNCE ” ( Oord et al. , 2018 ) . 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 : JMI = − 1 M2 M2∑ i=1 Ĩ ( L ( i ) φ ( x ) , Eφ ( x ) ) ( 3 ) 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 : DKL ( p ( x̃|y ) ||pθ ( x̃|y ) ) = Ex̃ , y∼p ( x̃ , y ) [ log p ( x̃|y ) − log pθ ( x̃|y ) ] = Ex̃ , y∼p ( x̃ , y ) [ log p ( x̃|y ) ] − Ex̃ , y∼p ( x̃ , y ) [ log pθ ( x̃|y ) ] ( 4 ) where the first item on RHS is a constant independent of the model parameters θ. Eq . 4 equals to maximize the expectation Ex̃ , y∼p ( x̃ , y ) [ log pθ ( x̃|y ) ] . In practice , we minimize the following loss JNLL , equivalent to empicically maximize the above expectation over { x̃i = Eφ ( xi ) , yi } Ni=1 : JNLL = −Ex̃ , y∼p ( x̃ , y ) [ log pθ ( x̃|y ) ] ≈ − 1 N N∑ i=1 log pθ ( x̃i|yi ) ( 5 ) Besides the introduction of supervision , SDIM differs from DIM in its way of enforcing the statistical constraints : DIM use adversarial learning ( Makhzani et al. , 2015 ) to 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 : JLM = 1 N · 1 C − 1 N∑ i=1 C∑ c=1 , c 6=yi max ( log p ( x̃i|y = c ) +K − log p ( x̃i|y = yi ) , 0 ) 2 ( 6 ) where K is a positive constant to control the margin . For each encoder output x̃i , the C − 1 truefalse class conditional gaps are squared1 , which quadratically increases the penalties when the gap becomes large , then are averaged . Putting all these together , the complete loss function we minimize is : JSDIM = α · JMI + β · JNLL + γ · JLM ( 7 ) 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 et al . ( 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 : p ( y|x ) = p ( x|y ) p ( y ) p ( x ) ∝ p ( x|y ) p ( y ) 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 : y∗ = argmax c= [ 1 ... C ] log p ( x∗|y = c ) . ( 8 ) 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 : { y∗ , if log p ( x∗|y∗ ) ≥ δy∗ Rejection , otherwise ( 9 ) 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 . This is also known as selective classification ( Geifman & El-Yaniv , 2017 ) or classification with reject option ( Nalisnick et al. , 2019 ) ( See Supp . A )

Summary: This paper studies classification problems via a reject option. A reject option could be useful in prediction problems to handle Out-of-distribution examples. The classification procedure studied in this paper builds on three components 1. An auto-encoder that obtains a latent low-dimensional representation of the data point 2.  A generative model that models the class-conditional probability model and 3.  a margin based loss function that learns a classifier that provides a large probability mass to the class-conditional distribution corresponding to the correct class.  The final decision procedure is to reject an input if the best class conditional probability is small and to use the class corresponding to the best class conditional probability otherwise.