Rethinking Generalization in Few-Shot Classification

As depicted in Figure 2, we encode the image patches $\mathcal{P}_{s}$ of the support set samples along with the query sample patches $\mathbf{p}_{q}$ via $f_{\theta}$ and obtain corresponding sets of tokens $\mathcal{Z}_{s}$ and $\boldsymbol{z}_{q}$, respectively³³3Note that some tensor shapes in the illustrations might differ from the equations for ease of visualization.. It is to be noted that while we choose to illustrate our method via the use of one single query sample, the classification of all query samples is computed at the same time in one single pass in practice. We retrieve our ‘prior’ correspondence map $\boldsymbol{S}$ expressing the token-wise similarity between the encoded semantic content of the local regions in the query sample and all patches of all support samples, allowing us to consider all available information jointly without incurring information loss due to averaging or similar operations. This ‘prior’ similarity map represents correspondences between regions of samples irrespective of their individual class, \xperiodafteri.ealso between entities that might not be relevant or potentially even harmful for correct classification. Using the annotated support set samples, we now infer a task-specific importance weight factor $v^{j}$ for each support token $z_{s}^{j}$ representing its contribution to correctly classify other samples in the support set via online optimization at inference time (Section 2.4). We then reweight the prior similarities to obtain our classification result for the query sample, jointly considering all available information.

To overcome the problem of supervision collapse induced by image-level annotations, we split the input images into smaller parts where each region has a higher likelihood of only containing one major entity and hence a more distinct semantic meaning. Since no labels are available for this more fine-grained data, we encode the information of each local region via an unsupervised method.

We build our approach around the recently introduced idea of using Masked Image Modeling (MIM) [3, 59] as a pretext task for self-supervised training of Vision Transformers. In contrast to previous unsupervised approaches [5, 7] which focused mainly on global image-level representations, MIM randomly masks a number of patch embeddings (tokens) and aims to reconstruct them given the remaining information of the image. The introduced token constraints help our Transformer backbone to learn an embedding space that yields semantically meaningful representations for each individual image patch. We then leverage the information of the provided labels through fine-tuning the pretrained backbone in conjunction with our inner loop token importance weighting described in the following sections while successfully avoiding supervision collapse (see experimental results in Section 3.2).

As illustrated in Figure 2, we split each input image $\boldsymbol{x}\in\mathbb{R}^{H\times W\times C}$ into a sequence of $M\!=\nicefrac{{H\cdot W}}{{P^{2}}}$ patches $\mathbf{p}=\{p^{i}\}^{M}_{i=1}$, with each patch $p^{i}\in\mathbb{R}^{P^{2}\times C}$. We then flatten and pass all patches of the support and query images as input to the Transformer architecture, obtaining the set of support tokens $\mathcal{Z}_{s}\!=\!f_{\theta}(\mathcal{P}_{s})$ with $\mathcal{Z}_{s}\!=\!\{{\boldsymbol{z}}_{s}^{nk}|n\!=\!1,\ldots,N,k\!=\!1,\ldots,K\},{\boldsymbol{z}}_{s}^{nk}\!=\!\{z_{s}^{nkl}|l\!=\!1,\ldots,L;z_{s}^{nkl}\in\mathbb{R}^{D}\}$ and query tokens ${\boldsymbol{z}}_{q}\!=\!f_{\theta}(\mathbf{p}_{q})$ with ${\boldsymbol{z}}_{q}\!=\!\{z_{q}^{l}|l\!=\!1,\ldots,L;z_{q}^{l}\in\mathbb{R}^{D}\}$. Vision Transformers like ViT [10, 46] satisfy $L=M$ whereas hierarchical Transformers like Swin [27] generally emit a reduced number of tokens $L<M$ due to internal merging strategies.

