META LEARNING LOW RANK COVARIANCE FACTORS
Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 RELATED WORK
text: 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).

heading: 3 APPROACH
text: 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)

heading: 3.1 LIMITATIONS OF EXISTING GENERATIVE CLASSIFIERS
text: 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) cannot 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.

heading: 3.2 META-LEARNING OF THE CLASS-WISE COVARIANCE
text: In order to remedy the limitations of empirical covariance, and capture a nontrivial structure of the class-conditional distribution even with a small support set, we propose to meta-learn the class-wise covariances over p(τ). Specifically, we meta-learn a set encoder gφ that takes a class set Sc as input and outputs a covariance matrix corresponding to the density p(z|y = c,S), for each class c = 1, . . . , C. We expect gφ to encode shared meta-knowledge gained through episodic training over tasks from p(τ), which, as we will demonstrate in section 4, fills a key shortcoming of applying existing methods such as DDU (Mukhoti et al., 2021) and SNGP (Liu et al., 2020a). We denote the set-encoder gφ for each class c as
Λc,Φc = gφ(Zc), Zc = {z̃− µc|z̃ = fθ(x̃) and x̃ ∈ Sc}. (4)
Algorithm 1 Proto Mahalanobis – Training 1: Input: Task distribution p(τ), initial θ and φ 2: Output: Meta-learned θ and φ 3: while not converged do 4: Sample a task τ = (S,Q) 5: for c = 1 to C do 6: µc ← 1|Sc| ∑ x̃∈Sc fθ(x̃)
7: Λc,Φc ← gφ(Zc) . Eq. 4 8: Σc ← Λc + ΦcΦ>c . Eq. 5 9: Compute Σ−1c and |Σc| . Eq. 8,9
10: end for 11: Lτ ← 1|Q| ∑ (x,y)∈Q− log pθ,φ(y|E[ω]) . Eq. 11 12: (θ, φ)← (θ, φ)− α∇θ,φLτ 13: end while
Algorithm 2 Proto Mahalanobis – Inference 1: Input: Task τ , meta-learned θ and φ 2: for c = 1 to C do 3: µc ← 1|Sc| ∑ x̃∈Sc fθ(x̃)
4: Λc,Φc ← gφ(Zc) . Eq. 4 5: Σc ← Λc + ΦcΦ>c . Eq. 5 6: Compute Σ−1c and |Σc| . Eq. 8,9 7: end for 8: Eval. pθ,φ(y|z,S) for (y,x) ∈ Q . Eq. 11
where Λc ∈ Rd×d is a diagonal matrix and Φc ∈ Rd×r is a rank-r matrix. Now, instead of the identity covariance matrix or empirical covariance estimation, we have the meta-learnable covariance matrix consisting of the strictly positive diagonal and low-rank component for each class c = 1, . . . , C.
Σc = Λc + ΦcΦ > c . (5)
It is easy to see that Σc is a valid positive semi-definite covariance matrix for non-negative Λc. Note that the covariance becomes diagonal when r = 0. A natural choice for gφ is the Set Transformer (Lee et al., 2019) which can model pairwise interactions between elements of the input set, an implicit requirement for covariance matrices.
Now, we let pθ,φ(z|y = c,S) = N (z;µc,Σc). From Bayes’ rule (see Appendix A.1), we compute the predictive distribution in the form of softmax function as follows,
pθ,φ(y = c|z,S) = pθ,φ(z|y = c,S)∑C
c′=1 pθ,φ(z|y = c′,S) (6)
= exp(− 12 (z− µc) >Σ−1c (z− µc)− 12 log |Σc|)∑C c′=1 exp(− 1 2 (z− µc′)>Σ −1 c′ (z− µc′)− 1 2 log |Σc′ |)
(7)
Covariance inversion and log-determinant. Note that the logit of the softmax function in Eq. (7) involves the inverse covariance Σ−1c and the log-determinant log |Σc|. In contrast to both DDU and SNGP which propose to calculate and invert an empircal feature covariance after training, the meta-learning setting requires that this inference procedure be performed on every iteration during meta-training, which may be cumbersome if a full O(d3) inversion is to be performed. Therefore, we utilize the matrix determinant lemma (Ding & Zhou, 2007) and the Sherman-Morrison formula in the following recursive forms for both the inverse and the log determinant in Equation 7.
(Σi + Φi+1Φ > i+1) −1 i+1 = Σ −1 i −
Σ−1i Φi+1Φ > i+1Σ −1 i 1 + Φ>i+1Σ −1 i Φi+1
(8)
det(Σi + Φi+1Φ>i+1)i+1 = (1 + Φ > i+1Σ −1 i Φi+1)det(Σi) (9)

