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 work suggests a *noise-contrastive loss* for variational information bottleneck networks to resolve the poor performance at separating correct and incorrect predictions in regularization methods. Standard regularization methods suffer *separability problems*: models indiscriminately regularize confidence to improve calibration perform much worse at differentiating correct from wrong predictions. The authors address this problem in the variational information bottleneck, suggesting a remedy that utilizes the distribution matching in the latent space. The authors have applied the proposed loss term to the two variational approximations of the information bottleneck, the deep variational information bottleneck (VIBN, Alemi et al., 2017) and information dropout (IDO, Achille & Soatto, 2018).