Uncertainty-Aware Few-Shot Image Classification

1 Introduction

The strong capability of deep learning in part relies on the using of a large amount of labeled data for training, while humans are more readily able to learn knowledge quickly given a few examples. Few-shot learning [Vinyals et al., 2016; Finn et al., 2017; Karlinsky et al., 2019; Dong and Xing, 2018; Mishra et al., 2018] strives to a further step to develop deep learning approaches which generalize well to unseen tasks/classes with just few labeled samples. Among them, few-shot image classification [Vinyals et al., 2016; Snell et al., 2017; Sung et al., 2018; Oreshkin et al., 2018; Finn et al., 2017; Nichol et al., 2018; Lee et al., 2019] has attracted much attention, which aims to classify data (query samples) of new categories given only a few labeled samples of these categories as examples (support samples).

∗This work was done when Zhizheng Zhang was a visiting student at Columbia University.

†Corresponding Author

For few-shot image classification, several latest research works reveal that good feature embedding is important to deliver favorable performance for the similarity-based classification [Chen et al., 2019; Chen et al., 2020; Huang and Tao, 2019; Tian et al., 2020]. They average the feature embeddings of the support samples of a category to be the prototype of this category. Thus, with each prototype being categoryspecific, the set of prototypes plays the role of similaritybased classifier to identify the category of a query sample by finding the nearest prototype over the support set in the embedding space. During training, for a query sample, the logit vector (classification probability) is obtained by feeding the similarities between the query sample and the prototypes to a SoftMax function. Hence, the reliability of the estimated similarity is vital to the classification performance.

The quality of network output has been investigated and modeled with uncertainty in regression, classification, segmentation, multi-task learning [Kendall et al., 2015; Kendall and Gal, 2017; Chang et al., 2020; Kendall et al., 2018] to benefit the optimization. Aleatoric uncertainty identifies the confidence level of the network on its output for a sample, which captures the noise inherent in the observations. For few-shot image classification, the network actually has different uncertainties on the calculated similarities of different query-prototype pairs. An observed similarity of a queryprototype pair, being a one time sampling, suffers from observation noise, where the higher of the uncertainty, the less reliable of the estimated similarity. For each few-shot task during the optimization, the number of experienced query-prototype pairs is limited and thus the side effect of observation noises due to the high uncertainty limits the optimization. Therefore, modeling uncertainty is especially vital for few-shot learning with the limited samples, but still under-explored.

In this paper, we propose an efficient uncertainty-aware few-shot image classification framework to model uncertainty and perform uncertainty-aware optimization. For each query-prototype pair, we convert the observed feature similarity between the query and the prototype from a deterministic scalar to a distribution, with the variance of this distribution characterizing the uncertainty of the observed similarity. The similarity-based classification losses are calculated based on the Monte Carlo sampling on the similarity distributions to alleviate the influence of observation noises. Since the classification probability of a query sample is determined based on the similarity of this query sample with its N prototypes in the support set, this is a joint determination process. Thus, we adopt a graph-based model to jointly estimate the uncertainties for the N query-prototype pairs, which facilitates the information passing among them for the joint optimization.

In summary, our contributions lie in the following aspects:

• We are the first to explore and model the uncertainty of the similarities of query-prototype pairs in few-shot image classification, where a better understanding of the limited data is particularly important.
• We perform uncertainty-aware optimization to make the few-shot learning more robust to observation noises.
• We design a model to jointly estimate the uncertainty of the similarities between a query image and the prototypes in the support set.

We conduct comprehensive experiments and demonstrate the effectiveness of our proposed method in the popular inductive setting, where each query sample is evaluated independently from other query samples. Our scheme achieves the state-of-the-art performance on multiple datasets.

2 RelatedWork

2.1 Few-shot Image Classification

Few-shot image classification aims to recognize novel (unseen) classes upon limited labeled examples. Representative approaches can be summarized into four categories.