heading: 3.3 OUT-OF-DISTRIBUTION CALIBRATION WITH SCALED ENERGY
text: Next, in order to tackle the overconfidence problem caused softmax shift invariance, we propose incorporating a positive constrained function of energy h(E) = max(0,− 1T log ∑ c exp(−Ec)), with temperature T , into the predictive distribution. Energy has been used for OOD detection (Liu et al., 2020b) and density estimation (Grathwohl et al., 2019), and the success of energy in these tasks implies that it can be used to calibrate the predictive distribution (example in Figure 1h). Results in Grathwohl et al. (2019) show improvements in calibration, but their training procedure requires a full input space generative model during training, adding unwanted complexity if the end goal does not require input space generation. Our method makes use of our logit values ω = (ω1, . . . , ωC) to parameterize the mean of a logit-normal distribution with the variance given by h(E). In this way,
the logit-normal distribution variance rises in conjunction with the energy, making predictions more uniform over the simplex for higher energies.
pθ,φ(ωc|z,S) = N (ωc; µ̃c, σ̃2), where µ̃c = − 1
2 (z− µc)>Σ−1c (z− µc)−
1 2 log |Σc|,
σ̃2 = − 1 T log ∑ c′ exp ( −1 2 (z− µc′)>Σ−1c′ (z− µc′) ) (10)
Intuitively, h(E) is dominated by minc(Ec) thereby acting as a soft approximation to the minimum energy, which only becomes large when the energy is high for all classes represented in the logits. Then, the predictive distribution becomes
pθ,φ(y = c|z,S) = ∫ p(y = c|ω)pθ,φ(ω|z,S)dω (11)
≈ 1 M M∑ m=1 exp(ω (m) c )∑ c′ exp(ω (m) c′ ) , ω(m)c ∼ p(ωc|z,S). (12)
Meta-training At training time, we do not sample ω and use the simple deterministic approximation pθ,φ(y|z,S) ≈ pθ,φ(y|E[ω]). Therefore, the loss for each task becomes Lτ (θ, φ) = 1 |Q| ∑
(x,y)∈Q− log pθ,φ(y|E[ω]). We then optimize θ and φ by minimizing the expected loss Ep(τ)[Lτ (θ, φ)] over the task distribution p(τ) via episodic training. Energy scaling. Utilizing equation 11 directly for inference may lead to underconfidence when using unscaled σ̃ from 10. Therefore, in order properly scale the variance to maintain ID performance, we tune the temperature parameter T after training. Specifically, we start with T = 1 and iteratively increase T by 1 until ED[− log p(y|z,S)] ≤ ED[− log p(y|E[ω])], where − log p(y|E[ω]) is the NLL evaluated by using only the deterministic logits E[ω].

heading: 3.4 SPECTRAL NORMALIZATION.
text: Lastly, we enforce a bi-Lipschitz regularization fθ by employing both residual connections and spectral normalization on the weights Liu et al. (2020b), such that Equation 13 is satisfied. Using features Z, the calculation of covariance (Z−µc)(Z−µc)> and subsequently, the mean and variance of 10 both implicitly utilize distance, therefore we require bi-Lipschitz regularization of fθ. We choose spectral normalization via the power iteration method (also known as the Von Mises Iteration (Mises & Pollaczek-Geiringer, 1929)) due to its low memory and computation overhead as compared to second order methods such as gradient penalties (Arjovsky et al., 2017). Specifically, for features at hidden layer h(·), at depth l, and for some constants α1, α2, for all zi and zj , we enforce:
α1||z(l)i − z (l) j || ≤ ||h(z (l−1) i )− h(z (l−1) j )|| ≤ α2||z (l) i − z (l) j ||. (13)

