Disentangled Feature Representation for Few-shot Image Classification

According to the meta-learning framework (Vinyals et al. 2016), there are mainly three types of few-shot learning methods. Firstly, the gradient-based methods utilize a good model initialization (Finn, Abbeel, and Levine 2017; Nichol, Achiam, and Schulman 2018) or optimization strategy (Ravi and Larochelle 2017; Rusu et al. 2019; Lee et al. 2019; Liu, Schiele, and Sun 2020) to quickly adapt to novel tasks. Secondly, the data augmentation-based methods focus on generating (Gidaris and Komodakis 2019; Li et al. 2020a) or gathering augmented data (Hariharan and Girshick 2017; Wang et al. 2018; Yang, Liu, and Xu 2021) to enable classification from limited samples. In this work, we focus on the third type, namely the metric learning-based methods, i.e., to learn the discriminative feature embedding for distinguishing different image classes. For example, ProtoNet (Snell, Swersky, and Zemel 2017) considered the class-mean representation as the prototype of each class and applied the Euclidean distance metric for classification. LCR (Tokmakov, Wang, and Hebert 2019) applied the subspace-based embedding for each class and DeepEMD (Zhang et al. 2020) adopted the earth mover’s distance as the metric function to compare the similarity between two feature maps in a structured way. FEAT (Ye et al. 2020) defined four kinds of set-to-set transformation including self-attention transformer (Jaderberg et al. 2015) to learn a task-specific feature embedding for few-shot learning. Based on prior knowledge, COMET (Cao, Brbic, and Leskovec 2021) mapped some high-level visual concepts into a semi-structured metric space and then learned an ensemble classifier by combining the outputs of independent concept learners. Tang et al. (Tang, Wertheimer, and Hariharan 2020) also uses a semi-structured feature space based on independent prior knowledge concepts to do pose normalization for fine-grained tasks. Our work does not intend to propose new metrics, but focuses on extracting the class-specific features from the variations distracting the metric learning, thus to improve few-shot classification.

Disentangled feature representation aims to learn an interpretable representation for image variants, which has been widely studied in tasks such as face generation (Chen et al. 2016), style translation (Lee et al. 2020; Liu et al. 2019), image restoration (Li et al. 2020b), video prediction (Hsieh et al. 2018) and image classification (Prabhudesai et al. 2021; Li et al. 2021). InfoGAN (Chen et al. 2016) applied an unsupervised method to learn interpretable and disentangled representations by maximized mutual information. DRIT (Lee et al. 2020) embedded images into a content space and a domain-specific attribute space and applied a cycle consistency loss for style translation. FDR (Li et al. 2020b) applied channel-wise feature disentanglement to reduce the interference between hybrid distortions for hybrid-distorted image restoration. Li et al. (Li et al. 2021) proposed a disentangled-VAE to excavate category-distilling information from visual and semantic features for generalized zero-shot learning.

It is noteworthy that the very recent D3DP (Prabhudesai et al. 2021) also adopted a feature disentangling scheme for few-shot detection and VQA, by dividing high-dimensional data (e.g., RGB-D) into individual objects and other attributes. However, our DFR significantly differs from D3DP in the following aspects: DFR is a general-purpose feature extractor for image classification, while D3DP only disentangles individual object to tackle more specific object detection and VQA in 3D scenes. Besides, our DFR works on real images from few-shot benchmarks while D3DP only works with synthetic scenes. Moreover, the DFR framework is much simpler comparing to D3DP with fewer parameters and more efficient algorithms. Thus DFR can serve an enhanced feature extractor for classic backbones that are widely used in most of the existing few-shot methods.

Given a training image set with the base classes $\mathcal{C}_{train}$, few-shot image classification task aims to predict the novel classes $\mathcal{C}_{test}$ from the testing set, i.e., $\mathcal{C}_{train}\cap\mathcal{C}_{test}=\emptyset$. Thus, the trained classifier from $\mathcal{C}_{train}$ needs to be generalized to $\mathcal{C}_{test}$ in the testing stage with only few labeled samples. In this paper, we follow the meta-learning strategy (i.e., the $N$-way $K$-shot setting) (Vinyals et al. 2016) to simulate meta-tasks in the training set that are similar to the few-shot setting at the testing stage, i.e., each meta-task $\mathcal{T}_{FS}$ contains a support set $\mathcal{S}$, and a query set $\mathcal{Q}$. The support set $\mathcal{S}$ contains $N$ classes with $K$ labeled samples ($N$ and $K$ are both very small) and a query set $\mathcal{Q}$ with unlabeled query samples from $N$ classes is used to evaluate the performance.

