PS-TTL: Prototype-based Soft-labels and Test-Time Learning for Few-shot Object Detection

Object detection aims to identify and localize objects within images, constituting a fundamental challenge in the fields of computer vision and multimedia. Recently, the success of deep learning has triggered numerous effective object detection methods. They can be roughly categorized into two main groups: single-stage and two-stage. Single-stage detectors (e.g. SSD (Liu et al., 2016) and RetinaNet (Lin et al., 2017)) predict bounding boxes and classification scores based on pre-defined anchors, exhibiting good real-time performance. The YOLO (Redmon et al., 2016; Wang et al., 2023) series, by continuously assimilating the latest advancements, such as label assignment and multi-scale feature fusion, has delivered high-precision real-time object detection. Although the structure of single-stage detectors is straightforward and efficient, their integrated design makes them less adaptable to FSOD tasks. In contrast, two-stage detectors (e.g. Faster R-CNN (Ren et al., 2016)) usually first use a Region Proposal Network (RPN) to generate potential proposals and then refine them by other modules. Compared to single-stage detectors, two-stage ones report higher detection performance and are commonly used for FSOD.

FSOD methods enable detectors to rapidly adapt to new objects with minimal data while preserving competitive performance. There are mainly two paradigms: meta-learning based and fine-tuning based. Meta-learning based methods (Wang et al., 2019; Li et al., 2021; Li and Li, 2021; Yin et al., 2022) employ a series of $N$-way $K$-shot detection tasks for training, aiming to well generalize to novel tasks with limited samples. FSRW (Kang et al., 2019) proposes a feature reweighting strategy, which extracts class-specific representations from support images and utilizes them to reweight the importance of query features. Similarly, Meta R-CNN (Yan et al., 2019) combines a two-stage detector and reweights RoI features in the detection head. Attention-RPN (Fan et al., 2020) exploits matching relationship between the few-shot support set and query set with a contrastive training scheme, which can then be applied to detect novel objects without retraining and fine-tuning. QA-FewDet (Han et al., 2021) uses a graph model to capture multi-relations among the proposal and class nodes. While theoretical promising, their training and inference processes are highly complex, making them difficult to deploy in real-world scenarios.

Fine-tuning based methods leverage a two-stage training procedure, i.e., first base training and then few-shot fine-tuning, which expects to transfer prior knowledge from base classes to novel ones. LSTD (Chen et al., 2018) is the earliest attempt in this domain, incorporating detection knowledge transfer and background depression regularization. TFA (Wang et al., 2020) simply freezes the backbone and only fine-tunes the detection head with novel classes, which largely improves the performance. Subsequent research refines the TFA method and integrates it with other techniques to further enhance the FSOD performance (Sun et al., 2021; Nguyen et al., 2022; Pei et al., 2022; Wu et al., 2022; Kim et al., 2021; Zhu et al., 2021; Lu et al., 2022; Qiao et al., 2021; Fan et al., 2021; Yang et al., 2022). FSCE (Sun et al., 2021) introduces contrastive learning to learn discriminative object proposal representations, alleviating the misclassification issue in novel classes. DeFRCN (Qiao et al., 2021) employs the gradient decoupled layer for multi-stage decoupling and the prototypical calibration block to align original classification scores. Although fine-tuning based methods demonstrate satisfactory results, the limited samples of novel classes still make it challenging for detectors to capture accurate data distributions.

To alleviate the issue above, HallucFsDet (Zhang and Wang, 2021) introduces a hallucinator network trained on base classes to synthesize additional training examples for novel classes. Zhao et al. (Zhao et al., 2022) assume that the features of both base and novel classes follow a Gaussian distribution and generate samples for novel classes according to the variance of similar base classes. However, these hallucination methods relying on base classes tend to result in biased synthetic novel samples. The other semi-supervised FSOD methods (Cao et al., 2022; Kaul et al., 2022; Tang et al., 2023) investigate mining implicit novel instances from training data with the assumption that novel instances do appear in base data. Kaul et al. (Kaul et al., 2022) present a simple pseudo-labelling strategy to detect potential novel instances in base datasets while MINI (Cao et al., 2022) comprises an offline and online mechanism, facilitating novel instance mining. However, as such an assumption does not always hold (e.g. in rare disease or animal detection), they are problematic in practical cases. Considering the accessibility of novel instances in test data, we are motivated to explore fine-tuning detectors at test-time.