heading: 4 EXPERIMENTS
text: The goal of our experimental evaluation is to answer the following questions. 1) What is the benefit of each component of our proposed model. 2) Does gφ produce more expressive covariances than
empirical features? 3) How does the ID/OOD calibration and accuracy compare with other popular baseline models.
Datasets. For few shot learning, we evaluate our model on both the Omniglot (Lake et al., 2015) and MiniImageNet (Vinyals et al., 2017) datasets. We utilize corrupted versions (Omniglot-C and MiniImageNet-C) which consists of 17 corruptions at 5 different intensities (Hendrycks & Dietterich, 2019). We follow the precedent set by Snell et al. (2017) and test Omniglot for 1000 random episodes and MiniImageNet for 600 episodes. For corruption experiments, the support set is uncorrupted, and corruption levels 0-5 are used as the query set (0 being the uncorrupted query set). We also experiment with multiple toy datasets which include half-moons, and concentric circles for binary classification and random 2D multivariate Gaussian distributions for multiclass classification (Figure 1). On the toy datasets, we create task distributions by sampling random tasks with biased support sets, applying random class shuffling and varying levels of noise added to each task. Randomly biasing each task ensures that no single task contains information from the whole distribution and therefore, the true distribution must be meta-learned through the episodic training over many such tasks. For a detailed explanation of the exact toy dataset task creation procedure, see the appendix section A.2.
Baselines. We compare our model against Protonets (Snell et al., 2017), A spectral normalized version of Protonets (Protonet-SN), MAML (Finn et al., 2017), Reptile (Nichol et al., 2018), Deep
Kernel Transfer (DKT) (Patacchiola et al., 2020), and straightforward few-shot/protonet adaptations of Spectral Normalized Neural Gaussian Processes (ProtoSNGP) (Liu et al., 2020a) and Deep Deterministic Uncertainty (ProtoDDU) (Mukhoti et al., 2021). These models represent a range of both metric based, gradient based, and covariance based meta learning algorithms. All baseline models are temperature scaled after training, with the temperature parameter optimized via LBFGS for 50 iterations with a learning rate of 0.001. This follows the temperature scaling implementation from Guo et al. (2017).
Calibration Error. We provide results for ECE on various types of OOD data in Figures 4 and 5 as well as Table 1. Accuracy and NLL are reported in Appendix A.8. Meta learning generally presents a high correlation between tasks, but random classes from different tasks which are not in the current support set S should still be treated as OOD. In Table 1 we provide results where the query set Q consists of random classes not in S . ProtoMahalanobis models perform the best in every case except for Omniglot 20-way/1-shot, where Reptile showed the lowest ECE. The reason for this can be seen in Figure 4, where Reptile shows poor ID performance relative to all other models. Under-confidence on ID data can lead to better confidence scores on OOD data, even though the model is poorly calibrated. Likewise we also evaluate our models on Omniglot-C and MiniImageNet-C in Figure 4 and 5. As the corruption intensity increases, ProtoMahalanobis models exhibit lower ECE in relation to baseline models while maintaining competitive ID performance. There are instances where Reptile and DKT appear to show strong OOD performance, but in each of these cases, they both show far worse performance on ID data, indicating that they are underconfident and very miscalibrated on ID data.
A compelling hypothesis for this behavior from Reptile is that it is related to the transductive BatchNorm used by both MAML and Reptile baselines. The query sets used in each corruption are larger than seen during training which led to lower MAML and Reptile accuracies in Figure 19 and 21, may have conversely had an opposite effect on calibration error in Figure 4 and 5. We are not aware of any work which analyzes the effects of transductive BatchNorm on ECE, which may be an interesting line of future work.
Eigenvalue Distribution. In Figure 6, we evaluate the effectiveness of meta learning the low rank covariance factors with gφ by analyzing the eigenvalue distribution of both empirical covariance (from DDU/SNGP baselines), and the covariance from gφ (eq. (5)). The empirically calculated covariance exhibits low diversity in eigenvalues, which implies that the learned Gaussian distribution is near spherical for every class. ProtoMahalanobis models, on the other hand, exhibit a diverse range of eigenvalues, leading to non-trivial ellipsoid distributions. We also note that in addition to
more diverse range of eigenvalues, the differences between the distributions of each class in S are also amplified in ProtoMahalanobis models, indicating a class specific variation between learned covariance factors. Extra figures are reported in the Appendix A.7, where it can be seen that the eigenvalue distribution becomes less diverse for ProtoMahalanobis models on the one-shot setting.
Architectures. For both Omniglot and MiniImageNet experiments, we utilize a 4 layer convolutional neural network with 64 filters, followed by BatchNorm and ReLU nonlinearities. Each of the four layers is followed by a max-pooling layer which results in a vector embedding of size 64 for Omniglot and 1600 for MiniImageNet. Exact architectures can be found in Appendix A.9. Protonet-like models use BatchNorm with statistics tracked over the training set, and MAML-like baselines use transductive BatchNorm as the original works do. As spectral normalized models require residual connections to maintain the lower Lipschitz bound in equation 13, we add residual connections to the CNN architecture in all Protonet based models.

