Variational Deterministic Uncertainty Quantification

1 INTRODUCTION . Deploying machine learning algorithms as part of automated decision making systems , such as self driving cars and medical diagnostics , requires implementing fail-safes . Whenever the model is presented with a novel or ambiguous situation , it would not be wise to simply trust its prediction . Instead , the system should try to get more information or simply withhold or defer judgment . While significant progress has been made towards estimating predictive uncertainty reliably in deep learning ( Gal & Ghahramani , 2016 ; Lakshminarayanan et al. , 2017 ) , there is no single method that is shown to work on large datasets in classification and regression without significant computation overheads , such as multiple forward passes . We propose Variational Deterministic Uncertainty Quantification ( vDUQ ) , a method for obtaining predictive uncertainty in deep learning for both classification and regression problems in only a single forward pass . In previous work , van Amersfoort et al . ( 2020 ) show that combining a distance aware decision function with a regularized feature extractor in the form of a deep RBF network , leads to a model ( DUQ ) that matches a softmax model in accuracy , but is competitive with Deep Ensembles for uncertainty on large datasets . The feature extractor is regularized using a two-sided gradient penalty , which encourages the model to be sensitive to changes in the input , avoiding feature collapse , and encouraging generalization by controlling the Lipschitz constant . This model , however , has several limitations ; for example the uncertainty ( a distance in feature space ) can not be interpreted probabilistically and it is difficult to disentangle aleatoric and epistemic uncertainty . Additionally , the loss function and centroid update scheme are not principled and do not extend to regression tasks . A probabilistic and principled alternative to deep RBF networks are Gaussian Processes ( GPs ) in combination with Deep Kernel Learning ( DKL ) ( Hinton & Salakhutdinov , 2008 ; Wilson et al. , 2016b ) . DKL was introduced as a “ best of both worlds ” solution : apply a deep model on the training data and learn the GP in feature space , ideally getting the advantages of both models . In practice , however , DKL suffers from the same failure as Deep RBF networks : the deep model is free to map out of distribution data close to the feature representation of the training data , removing the attractive properties of GPs with distance sensitive kernels . Using insights from DUQ , we are able to mitigate the problems of uncertainty collapse in DKL . In particular , we use direct spectral normalization ( Gouk et al. , 2018 ; Miyato et al. , 2018 ) in combination with a ResNet ( He et al. , 2016 ) , a variation that was suggested in Liu et al . ( 2020 ) . The spectral normalization enforces smoothness , while the residual connections enforce sensitivity of the feature represenation to changes in the input , obtaining a similar effect as the gradient penalty of DUQ . We use an inter-domain inducing point variational approximation of the GP predictive distribution ( Lázaro-Gredilla & Figueiras-Vidal , 2009 ; Hensman et al. , 2015 ) , which places inducing points in feature space leading to needing fewer inducing points than previous work ( Wilson et al. , 2016a ) . These two techniques combined speed up inference in the GP model and decouple it from the dataset size . We release our code1 and hope that it will become a drop in alternative for softmax models with improved uncertainty . In Figure 1 , we show how vDUQ and Deep Ensembles ( Lakshminarayanan et al. , 2017 ) , the current state of the art for uncertainty quantification ( Ovadia et al. , 2019 ) , perform on simple 1D regression . This task is particularly hard for deep networks as shown in Foong et al . ( 2019 ) . vDUQ shows the desired behavior of reverting back to the prior away from the data , while the Deep Ensemble extrapolates arbitrarily and confidently . In between the two sinusoids , the Deep Ensemble is certain while vDUQ increases its uncertainty . In summary , our contributions are as follows : • We improve training a DKL model and for the first time match the accuracy and speed of training a deep network using regular softmax output on standard vision benchmarks . • We demonstrate excellent uncertainty quantification in classsification which matches or exceeds the state of the art on CIFAR-10 , including ensembling approaches . • We show state of the art performance on causal inference for personalized medicine , an exciting real world application . This task requires calibrated uncertainty in regression to be able to defer treatment to an expert when uncertainty is high . 2 BACKGROUND . Gaussian Processes ( GPs ) provide an interpretable , explainable and principled way to make predictions , and can work well even with little training data due to their use of Bayesian inference . In contrast to deep neural networks , GPs have high uncertainty away from the training data and on noisy inputs . There are however two main issues with the standard GP setup : poor performance on high dimensional inputs and inefficient computational scaling with large datasets . The poor performance on high dimensional inputs is due to the fact that most standard shift-invariant kernels are based on 1Available at : anonymized-for-review the Euclidean distance , which is a poor metric for highly structured data like images ( Theis et al. , 2016 ) . While kernels exist that better address this ( Van der Wilk et al. , 2017 ; Jacot et al. , 2018 ) , these are more computationally expensive and typically still underperform standard convolutional neural networks . Deep Kernel Learning ( Calandra et al. , 2016 ; Hinton & Salakhutdinov , 2008 ; Wilson et al. , 2016b ) is a way to combine the expressiveness of deep neural networks with the attractive properties of Gaussian Processes . The core idea is to use a deep neural network inside the kernel of a GP , k ( xi , xj ) → k ( fθ ( xi ) , fθ ( xj ) ) , where fθ ( · ) is a deep neural network , such as a Wide ResNet ( Zagoruyko & Komodakis , 2016 ) up to the last linear layer , parametrized by θ . The kernel k ( · , · ) can be any of the standard kernels , such as the RBF or Matérn kernel . With the deep network it becomes possible to train the GP on datasets that contain high dimensional points , such as images . As the deep network is unconstrained , uncertainty away from the data collapses ( Bradshaw et al. , 2017 ) , and the model can be arbitrarily confident while no data has been observed in that region during training . This has been a major drawback to using DKL in practice . In Section 3 we provide a detailed discussion of this effect and ways to mitigate it . While DKL is a potentially powerful solution to overcome the problem of inexpressive kernels , it does not address the other scaling problem of Gaussian Processes : large datasets . The poor computational scaling comes from the fact that making predictions with an exact GP requires solving a linear system the size of the training data , namely computing K ( X , X ) −1K ( X , x∗ ) where K ( X , X ) is the kernel matrix evaluated on the training dataset X and K ( X , x∗ ) the kernel between the training dataset and a test point x∗ . This is a significant computational problem in its own right . A powerful and principled alternative to exact GP inference is inducing point GP approximate inference , where only a small set of inducing points in the input space U are used to represent the entire training set , reducing the linear system to be solved to K ( u , u ) −1K ( u , x∗ ) . The new linear system is m by m , where m is the number of inducing points , rather than N by N . Finding the optimal set of points , which are not necessarily part of our training set , can be done by treating them as variational parameters and maximizing a lower bound on the marginal log-likelihood ( Titsias , 2009 ) . We follow the variational inducing point formulation of SVGPC ( Hensman et al. , 2015 ) as implemented in GPyTorch ( Gardner et al. , 2018 ) . 3 METHOD . vDUQ is an instance of Deep Kernel Learning , combining a constrained deep neural network with an inducing point GP . The model is learned end-to-end using variational inference by means of gradient Algorithm 1 Algorithm for training vDUQ Initialization : - Residual NN fθ : x → Rd with feature space dimensionality d and parameters θ initialized using He et al . ( 2015 ) . - Approximate GP with variational parameters φ and number of inducing points m. Using a random subset of k points of our training data , X init ⊂ X , compute : Initial inducing points : K-means on f ( X init ) with K = m. Use found centroids as initial inducing point locations in GP . Initial length scale : l = 1 ( k2 ) ∑k i=0 ∑k j=i+1 |f ( X initi ) − f ( X initj ) |2 . 1 : for minibatch Bx = { x1 , ... , xb } from X and By = { y1 , ... , yb } from Y do 2 : θ′ ← spectral normalization ( θ ) 3 : ψ ← fθ′ ( Bx ) . Compute feature representation 4 : p ( Y ′|Bx ) ← GPφ ( ψ ) . Compute posterior over labels , implemented in GPyTorch 5 : L← ELBOφ ( p ( Y ′|Bx ) , By ) . Compute loss , implemented in GPyTorch 6 : φ , θ ← φ , θ + η ∗ ∇φ , θL . η the learning rate , alternatively use ADAM 7 : end for descent . In this section we discuss how vDUQ overcomes previous problems with collapsing uncertainty in DKL ( Bradshaw et al. , 2017 ) and explain how to learn the model with no pre-training , few inducing points , and a standard minibatch size , avoiding some shortcomings of Wilson et al . ( 2016a ) . Without restrictions , the deep network inside the kernel is free to map input points that are far away from the training distribution to a feature representation that resembles those of data points in the training distribution . This behavior is also exhibited in standard neural networks ( Smith & Gal , 2018 ) and sometimes referred to as feature collapse . We visualize this in Figure 2 , where we map a 2D input space into feature space . With the feature space of the unconstrained model it is impossible to recover uncertainty : many points are collapsed on top of each other . In DUQ ( van Amersfoort et al. , 2020 ) , the authors show that it is possible to reduce feature collapse by enforcing two constraints on the model : sensitivity and smoothness . Sensitivity implies that when the input changes the feature representation also changes : this means the model can not simply collapse feature representations arbitrarily . Smoothness means small changes in the input can not cause massive shifts in the output . This appears to help optimization , and ensures the feature space accords with the implicit assumptions that for example RBF kernels make about the data . We discuss a connection between these two requirements , sensitivity and smoothness , and bi-Lipschitz functions in Section 6 . There is a number of methods proposed in the literature that attempt to satisfy these constraints and each comes with different trade-offs : • Two-sided gradient penalty : In DUQ ( van Amersfoort et al. , 2020 ) , the properties are achieved by regularizing using a two-sided gradient penalty , that penalizes the squared distance of the gradient from a fixed value at every input point . This approach is easy to implement , but is not guaranteed to work , and in practice both the stability of training and its effectiveness as a regularizer can be fairly sensitive to the weighting of this penalty . • Direct spectral normalization and residual connections : spectral normalization ( Miyato et al. , 2018 ; Gouk et al. , 2018 ) on the weights leads to smoothness and it is possible to combine this with an architecture that contains residual connections for the sensitivity constraint ( Liu et al. , 2020 ) . This method is faster than the gradient penalty , and in practice a more effective way of mitigating feature collapse . • Reversible model : A reversible model is constructed by using reversible layers and avoiding any down scaling operations ( Jacobsen et al. , 2018 ; Behrmann et al. , 2019 ) . This approach can guarantee that the overall function is bi-Lipschitz , but the resulting model consumes considerably more memory and can be difficult to train . In this work , we use direct spectral normalization and residual connections , as we find it to be more stable than a direct gradient penalty and significantly more computationally efficient than reversible models . In Figure 2 , we show that a constrained model is unable to collapse points on top of each other in feature space , enabling the GP to correctly quantify uncertainty . Using a stationary kernel , the model then reverts back to the prior away from the training data just like a standard GP . The regularized feature extractor allows us to offload computational complexity from the GP onto the deep model without sacrificing uncertainty . An expressive deep model is able to find a feature representation which generalizes across all intra class variation , and thus cluster inputs of the same class in feature space . We define the inducing points in feature space instead of input space ( also known as an interdomain approximation ) ( Lázaro-Gredilla & Figueiras-Vidal , 2009 ; Hensman et al. , 2017 ) , which takes advantage of this clustering to reduce the number of inducing points required . In practice , we find that only very few inducing points , the number of classes in the case of classification , are necessary to obtain a well performing GP and we found similarly low numbers ( in the 10 to 100 range ) work well for regression , which means that solving the linear system is fast and the GP has minimal overhead compared to a softmax output . The fact that we can use few inducing points is not true in general for DKL ( e.g . Wilson et al . ( 2016a ) use several hundred points in input space ) , but requires the extra restrictions we have placed on our feature extractor to avoid the pathological feature collapse . We also find we can train with standard minibatch sizes of 128 , while previous work used 5,000 on CIFAR-10 to combat gradient variance ( Wilson et al. , 2016a ) . It is important to note that in DKL and also in this paper we are not being Bayesian about our neural network parameters θ , so we only need a single pass through the feature extractor both when training and during inference .

Summary: This paper proposes variational deterministic uncertainty quantification (vDUQ), which adopts the stochastic (sparse) variational deep kernel learning (DKL) method to enable uncertainty estimations for deep models. To avoid uncertainty collapse, the deep neural network in the GP kernel is regularized with spectral normalization, which ensures a bi-Lipschitz constraint. Experiments show that vDUQ is effective in uncertainty quantification tasks.