Figure 2 is an overview of the proposed DFR framework. With a few-shot task $\mathcal{T}_{FS}$ of support set $\mathcal{S}$ and query set $\mathcal{Q}$, the objective is to extract discriminative features for classification from the excursive information of each image $x_{i}$. The proposed DFR consists of two branches with two encoders, i.e., $E_{cls}$ and $E_{var}$ for the classification and variation branches, respectively, and one decoder $D$, as well as a discriminator with a gradient reverse layer and a relation module.

To be specific, the output feature of the classification branch is fed to the MLP module $g$ to extract class-specific information $(\mu,\sigma)$ of each sample for scaling the feature of the variation branch in the follow-up decoder. The decoder can reconstruct or translate the source image based on different sources of $(\mu,\sigma)$, shown in Figure 2, as

$$\hat{x}_{i}=D(E_{var},g(X)),$$

(3)

where $X$ can be the feature of the $i$-th sample itself for self-reconstruction, or the mean feature of class $y_{i}$ for class-reconstruction or another class $y_{j}$ with $j\neq i$ for class-translation.

The objective function consists of the discriminative loss $L_{dis}$, cross-entropy loss $L_{cls}$, reconstruction loss $L_{rec}$ and translation loss $L_{tran}$.

Moreover, the perceptual loss is also adapted to measure perceptual differences between the output image $\hat{x}_{i}^{c_{j}}$ and the support set of the $j$-th class for class-translation to achieve feature disentanglement as

$$L_{tran}=\frac{1}{N}\sum_{i=1}^{M}\sum_{l=1}^{K}\|\phi(x_{l}^{\mathcal{S}_{j}})-\phi(\hat{x}_{i}^{c_{j}})\|_{1}.$$

(7)

The total loss for training DFR can be formulated as

$$L_{total}=\lambda_{1}\cdot L_{dis}+\lambda_{2}\cdot L_{rec}+\lambda_{3}\cdot L_{tran}+L_{cls},$$

(8)

where $\lambda_{1}$, $\lambda_{2}$ and $\lambda_{3}$ denote the weights parameters of $L_{dis}$, $L_{rec}$ and $L_{tran}$ relative to $L_{cls}$, respectively.

DFR framework aims to extract only class-related information for classification. Different from other attempts towards more adaptive embedding using attention mechanism (Hou et al. 2019; Li et al. 2019; Ye et al. 2020), our classification and variation branches play as the adversaries by minimizing $L_{cls}$ and $L_{dis}$ simultaneously. In practice, the classification and variation features of an image are always complementary, thus the image reconstruction quality is enforced after fusion by minimizing $L_{rec}$. It is essential to preserve the image representation in DFR for few-shot classification: as the class-specific features can be task-varying thus hard to be generalized, any information loss throughout the inter-state flow may potentially limit the model performance. Such design is contrast to the classic feature embedding for few-shot learning, in which image features are always projected onto the lower-dimensional manifolds (Simon et al. 2020). Our $E_{cls}$ feature has much lower dimension comparing to the $E_{var}$ feature, as the class-irrelevant information (e.g., image style and background) are typically excessive. To this end, a more restrictive classification feature will significantly reduce the model bias, thus to enhance its generalizability in few-shot tasks.

We visualize the feature distributions w/o and w/ DFR using t-SNE (Van der Maaten and Hinton 2008) to verify our intuition. Figure 3 (a) shows that the learned features extracted from the ResNet-12 backbone are less discriminative without using the DFR framework. While when applying the DFR framework, the classification branch clusters in Figure 3 (b) are more separable from each other, and the output features of the variation branch in Figure 3 (c) contains more class-irrelevant information that meets our expectations.

We use ResNet-12 network (Lee et al. 2019) as our backbone $E_{cls}$ for classification branch and set the number of channels as $[64,160,320,640]$, which are similar to the competing methods. The encoder $E_{var}$ consists of four convolutional blocks and two residual blocks. The decoder contains a two-layer MLP block and a decoder network $D$ with residual blocks and upscale convolutional blocks. The I/O channel numbers of variation encoder and decoder are all set to $128$. The level ratio $\lambda$ of GRL layer is set to $1$.

Data augmentation including resizing, random cropping, color jitter and random flipping following (Ye et al. 2020) are applied for all methods in training. Our models are all trained using SGD optimizer, with the weight decay as $5e{-4}$, and the momentum as $0.9$.