heading: 4.1 IMPLEMENTATION DETAILS
text: ProtoSNGP & ProtoDDU Both ProtoSNGP and ProtoDDU baselines are adapted to meta learning by using the original implementation plus the addition of one meta parameter which represents the first diagonal term in Equation 8. This both acts as a meta parameter and a necessary first step in applying the recursive formula for inversion to make a prediction on each query set seen during during training.
Covariance Encoder gφ. We utilize the Set Transformer (Lee et al., 2019), as the self-attention performed by the transformer is an expressive means to encode pairwise information between inputs. We initialize the seeds in the pooling layers (PMA), with samples from N (0, 1). We do not use any spectral normalization in gφ, as it should be sufficient to only require that the input to the encoder is composed of geometry preserving features. Crucially, we remove the residual connection Q+ σ(QK>)V as we found that this led to the pooling layer ignoring the inputs and outputting an identical covariance for each class in each task. In the one-shot case, we skip the centering about the centroid in Equation 4 because it would send all class centroids to the origin.

heading: 5 CONCLUSION
text: It is widely known that DNNs are miscalibrated especially for OOD data. We have shown that existing covariance based methods fail to calibrate well when given a limited amounts of data for class-specific covariance construction. In this work, we have proposed a novel method which metalearns a diagonal or diagonal plus low rank covariance matrix which can be used for downstream tasks such as uncertainty calibration. Additionally, we have proposed an inference procedure and energy tuning scheme which can overcome miscalibration due to the shift invariance property of softmax. We further enforce bi-Lipschitz regularization of neural network layers to preserve relative distances between data instances in the feature spaces. We validated our methods on a synthetic data and two benchmark few-shot learning datasets, and showed that the final predictive distribution of our method is well calibrated under a distributional dataset shift when compared with relevant baselines.

heading: A APPENDIX

heading: A.1 LOSS DERIVATION (EQUATION 7)
text: The full derivation of Equation 7 can be achieved by first applying Bayes’ Rule, assuming a simple uniform prior over the class labels, p(yi|xi) can be proportionately expressed as,
p(yi|xi) = p(xi|yk)p(yk)
p(xi) ∝ p(xi|yk)p(yk) (14)
In which case the objective of the model becomes raising the class conditional p(xi|yi), while simultaneously lowering p(xi|yj) ∀ j 6= i. This is in fact equivalent to a softmax + cross entropy loss over the class conditional densities which are output from our model. In the softmax case, maximizing p(yi|xi) for a given class can be done by,
p(yi|xi) = ezi∑ z′ e z′ (15)
Which them implies that the loss to be minimized is the following, commonly known as the negative log likelihood of the data, or the empirical cross entropy between the true data distribution and the predictive distribution of the model.
LNLL = ED[− log p(y|x)]
= 1
N N∑ i=0 − log p(yi|xi)
= 1
N N∑ i=0 ( − zj + log ∑ z′j exp(z′j) ) i
LCE = − ∫ x px(y|x) log pθ(y|x)dx
≈ 1 N N∑ i=0 − log pθ(y|x)
= 1
N N∑ i=0 ( − zj + log ∑ z′j exp(z′j) ) i
(16)
In our case, assuming a uniform prior over the classes, we can analogously formulate the loss as,
L = ED[− log p(x|y)p(y)]
= 1
N N∑ i=0 ( − log p(xi|yk)p(yk) + log ∑ x′ p(x′i|yk)p(yk) )
= 1
N N∑ i=0 ( − log p(xi|yk)− log p(yk) + log p(yk) + log ∑ x′ p(x′i|yk) )
= 1
N N∑ i=0 ( − log p(xi|yk) + log ∑ x′ p(x′i|yk) )
(17)

