Adaptive Distribution Calibration for Few-Shot Learning with Hierarchical Optimal Transport

Optimal Transport (OT) is a powerful tool for the comparison of probability distributions, which has been widely used in various machine learning problems, such as generative models [15], text analysis [16, 17], adversarial robustness [18], and imbalanced classification [19]. Here we limit our discussion to OT for discrete distributions and refer the reader to Peyré and Cuturi [20] for more details. Denote $p=\sum_{i=1}^{{n}}a_{i}\delta_{x_{i}}$ and $q=\sum_{j=1}^{m}b_{j}\delta_{y_{j}}$ as two $n$ and $m$ dimensional discrete probability distributions, respectively. In this case, $\boldsymbol{a}\in\Delta^{n}$ and $\boldsymbol{b}\in\Delta^{m}$, where $\Delta^{m}$ denotes the probability simplex of $\mathbb{R}^{m}$. The OT distance between $p$ and $q$ is defined as

where $\langle\cdot,\cdot\rangle$ denotes the Frobenius dot-product; $\mathbf{C}\in\mathbb{R}_{\geq 0}^{n\times m}$ is the transport cost function with element $C_{ij}=C(x_{i},y_{j})$; $\mathbf{T}\in\mathbb{R}_{>0}^{n\times m}$ denotes the doubly stochastic transport probability matrix such that $\Pi(p,q):=\{\mathbf{T}\,|\,\sum_{i}^{n}T_{ij}=b_{j},\sum_{j}^{m}T_{ij}=a_{i}\}$, meaning that $\mathbf{T}$ has to be one of the joint distribution of $p$ and $q$. As directly optimising Equation (1) can be time-consuming for large-scale problems, the entropic regularization, $H\!=\!-\sum_{ij}T_{ij}\ln T_{ij}$, is introduced in Cuturi [21], resulting in the widely-used Sinkhorn algorithm for discrete OT problems with reduced complexity.

Following a typical few-shot learning problem, we divide the whole dataset with labeled examples into a base dataset $\mathbf{D}_{\text{base }}$ with $B$ base classes and a novel dataset $\mathbf{D}_{\text{novel }}$ with $N_{all}$ novel classes, each with a disjoint set of classes. To build a commonly-used $N$-way-$K$-shot task [8, 14], we randomly sample $N$ classes from $N_{all}$ novel classes, and in each class, we only pick $K$ ($e.g.$, 1 or 5) samples for the support set $\mathcal{S}\!=\!\left\{\left(x_{i},y_{i}\right)\right\}_{i=1}^{N\times K}$ to train or fine-tune the model and sample $q$ instances for the query set $\mathcal{Q}\!=\!\left\{\left(x_{i},y_{i}\right)\right\}_{i=N\times K+1}^{N\times K+N\times q}$ to evaluate the model. By averaging the accuracy on the query set of multiple tasks from the novel dataset, we can evaluate the performance of a model.

Distribution calibration [14] uses the statistics of base classes to estimate the statistics of novel samples in the support set and generate more samples. Specifically, for the $b$th base class, its samples are assumed to be generated from a Gaussian distribution, whose mean and covariance matrix are:

where $b\!\in\![1,B]$, $\boldsymbol{x}_{j}\in\mathbb{R}^{V}$ is the $V$-dimensional feature of sample $j$ extracted from the pre-trained feature encoder, $J_{b}$ the number of samples in class $b$, and $\{\mu_{b},\boldsymbol{\Sigma}_{b}\}$ are the statistics of base class $b$.

The samples of a novel class are also assumed to be generated from a Gaussian distribution with mean $\boldsymbol{\mu}^{\prime}$ and covariance $\boldsymbol{\Sigma}^{\prime}$. As the novel class only has one or a few labeled samples, it is hard to accurately estimate $\boldsymbol{\mu}^{\prime}$ and $\boldsymbol{\Sigma}^{\prime}$. Thus, the key idea is to transfer the statistics of the base classes to calibrate the novel class’s distribution. Once the distribution of the novel class is calibrated, we can generate more samples from it, which are useful for training a good classifier. As a result, how to effectively transfer the statistics from the base classes is critical to the success of distribution calibration-based methods for few-shot learning. Accordingly, Free-Lunch [14] designs a heuristic approach that calibrates the Gaussian parameters of a novel distribution with one data sample $\boldsymbol{x}$:

