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: In this paper, the authors proposed methodologies to estimate the covariance matrix of embedded features under the meta-learning setting. The authors first showed that the empirical covariance matrix is inappropriate in this setting and that due to the shift invariant property of softmax function, ProtoNet could produce overly confident predictions.   The authors instead propose to learn the covariance matrix through a Set Transformer structure, where the covariance matrix can be decomposed into the sum of a diagonal component and a low rank component. The authors used scaled energy and bi-Lipschitz regularization to make the methodology more performant.   Numerical studies are performed to evaluate the calibration performance as well as the eigenvalue distribution of several methods.