Classification-based methods train both a feature extractor and classifiers with meta-learning and learn a new classifier (e.g., linear softmax classifier) for the novel classes during meta-testing [Chen et al., 2019; Ren et al., 2019; Lee et al., 2019]. They in general need to update/fine-tune the network given a few-shot testing task.

Optimization-based methods exploit more efficient metalearning optimization strategies for few-shot learning [Finn et al., 2017; Grant et al., 2018; Nichol et al., 2018; Lee et al., 2019]. MAML [Finn et al., 2017] is a representative work, which trains models by aligning the gradients of several learning tasks, such that only a few iterations are needed to achieve rapid adaptation for a new task.

Hallucination-based methods enhance the performance by augmenting data, which alleviates data deficiency by using generated data. The generators transfer appearance variations [Hariharan and Girshick, 2017; Gao et al., 2018] or styles [Antoniou et al., 2018] from base classes to novel classes. Wang et al. create an augmented training set through a generator which is trained end-to-end along with the classification [Wang et al., 2018b].

Similarity-based/Metric-based methods classify an unseen instance into its nearest class based on the similarities with a few labeled examples. They learn to project different instances to a metric space wherein the similarities among instances of the same category are larger than that of instances of different categories. Matching networks [Vinyals et al., 2016] propose to assign an unlabeled data with the label of its nearest labeled data in one-shot learning. Furthermore, the prototypical network [Snell et al., 2017] averages the features of several samples of the same category as the prototype of this class and classify unseen data by finding the nearest prototype in the support set, where a prototype is taken as the representation of a class. Relation networks [Sung et al., 2018] propose an additional relation module which learns to determine whether a query and a prototype belongs to the same class and is jointly trained with deep feature representations.

Some recent metric-based works reveal the importance of learning good feature representations via classificationbased pre-training for few-shot tasks [Chen et al., 2019; Tian et al., 2020; Chen et al., 2020]. They propose strong baselines by training the network with classification supervision over all base categories and further performing metalearning on sampled categories to simulate the few-shot setting. Chen et al. confirmed that the classification-based pretraining can provide extra transferability from base classes to novel classes in the meta-learning stage [Chen et al., 2020].

In this paper, on top of such a strong baseline, we look into an important but still under-explored issue, i.e., modeling and exploiting the inherent uncertainty of the calculated similarities for more effective optimization.

2.2 Uncertainty in Deep Learning

There are two main types of uncertainty studied for deep neural networks: aleatoric uncertainty, and epistemic uncertainty [Kendall and Gal, 2017; Gal, 2016; Kendall et al., 2018]. Epistemic uncertainty captures model-related stochasticity, including parameters choices, architecture, etc. Aleatoric (Data-dependent) uncertainty captures the noise inherent in the observations [Kendall and Gal, 2017; Gal, 2016; Kendall et al., 2018], which depends on the input sample of the model. A model usually has high uncertainty on noisy input or rarely seen input.

For few-shot classification, the "hardness" of a few-shot episode is quantified with a metric in [Dhillon et al., 2020] but is only used to evaluate few-shot algorithms rather than improve the optimization. In [Wang et al., 2020], a statistical approach is adopted to measure the credibility of predicted pseudo-labels for unlabeled samples, where the most trustworthy pseudo-labeled instances are selected to update the classifier in the semi-supervised or transductive few shot learning setting. There is still a lack of exploration of the uncertainty for more general inductive few-shot setting, where the network does not use/need testing samples for updating. In this paper, we model the data-dependent uncertainty of the calculated similarities between a query sample and the prototypes, and leverage it to better optimize the network. To our best knowledge, we are the first to explore the uncertainty of pair-similarity in few-shot classification towards more efficient optimization.

3 ProposedMethod

We propose a Uncertainty-Aware Few-Shot (UAFS) image classification framework (see Figure 1), which models uncertainty of the query-prototype similarity and perform uncertainty-aware optimization. We first present the problem formulation and a strong baseline in Section 3.1. Then, we elaborate our proposed uncertainty-aware few-shot image classification in Section 3.2.