We conduct experiments under both $5$-way $1$-shot and $5$-way $5$-shot settings with $15$ query images each class, i.e., $5\times(1\,and\,5)+5\times 15$ samples for $1$-shot and $5$-shot tasks, respectively. We report the mean accuracy of randomly sampled $10k$ tasks as well as the $95\%$ confidence intervals on the testing set as mentioned in (Ye et al. 2020; Zhang et al. 2020). To verify the effectiveness of our proposed DFR framework, we combined DFR with three few-shot algorithms: a commonly used baseline ProtoNet (Snell, Swersky, and Zemel 2017), two state-of-the-art methods DeepEMD (Zhang et al. 2020) and FEAT (Ye et al. 2020).²²2We utilized the official codes released by the authors, for implementations of ProtoNet, DeepEMD and FEAT and the corresponding DFR models. The results are all obtained by following the unified setting for fair comparison, which may not exactly match with the results reported in their original papers. Note that we only adopt the FCN version of DeepEMD for comparison over all datasets.

We first conduct experiments on two general few-shot benchmarks: mini-ImageNet and tiered-ImageNet.

The classification results are shown in Table 1. It is clear that FEAT+DFR achieves state-of-the-art results on mini-ImageNet benchmark, while DeepEMD+DFR achieves state-of-the-art results on tiered-ImageNet benchmark. Moreover, we observe that the improvements by DFR remain inconsistent for all baselines. By adopting the DFR framework, the 5-way 1-shot accuracies by ProtoNet are increased by $3.0\%$ and $3.4\%$ on mini-ImageNet and tiered-ImageNet, respectively, which are even comparable to more sophisticated methods. For the other two methods DeepEMD and FEAT, which are the current state-of-the-art FS methods, their FS classification results can still be further boosted by $1\%$ in average, after applying the DFR framework.

Domain generalization (DG) aims to learn a domain-agnostic model from multiple sources that can classify data from any target domain. DG tasks become more challenging when there exists a class gap (besides domain gap) between the training and testing sets, i.e., DG under the few-shot setting. General few-shot learning does not consider the influence caused by the domain gap, thus the DG models can hardly be generalized to unseen domains. In this work, we consider a more challenging Few-Shot Domain Generalization (FS-DG) problem, i.e., both domain and class gaps exist between the training (source) and testing (target) sets. In our FS-DG experiments (under both the Setting A and B), only the training samples from the source domains are selected in training. Specifically, an $N$-way $K$-shot FS-DG task contains support and query samples from $N$ classes on the source domain in the meta-training step, and then the trained model is to predict the query data label out of the testing classes on the target domain. Here we propose two FS-DG evaluation settings based on different domains of support set $\mathcal{S}$ as: (1) Setting A: Support set is only from the target domain and (2) Setting B: Support set is only from the source domain. Both settings can evaluate the generalizability of the model, i.e., ability to extract domain-invariant and class-specific features. Recent works (Ye et al. 2020; Du et al. 2021) also attempted simple FS-DG tasks to evaluate their proposed FS models. However, only preliminary results are reported following the simple setting (i.e., Setting B in Table 2) without comprehensive investigation of the effect of domain gap on novel classes (test class set). We further conduct experiments with full evaluation settings to validate the proposed DFR for FS-DG tasks using a novel FS-DomainNet Benchmark.

We further evaluate DFR on a fine-grained benchmark, i.e., Caltech-UCSD Birds 200-2011 (CUB) (Wah et al. 2011) which was initially proposed for fine-grained image classification, which contains 200 different birds with 11788 images. Following the split in (Chen et al. 2019; Hilliard et al. 2018), 200 classes are divided into 100, 50 and 50 for training, validation and testing, respectively. We also pre-process the data by cropping each image with the provided bounding box according to the prior work (Ye et al. 2020; Wertheimer, Tang, and Hariharan 2021).

Table 4 reports the fine-grained few-shot classification results with both 5-way 1-shot and 5-way 5-shot tests. Comparing to the general and multi-domain few-shot benchmarks that contain significant differences between the categories, fine-grained classification only includes minor intra-class differences. The domain information in the fine-grained dataset may contribute to the category, making it a challenging task. It is clear that the proposed DFR can also significantly and consistently boost all FS baselines, with $0.5\%$ to $1.9\%$ additional improvements on CUB dataset. It demonstrates that DFR can effectively remove the excursive features, and thus highlight the subtle traits which are critical for fine-grained FS classification.

We investigate the weights in our formulation and incorporate FEAT as the baseline method. Table 5 shows that FEAT+DFR achieves the best performance when weighting parameters are all set to $1.0$. Compared with $L_{rec}$ and $L_{tran}$, the discriminative loss $L_{dis}$ has a more significant impact on performance as it affects the class-specific information removed from the variation branch, which is directly related to the classification ability of the classification branch. Overall, we find that the performance is minimally affected by loss weight which also shows the robustness of our framework.