heading: A.2 TOY DATASETS
text: To add bias to each samples task from our 2D toy datasets, we first randomly choose an axis (X or Y) for each class and then slice the datapoints in half randomly. We then sample the support set from the
chosen biased subset and leave the rest of the remaining points for the query set. Each sampled task calculates the mean and variance from the support set, which are then used to normalize all instances in S and Q.
Dataset N-Way K-Shot
Circles 2 5 Moons 2 5 Gaussians 10 10

heading: A.2.1 META MOONS
text: For the Meta Moons dataset, we randomly invert the classes to make sure that the class indices appear in a random order for each task. We add a random amount of Gaussian noise to each moon with a uniform standard deviation in the range of (0, 0.25].

heading: A.2.2 META CIRCLES
text: For the Meta Circles dataset, we randomly invert the order of the classes so that the inner circle and the outer circle are not guaranteed to appear in the same order on every task. We inject a random amount of Gaussian noise into the data, with a uniformly random standard deviation in the range of (0, 0.25]. We also randomly choose the scale factor between the size of the inner circle and the outer circle, which is uniformly random in the range of (0, 0.8]

heading: A.2.3 META GAUSSIANS
text: The task construction of the Meta Gaussians dataset requires that we construct random positive semidefinite covariance matrices for each class. We first uniformly sample N 2 × 2 matrices in the range U(−1, 1) and perform a QR decomposition to extract orthonormal matrices Q. We then sample a random diagonal D ∼ U(0, 1), and construct the final matrix as QDQ> which is positive
semi-definite. This leads to the distribution of each class being an elliptical multivariate Gaussian distribution.

heading: A.3 EXTRA RESULTS
text: We provide extra results on the MiniImageNet-C and Omniglot-C dataset here. Tables 2 and 3 contain results averaged over the whole corrupted dataset, including the natural test set and all 5 levels of corruption

heading: A.4 SET ENCODING RELATED WORKS
text: Set encoding functions require special end-to-end design considerations such as obeying permutation invariance w.r.t. the input set f({X1,X2, ...,Xn}) = f({Xπ(1),Xπ(2), ...,Xπ(n)}) for any random permutation of indices π(.). Likewise, the intermediate latent representations must satisfy permutation equivariance such that f({Xπ(1),Xπ(2), ...,Xπ(n)}) = {fπ(1)(X), fπ(2)(X), ..., fπ(1)(X)}. Deepsets (Zaheer et al., 2017) first proposed basic adaptations of linear and convolutional neural networks which obey the above required properties and have the addition of a sum decomposable (permutation invariant) pooling function and decoder to match the requirements of the given task. As sets can have complex interactions between elements, it may be beneficial to model pairwise interactions between set elements. The Set Transformer (Lee et al., 2019) uses a transformer architecture with self attention to model such pairwise interactions between set elements. As
transformers have a quadratic complexity w.r.t. input set length, it may not be possible to process a large set with a transformer and maintain permutation invariance, if the set will not fit into memory. Therefore, recent works have also further explored how to make an attentive set encoder which can process sets in batches (Andreis et al., 2021) while maintaining the above requirements of set functions.
For our model, we chose to use the set transformer architecture, as it models pairwise interactions between elements which is an implicit requirement of construction a Gaussian covariance matrix. Therefore, it has the proper inductive biases needed to satisfy our requirement of predicting low rank covariance factors given an input set of features.

heading: A.5 EXTRA TOY RESULTS
text: In Figures 10, 11, 12, 13, 14, 15, and 16 we provide extra qualitative results on toy dataset covariances and entropy surfaces. In Tables 6, 7, and 8 we provide tabular results of all toy experiments, showcasing the differences between in-distribution data and random uniform OOD noise. While the best ID performanc is sometimes mixed between the