We follow the standard FSOD setting introduced in (Wang et al., 2020) with two disjoint training sets: a base dataset $D^{b}=\{x^{b}_{i},y^{b}_{i}\}$ with exhaustively annotated instances for each base class $C^{b}$ and a novel dataset $D^{n}=\{x^{n}_{i},y^{n}_{i}\}$ with only $K$ instances (usually less than 30) for each novel class $C^{n}$, where $x_{i}$ and $y_{i}$ refer to the input image and the Ground Truth (GT) label, respectively. It is worth noting that there is no intersection between base and novel classes, i.e., $C^{b}\cap C^{n}=\emptyset$. In this case, the ultimate goal of FSOD is to train a robust detector based on $D^{b}$ and $D^{n}$ to deal with the objects in the test set that contains both types of instances in $C^{b}\cup C^{n}$.

DeFRCN (Qiao et al., 2021) is a state-of-the-art fine-tuning based few-shot object detector, built through two training stages. In the first phase, Faster-RCNN is trained on base classes $C^{b}$ with sufficient samples, while in the second phase, transfer learning is conducted by fine-tuning Faster-RCNN on base classes and novel classes $C^{b}\cup C^{n}$ with $K$ instances per class. Fine-tuning on a balanced set $D^{balanced}$ containing training samples for both base and novel classes helps to preserve the performance on base classes. The entire procedure is summarized as follows:

where $M_{init}$, $M_{base}$, and $M_{novel}$ indicate the detector in the initialization, base training, and novel fine-tuning stages, respectively.

Different from previous fine-tuning based methods, which only fine-tune a small number of parameters of Faster-RCNN, such as the prediction head, to avoid overfitting of the detector, DeFRCN introduces a gradient decoupled layer during fine-tuning to stop the gradient between RPN and backbone while scaling the gradient between RCNN and backbone. This allows the detector to sufficiently learn from novel data while preventing overfitting, making DeFRCN remarkably superior to other existing counterparts.

Despite the significant progress made by the fine-tuning based methods, given only $K$ novel instances, they fail to accurately capture the data distribution. To overcome this obstacle, we propose Prototype-based Soft-labels and Test-Time Learning (PS-TTL) to mine novel instances in test data. The overall architecture of the model is illustrated in Fig. 2.

Self-training has shown promising performance for semi-supervised object detection (Lee, 2013; Tarvainen and Valpola, 2017; Liu et al., 2021). It typically predicts pseudo-labels for unlabeled data, where high-confidence pseudo-labels are used to supervise detector training.

In this work, we aim to fully leverage novel instances in test data, especially in the scenario of online learning, called Test-Time Learning (TTL). To this end, we employ a mean-teacher self-training paradigm (Tarvainen and Valpola, 2017), which mainly consists of two architecturally identical detectors, i.e. the student network and the teacher network. The teacher network first works on test data and renders pseudo-labels from its detection results through post-processing procedures (e.g., Non-Maximum Suppression (NMS) and filtering using a confidence threshold). The high-quality pseudo-labels are screened out to supervise the student network, enhancing its capability.

Since the teacher network is expected to produce reliable pseudo-labels of novel classes in test data, the few-shot detector $M_{novel}$ which has been fine-tuned on novel data is used to initialize both the student and teacher networks. However, the self-training paradigm inevitably generates noisy pseudo-labels, particularly for novel classes. Considering that excessively noisy pseudo-labels deteriorate the performance of the student network as training progresses, to address this issue, we first apply NMS for each class to remove duplicate detection boxes. Then, we set a high confidence threshold $\delta_{upper}$ to exclude uncertain labels. Finally, we optimize the student network using the remaining high-quality pseudo-labels with the loss function as follows:

where $X^{t}$ is the input test image and $\hat{Y^{t}}$ denotes the filtered pseudo-label. Note that the unsupervised loss is only applied to the classification heads of RPN and RoI.