where: 1) $\tilde{\boldsymbol{x}}$ is the transformed data by Tukey’s Ladder of Powers transformation (TLPT) [22], $i.e.$, $\tilde{\boldsymbol{x}}\!=\!\boldsymbol{x}^{\lambda}\text{ if }\lambda\!\neq\!0$ and $\tilde{\boldsymbol{x}}\!=\!\log(\boldsymbol{x})\text{ if }\lambda\!=\!0$, for reducing the skewness of distributions and make distributions more Gaussian-like. 2) $\mathbb{S}_{d}\!=\!\left\{-\left\|\boldsymbol{\mu}_{b}\!-\!\tilde{\boldsymbol{x}}\right\|^{2}\mid b\!=\![1,B]\right\}$ is a distance set defined by the Euclidean distance between the transformed feature $\tilde{\boldsymbol{x}}$ of a novel sample in support set and the mean $\boldsymbol{\mu}_{b}$ of the base class $b$. 3) $\operatorname{topk}(\mathbb{S}_{d})$ is the operation that selects the top-$k$ closest base classes from the set $\mathbb{S}_{d}$. 4) $\alpha$ determines the degree of dispersion of features sampled from the calibrated distribution.

Although effective, Free-Lunch represents a base class by the unweighted average over all its samples when computing its closeness with a novel sample, which ignores the fact that each sample of a base class may contribute to the classification boundary differently. In addition, the Euclidean distance in the feature space may not well capture the relations between a base class and a novel sample. Moreover, each selected top-$k$ base class has an equal weight ($i.e.$, $1/k$), which may not reflect the different contributions of the base classes and omit useful information in unselected base classes.

In this work, we propose a novel adaptive distribution calibration framework, a holistic method for few-shot classification. Compared to the novel classes, which only have a limited number of labeled samples, the base classes typically have a sufficient amount of data, allowing their statistics to be estimated more accurately. Due to the correlation between novel and base classes, it is reasonable to use the statistics of base classes to revise the distribution of the novel sample. Therefore, the key is how to transfer the statistics from the base classes to a novel class to achieve the best calibration results, which is the focus of this paper. Here, we develop the H-OT to learn the Transport Plan matrix between base classes and novel samples, where each element of the transport plan measures the importance of each base class for each novel sample and more relevant classes usually have a larger transport probability. Computing the high-level OT requires the specification of the cost function between one base class and one novel class, which leads to the introduction of a low-level OT problem. By viewing the learned transport plan as the adaptive weight matrix, we provide an elegant and principled way to transfer the statistics from the base classes to novel classes.

Moving beyond the Free-Lunch method [14], which uses the Euclidean distance between a novel sample and a base class to decide their similarity and endow the chosen base classes with the equal importance, we aim to capture the correlations between the base class and novel samples at multiple levels and transfer the related statistics from base classes to novel samples. We learn the similarity by minimizing a high-level OT distance between base and novel classes and build the cost function used in high-level OT by further introducing a low-level OT distance. To formulate the task as the high-level OT problem, we model $P$ as a following discrete uniform distribution over $B$ base classes:

where $R_{b}$ represents the base class $b$ in the $V$-dimensional feature space, which will be introduced later. Taking the $N$-way-$1$-shot task as the example, where each novel class has one labeled sample $\boldsymbol{x}$, we represent $Q$ as a discrete uniform distribution over $N$ novel classes from support set:

where $\tilde{\boldsymbol{x}}_{n}$ is the transformed feature from $\boldsymbol{x}$ following Yang et al. [14] and detailed below Equation 3. The OT distance between $P$ and $Q$ is thus defined as $\text{OT}(P,Q)=\min_{{\bf T}\in\Pi(P,Q)}\langle{\bf T},{\bf C}\rangle$. We adopt a regularised OT distance with an entropic constraint [21] and express the optimisation problem as:

where $\epsilon\textgreater 0$, ${\bf C}\in\mathbb{R}_{\geq 0}^{B\times N}$ is the transport cost matrix, and $C_{bn}$ indicates the cost between base class $b$ and novel sample $n$. Importantly, the transport probability matrix ${\bf T}\in\mathbb{R}_{>0}^{B\times N}$ should satisfy $\Pi(P,Q):=\left\{\sum_{n}^{N}T_{bn}=\frac{1}{B},\sum_{b}^{B}T_{bn}=\frac{1}{N}\right\}$ with element $T_{bn}=T(R_{b},\tilde{\boldsymbol{x}}_{n})$, which denotes the transport probability between the $b$th base class and the $n$th novel sample and is an upper-bounded positive metric. Therefore, $T_{bn}$ provides a natural way to weight the importance of each base class which can be used as the class similarity matrix when calibrating the novel distribution. Hence, the transport plan is the main thing that we want to learn from the data.

To compute the OT in Equation 6, we need to define the cost function ${\bf C}$, which is the main parameter defining the transport distance between probability distributions and thus plays the paramount role in learning the optimal transport plan. In terms of the transport cost matrix, a naive method is to specify the ${\bf C}$ with Euclidean distance or cosine similarity between the feature space of the novel sample and mean of the features from the base classes, such as $C_{bn}\!=\!1\!-\!\cos\left(\tilde{\boldsymbol{x}}_{n},\boldsymbol{\mu}_{b}\right)$. However, these manually chosen cost functions might have the limited ability to measure the transport cost between a base class and a novel sample. Besides, representing the base class only with the average of all features in class $b$ might ignore the contributions of different samples for this class. Hence the optimal transport plan for these cost functions might be inaccurate. To this end, we further introduce a low-level OT optimization problem to automatically learn the transport cost function ${\bf C}$ in (6). Specifically, we further treat each base class $b$ as an empirical distribution $R_{b}$ over the features within this class:

where $p_{j}^{b}$ is the weight of data $\boldsymbol{x}_{j}^{b}$ to base class $b$ and captures the importance of this sample and will be described in short order. Specifically, we train a classifier parameterized by $\phi$ with the samples in the base classes, which predicts which base class a sample is in. The predicted probability of sample $j$ belonging to the base class $b$ is denoted by $s^{b}_{j}$ and then $[p_{1}^{b},\dots,p_{J_{b}}^{b}]$ is obtained by normalizing the vector $[s_{1}^{b},\dots,s_{J_{b}}^{b}]$ with the Softmax function. We further define the low-level OT distance between each distribution $R_{b}$ and distributions $Q$ with an entropic constraint as

where the transport probability matrix ${\bf M}^{b}\in\mathbb{R}_{>0}^{J_{b}\times N}$ should satisfy $\Pi(R_{b},Q):=\left\{\sum_{j}^{J_{b}}M_{jn}^{b}p_{j}^{b}=\frac{1}{N},\sum_{n}^{N}M_{jn}^{b}\frac{1}{N}=p_{j}^{b}\right\}$. The element $D_{jn}^{b}$ of cost function ${\bf D}^{b}\in\mathbb{R}_{>0}^{J_{b}\times N}$ is the distance between the $n$-th novel sample and the $j$-th sample from base class $b$, where we naturally use the distance between their features, i.e., $\tilde{\boldsymbol{x}}_{n}$ and $\boldsymbol{x}_{j}^{b}$. Empirically, we find that the cosine similarity $D_{jn}^{b}=1-\text{cos}(\tilde{\boldsymbol{x}}_{n},\boldsymbol{x}_{j}^{b})$ works well in practice, although other choices are possible.

Minimizing the low-level OT loss in Equation (8) can learn an optimal transport probability matrix ${\bf M}^{b}$ for $b\in B$, where $M^{b}_{jn}$ tells us the transport weight between the $n$th novel class and $j$th sample in $b$ base class. Back to the cost function $C_{bn}$ for the high-level OT, instead of using the manually chose cost functions, we further adopt the learned total transport cost between novel sample $\boldsymbol{x}_{n}$ and all samples from base class $b$ in the low-level OT to represent the cost function $C_{bn}$ in high-level OT:

where $C_{bn}$ is fed into (6) for learning the transport plan between base classes and novel samples. Defined in this way $C_{bn}$ is an adaptive cost between base class $b$ and novel sample $n$ in support set, taking full advantage of all samples in a base class. With the proposed low-level OT distance, the distance between a novel sample and a base class is no longer representing a base class by the unweighted average over all its samples and then using the Euclidean distance. In our method, the weights of the samples are considered in a principled way.

Once we obtain the transport plan matrix by minimizing the OT problem in (6), we can compute the statistics of novel samples by following Yang et al. [14]. For the $n$-th transformed feature $\tilde{\boldsymbol{x}}_{n}$ in the support set, we calibrate its mean and covariance as follows:

where $\mu_{b}$ and $\boldsymbol{\Sigma}_{b}$ are the statistics of base class $b$ by (2), respectively; $T_{bn}$ provides an adaptive way to weight the importance of base class $b$ for novel sample $n$; $N$ is used to scale $T_{bn}$ since $\sum_{b\in B}T_{bn}=1/N$; $\alpha$ is a hyper-parameter explained in Equation  (3).

For $N$-way-$K$-shot task with $K\textgreater 1$, the $N$ in aforementioned Equations (6) and (8) should be replaced with $N*K$, and the distribution calibration in Equation (10) should be undertaken $K$ times for each novel class. Notably, this way can potentially achieve more diverse and accurate distribution estimation for avoiding the bias provided by one specific sample. Thus, for a class $y\in N$ in novel task, we denote the calibrated distribution with a set of statistics as

Based on set $\mathbb{S}_{y}$ for novel class $y$, we can sample from the calibrated Gaussian distributions to generate a set of feature vectors with label $y$: $\mathbb{D}_{y}=\left\{(x,y)\mid x\sim\mathcal{N}(\mu,\Sigma),\forall(\mu,\Sigma)\in\mathbb{S}^{y}\right\}$. Given the transformed support set $\tilde{\mathcal{S}}$ with TLPT and the generated features $\mathbb{D}_{y}$, we only need to train the classifier for a task $\mathcal{T}$ ($under$, N-way-K-shot) by minimizing the cross-entropy loss:

where $\mathcal{Y}^{\mathcal{T}}$ is the set of classes for the task $\mathcal{T}$ and the classifier is parameterized by $\theta$. We describe our proposed framework in Algorithm 1.

Datasets We evaluate our proposed method on several standard few-shot classification datasets with different levels of granularity, including miniImageNet [33], tieredImageNet [34], CUB [35], and CIFAR-FS [36]. Here, miniImageNet, tieredImageNet, and CIFAR-FS have a broad range of classes including various animals and objects while CUB is a more fine-grained dataset that includes various species of birds. The miniImageNet, which is randomly chosen from ILSVRC-12 dataset [37], contains 100 diverse classes with 600 samples per class. The image size is $84\times 84\times 3$. Following the previous work [33], we split the dataset into $64$ base classes, $16$ validation classes, and $20$ novel classes. The tieredImageNet is a larger subset of ILSVRC-12 dataset [37]. It contains 608 classes sampled from a hierarchical category structure, where each class belongs to one of 34 higher-level categories from the high-level nodes in the ImageNet. The average number of images in each class is 1281. We adopt 351, 97, and 160 classes for training, validation, and test, respectively. The CUB is a fine-grained few-shot classification benchmark, which contains 200 different classes of birds with a total of 11,788 images, and we resize them as $84\times 84\times 3$. We follow previous works [3] and split the dataset into 100 base classes, 50 validation classes, and 50 novel classes. The CIFAR-FS dataset is a recently proposed few-shot image classification benchmark, consisting of all 100 classes from CIFAR-100 [38]. The classes are randomly split into $64$, $16$, and $20$ for meta-training, meta-validation, and meta-testing, respectively. Each class contains 600 images of size $32\times 32\times 3$.