Having obtained all patch embeddings, we establish semantic correspondences by computing the pair-wise patch similarity matrix between the set of support tokens $\mathcal{Z}_{s}$ and query tokens $\boldsymbol{z}_{q}$ as $\boldsymbol{S}\in\mathbb{R}^{N\cdot K\cdot L\times L}$, where each element in $\boldsymbol{S}$ is obtained by $s_{nk}^{l_{s},l_{q}}={\mathrm{sim}}(z_{s}^{nkl_{s}},z_{q}^{l_{q}})$, where $l_{s}\!=\!1,\ldots,L$ and $l_{q}\!=\!1,\ldots,L$. Note that local image regions representing similar entities exhibit higher scores. While any distance metric can be used to compute the similarity ($\mathrm{sim}$), we found cosine similarity to work particularly well for this task. We then use the task-specific token importance weights $\boldsymbol{v}\in\mathbb{R}^{N\cdot K\cdot L\times 1}$ inferred via online optimization based on the annotated support set samples (see Section 2.4) to reweight the similarities through column-wise addition and obtain our task-specific similarity matrix as ${\boldsymbol{\tilde{S}}}=\boldsymbol{S}+\left[\boldsymbol{v}\cdot\mathbbm{1}^{1\times L}\right]$, with elements $\tilde{s}_{nk}^{l_{s},l_{q}}$. Note that this addition of our reweighting logits corresponds to multiplicative reweighting in probability space. We temperature-scale the adapted similarity logits with $\nicefrac{{1}}{{\tau_{S}}}$ and aggregate the token similarity values across all elements belonging to the same support set class via a LogSumExp operation, \xperiodafteri.eaggregating $K\cdot L^{2}$ logits per class followed by a softmax – resulting in the final class prediction $\hat{\boldsymbol{y}}_{q}$ for the query sample $\boldsymbol{x}_{q}$ as

$$\hat{\boldsymbol{y}}_{q}=\mathrm{softmax}\left(\left\{\hat{y}_{q}^{n}\right\}_{n=1}^{N}\right)=\mathrm{softmax}\left(\left\{\log\sum_{k=1}^{K}\sum_{l_{q}=1}^{L}\sum_{l_{s}=1}^{L}\mathrm{exp}\left(\nicefrac{{\tilde{s}_{nk}^{l_{s},l_{q}}}}{{\tau_{S}}}\right)\right\}_{n=1}^{N}\right)\,.$$

(1)

We use all samples of a task’s support set together with their annotations to learn the importance for each individual patch token via online optimization at inference time. As depicted in Figure 3, we formulate the same classification objective as in the previous section but aim to classify the support set samples instead of a query sample. In other words, we use two versions of the support set samples: one with labels as ‘support’ $\mathcal{Z}_{s}$ and one without as ‘pseudo-query’ $\mathcal{Z}_{sq}$, and obtain the similarity matrix $\boldsymbol{S}_{s}\in\mathbb{R}^{N\cdot K\cdot L\times N\cdot K\cdot L}$. The token importance weights are initialized to $\boldsymbol{v}^{0}=\boldsymbol{0}\in\mathbb{R}^{N\cdot K\cdot L\times 1}$ and column-wise added to form $\tilde{\boldsymbol{S}}_{s}=\boldsymbol{S}_{s}+\left[\boldsymbol{v}^{0}\cdot\mathbbm{1}^{1\times N\cdot K\cdot L}\right]$.

The goal is now to determine which tokens of $\mathcal{Z}_{s}$ are most helpful in contributing towards correctly classifying $\mathcal{Z}_{sq}$, and which ones negatively affect this objective. To prevent tokens from simply classifying themselves, we devise the following strategy for our $N$-way $K$-shot tasks. For scenarios with $K>1$ samples per class, we apply block-diagonal masking with blocks of size $L\times L$ to the similarity matrix $\tilde{\boldsymbol{S}}_{s}$ – meaning that we enforce classification of each token in $\mathcal{Z}_{sq}$ exclusively based on information from other images. Since there are no other in-class examples available in $1$-shot scenarios, we slightly weaken the constraint and apply local masking in an $m\times m$ window around each patch, forcing the token to be classified based on the remaining information in the image. We found a local window of $m=5$ to work well throughout our experiments for both architectures.

We use temperature-scaling and aggregate all modified similarity logits across all elements belonging to the same support set class of the annotated $\mathcal{Z}_{s}$ for each element $\boldsymbol{z}_{sq}^{nk}$, apply a softmax operation and obtain the predicted class probabilities $\hat{\boldsymbol{y}}^{nk}_{s}$ for each support set sample (cf. Equation 1). Given that the prediction is now dependent on the initialized token similarity weights $\boldsymbol{v}$, we can formulate an online optimization objective by using the support set labels $\boldsymbol{y}^{nk}_{s}$ as

$$\operatorname*{arg\,min}_{\boldsymbol{v}}\sum_{n=1}^{N}\sum_{k=1}^{K}\mathcal{L}_{\mathrm{CE}}\left(\boldsymbol{y}^{nk}_{s},\;\hat{\boldsymbol{y}}^{nk}_{s}\left(\boldsymbol{v}\right)\right)\,.$$

(2)