On the other side, as the few-shot detector $M_{novel}$ is not strong enough, pseudo-labels still contain noise even after filtering out low-confidence predictions. Therefore, to alleviate the degradation of the few-shot detector during test-time learning, we propose to take $N$-way $K$-shot data $D^{balanced}$ as supervision signals. The supervised loss for training the student network can thus be defined as:

where $\{X^{s},Y^{s}\}\in D^{balanced}$. Both the RPN and RoI heads adopt the classification loss and bounding box regression loss.

Following the mean-teacher framework (Tarvainen and Valpola, 2017), we update the weights of the student network using both $L_{unsup}$ and $L_{sup}$. To obtain high-quality pseudo-labels from test data, we update the weights of the teacher network via Exponential Moving Average (EMA) of those of the student as below:

where $\theta_{t}$ and $\theta_{s}$ are the network parameters of the teacher and the student, respectively. $\alpha$ is the EMA momentum coefficient.

The mean-teacher self-training framework presented in Section 3.3 for TTL on test data promotes the detection performance, and a large threshold $\delta_{upper}$ to filter the generated pseudo-labels is more beneficial (please see the experiments). However, this leads to severe detection missing, and many test images do not have any pseudo-label. Different from semi-supervised object detection, where multiple rounds of fine-tuning can be conducted on unlabeled data, under the TTL setting, we can only perform one epoch of training on test data, making it challenging to fully utilize every input test image.

As shown in Fig. 3, we observe that relatively low-confidence pseudo-labels, despite suffering from classification confusion, mostly recall foreground. Based on this phenomenon, we propose a Prototype-based Soft-labels (PS) strategy to replace the hard labels of these implicit foreground predictions with soft-labels for fully unleashing the potential of low-quality pseudo-labels.

Specifically, we first introduce a lower bound confidence threshold $\delta_{lower}$, and the predicted results between $\delta_{lower}$ and $\delta_{upper}$ are also assigned as foreground. Due to unavoidable classification confusion in these implicit foreground predictions, it fails to effectively remove redundant boxes by employing class-specific NMS in the teacher network. Instead, after removing the hard labels of these implicit foreground predictions, we apply class-agnostic NMS to them using every high-confidence pseudo prediction (i.e., whose confidence score is greater than $\delta_{upper}$) to filter the redundant ones.

We then generate soft-labels for these implicit foreground predictions by measuring their similarities to each class. Formally, given an implicit foreground prediction $r$, we define its similarity to a class $c$ as the cosine distance between its RoI feature $F_{r}$ and the prototype $P^{c}$ of the class $c$:

Finally, $s_{r}=[s^{1}_{r},s^{2}_{r},...,s^{N}_{r}]$ follows a softmax function to generate $u_{r}$, which represents the soft-label of the implicit foreground prediction $r$. We minimize the Kullback-Leibler (KL) divergence between the soft-label and the class probability of each implicit foreground prediction $r$:

where $v_{r}$ is the class probability, $N+1$ denotes $N$ foreground classes and one background class. Hence, we set $u^{N+1}_{r}=0$.

To leverage soft-labels of implicit foreground predictions at the early stage during TTL, we initialize the class prototypes with $N$-way $K$-shot data:

where $F^{c}_{i}$ is the RoI feature of the $i$-th instance for class $c$. Because the $N$-way $K$-shot data cannot accurately represent the class prototypes, we propose to dynamically update them using both the supervised $N$-way $K$-shot data and test data with high-confidence pseudo-labels, enabling the class prototypes to converge to the true representations as training progresses. In this case, we update the class prototypes using the following formula:

where ${F^{c}_{avg}}$ is the averaged RoI feature of high-confidence predictions for class $c$ during TTL, and $sim(\cdot,\cdot)$ denotes the normalized cosine similarity function as:

During TTL, the total loss to optimize is as follows:

consisting of the supervised loss of $N$-way $K$-shot data, the unsupervised loss of pseudo-labels on test data, and the KL loss of soft-labels on test data. Here, $\lambda_{1}$ and $\lambda_{2}$ are hyper-parameters to balance such losses.

