M ETA L EARNING L OW R ANK C OVARIANCE F ACTORS FOR E NERGY -B ASED D ETERMINISTIC U NCERTAINTY
A
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 efﬁciently construct task speciﬁc covariance matrices.
Additionally, we propose an inference procedure which utilizes scaled energy to achieve a ﬁnal predictive distribution which is well calibrated under a distributional dataset shift
I NTRODUCTION
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 miscalibrated 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 conﬁdence 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 in the singularity, thereby destroying 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, enforcing both an upper and lower Lipschitz bound on each function 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 quantiﬁcation.
Broadly speaking, previous works have done this by constructing covariance matrices from the resulting features in order to aid in uncertainty quantiﬁcation (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 meta-learning is to learn meta-knowledge over a task distribution, but as we show, empirical covariance estimation methods are not able to effectively encode such knowledge, 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 feature covariances to each task.
Another obstacle, highlighted by Mukhoti et al. (2021), is that plain softmax classiﬁers cannot accurately model epistemic uncertainties.
We identify a contributing factor to this, which is the shift [Figure] invariance property of the softmax function.
Speciﬁcally, even if an evaluation point comes from an
OOD area and is assigned low logit values (high energy), this alone is insufﬁcient for a well calibrated prediction.
Small variations in logit values can lead to arbitrarily conﬁdent 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 classiﬁers as well.
In the following work, we ﬁrst propose a method of meta-learning class-speciﬁc covariance matrices that is transferable across the task distribution.
Speciﬁcally, we meta-learn a function that takes a set of class examples as an input and outputs a class-speciﬁc 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 minimum energy increases, and vice versa.
In this way, the softmax prediction can become progressively more uniform between ID and OOD data, after marginalizing the logit-normal distribution (see example in Figure 1).
By combining those two components, we have an inference procedure which 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 better calibrated uncertainty on OOD data.
R ELATED W ORK
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 classiﬁers 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 classiﬁer.
The method outlined in Lee et al.
(2018) provides a solid groundwork for our method, which also utilizes [Figure]
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 (temperature) before applying 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 single tasks, these post-processing methods make less expressive covariances in the meta learning setting, as demonstrated in Figure 1.
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 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 separability.
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 over 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; Nichol et al., 2018).
In this work, our model utilizes a metric-based approach as they are closely related to generative classiﬁers, which have been shown to be important for epistemic uncertainty (Mukhoti et al., 2021).
A PPROACH
We start by introducing a task distribution p ( τ ) which randomly generates tasks containing a support set S
= { (˜ x i , ˜ y i ) }
N s i =1 and a query set Q
= { ( x i , y i ) } N q i =1 .
Then, given randomly sampled task τ
= ( S , Q ) , we meta-learn a generative classiﬁer 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 classiﬁer 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.
[Equation]
3. L IMITATIONS OF E XISTING G ENERATIVE C LASSIFIERS
Possibly one of the simplest forms of deep generative classiﬁer is Prototypical Networks (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 |S c | (cid:80) ˜ x ∈S c f θ (˜ x ) is the class-wise embedding mean computed from S c , 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.
[Equation]
Limitations of ﬁxed 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 ﬁxed covariance I for all classes, such that p θ ( z | y
= c, S )
=
N ( z ; µ c , I ) .
For this reason, many-shot models such as SNGP (Liu et al., 2020b) and DDU (Mukhoti et al., 2021) calculate empirical covariances from data after training to aid in uncertainty quantiﬁcation.
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 can provide unreliable estimates of the true class covariance.
Unreliable covariance leads to poor estimation of Mahalanobis distances and therefore unreliable uncertainty estimation. [Equation] [Figure]
[Equation]
𝒛
Shift invariant property of softmax and OOD calibration.
Another critical limitation of Eq.
(2) is that it produces overconﬁdent predictions in areas distant from the class prototypes.
The problem can arise from the shift invariance property of the softmax function σ
[Equation]
ω
[Equation]
= e [Equation] [Equation] e ω (cid:48) with ω denoting the logits, such that σ
[Equation]
(cid:80)
[Equation]
e ω
[Equation]
= σ ( ω ) for any shift s .
More speciﬁcally, suppose we have two classes c
= 1 , 2 , and z moves along the line extrapolating the prototypes µ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:
[Equation]
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 µwith extreme values of c , we still have equally conﬁdent predictions.
See Figure 3 for illustration.
3. M ETA - LEARNING OF THE C LASS - WISE C OVARIANCE
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 ( τ ) .
Speciﬁcally, we meta-learn a set encoder g φ that takes a class set S c 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, ﬁlls 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
[Equation]
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 |S c | (cid:80) ˜ x ∈S c f θ (˜ x ) 7: Λ c , Φ c ← g φ ( Z c ) (cid:46) Eq.8: Σ c ← Λ c + Φ c Φ (cid:62) c
(cid:46) Eq.9:
Compute Σ − 1 c and | Σ c | (cid:46) Eq.
8,9 10: end for 11:
L τ ← 1 |Q| (cid:80) ( x ,y ) ∈Q − log p θ,φ ( y | E [ ω ]) (cid:46) Eq.12: ( θ, φ ) ← ( θ, φ ) − α ∇ θ,φ L τ 13: end while
1: Input: Task τ , meta-learned θ and φ 2: for c = 1 to C do 3: µ c ← 1 |S c | (cid:80) ˜ x ∈S c f θ (˜ x ) 4: Λ c , Φ c ← g φ ( Z c ) (cid:46) Eq.5: Σ c ← Λ c + Φ c
Φ (cid:62) c (cid:46) Eq.6:
Compute Σ − 1 c and | Σ c | (cid:46) Eq.
8,9 7: end for 8: Eval.
p θ,φ ( y | z , S ) for ( y, x ) ∈ Q (cid:46)
Eq.where Λ c ∈ R d
× d is a diagonal matrix and Φ c ∈ R d × 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 .
[Equation]
It is easy to see that Σ c is a valid positive semi-deﬁnite covariance matrix for positive Λ c .
Note that the covariance becomes diagonal when r = 0 .
A natural choice for g φ is the Set Transformer (Lee et al., 2019) which models 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,
[Equation]
Covariance inversion and log-determinant.
Note that the logit of the softmax function in Eq.
(7) involves the inverse covariance
Σ − 1 c and the log-determinant log | Σ c | .
In contrast to both DDU and SNGP which propose to calculate and invert an empirical feature covariance during post-processing, 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 ( d 3 ) 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.
[Equation]
3. O UT - OF -D ISTRIBUTION C ALIBRATION WITH S CALED E NERGY
Next, in order to tackle the overconﬁdence problem caused softmax shift invariance (Figure 3), we propose incorporating a positive constrained function of energy h ( E ) = max( (cid:15), − 1 T log (cid:80) c exp( − E c )) , 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
[Table]
04 ± 0.
49.28 ±
[Table] [Table]
±
[Table] [Table]
±
[Table]
way, the logit-normal distribution variance rises in conjunction with the energy magnitude, making predictions more uniform over the simplex for higher magnitude energies. p
[Equation]
z ,
[Equation]
) = N [Equation] ω c ; ˜ µ c ,
[Equation]
σ ) , where ˜ µ [Equation] Intuitively, h ( E ) is dominated by min c ( | E c | ) thereby acting as a soft approximation to the minimum energy magnitude (shortest Mahalanobis distance), which only becomes large when the energy is high for all classes represented in the logits.
Then, the predictive distribution becomes
[Equation]
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 τ ( θ, φ )
=|Q| (cid:80) ( x ,y ) ∈Q − log p θ,φ ( y | E [ ω ]) .
We then optimize θ and φ by minimizing the expected loss
E p ( τ ) [ L τ ( θ, φ )] over the task distribution p ( τ ) via episodic training.
Energy scaling.
Inference with equation 11 can still beneﬁt from temperature scaling of ˜ σ in 10.
Therefore, in order properly scale the variance to avoid underconﬁdent ID performance, we tune the temperature parameter T after training.
Speciﬁcally, we start with T = 1 and iteratively increase T by 1 until E D [ − log p ( y | z , S )] ≤ E D [ − log p ( y | E [ ω ])] , where − log p ( y | E [ ω ]) is the NLL evaluated by using only the deterministic logits E [ ω ] .
3. S PECTRAL N ORMALIZATION .
Lastly, we enforce a bi-Lipschitz regularization f θ by employing both residual connections and spectral normalization on the weights (Liu et al., 2020a), such that Equation 13 is satisﬁed.
Using features Z , the calculation of covariance ( Z − µ c )( Z − µ c ) (cid:62) and the subsequent 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).
Speciﬁcally, for features at hidden layer h ( · ) , at depth l , and for some constants α 1 , α 2 , for all z i and z j , we enforce:
[Equation]
E XPERIMENTS
The goal of our experimental evaluation is to answer the following questions.
What is the beneﬁt of each component of our proposed model? 2)
Does g φ produce more expressive covariances than
[Figure]
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 [Figure] 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 classiﬁcation and random 2D multivariate Gaussian distributions for multiclass classiﬁcation (Figure 1).
On the toy datasets, we create task distributions by sampling random tasks with biased support sets, applying random class shufﬂing 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), and
[Figure]
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.
Calibration Error.
We provide results for Expected Calibration Error (ECE) (Guo et al., 2017) 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-conﬁdence on ID data can lead to better conﬁdence scores on OOD data, even though the model is poorly calibrated.
Likewise we also evaluate our models on Omniglot-C and MiniImageNet-C in Figures 4 and 5.
As the corruption intensity increases, ProtoMahalanobis models exhibit lower ECE in relation to baseline models while maintaining competitive ID performance.
Overall, Reptile shows the strongest calibration of baseline models although it can be underconﬁdent on ID data as can be seen in Figure 4.
In our experiments, transductive batch normalization used in MAML/ Reptile led to suboptimal results, as the normalization statistics depend on the query set which is simultaneously passed through the network.
Passing a large batch of corrupted/uncorrupted samples caused performance degradation on ID data and presented an unrealistic setting.
We therefore utilized the normalization scheme proposed by Nichol et al. (2018) which creates batch normalization statistics based on the whole support set plus a single query instance.
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 and the encoded covariance from g φ (Equation 5).
The empirically calculated covariances exhibit lower diversity in eigenvalues, which implies that the learned Gaussian distribution is more spherical and uniform for every class.
ProtoMahalanobis models, on the other hand, exhibit a more 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 ampliﬁed in ProtoMahalanobis models, indicating a class speciﬁc variation between
[Figure]
learned covariance factors.
Extra ﬁgures are reported in the Appendix A.7, where it can be seen that the eigenvalue distribution becomes less diverse for ProtoMahalanobis models in the one-shot setting.
Architectures.
For both Omniglot and MiniImageNet experiments, we utilize a 4 layer convolutional neural network with 64 ﬁlters, 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 Reptile Norm (Nichol et al., 2018).
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.
4. I MPLEMENTATION D ETAILS
ProtoSNGP & ProtoDDU Both ProtoSNGP and ProtoDDU baselines are adapted to meta learning by using the original backbone implementation plus the addition of a positive constrained meta parameter for the ﬁrst diagonal term in Equation 8 which is shared among all classes.
This provides meta knowledge and a necessary ﬁrst step in applying the recursive formula for inversion to make predictions 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 sufﬁcient to only require that the input to the encoder is composed of geometry preserving features.
Crucially, we remove the residual connection
Q + σ ( QK (cid:62) )
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 place all class centroids at the origin.
C ONCLUSION
It is widely known that DNNs can be miscalibrated for OOD data.
We have shown that existing covariance based uncertainty quantiﬁcation methods fail to calibrate well when given a limited amounts of data for class-speciﬁc covariance construction for meta learning.
In this work, we have proposed a novel method which meta-learns 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 both synthetic data and two benchmark few-shot learning datasets, showing that the ﬁnal predictive distribution of our method is well calibrated under a distributional dataset shift when compared with relevant baselines.  A CKNOWLEDGEMENTS
This work was supported by the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2019-0-00075, Artiﬁcial
Intelligence Graduate School Program(KAIST)), the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF2018R1A5A1059921), the Institute of Information & communications Technology Planning &
Evaluation (IITP) grant funded by the Korea government (MSIT)
No. 2021-0-02068 (Artiﬁcial Intelligence Innovation Hub) and the National Research Foundation of Korea (NRF) funded by the Ministry of Education (NRF2021R1F1A1061655). A A PPENDIX A. L OSS D ERIVATION (E QUATION 7)
The full derivation of Equation 7 can be achieved by ﬁrst applying Bayes’ Rule, assuming a simple uniform prior over the class labels, p ( y i | x i ) can be proportionately expressed as,
[Equation]
In which case the objective of the model becomes raising the class conditional p ( x i | y i ) , while simultaneously lowering p ( x i | y j )
∀ j (cid:54)
=
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 ( y i | x i ) for a given class can be done by,
[Equation]
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.
[Equation]
In our case
[Equation]
assuming a uniform prior over the classes
[Equation]
we can analogously formulate the loss as,
[Equation]
A. T OY D ATASETS
To add bias to each samples task from our 2D toy datasets
[Equation]
we ﬁrst 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
Moons
Gaussians
[Figure]
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] .
[Figure]
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]
[Figure]
The task construction of the Meta Gaussians dataset requires that we construct random positive semideﬁnite covariance matrices for each class.
We ﬁrst 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 ﬁnal matrix as QDQ (cid:62) which is positive
[Table]
semi-deﬁnite.
This leads to the distribution of each class being an elliptical multivariate Gaussian distribution.
A. E XTRA R ESULTS
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
A. S ET E NCODING R ELATED W ORKS
Set encoding functions require special end-to-end design considerations such as obeying permutation invariance w.r.t.
the input set f
( { X 1 , X 2 , ..., X n } )
=
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)
ﬁrst 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 beneﬁcial 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
[Table]
[Table] [Table]
transformers have a quadratic complexity w.r
[Table] 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 ﬁt 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
[Table]
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.
A. E XTRA T OY R ESULTS
In Figures 10, 11, 12, 13, 14,
[Table]
and 16 we provide extra qualitative results on toy dataset covariances and entropy surfaces.
In Tables 6,
[Table]
and 8 we provide tabular results of all toy experiments, showcasing the differences between in-distribution data and random uniform OOD noise.
[Table]
[Figure] [Figure] [Table]
[Table] [Table]
[Figure] [Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure]
F URTHER I MPLEMENTATION D ETAILS
Positive Diagonal Constraint
In order to constrain the diagonal Λ of Proto Mahalanobis models (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 ﬁnal 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 Λ + ΦΦ (cid:62) , where Λ is a positive constrained (via softplus) meta learned parameter.
Λ can be seen as a shrinkage estimation ( δ Λ + (1 − δ ) ΦΦ (cid:62) ) for low sample size, with a meta learned mixing coefﬁcient δ .
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.
Empirically, we found that using the SNGP logit-normal inference procedure led to a severe performance decrease, therefore our results utilized Mahalanobis distance instead.
OOD AUPR/AUROC
In order to evaluate the OOD AUPR/AUROC metrics in the supplementary tables, we utilize the method proposed by Liu et al.
(2020b).
Speciﬁcally, we use the total energy in the logits log (cid:80) i exp( z i ) as the score when evaluating AUPR/AUROC.
Optimizers All models are trained with the Adam (Kingma & Ba, 2014) optimizer
A. A DDITIONAL E IGENVALUE D ISTRIBUTIONS
The eigenvalue distributions highlighted in section 4 exhibit the most diverse case of eigenvalues.
However, the eigenvalues of ProtoMahalanobis precision matrices become less diverse in the one-shot setting which is also where we are unable to mean center the respective features by class.
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Table]
[Table]
A. A DDITIONAL B OXPLOT R ESULTS
Figures 18, 20 show extra boxplot results for accuracy while Figures 19, 21 show negative log likelihood.
A. A RCHITECTURE D ETAILS
Tables 9, and 10 show the backbone architectures for MAML/Reptile or Omniglot and MiniImageNet respectively.
Table 11 shows the backbone architecture for all Protonet based models.
A. R UNTIME A NALYSIS
In Tables [Table] and 13 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.
SNGP/DDU also utilize the matrix inversion outlined in Equation 8.
Mahalanobis models show slightly better (Linear) or similar (CNN) latency to SNGP/DDU for diagonal and rank-1 variants.
Latency increases as the rank goes higher due to more factors and more iterations required for inversion and log-determinant calculations.
Comparing Protonet, Protonet-SN, and other variants which need to construct a covariance, we can see that constructing the covariance matrix adds a cost which is roughly equivalent to spectral normalization.
[Table]
[Table]
[Table]
A. F URTHER E IGENVALUE E XPERIMENTS
In Table [Table] we perform further experiments and analysis into the behavior of the low rank covariance encoder outlined in Section 3, we analyze the signiﬁcance of the eigenvalues of the precision matrix.
In this experiment, we ﬁrst obtain the predicted precision matrix and perform an eigendecomposition
A = Q Λ Q − 1 ∈ R N × N .
We then construct a set of alternate precision matrices S
=
{ A (cid:48) } Ni =1 , where each set element is a recomposition A (cid:48) i = Q Λ (cid:48)
i Q − 1 , where Λ (cid:48) i has one eigenvalue reset to 1.
We then compute the ﬁnal 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. This experiment utilizes Omniglot 5-way/5-shot and the ProtoMahalanobis Rank-1 variant. [Figure]
[Table]