It is to be noted that by using column-wise addition of $\boldsymbol{v}$, we share the importance weights of each support token across all ‘pseudo-query’ tokens and thus constrain the optimization to jointly learn token importance with respect to all available information. In other words, task-specific intra-class similarities will be strengthened while inter-class ones will be penalized accordingly. We further do not require any second order derivatives like other methods during meta fine-tuning, since the optimization of $\boldsymbol{v}$ is decoupled from the network’s parameters.

Our training strategy is divided into two parts: self-supervised pretraining followed by meta fine-tuning. It is to be noted that each architecture is exclusively trained on the training set data of the respective dataset that is to be evaluated, and no additional data is used.

Datasets. We train and evaluate our methods using four popular few-shot classification benchmarks, namely miniImageNet [48], tieredImageNet [37], CIFAR-FS [4] and FC-100 [34].

Architectures. We compare two different Transformer architectures in this work, the monolithic ViT architecture [10, 46] in its ‘small’ form (ViT-small aka DeiT-small) and the multi-scale Swin architecture [27] in its ‘tiny’ version (Swin-tiny).

Self-supervised pretraining. We employ the strategy proposed by [59] to pretrain our Transformer backbones and mostly stick to the hyperparameter settings reported in their work. Our ViT and Swin architectures are trained with a batch size of $512$ for $1600$ and $800$ epochs, respectively. We use 4 Nvidia A100 GPUs with 40GB each for our ViT and 8 such GPUs for our Swin models.

Meta fine-tuning. We use meta fine-tuning to further refine our embedding space by using the available image-level labels in conjunction with our token-reweighting method. We generally train for up to $200$ epochs but find most architectures to converge earlier. We evaluate at each epoch on 600 randomly sampled episodes from the respective validation set to select the best set of parameters. During test time, we randomly sample 600 episodes from the test set to evaluate our model.

Token importance reweighting and classifier. We use cosine similarity to compute $\boldsymbol{S}$. While the temperature $\tau_{S}$ used to scale the logits can be learnt during meta fine-tuning, we found $\tau_{S}=\nicefrac{{1}}{{\sqrt{d}}}$ to be a good default value (or starting point if learnt, see supplementary material). We use SGD as optimizer with a learning rate of 0.1 for the token importance weight generation.

Please refer to supplementary material of this paper for a more detailed discussion regarding datasets, implementation and hyperparameters.

In this section, we investigate the influence of self-supervised pretraining compared to its supervised counterpart using the provided image-level labels (de-facto standard in most current state-of-the-art methods). We only vary the training strategy of our backbone and use our introduced token similarity-based classifier with 15 inner-loop reweighting steps for all experiments. Figure 4 illustrates that both ViT-small and Swin-tiny with self-supervised pretraining alone (w/o meta fine-tuning) learn more generalizable features than their respective supervised versions for both $1$-shot (left-hand side) and $5$-shot (right-hand side) scenarios. In fact, across all cases w/o meta fine-tuning, the self-supervised pretrained Transformers outperform their supervised counterparts by more than $10\%$ and up to a significant $18.19\%$ in the case of Swin-tiny. This clearly demonstrates that our self-supervised pretraining strategy explores information that is richer and beyond the labels. Another interesting insight is that the meta fine-tuning stage does not improve supervised backbones as much as their unsupervised counterparts. Specifically after unsupervised pretraining, meta fine-tuning of the Transformers in conjunction with our token-reweighting strategy is able to boost the performance by $6.77\%$ (ViT-small) and $13.01\%$ (Swin-tiny) for $1$-shot, and $9.94\%$ (ViT-small) and $10.10\%$ (Swin-tiny) for $5$-shot. While such significant improvements cannot be observed across the supervised networks, the Swin versions seem to generally start off lower after pretraining but benefit more from the fine-tuning than ViT. The observed results clearly indicate that our token-reweighted fine-tuning strategy is able to further improve the generalization of self-supervised Transformers, thus performing better on the novel tasks of the unseen test set. Figure 5 additionally depicts projected views of the tokens of 5 instances from a novel class as well as the entire novel support set in embedding space. Representations obtained with our classifier seem to retain the instance information (‘w/o $\boldsymbol{v}$’) and separation is improved when using token importance reweighting (‘w/ $\boldsymbol{v}$’). While the projected tokens of the entire support set show partial overlap between classes as is expected due to commonalities like \xperiodaftere.gsimilar background, our reweighting clearly determines the class-characteristic tokens (displayed in their original class-respective color). These results indicate that our similarity-based classifier coupled with task-specific token reweighting is able to better disentangle the embeddings of different instances from the same class as well as other classes, which prevents the network from supervision collapse and achieves the higher performance observed on the benchmarks. They further show that self-supervised pretraining is helpful but not sufficient to achieve well-separated representations without supervision collapse that are suitable for few-shot classification.