heading: A.6 FURTHER IMPLEMENTATION DETAILS
text: Positive Diagonal Constraint In order to constrain the diagonal Λ in Equation 5 to be positive as mentioned in Section 3.2, we utilize a truncated sigmoid function Λ = max(0.1, σ(z)). We truncate the values in order to avoid extreme values during the inversion.
SNGP & DDU Both SNGP (Liu et al., 2020a) and DDU (Mukhoti et al., 2021) were originally designed under the assumption that an entire dataset would be used in the final pass to construct a feature covariance matrix. Given that few-shot-learning contains a limited number of samples for each task, we compose the feature covariance as a diagonal + low-rank factor Λ + ΦΦ>, where Λ is a meta learned parameter. Λ can be seen as a shrinkage estimation (δΛ + (1− δ)ΦΦ>) for low sample size, with a meta learned mixing coefficient δ.
In order to extend SNGP to work in the few shot learning scenario under the prototypical network Snell et al. (2017) framework, we had to modify the original algorithm by replacing the last linear layer with the embedding layer and centroids used by prototypical networks. This means that SNGP can use the pretrained ProtoNet-12-SN model as a backbone and we can apply the Random Features and construct the Gaussian process at meta test time for each sampled task.
OOD AUPR/AUROC In ordre to evaluate the OOD AUPR/AUROC metrics in the supplementary tables, we utilize the method proposed by Liu et al. (2020b). Specifically, we use the total energy in the logits log ∑ i exp(zi)) as the score when evaluating AUPR/AUROC.

heading: A.7 ADDITIONAL EIGENVALUE DISTRIBUTIONS
text: The eigenvalue distributions highlighted in section 4 exhibit the most diverse case of eigenvalues. The eigenvalues of ProtoMahalanobis models become less diverse in the one-shot setting which is also where the calibration performance is more comptetitive with the baseline models in our experiments.

heading: A.8 ADDITIONAL BOXPLOT RESULTS

heading: A.9 ARCHITECTURE DETAILS

heading: A.10 RUNTIME ANALYSIS
text: Here, we provide a runtime analysis of different variants of our models and baselines. Linear models are evaluated by using the mean and standard deviations from 50 iterations of both training and inference on the MetaMoons dataset. Convolutional models are likewise evaluated on 50 iterations of the Omniglot dataset. All models were evaluated on a single GeForce GTX 1080 Ti GPU.
Mahalanobis models show slightly better (linear) or similar (CNN) speeds for Diagonal and rank-1 variants. Speeds decrease as the rank goes higher due to more factors and more iterations required for inversion and log-determinant calculations.

heading: A.11 FURTHER EIGENVALUE EXPERIMENTS
text: To perform further experiments and analysis into the behavior of the low rank covariance encoder outlined in Section 3, we analyze the significance of the eigenvalues of the precision matrix. In this experiment, we obtain the predicted precision matrix and perform an eigendecomposition A = QΛQ−1 ∈ RN×N . We then construct a set of alternate precision matrices S = {A′}Ni=1, where each set element is a recomposition QΛ′Q−1, where Λ′ has one eigenvalue reset to 1. We then compute the final Accuracy, NLL, and ECE once for each matrix in S. If the predicted eigenvalues are due to arbitrary error or noise, then we would expect to see that the test statistics would arbitrarily improve for some precision matrices in S.
Instead, in Table 14 we see that the precision matrix which is predicted from the set transformer gives the best results on the test set in all cases, showing that all of the predicted values are necessary for the given solution.

Review summaries:
This submission focus on improving the performance of uncertainty quantification and OOD detection using bi-Lipschitz regularized neural networks. There are two involved challenges here: (i) we are required to estimate the class-conditional covariance matrix of the latent features. This estimation is challenging under low sample size. (ii) generative classifiers using softmax activation function suffer from shift invariance, and thus may become over-confident for OOD samples.  To resolve these issues, the paper focus on two innovations:  (i) a meta-learning model to estimate the class-conditional covariance matrices, (ii) a modification of the inference procedure with the energy function to detect OOD samples.  The paper concludes with numerical results showing the performance of the proposed solution approach.

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.

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.