3.1 Preliminaries and a Strong Baseline

Problem Formulation. For few-shot image classification, all the categories/classes of the dataset are divided into base classes Cbase for training and novel classes Cnovel for testing without class overlapping [Snell et al., 2017; Chen et al., 2019; Dhillon et al., 2020]. In few-shot learning, episodic training is widely used. Each N-way K-shot task randomly sampled from base classes is defined as an episode, where the support set S includes N classes with K samples per class, and the query set Q contains the same N classes with M samples per class. Our method is applied to the metric-based few-shot learning which performs similarity-based classification based on matching the prototypes in the support set for a given query sample.

Given an image xi, we feed it to a Convolution Neural Network (CNN) to obtain a feature map, and perform global averaging pooling to get a D-dimensional feature vector zi = fφ(xi) ∈ RD as the embedding result of xi, where fφ(·) denotes the feature extractor. The prototype of a class indexed by k is calculated by averaging the feature embeddings of all the samples of this class in the support set as:

where xi denotes the sample indexed by i and yi denotes its label, Sk denotes the set of samples of the k-th class in the support set S = {S1, · · · , SN}, |Sk| denotes the number of samples in Sk. For a query image x, the probability of classifying it into the class k is:

where the temperature parameter τ is a hyper-parameter, and N is the number of classes in the support set (i.e., N-way).

The s(fφ(xi), ck) represents the similarity between the given query-sample xi and the prototype ck of the class k. Here, the N prototypes of N classes in the support set construct a similarity-based classifier.

Baseline. Many works [Vinyals et al., 2016; Snell et al., 2017; Chen et al., 2019; Li et al., 2019; Chen et al., 2020] reveal that the core of metric-based few-shot learning is to learn a good feature embedding function fφ(·) from base classes Cbase which could generalize well on unseen classes. We follow the latest work [Chen et al., 2020] to build a strong baseline with two-stage training:

Stage-1: Classification-based pre-training. We employ a Fully-Connected (FC) layer as the classifier with cosine similarity over all base classes to perform pre-training. We train the feature extractor fφ(·) (backbone) and the classifier in an end-to-end manner using a standard cross-entropy loss.

Stage-2: Meta-learning. For each episode/task, we sample N-way K-shot from base classes for training, which can be understood as a simulation of few-shot test process. In each task, N × K images and N × M images are sampled to constitute the support set and the query set, respectively. We adopt the cross-entropy loss formulated by L =

- ∑N×M i=1 yi · log(p(yi|xi)), where p(yi|xi) is calculated in Eq. (2) with s(·, ·) instantiated by cosine similarity.

3.2 Uncertainty-aware Few-shot Learning

A neural network has different uncertainties regarding its outputs for different input data and there will be observation noises for an instantiated output [Kendall and Gal, 2017]. For few-shot image classification, the model also has uncertainty on the estimated similarities for the query-prototypes pairs. The small number of samples in each few-shot episode exacerbate the side effect of observation noises from uncertainty. To address this, we model the data-dependent uncertainty of similarities and leverage it for better optimizing the network.

Probabilistic Similarity Representation. To characterize the uncertainty of similarity of a query-prototype pair, we convert the similarity representation of this query-prototype pair from a deterministic scalar to a probabilistic representation, i.e., a parameterized distribution.

Given a query sample xi (with its feature zi = fφ(xi) in the latent space) and a prototype cj obtained by Eq. (1), instead of using a deterministic scalar, we model the similarity as a probability distribution. Similar to other works [Kendall and Gal, 2017; Gal, 2016], we simply model the distribution of the similarity between this query and the prototype cj by a Gaussian distribution with mean µij and variance σ2 ij :

where C = {c1, · · · , cN} denotes the N prototypes in the support set, cj is the j-th prototype within it.