Figure 6 shows a visualization of the patch importance weights $\boldsymbol{v}$ that are learned at inference time during the inner loop adaptation for the support set images. Brighter regions represent a higher importance weight, meaning that these patches will contribute most to the classification of query samples if matches with high similarity can be found. Judging from the visualized weights, FewTURE seems to consistently select characteristic regions of the depicted objects, \xperiodaftere.g, the rim of the bowls, strings of the guitar or the dogs’ facial area, and to exclude unimportant or out-of-task information.

We further investigate the influence of the number of optimization steps in our inner loop token importance weighting method using ViT-small on miniImageNet. The results in Figure 7 indicate that increasing the steps up to 20 aligns with increased performance, both during validation and testing. While the initial increase in test accuracy when using our token reweighting (steps $>0$) is rather significant with $1.15\%$, the contribution of higher step numbers comes at the cost of higher computational complexity, and we generally found anything between 5 and 15 steps to be a good performance \xperiodaftervsinference-time trade-off (see supplementary material for further details).

We conduct experiments using the few-shot settings established in the community, namely $5$-way $1$-shot and $5$-way $5$-shot – meaning the network has to distinguish samples from 5 novel classes based on a provided number of 1 or 5 images per class. We evaluate our method FewTURE using two different Transformer backbones and compare our results against the current state of the art in Table 1 for the miniImageNet and tieredImageNet, and in Table 2 for the CIFAR-FS and FC100 datasets. It is to be noted that in contrast to previous works, we do not employ the help of any convolutional backbone but instead (and as far as we are aware for the first time) use a Transformer backbone together with our previously introduced token importance reweighting method to achieve these results. We are able to set new state of the art results across all four datasets in both $5$-shot and $1$-shot settings, improving particularly the $1$-shot results on miniImageNet and tieredImageNet by significant margins of $3.03\%$ and $1.74\%$, respectively.

To further improve the computational efficiency of our method, we investigate its behaviour when limiting the number of tokens that are considered to establish patch-wise correspondences to only a subset – allowing FewTURE to scale to potentially large many-way many-shot settings. We use the attention maps inherent in our approach (averaged over all heads) to prune the number of patch tokens by only using the ones within the top-k attention values to compute the similarity matrix $\boldsymbol{S}$. The results obtained with our ViT-small backbone for pruning the number of tokens to $75\%$, $50\%$, $25\%$ and $10\%$ of the original token number indicate that such pruning might be an interesting avenue for future work, with our method still achieving $96.4\%$ of its original performance when only retaining $10\%$ of the number of tokens.

We train a linear classifier as well as prototypical approach using our pre-trained ViT-small backbone for a 5-way 5-shot setting on miniImagenet [48] to investigate the influence of the choice of classifier (Table 4). We were able to obtain a test accuracy of $82.80\%$ for the prototypical network after optimising the pre-trained backbone with meta-finetuning, which is competitive to the results we obtain for our method without reweighting (’0 step’ in Figure 7(7(b))), but is clearly outperformed by our reweighting-based approach. To provide a fair comparison, we optimize the linear classifier at inference time to adapt to the support set and obtained a maximum test accuracy of $82.37\%$. Both results indicate the quality of embedding our backbone is able to produce but also demonstrate the importance of our task-specific reweighting-based approach. For further ablation studies, please refer to the supplementary material.

Our introduced method is arguably more powerful for cases where multiple examples of the same class are provided, \xperiodafteri.ein $N$-way $K$-shot settings with $K>1$. While FewTURE works well in $1$-shot settings (Table 1 and 2), the inner loop adaptation procedure still aims to exclude cross-class similarities that hurt classification performance, but has less diverse information for selecting the most helpful regions due to the lack of other in-class comparison samples – which might thus yield slightly less-refined token selections compared to multi-shot scenarios (supplementary material).

While we did not face significant problems regarding the comparably very small size of our datasets (\xperiodaftere.gminiImageNet with 38K compared to the usually used ImageNet with 1.28M training images), more specialized applications with highly limited training data might be negatively impacted and successfully training our method due to the reduced inductive bias present in the Transformer architecture could prove challenging.