During this procedure, few-shot detectors are able to learn from test data. When a mini-batch of test samples arrive, we update the weights of the model through the total loss $L_{total}$. A detailed description is provided in Algorithm 1.

PASCAL VOC. For PASCAL VOC (Everingham et al., 2015), the overall 20 classes are divided into 15 base classes and 5 novel classes. Following TFA (Wang et al., 2020), we utilize three different class splits, namely split 1, 2, and 3. For each split, base classes are exhaustively annotated, but novel classes only have $K=1,2,3,5,10$ annotated instances per class. Both base and novel instances are sampled from the PASCAL VOC (07+12) trainval set, and the model is tested on the PASCAL VOC07 test set. We report AP50 for novel classes during evaluation.

MS COCO. MS COCO (Lin et al., 2014) has 80 classes, and we select the 20 classes that overlapped with PASCAL VOC as novel classes and the remaining 60 classes as base classes. In this case, we evaluate our method with $K=10,30$ shots for each novel class. We report mAP and AP75, respectively.

Our method can be combined with the majority of fine-tuning based few-shot object detectors. For simplicity, we choose the most representative SOTA method, i.e. DeFRCN (Qiao et al., 2021), as our baseline. DeFRCN adopts Faster-RCNN (Ren et al., 2016) as the detection model and uses the ResNet-101 backbone pre-trained on ImageNet. We use DeFRCN, which has been pre-trained on base classes and fine-tuned on novel classes, to initialize our model, and then fine-tune it on test data. During TTL, we fine-tune our model with a mini-batch of 2 on a single GPU, which simulates the real inference process of the few-shot detector. Besides, we follow a one-epoch setting, where we fine-tune on test data for only one epoch. We also utilize the $N$-way $K$-shot data for novel fine-tuning during the testing process. Due to the unsatisfactory performance of the few-shot detector, we apply weak data augmentation to both the $N$-way $K$-shot data and the test data, including random resizing and random horizontal flipping. For the hyper-parameters, we set $\lambda_{1}=0.5$ and $\lambda_{2}=0.1$ for all the experiments. We set the thresholds $\delta_{upper}=0.9$ and $\delta_{lower}=0.7$. We optimize the network using Stochastic Gradient Descent (SGD) and set the learning rate to 0.00125. The momentum coefficient of EMA for the teacher network is set to 0.9996.

PASCAL VOC. Experimental results on the PASCAL VOC dataset are shown in Table 1. We use DeFRCN as our baseline, which incorporates an additional Prototypical Calibration Block (PCB) to refine predictions. However, we find that the $N$-way $K$-shot data utilized by PCB may not align with that used at the novel fine-tuning stage. Therefore, we exclude PCB and present our re-implementation results DeFRCN* in Tables 1 and 2. It can be observed that our method achieves decent improvements across various splits and different shots on the PASCAL VOC benchmark. Our method outperforms HallucFsDet (Zhang and Wang, 2021) and LVC (Kaul et al., 2022), which represent synthetic novel class data and semi-supervised learning on base data, respectively. Meanwhile, we find that the performance gain obtained from TTL becomes more significant as the shot number decreases, especially in the 1-shot scenario.

MS COCO. Table 2 shows the detection results on MS COCO. The MS COCO dataset conveys more categories, and typically, a single image contains multiple instances. few-shot detectors generally do not perform well on MS COCO due to these factors, which also undermine the performance of our method. However, we observe that our method still achieves a significant improvement compared to the baseline, especially in the mAP75 metric. There is a 5.0% improvement in AP75 at 10 shots and a 9.2% improvement in AP75 at 30 shots. LVC (Kaul et al., 2022) demonstrates a noticeable improvement on the MS COCO dataset, because the base data in the MS COCO benchmark include a large number of implicit novel instances. However, an issue arises from its setting of few-shot detection, which does not match the real-world scenario. In contrast, under the TTL setting, we only have 5,000 images available for mining implicit novel instances.

Ablation experiments are carried out on novel split 1 of the PASCAL VOC benchmark to highlight the necessity of the main components of the proposed method.