We estimate the mean of the similarity by their inner product, i.e., µij = 〈zi, cj〉, which denotes the most likely value of the similarity between zi and ci. The variance σ2 ij is used to characterize the uncertainty of this pairwise similarity. We estimate it by a graph neural network, which will be elaborated in the subsection "Similarity Uncertainty Estimation".

Uncertainty-aware Optimization. With the similarity of a query-support pair being represented by a distribution rather than a scalar, we incorporate such probabilistic similarity representation in the classification loss for better optimization, which could be more robust to the observation noises.

The analytic solution of integrating over these distributions for the optimization of classification loss function is difficult. Following [Kendall and Gal, 2017], we thus approximate the optimization objective by Monte Carlo integration. Particularly, Monte Carlo sampling is performed on the N similarity distributions. For a given query image xi and N prototypes cj, j = 1, · · · ,N which form N similarity pairs, we repeat T random sampling over the similarity distributions for each query-prototype pair sij (see (4) with t = 1, · · · , T ) to obtain statistical results. We enable a differentiable representation of the probabilistic similarity for a query and prototype pair sij by re-parameterizing it as:

For a given query image xi with groundtruth label yi = k, we obtain its corresponding classification loss as:

Similarity Uncertainty Estimation. We characterize the uncertainty by variance σ2 ij as mentioned before. Hereafter, we elaborate on how to estimate the uncertainty.

For a N-way similarity-based classification, the given query is compared with all the N prototypes and the class of the prototype with the highest similarity value is taken as the predicted class, which is a joint determination process.

The joint determination nature inspires us to estimate uncertainty parameter σij for each pair after taking a global glance of the N pairs with j = 1, · · · ,N. Therefore, we design a graph-based model as the similarity uncertainty estimator to jointly infer all the σij for a given query xi and N prototypes cj. This enables the exploration of global scope information for the current task to jointly estimate the similarity uncertainties. Besides, the graph-based design has another advantage, i.e., its scalability to the number of categories. This allows us to pre-train the parameters of the graph-based model it in Stage-1. We will demonstrate its effectiveness by comparing it with other alternative designs: 1) using Convolutional Neural Network (CNN) to infer σij independently for each pair; 2) adopting Fully-Connected (FC) based model that is able to exploit global information for N pairs but cannot handle the consistency classes in different episodics well.

Given a query and N prototypes as illustrated in Fig. 1,

we build a graph of N nodes, with the j th node denoted by the information of the query and the j th prototype. The output of the node is the predicted uncertainty of the similarity of this query-prototype pair. We feed a query image xi to the feature extractor to get its embedding vector zi = fφ(xi) ∈ RD . Similarly, we get the N prototypes of the support set as cj ∈ RD , j=1,· · ·,N. We use the group-wise relations/similarities between the query and the j th prototype feature to represent the information of this query-prototype pair. Specifically, we split both zi and cj into L groups along their channel dimensions to have {zl i |l = 1, · · · ,L} and {clj |l = 1, · · · ,L}, where L ≤ D. We calculate the cosine similarity of the lth group as rlij = zl i clj T ‖zl i ‖‖clj ‖ , and stack them to be a relation feature vector vij = [r1 ij , · · · , rL ij ] ∈ RL. Note that the relations/similarities observed from multiple groups provide valuable hints of the uncertainty since they reflect the similarity from different perspectives.

The N nodes in the graph are represented by a matrix Vi = (vi1; · · · ; viN) ∈ RN×L. We use graph convolutional network (GCN) to learn the similarity uncertainty for each node. Similar to [Shi et al., 2019; Wang et al., 2018a; Wang and Gupta, 2018], the edge from the j th node to the j ′ th node is modeled by their affinity in the embedded space, which is formulated as:

where ϕ1 and ϕ2 denote two embedding functions implemented by FC layers. All edges constitute an adjacent matrix denoted by Ei ∈ RN×N . We normalize each row of Ei with SoftMax function so that all the edge values connected to a target node is 1, yielding a numerically-stable message passing through the modeled graph. Similar to [Wang et al., 2018a], we denote the normalized adjacent matrix by Gi and update the nodes through GCN as:

Wυ ∈ RL×L and Wy ∈ RL×L are two learnable transformation matrices. Wυ is implemented by a 1 × 1 convolutional layer. Wy is implemented by two stacked blocks. Each block consists of a 1 × 1 convolution layer, followed by a Batch Normalization (BN) layer and an LeakyReLU activation layer. We thus infer the similarity uncertainty vector ui = [σi1, · · · , σiN] ∈ RN for the N pairs from the updated graph nodes, which can be formulated as:

where the j th -dimension of ui is the aforementioned σij, denoting the similarity uncertainty for the query xi and the j th prototype. Wu1 ∈ RL×L and Wu2 ∈ RL×1 are transformation matrices both implemented by a 1 × 1 convolutional layer. "BN" refers to the Batch Normalization layer and α(·) refers to the LeakyReLU activation function.

Discussion: We model pairwise uncertainty for the similarity instead of for a sample itself like in other tasks [Kendall and Gal, 2017]. This is because the classification of a query image is influenced by not only the query itself but also the prototypes in the support set. We further conduct extensive experiments to compare the effectiveness of modeling feature uncertainty (query itself) vs. similarity uncertainty (query and prototypes) for few-shot image classification in our experiments. Unlike the previous work [Garcia and Bruna, 2018], which adopts graph-based model to propagate information from labeled samples to unlabeled ones thus delivers a transductive solution, we use graph-based model to achieve the joint estimation of uncertainty of N query-prototype pairs for a given query (non-transductive), where the uncertainty is used to better optimize the network. We do not introduce any complexity increase in the testing where the GCN-based uncertainty estimator is discarded. In contrast, the GCN in [Garcia and Bruna, 2018] is a part of their network and is needed during testing.

Joint Training. Similar to the baseline configuration described in Section 3.1, we train the entire network in two stages: classification-based pre-training and meta-learning.

In the classification-based pre-training stage, the FC coefficients can be considered as the prototypes and the number of input nodes in GCN is the total number of classes in Cbase. In the meta-learning stage, the number of input nodes in GCN is N. Because our graph-based uncertainty estimator has the scalability to the number of involved categories for classification, the parameters in GCN are shared in the two stages. We fine-tune the backbone and the similarity uncertainty estimator in the meta training stage.

4 Experiments

4.1 Dataset and Implementation

Datasets. For few-shot image classification, we conduct experiments on four public benchmark datasets: miniImageNet [Vinyals et al., 2016], tiered-ImageNet [Ren et al., 2018], CIFAR-FS [Bertinetto et al., 2018], and FC100 [Oreshkin et al., 2018].

Networks and Training. Following recent common practices [Oreshkin et al., 2018; Lee et al., 2019], we adopt a wider ResNet-12 with more channels as the backbone in our work unless claimed otherwise. We build our strong baseline (see Section 3.1) with two-stage training by following [Chen et al., 2020]. For all our models, the classification temperature τ in (2) is a trainable parameter which is initialized with 10. For the parameter L, we found the performance is very similar when it is in the range of 16 to 64 and we set it to 32.

Evaluation. We conduct all experiments in the inductive (non-transductive) setting (i.e., each query sample is evaluated independently from other query samples). We create each few-shot episode by uniformly sampling 5 classes (i.e.N=5) from the test set and further uniformly sampling support and query samples for each sampled class accordingly. Consistent with other works, we report the mean accuracy and the standard deviation with a 95% confidence interval over 1000 randomly sampled few-shot episodes.

4.2 Ablation Study

Effectiveness of Proposed UAFS. Table 1 shows the ablation study on the effectiveness of our proposed method. Model-3 corresponds to our strong baseline Meta-Base (see Table 3) with two-stage training. Model-6 corresponds to our final scheme Meta-UAFS with the proposed similarity uncertainty modeling and optimization (UMO).