Implementation Details Since Free-Lunch is the SOTA distribution calibration method, we mainly compare with Free-Lunch. For miniImageNet and CUB, we directly use pre-extracted features from Free-Lunch’s code [14], which are produced by WideResNet trained from scratch on the base classes by following Mangla et al. [39] to optimize a classification loss and a self-supervised rotation loss jointly, denoted as WRN28+RTLoss. Free-lunch did not share the features on tieredImageNet and CIFAR-FS, thus we have taken the pre-trained WRN28+RTloss backbone of Mangla et al. [39] to extract the features, where we use this same backbone to implement Free-lunch with its official code and we have explored and selected an appropriate $w$ for Free-lunch. To confirm the validity of our proposed method on the commonly-used backbone, we also take the ResNet12 as the example, which is trained from scratch by minimizing a standard classification loss (without the rotation loss). To make sure the log transform in TLPT is valid to its inputs, we need to guarantee the non-negativity of features. Therefore, we extract the features of given samples from the penultimate layer (with a ReLU activation function) of the feature extractor. To demonstrate the effectiveness of our method, we consider the simple Logistic Regression (LR) classifier for $\theta$ and $\phi$ in Algorithm 1, where the LR implementation of scikit-learn [40] with the default settings is adopted. We use the same hyper-parameter value for all datasets except for $\lambda$. Specifically, the number of generated features is $750$, the $\epsilon$ in Sinkhorn algorithm is $0.01$, the $\alpha$ in (10) is $0.21$; and $\lambda$ is $0.5$, $1$, $1$ and $0.8$ for miniImageNet, tieredImageNet, CUB, and CIFAR-FS, respectively, selected by a grid search using the validation set. The maximum iteration number in Sinkhorn algorithm is set as $200$. The reported results are the averaged classification accuracy over 10,000 tasks.

We now conduct experiments on the most common setting in few-shot classification, 1-shot and 5-shot classification, where the results of different models on miniImageNet, tieredImagenet, CUB and CIFAR-FS are shown in Tables 1 and 2. To validate the effectiveness of our proposed model, we compare it with three main groups of the few-shot learning method, including optimization-based, metric-based, and distribution calibration-based. There are several observations. First, we can find that our H-OT equipped with WRN28+RTLoss and simple linear classifier outperforms all previous competing methods on 5-way-1-shot and 5-way-5-shot settings, which illustrates the superiority of the proposed framework. Second, our H-OT equipped with ResNet12 is better than other baselines except for the 5way1shot task on miniImageNet, where in this case our H-OT is still competitive with DeepEMD [30]. This observation proves the effectiveness of distribution calibration and transferring the statistics from base classes to novel classes using OT. Third, the comparison between our model and the strong baseline of distribution calibration-based methods, i.e., Free-Lunch [14], confirms the validity of incorporating the adaptive weight matrix when transferring the statistics from the base classes to the novel classes. Although both methods utilize distribution calibration to perform few-shot classification, our model can more effectively transfer the knowledge of the base classes to estimate the distribution of novel samples more accurately with the help of hierarchical OT (H-OT). Fourth, in terms of 5-way 1-shot tasks on all datasets, the performance gain from our model over Free-Lunch on the CUB is larger than that of other datasets. However, for 5-way 5-shot tasks on all datasets, the improvement gain from our model over Free-Lunch on the CUB is not as significant as that on other datasets. The reason behind this might be that CUB is a fine-grained image classification dataset, which makes it more effective for statistics transfer compared to other datasets under the 1-shot setting. Furthermore, with the development of label examples in the novel classes (5-shot), our method can perform more effective knowledge transfer from base classes to novel samples in the coarse-grained image classification dataset. The results reveal that the granularity of the dataset is an important factor in the current few-shot classification setting. Our model avoids the costly fine-tuning for backbone and it only takes time at the testing stage. We defer the details on comparison of computational cost to Appendix A and summary of test results to Appendix B.

