Noise-Contrastive Variational Information Bottleneck Networks

1 INTRODUCTION Deep neural networks ( DNNs ) have become the standard tool for challenging classification tasks , e.g . image classification or semantic segmentation , due to their excellent predictive accuracy . However , predictions of DNNs often tend to be overconfident , leading to miscalibration ( Guo et al. , 2017 ) . This problem is amplified in the presence of distributional shift in the test data , such as from image corruptions ( Ovadia et al. , 2019 ) . Multiple methods to regularize the output distribution of DNNs during training have been proposed ( Joo & Chung , 2020 ; Joo et al. , 2020 ; Pereyra et al. , 2017 ; Szegedy et al. , 2016 ) to obtain well-calibrated models . We note , however , that evaluating in-domain uncertainty quantification in terms of model calibration is not sufficient , as it does not indicate how well correct and incorrect predictions can be discerned based on the predictive uncertainty ( Ding et al. , 2020 ) . As Fig . 1 shows , methods that indiscriminately regularize the confidence to improve calibration perform significantly worse at separating correct from incorrect predictions . To address this , we turn to deep variational information bottleneck networks ( VIBN ; Alemi et al. , 2017 ) . Similarly to output regularization methods , they benefit generalization and calibration ( Alemi et al. , 2018 ) but , as we show empirically , suffer from the same separability problem . However , unlike these methods , VIBNs allow us to utilize the distribution matching of the variational approximation to define a noise-contrastive loss term , which can overcome the problem of insufficient separation of correct and incorrect predictions based on the uncertainty while retaining the benefits of the VIBN in terms of generalization and calibration . To this end , we propose a novel model , the noise-contrastive variational information bottleneck network ( NC-VIBN ) , which builds upon the VIBN to improve uncertainty estimation . Instead of using distribution matching as the primary source of regularization , our model utilizes it to define a loss term that explicitly encourages high predictive entropy only for uninformative samples from the latent prior . Additionally , we account for weight uncertainties in the decoder and use L2-normalization before computing the latent embeddings to further alleviate the described problems while improving the calibration and generalization capabilities of the model . We make the following contributions : ( i ) We empirically show that models that explicitly regularize the prediction confidence make it harder to distinguish between correct and incorrect predictions based on the estimated uncertainty . ( ii ) We link the VIBN to these methods and find that it suffers from the same behavior due to the implicit L2-regularization through the latent KL-divergence term . ( iii ) We circumvent these ill effects regarding separability of correct and incorrect predictions by proposing a noise-contrastive loss term that utilizes distribution matching in the latent space and , combined with architectural refinements , leads to improved separability , calibration , and accuracy . ( iv ) Our results show that our proposed model also leads to improved accuracy and calibration in the presence of distributional shift introduced by image corruptions . 2 RELATED WORK . Information Bottleneck . The information bottleneck principle has been proposed by Tishby et al . ( 1999 ) as means to analyze generalization for deep neural networks ( Tishby & Zaslavsky , 2015 ) and relies on computing the mutual information of input , output , and intermediate variables . Since these quantities are in general intractable , variational approximations of the mutual information terms have been introduced ( Alemi et al. , 2017 ; Achille & Soatto , 2018 ) . These variational approximations are also able to overcome problems of the information bottleneck objective for deterministic representations ( Amjad & Geiger , 2019 ) and share some of the benefits of Bayesian models ( Alemi et al. , 2020 ) . Further , as shown by Alemi et al . ( 2018 ) , variational information bottleneck networks can improve calibration , which we want to further improve with our method . Noise-contrastive estimation . Noise-contrastive estimation ( Gutmann & Hyvärinen , 2010 ) is an estimation method for parameterized densities , which is based on discriminating between real and artificial data by logistic regression using the log-density functions . Inspired by noise-contrastive estimation , Hafner et al . ( 2020 ) propose noise-contrastive priors as data space priors that encourage uncertain predictions at the boundary of the training data for regression tasks by minimizing the KL-divergence between a high-variance Gaussian and the predicted output distribution for perturbed training data points . We focus on classification instead . Uncertainty estimation . Bayesian neural networks ( BNNs ) are a popular and theoretically wellfounded tool for uncertainty estimation . A multitude of methods have been proposed to approximate the intractable weight posterior , including variational inference ( Blundell et al. , 2015 ; Graves , 2011 ) , Markov chain Monte Carlo methods ( Welling & Teh , 2011 ) , Laplace approximation ( MacKay , 1992 ; Ritter et al. , 2018 ) , and assumed density filtering ( Hernández-Lobato & Adams , 2015 ) , as well as approximate variational inference methods based on the dropout ( Srivastava et al. , 2014 ) regularization scheme ( Gal & Ghahramani , 2016 ) . In the last years , these methods have been successfully scaled to larger models ( Dusenberry et al. , 2020 ; Heek & Kalchbrenner , 2019 ; Maddox et al. , 2019 ; Osawa et al. , 2019 ; Zhang et al. , 2020 ) . Deep ensembles ( Lakshminarayanan et al. , 2017 ) have also been used for uncertainty estimation and can be interpreted as Bayesian model averaging . Since model averaging with respect to the approximate posterior requires multiple forward passes , BNNs incur a substantial computational overhead . To lighten this , there has been an interest in Bayesian last layer approaches ( van Amersfoort et al. , 2021 ; Kristiadi et al. , 2020 ; Liu et al. , 2020 ; Riquelme et al. , 2018 ; Snoek et al. , 2015 ; Wilson et al. , 2016 ) . Our approach similarly employs a Bayesian treatment only for the last ( few ) layers , however combining it with the information bottleneck principle and a noise-contrastive loss . An alternative approach for estimating uncertainty in classification networks is parameterizing the more expressive Dirichlet distribution ( Gast & Roth , 2018 ; Joo et al. , 2020 ; Malinin & Gales , 2018 ; Sensoy et al. , 2018 ) instead of the categorical distribution at the output layer . We instead gain additional expressiveness by modelling the latent space distributions . Related methods in out-of-distribution detection . Lee et al . ( 2018 ) proposed to train confidencecalibrated classifiers by using a generative adversarial network ( GAN ) that learns to generate samples at the data boundary . The generator is trained to generate data points , which are hard to separate from in-distribution data by the discriminator while given an almost uniform labeling by the classifier . In contrast , Sricharan & Srivastava ( 2018 ) train the generator to produce low-entropy in-distribution samples while requiring the classifier to maximize the entropy of those samples . 3 VARIATIONAL APPROXIMATION OF THE INFORMATION BOTTLENECK . We begin by recapitulating two variational approximations of the information bottleneck , the deep variational information bottleneck ( Alemi et al. , 2017 ) and information dropout ( Achille & Soatto , 2018 ) , which we will use to explain the behavior of such models and to build our own model upon . Deep variational information bottleneck . The information bottleneck was first introduced ( Tishby et al. , 1999 ; Tishby & Zaslavsky , 2015 ) to find a low-complexity representation Z depending on a feature vector X that maximizes the mutual information with a target variable Y . To constrain the complexity of Z , the mutual information between X and Z is bounded , resulting in a maximization problem with inequality constraints . Alemi et al . ( 2017 ) proposed to use a variational approximation of the mutual information terms of the Lagrangian with Lagrange multiplier β , resulting in the objective min ϕ , ψ 1 N N∑ n=1 Epϕ ( z|xn ) [ − log qψ ( yn|z ) ] + βDKL [ pϕ ( z|xn ) ∥∥r ( z ) ] , ( 1 ) where the stochastic encoder pϕ ( z|x ) and decoder qψ ( y|z ) are modeled by neural networks , parameterized by ϕ and ψ respectively , and r ( z ) is a variational approximation of the marginal distribution p ( z ) = ∫ pϕ ( z|x ) p ( x ) dx of z . To draw a parallel to the variational inference literature ( Kingma & Welling , 2014 ) , r ( z ) is also referred to as the latent prior . Alemi et al . ( 2017 ) assume r ( z ) to be a standard Gaussian and model the distribution of the latent encodings as Gaussians with diagonal covariance , resulting in DKL [ pϕ ( z|xn ) ∥∥r ( z ) ] = 1 2 ∑ i [ − log σ2zi|xn + σ 2 zi|xn + µ 2 zi|xn − 1 ] , ( 2 ) where µzi|xn and σ 2 zi|xn are the component-wise mean and variance of the latent embedding pϕ ( z|xn ) of xn , estimated by the encoder network . The decoder is a softmax classification network hψ , predicting class probability vectors from z , i.e . qψ ( y|z ) = Cat ( y|hψ ( z ) ) . Information dropout ( IDO ) . Achille & Soatto ( 2018 ) proposed an alternative approach , where the encoder network gϕ predicts a non-negative vector µx as well as αx , parameterizing the log-normal distribution logN ( 0 , α2x ) . The distribution of the latent encodings is modeled as z ∼ µx ⊙ ϵ with ϵ ∼ logN ( 0 , α2x ) . They show that if µx is the output of ReLU units and the latent prior r ( z ) is chosen to be a mixture of the delta distribution at 0 and a log-uniform distribution , i.e . r ( z ) ∝ qδ0 ( z ) + c/z , the KL-divergence is given as DKL [ pϕ ( z|x ) ∥∥r ( z ) ] = { − log q µx = 0−H [ pαx ( log ϵ ) ] + log c µx > 0 , ( 3 ) where the entropy term H [ pαx ( log ϵ ) ] is given by logαx for ϵ ∼ logN ( 0 , α2x ) up to an additive constant . In the original formulation of Achille & Soatto ( 2018 ) , the mean of ϵ grows with αx , resulting in a higher level of saturation of the softmax outputs , hence in overconfidence . Note that if ϵ is log-normal distributed , log ϵ is normal distributed and the entropy does not depend on its mean . Therefore , we here instead employ the mean-corrected log-normal distribution logN ( − 12α 2 x , α 2 x ) so that Epαx ( ϵ ) [ µx ⊙ ϵ ] = µx without changes to the KL-divergence . 4 UNCERTAINTY QUANTIFICATION UNDER OUTPUT DISTRIBUTION REGULARIZATION . A frequently used metric to assess uncertainty estimation is the expected calibration error ( Guo et al. , 2017 ) or related calibration metrics , which measure how well the prediction confidence coincides with the prediction accuracy . Methods that achieve better calibration by output distribution regularization include label smoothing ( Müller et al. , 2019 ; Szegedy et al. , 2016 ) , regularization of the predictive entropy ( Pereyra et al. , 2017 ) for the categorical distribution , or evidential deep learning ( Sensoy et al. , 2018 ) and belief matching ( Joo et al. , 2020 ) for the Dirichlet distribution . Alternatively , the model can be regularized in function or distribution spaces prior to the final softmax layer ( Joo et al. , 2020 ) , and in the simplest case results in norm penalties for the predicted logits . However , these calibration metrics actually do not indicate how well correct predictions can be separated from incorrect ones based on some uncertainty estimate like the predictive entropy . In fact , a model predicting with a confidence equal to the expected accuracy on in-domain data for both correct classifications and misclassifications would be perfectly calibrated according to the expected calibration error , but the confidence of a prediction would give us no information about an example being correctly or incorrectly classified ( Ding et al. , 2020 ) . A suitable method to examine how easily correct and incorrect predictions can be discriminated , is to analyze how the predictive performance behaves under the rejection of uncertain predictions ( Ding et al. , 2020 ; Nadeem et al. , 2009 ) . As can be seen in Fig . 1 , models that are trained with an objective that regularizes the output distribution perform significantly worse at assigning incorrect predictions a comparatively higher uncertainty and rejecting them early than the deterministic model . We argue that this is because these methods indiscriminately regularize the confidence of correct and incorrect predictions . More experimental evidence is presented in Sec . 6.1 .

Summary: This paper attempts to improve uncertainty estimates for neural network-based classification. To do so, they combine various techniques: the information bottleneck, a variational distribution over neural net weights, a noise contrastive loss term, and an L2-normalisation step for the last part of encoder. They present experimental results justifying these choices (including, notably, significantly improved NLL compared to the baselines).