We have the following observations. 1) Our proposed Meta-UAFS (Model-6) improves the strong baseline MetaBase (Model-3) by 1.12%, 1.41%, 1.72% and 1.76% in 1-shot classification on four datasets respectively, which demonstrates the effectiveness of our UMO. 2) Meta-Base (Model-

3) obviously outperforms Model-1 that has only Stage1 training, demonstrating the effectiveness of two-stage training. 3) Enabling UMO in the classification-based pre-training stage (Stage1) brings obvious improvement for one-stage training (Model-2 vs. Model-1), but only slight gains for two-stage training (Model-4 vs. Model-3). The higher the performance, the more difficult it is to obtain gain. On top of the strong baseline, the UMO is not helpful for Stage1 which does not suffer from limited training data as Stage2. 4) Adopting UMO only in Stage2 (Model-5) does not always bring improvements (when compared with Meta-Base), while applying UMO on both stages (Meta-UAFS) brings significant improvements. For Model-5, the uncertainty estimator is trained only in Stage2 (from scratch without pre-training in Stage1) and may be difficult to optimize. Our graph-based uncertainty estimator has the scalability to the number of involved categories for classification, which allows us to pre-train the parameter-sharable uncertainty estimator in Stage1 and finetune it in Stage2. When we enable the pre-training of uncertainty estimator in Stage1, our final scheme (Meta-UAFS) consistently delivers the best results. Note that the backbone network is trained in both stages.

Similarity Uncertainty Estimator Designs. Since the query is classified in a joint determination by comparing it with the N prototypes simultaneously, we propose to use a GCN to jointly estimate the uncertainties. This design is superior because it 1) enables the joint estimation of similarity uncertainties and 2) has the scalability to the number of categories which facilitates the pre-training of the uncertainty estimator in Stage-1. Table 2 shows that 1) our Graph-based design (last row) is superior to Conv-based design (the penultimate row) which uses 1 × 1 convolutional layers to independently infer the uncertainty for each query-prototype pair. FC-based design is inferior to our proposed design.

Sample Uncertainty vs. Pair Similarity Uncertainty. Table 2 shows that the scheme B+SimiU with Graph-based uncertainty estimator which models the uncertainty of the query-prototype similarity is more effective than the scheme B+SampleU which models the uncertainty of a sample. That is because in the few shot classification, for a given query, its estimated class depends on both the query and the N prototypes. Therefore, we model the uncertainty of the similarity between the query and the prototypes.

4.3 Comparison with State-of-the-arts

In Table 3, we compare our proposed Meta-UAFS with the state-of-the-art approaches. Note that our baseline Meta-Base is the same as MetaBaseline [Chen et al., 2020] and they have the similar performance. Compared to Meta-Base [Chen et al., 2020], our Meta-UAFS achieves significant improvement of 1.12%, 1.41%, 1.72%, and 1.76% in 1-shot accuracy on mini-ImageNet, tiered-ImageNet, CIFAR-FS, and FC00, respectively for 1-shot classification. Our Meta-UAFS achieves the state-of-the-art performance. CAN [Hou et al., 2019] introduces cross attention module to highlight the target object regions, making the extracted feature more discriminative. Our uncertainty modeling and optimization which reduces the side effects of observation noises is complementary to their attention design and incorporating their attention design into our framework would lead to superior performance.

5 Conclusion

In this paper, we propose an Uncertainty-Aware Few-Shot image classification approach where data-dependent uncertainty modeling is introduced to alleviate the side effect of observation noise. Particularly, we convert the similarity between a query sample and a prototype from a deterministic value to a probabilistic representation and optimize the networks by exploiting the similarity uncertainty. We design a graph-based uncertainty estimator to jointly estimate the similarity uncertainty of the query-support pairs for a given query sample. Extensive experiments demonstrate the effectiveness of our proposed method and our scheme achieves state-of-the-arts performance on all the four benchmark datasets.

Acknowledgements

This work was supported in part by the National Key Research and Development Program of China 2018AAA0101400 and NSFC under Grant U1908209, 61632001.