To explore the cross-domain generalization ability of H-OT for the few-shot classification task, we consider a practical evaluation setting following Chen et al. [3], where there exists domain shift between base and novel classes (e.g., sampling base classes from a coarse-grained dataset and novel classes from a fine-grained dataset). We design two cross-domain scenarios: miniImageNet $\rightarrow$ CUB, CIFAR-FS $\rightarrow$ CUB, where CUB is the target domain while both miniImageNet and CIFAR-FS are the source domains. Compared with collecting images from a general class, collecting images from fine-grained classes might be more difficult, so these two scenarios are very meaningful in reality. We consider the comparison with Free-Lunch, which is closely related to ours. As shown in Fig. 1, both H-OT and Free-Lunch achieve acceptable performance in these two scenarios, proving the effectiveness of the distribution calibration framework in the cross-domain case. Besides, H-OT shows its clear advantage over Free-Lunch. It indicates that H-OT can better utilize the statistics from the base classes of the source domain with the adaptive weight matrix and transfer them to the novel classes of the target domain. Moreover, H-OT under the CIFAR-FS $\rightarrow$ CUB (blue bar in the 3rd column) scenario still outperforms Free-Lunch under CUB (orange bar in the 1st column). These results show that H-OT is less affected by domain shift between the base and novel classes and possesses the desired cross-domain generalization ability.

Impact of adaptive cost from low-level OT Recall that we formulate our general goal as a high-level OT problem and we propose a low-level OT for computing the cost function for the high-level one in an adaptive way. The low-level OT captures the distance between a base class and a novel class. Here we consider two commonly-used cost functions to replace the adaptive cost, including Euclidean distance and cosine similarity between the feature space of the novel classes and the mean of the features from the base classes, denoted as High-level OT+Euc and High-level OT+cos. Besides, to exclude the possibility that the gain of our H-OT comes only from the effect of introducing weighted mean in the low-level OT, we also consider use the weighted mean with (7) to implement High-level OT, denoted as High-level OT+Euc+weighted and High-level OT+cos+weighted. As summarized in Table 6, our proposed H-OT achieves better performance than these degraded variants, which confirms the validity of adaptive cost from low-level OT in few-shot classification. And introducing the weighted mean into the cost function of high-level OT does not improve the performance. It indicates that the performance gain comes mainly from the low-level OT rather than the effect of weighted mean. The reason behind this might be that our proposed low-level OT considers all samples within a base class when deciding the cost between novel samples and this base class, instead of using the mean of the features from the base class. Besides, even with the given cost functions, our proposed method still outperforms the competing baseline (i.e., Free-Lunch), indicating the usefulness of the transport plan (i.e., adaptive matrix) learned from high-level OT.

Different backbones Since our framework is agnostic with the choice of backbones, following Yang et al. [14], we further employ other commonly used backbones, such as networks with four or six convolutional layers (Conv4 or Conv6), Resnet10 and Resnet18, which are trained from scratch following [3]. Table 4 illustrates the classification performance (5-way-1-shot) comparison between our H-OT and strong baseline (Free-Lunch) on miniImageNet with different backbones. We find that H-OT always outperforms its baseline given the same backbone. This result reveals the superiority of learning an adaptive weight matrix to select statistics of base classes on different backbones. As H-OT is agnostic to the feature extractor and does not require costly fine-tuning, it can also conveniently work with more advanced feature extractors to get better performance.

Number of generated features Fig. 3 shows the variation of the classification performance of distribution calibration models with the increasing number of generated features, where we consider our H-OT (red) and Free-Lunch (green) in a 5-way-1-shot classification setting for miniImageNet. We can see that increasing the number of generated samples results in consistent improvement in both cases. However, given the same number of generated samples, H-OT outperforms its baseline, which only selects the $top-k$ base classes for each novel sample using Euclidean distance [14]. Interestingly, our model achieves around 68.5 accuracy with only 50 generated samples, while Free-Lunch takes 750 generated samples to achieve a similar result. Hence, the adaptive weight matrix learned by H-OT can contribute to the distribution calibration strategy for sampling more representative features.

Learned transport plan To explore whether the learned transport plan can capture the correlations between the base classes and one given sample from the novel class (“plain”), we show a list of base classes with decreasing ranks in Fig. 3. The top-5 base classes are highly related to the given novel sample about their semantic meanings. It is reasonable since the images from wild animal-related classes are usually taken on the plain, such as “elephant” and “bear”. It indicates that the transport plan can be successfully optimized by our H-OT and provides an elegant way to measure the closeness between base classes and novel samples. Different from Free-Lunch that only selects top-$k$ base classes ($k$=2), our model takes full advantage of all base classes for learning adaptive weight for each base class. Additional quantitative results and qualitative are deferred to the Appendix B.