Meta Learning Low Rank Covariance Factors for Energy Based Deterministic Uncertainty

Numerous recent works utilize bi-Lipschitz regularization of neural network layers to preserve relative distances between data instances in the feature spaces of each layer . This distance sensitivity with respect to the data aids in tasks such as uncertainty calibration and out-of-distribution ( OOD ) detection . In previous works , features extracted with a distance sensitive model are used to construct feature covariance matrices which are used in deterministic uncertainty estimation or OOD detection . However , in cases where there is a distribution over tasks , these methods result in covariances which are sub-optimal , as they may not leverage all of the meta information which can be shared among tasks . With the use of an attentive set encoder , we propose to meta learn either diagonal or diagonal plus low-rank factors to efficiently construct task specific covariance matrices . Additionally , we propose an inference procedure which utilizes scaled energy to achieve a final predictive distribution which can better separate OOD data , and is well calibrated under a distributional dataset shift . 1 INTRODUCTION . Accurate uncertainty in predictions ( calibration ) lies at the heart of being able to trust decisions made by deep neural networks ( DNNs ) . However , DNNs can be mis-calibrated when given out-ofdistribution ( OOD ) test examples ( Ovadia et al. , 2019 ; Guo et al. , 2017 ) . Hein et al . ( 2019 ) show that the problem can arise from ReLU non-linearities introducing linear polytopes into decision boundaries which lead to arbitrary high confidence regions outside of the domain of the training data . Another series of works ( van Amersfoort et al. , 2021 ; Liu et al. , 2020a ; Mukhoti et al. , 2021 ; van Amersfoort et al. , 2021 ) link the problem to feature collapse , whereby entire regions of feature space collapse into singularities which then inhibits the ability of a downstream function to differentiate between points after the collapse , which destroys any information which could be used to differentiate them . When these collapsed regions include areas of OOD data , the model loses any ability to differentiate between in-distribution ( ID ) and OOD data . A solution to prevent feature collapse is to impose bi-Lipschitz regularization into the network which enforces both an upper and lower Lipschitz bound on the functions operating in feature space ( van Amersfoort et al. , 2021 ; Liu et al. , 2020a ) , preventing feature collapse . Such features from biLipschitz regualarized extractors are then used to improve downstream tasks such as OOD detection or uncertainty quantification . Broadly speaking , previous works have done this by constructing covariance matrices from the resulting features in order to aid in the uncertainty quantification ( Liu et al. , 2020a ; Van Amersfoort et al. , 2020 ) or OOD detection ( Mukhoti et al. , 2021 ) . Intuitively , features from a Lipschitz regularized extractor make for more expressive covariances , due to the preservation of identifying information within different features . However , empirical covariance estimation is limited when there are few datapoints on hand , such as in few-shot learning . A key aspect of few-shot learning is to learn from a task distribution , but empirical covariance estimation methods are not able to effectively encode shared knowledge over the task distribution , even when the features used to calculate the covariance come from a meta-learned feature extractor ( see Figure 6 ) . As a result , the empirical covariance matrices are not expressive given limited data and thus the model loses its ability to effectively adapt to a task and make calibrated predictions . Another obstacle , highlighted by Mukhoti et al . ( 2021 ) , is that plain softmax classifier can not accurately model epistemic uncertainties . We identify a contributing factor to this , which is the shift invariance property of the softmax function . Specifically , even if an evaluation point comes from an OOD area and is assigned low logit values ( low energy ) , this alone is insufficient for a well calibrated prediction . Small variations in logit values can lead to arbitrarily confident predictions due to the shift invariance . From the perspective of Prototypical Networks ( Snell et al. , 2017 ) , we highlight this problem in Figure 3 , although it applies to linear softmax classifiers as well . In the following work , we first propose a method of meta-learning class-specific covariance matrices that is transferable across the task distribution . Specifically , we meta-learn a function that takes a set of class examples as an input and outputs a class-specific covariance matrix which is in the form of either a diagonal or diagonal plus low-rank factors . By doing so , the resulting covariance matrices remain expressive even with limited amounts of data . Further , in order to tackle the limitation caused by the shift invariance property of the softmax function , we propose to use scaled energy to parameterize a logit-normal softmax distribution which leads to better calibrated softmax scores . We enforce its variance to increase as the energy increases , and vice versa . In this way , the softmax prediction can become progressively more uniform between ID and OOD data , after marginalizing over logit-normal distribution ( see example in Figure 1 ) . By combining those two components , we have the final inference procedure to effectively separate OOD data , and achieves a well calibrated probabilistic model using a deterministic DNN . Our contributions are as follows : • We show that existing approaches fail to generalize to the meta-learning setting . • We propose a meta learning framework which predicts diagonal or low-rank covariance factors as a function of a support set . • We propose an energy-based inference procedure which leads to improved OOD detection , and better calibrated uncertainty . 2 RELATED WORK . Mahalanobis Distance . Mahalanobis distance has been used in previous works for OOD detection ( Lee et al. , 2018 ) which also showed that there is a connection between softmax classifiers and Gaussian discriminant analysis , and that the representation space in the latent features of DNN ’ s provides for an effective multivariate Gaussian distribution which can be more useful in constructing class conditional Gaussian distributions than the output space of the softmax classifier . The method outlined in Lee et al . ( 2018 ) provides a solid groundwork for our method , which also utilizes Mahalanobis distance in the latent space , and adds a deeper capability to learn meta concepts which can be shared over a distribution of tasks . Post Processing . We refer to post-processing as any method which applies some function after training and before inference in order to improve the test set performance . In the calibration literature , temperature scaling ( Guo et al. , 2017 ) is a common and effective post-processing method . As the name suggests , temperature scaling scales the logits by a constant before going through the softmax function . The temperature is tuned such that the negative log-likelihood ( NLL ) on a validation set is minimized . Previous works which utilize covariance ( Lee et al. , 2018 ; Mukhoti et al. , 2021 ; Liu et al. , 2020a ) have also applied post-processing methods to construct latent feature covariance matrices after training . While effective for large data , these post-processing methods make less expressive covariances in the limited data regime of few-shot learning , as we demonstrate in Figure 2 . Bi-Lipschitz Regularization . Adding a regularizer to enforce functional smoothness of a DNN is a useful tactic in stabilizing the training of generative adversarial networks ( GANs ) ( Miyato et al. , 2018 ; Arjovsky et al. , 2017 ) , improving predictive uncertainty ( Liu et al. , 2020a ; Van Amersfoort et al. , 2020 ) , and aiding in OOD detection ( Mukhoti et al. , 2021 ) . By imposing a smoothness constraint on the network , distances which are semantically meaningful w.r.t . the feature manifold can be preserved in the latent representations , allowing for downstream tasks ( such as uncertainty estimation ) to make use of the preserved information . ( Van Amersfoort et al. , 2020 ) showed that without this regularization , a phenomena known as feature collapse can map regions of feature space onto singularities ( Huang et al. , 2020 ) , where previously distinct features now become indistinguishable . For both uncertainty calibration and OOD detection , feature collapse can map OOD features onto the same feature spaces as ID samples , adversely affecting both calibration and OOD seperability . Meta Learning . The goal of meta learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) is to leverage shared knowledge which may apply across a distribution of tasks . In the few shot learning scenario , models leverage general meta-knowledge gained through episodic training on a task distribution ( Vinyals et al. , 2016 ; Ravi & Larochelle , 2017 ) , which allows for effective adaptation and inference on a task which may contain only limited amounts of data during inference . The current meta-learning approaches are roughly categorized into metric-based ( Vinyals et al. , 2016 ; Snell et al. , 2017 ) or optimization-based approaches ( Finn et al. , 2017 ) . In this work , we focus on metric-based approaches because they are closely related to generative classifiers , which have been shown to be important for epistemic uncertainty ( Mukhoti et al. , 2021 ) . 3 APPROACH . We start by introducing a task distribution p ( τ ) which randomly generates tasks containing a support set S = { ( x̃i , ỹi ) } Nsi=1 and a query set Q = { ( xi , yi ) } Nq i=1 . Then , given randomly sampled task τ = ( S , Q ) , we meta-learn a generative classifier that can estimate the class-wise distribution of query examples , p ( x|y = c , S ) conditioned on the support set S , for each class c = 1 , . . . , C. A generative classifier is a natural choice in our setting due to fact that it utilizes feature space densities which has been shown to be a requirement for accurate epistemic uncertainty prediction ( Mukhoti et al. , 2021 ) . Under the class-balanced scenario p ( y = 1 ) = · · · = p ( y = C ) we can easily predict the class labels as follows . p ( y = c|x , S ) = p ( x|y = c , S ) ∑C c′=1 p ( x|y = c′ , S ) . ( 1 ) 3.1 LIMITATIONS OF EXISTING GENERATIVE CLASSIFIERS . Possibly one of the simplest forms of deep generative classifier is ProtoNets ( Snell et al. , 2017 ) . In ProtoNets we assume a deep feature extractor fθ that embeds x to a common metric space such that z = fθ ( x ) . We then explicitly model the class-wise distribution p ( z|y = c , S ) of the embedding z instead of the raw input x . Under the assumption of a regular exponential family distribution for pθ ( z|y = c , S ) and a Bregman divergence d such as Euclidean or Mahalanobis distance , we have pθ ( z|y = c , S ) ∝ exp ( −d ( z , µc ) ) ( Snell et al. , 2017 ) , where µc = 1|Sc| ∑ x̃∈Sc fθ ( x̃ ) is the class-wise embedding mean computed from Sc , the set of examples from class c. In ProtoNets , d is squared Euclidean distance , resulting in the following likelihood of the query embedding z = fθ ( x ) in the form of a softmax function . pθ ( y = c|z , S ) = exp ( −‖z− µc‖2 ) ∑C c′=1 exp ( −‖z− µc′‖2 ) . ( 2 ) 1 . Limitations of fixed or empirical covariance . Unfortunately , Eq . ( 2 ) can not capture a nontrivial class-conditional distribution structure , as Euclidean distance in Eq . ( 2 ) is equivalent to Mahalanobis distance with fixed covariance I for all classes , such that pθ ( z|y = c , S ) = N ( z ; µc , I ) . For this reason , SNGP ( Liu et al. , 2020b ) and DDU ( Mukhoti et al. , 2021 ) calculate empirical covariances from all available training data after training . However , such empirical covariance estimations are limited especially when the dataset size is small . If we consider the few-shot learning scenario where we have only a few training examples for each class , empirical covariances will provide unreliable estimates of the true class covariance . Unreliable covariance leads to poor estimation of Mahalanobis distances and therefore unreliable OOD detection and uncertainty estimation . 2 . Shift invariant property of softmax and OOD calibration . Another critical limitation of Eq . ( 2 ) is that it produces overconfident predictions in areas distant from the class prototypes . The problem arises from the shift invariance property of the softmax function σ ( ω ) = eω/ ∑ ω′ e ω′ with ω denoting the logits , such that σ ( ω + s ) = eω+s/ ∑ ω′ e ω′+s = eω/ ∑ ω′ e ω′ = σ ( ω ) for any shift s. More specifically , suppose we have two classes c = 1 , 2 , and z moves along the line extrapolating the prototypes µ1 and µ2 such that z = µ1 + c ( µ2 − µ1 ) for c ≤ 0 or c ≥ 1 . Then , we can easily derive the following equality based on the shift invariant property of the softmax function : pθ ( y = 1|z , S ) = 1 1 + exp ( ±‖µ2 − µ1‖ ) ( 3 ) where ± corresponds to the sign of c. Note that the expression is invariant to the value of c except for its sign . Therefore , even if z is OOD , residing somewhere distant from the prototypes µ1 and µ2 with extreme values of c , we still have equally confident predictions . See Figure 3 for illustration .

Summary: The authors explore the uncertainty calibration problem in meta-learning few-shot classification, which is a relatively new topic that was seldom discussed before.  To solve the challenges that the few samples in few-shot learning are usually insufficient for providing reliable estimation to the covariance matrices, the authors introduce meta-learning covariance matrices parametrized by a set encoder (implemented as set-transformer). And the authors further introduce scaled energy to parameterize a logit-normal softmax distribution for improving the uncertainty calibration of the softmax classification layer.