Improving Object Detection by Label Assignment Distillation

Modern object detectors can be single or multi-stages. Single-stage is simpler and more efficient, while multi-stages are more complex but predict with higher precision.

In single stage detectors, classification and localization are predicted concurrently for each pixel. There are anchor-based or anchor-free methods. For anchor-based detectors, such as SSD [31], RetinaNet [28], an anchor is typically considered as positive if its IoU with the nearest ground truth box is at least 0.5, negative if less than 0.4, and ignored otherwise. ATSS [48] improves RetinaNet by selecting 9 nearest anchors from the object center in each feature level, and use the mean plus standard deviation (std) of IoU from this set as the assignment threshold. Anchor-free methods instead resort objects by points. FCOS [42] assigns a point as positive if it is in a ground truth box, negative otherwise. If the point falls into two or more ground truth boxes, the smallest box is chosen. CornerNet [22] represents an object by two corners of the bounding box, and only the corner points are assigned as positive. During training, the farther the points from the corner center, the larger weight it contributes to the loss function as negative samples.

In multi-stage detectors, the first stage network proposes candidate regions, by which the features are cropped and fed to the following stages for the box refinement and class prediction. The process can continue in a cascade manner as many stages as needed. Anchor assignments are also different in each stage. In the first stage, i.e. proposal network, an anchor is typically considered as positive if its IoU with any ground-truth is greater than $0.7$, negative if less than $0.3$, and ignored otherwise. For the second stage, $0.5$ IoU threshold is used to separate positive and negative anchors. Cascade R-CNN [3] improves Faster R-CNN by adding more stages. Setting the IoU threshold too high leads to extremely few positive samples, thus insufficient to train the network. Therefore, Cascade R-CNN re-assign the samples with IoU thresholds progressively increasing from $0.5$ to $0.7$ after each stage.

In summary, label assignments in the aforementioned methods are static and heuristically designed.

Learning label assignment approaches update the network weights and learn the assignment iteratively towards final optimization. There are two main approaches.

In hard assignment, samples can only be either positive or negative. The pipeline generally includes three steps. We first select and pre-assign a set of potential samples as positive, then evaluate their performance by a cost function. Finally, we select a threshold to separate the samples having lower cost as true positive, while the rest are reassigned to negative or ignored. Table 1 compares different methods of learning label assignment in these aspects for single stage detector. For two-stage detectors, Dynamic-RCNN [47] progressively increases the IoU threshold in the second stage by mean of the predicted IoU of the proposal network, and modifies the hyper-parameter of SmoothL1 Loss function according to the statistics of regression loss.

In soft assignment, samples can have both positive and negative semantic characteristics. There are two main approaches: soft-weighted loss and soft-targets supervision. In the former, a set of positive candidates are first selected by a relaxed criterion, e.g. low IoU threshold, then their importance scores are evaluated and used as weights to sum the loss values of the samples. FreeAnchor [50] first constructs a bag of top-k anchors having highest IoU with object, then calculates the importance scores from the classification and localization confidence using Mean-Max function. Following this framework, NoisyAnchor [23] derives the importance scores, i.e. cleanliness, from the regression and classification losses of the samples. SAPD [54] utilizes distance to object center as the metrics to measure the importance of samples. Furthermore, it leverages FSAF [55] module to assign soft-weights to samples in different feature levels. Alternatively, soft-target supervision decides how much a sample belonging to positive or negative based on an estimated quality distribution. AutoAssign [53] improves FCOS [42] by estimating object’s spatial distribution, e.g. a 2D-Gaussian, for each category. Then, a point in the ground truth box is supervised by both positive and negative losses weighted by this distribution. Similarly, IQDet [33] builds a mixture of Gaussian models to capture IoU quality. Differently, positive points are randomly sampled by this distribution and via a bi-linear interpolation. The quality distribution is also used as the soft-label for classification supervision.

Distillation is first introduced in the seminal work of Hilton et al. [19] for image classification. By using the prediction of a teacher network as soft training targets, we can distill the knowledge from the teacher to a student. Distillation then becomes the principle for many other related problems, such as self-supervised learning. However, applying distillation to object detection is not straightforward, since classification and localization must be learned simultaneously. Zheng et al. [51] hypothesize that a teacher is better in estimating the localization uncertainty, which can be used as the dark knowledge for distillation. However, the method is only applicable to Generalized Focal detector [26] due to its formulation problem. Other works focus on distillation by feature mimicking. Chen et al. [8] combine feature mimicking and soft-label distillation on Faster-RCNN. Other researches [46, 43, 24] believe that background features are noisy, thus propose applying semantic masks to focus attention on the meaningful regions, especially in and near foregrounds. In contrast, Guo et al. [17] consider that background can capture the object’s relationship and suggest decoupling the background features during distillation.

Although these methods exploit different ways for distillation, they all enforce a direct mimicking the teacher’s final or intermediate outputs. This may be too restrictive if their architectures or performance are significantly different. Therefore, several papers [34, 51] propose to add a teacher assistant network to bridge their gaps. Nevertheless, this introduces more hyperparameters, larger memory, longer training, and complicates the process.

Our proposed LAD extends the distillation concept and is orthogonal with these methods. Generally, they can be combined to further improve the performance.

In network distillation, the teacher’s prediction is referred as soft-label, which is believed to capture the network’s dark knowledge and provide richer information for training a student. Adopting the convention from the classification problem [19], we use the Kullback-Leibler (KL) divergence loss. Follow Focal loss [28], we also add the focal term to deal with the extreme imbalance problem

where $p_{t}$ and $p_{s}$ denote the teacher and student prediction (i.e. after Sigmoid), $c$ is the class index of total $C$ classes in the dataset, and $w^{c}\!=\!|p_{t}^{c}\!-\!p_{s}^{c}|^{\gamma}$ is the focal term. Note, when replacing the soft-label $p_{t}$ with an one-hot vector, (1) becomes the standard Focal loss. Alternatively, we can also use $\mathcal{L}_{1}$ or $\mathcal{L}_{2}$ losses. Different from the classification losses, localization loss is only computed for positive samples. Hence, we select the predicted boxes having $IoU\!>\!0.5$ w.r.t. its nearest ground truth box for localization distilling.

LAD is very general and can be applied to many learning label assignment methods, such as those described in Sec. 2.2. For a concrete example, we adopt PAA to implement LAD. In PAA, for each object, we select a set of anchors $\{a_{i}\}$ that have $IoU\geq 0.1$, and evaluate the assignment cost

as if they are positive anchors, where $FL$ is the Focal loss, $(p_{i},b_{i})$ are the predicted class probability and bounding box of $a_{i}$, $(C,B)$ are the ground truth class and bounding box of the object. Then, we build a distribution model for $\{c_{i}\}$ by a mixture of two Gaussian: $G(\mu_{1},\sigma_{1})+G(\mu_{2},\sigma_{2}),\mu_{1}<\mu_{2}$. We then assign anchors with $c_{i}<\mu_{1}$ as the true positive, and reassign the other as negative, as illustrated in Fig. 2(a).

To convert PAA to LAD, we simply use the teacher’s prediction $(p_{T}^{i},b_{T}^{i})$ instead of $(p_{i},b_{i})$ in (2), and proceed exactly the steps, as illustrated in Fig. 2(b). Consequently, LAD presents a new distillation method for object detection that works without directly mimicking the teacher’s outputs.

Our motivation is to seek a method that utilizes the knowledge from teacher network, and this can be achieved either by soft-label or LAD, as shown in Fig. 1. Therefore, we conduct a pilot study to understand their properties.

The pilot study in Sec. 3.3 proves that the proposed LAD is a simple and effective solution to improve network training for object detection. However, like many other distillation methods, it requires a teacher pretrained in advance. Learning label assignment methods [6, 21, 53, 15] on the other hand, use a bootstrap mechanism to self-distill the labels. However, a potential drawback is that it can be trapped in local minima. This is because it only exploits the highest short-term reward, i.e. lowest assignment cost, without any chance for exploration.

Therefore, to avoid these drawbacks and combine their advantages, we propose the Co-learning Label Assignment Distillation (CoLAD). CoLAD does not need a beforehand teacher, while potentially avoiding falling into local minima. The framework is illustrated in Fig. 3. Concretely, we use two separate networks similarly to LAD, but none of them must be pre-trained. Instead, we train them concurrently, and dynamically switch the role of teacher and student based on their performance indicator $\rho$. We propose two criteria for $\rho$, namely Std/Mean score ($\rho_{\sigma/\mu}$) and Fisher score ($\rho_{\rm Fisher}$). Let $\{c_{i}\}$ be the set of the assignment costs evaluated by a network on the pre-assigned positive anchors, the criteria are respectively defined as:

where $(\mu,\sigma)$ are the mean and std. of $\{c_{i}\}$, and in the Fisher Score, we approximate the distribution of $\{c_{i}\}$ by a mixture of two Gaussian models $G(\mu^{+},\sigma^{+})+G(\mu^{-},\sigma^{-})$.

At each iteration, the network with higher $\rho$ can be selected as the teacher. Therefore, by dynamically switching the teacher-student, the process no longer relies on a single network, stochastically breaking the loop of training-self assignment. This also brings more randomness at the initial training due to network initialization, hence potentially reducing the risk of local minimal trapping.

Finally, we can also apply the dynamic switch (3) to LAD, in which the teacher is pretrained and fixed. This allows the student learn from the teacher at the early training stage. However, when the student is gradually improved and potentially surpasses its teacher, it can switch to the self-distillation mode. This makes LAD more adaptive.

We use PAA detector with different ResNet backbones R18, R50, R101 [18] with FPN [27] for the baseline. The backbones are initialized with weights pre-trained on ImageNet [12]. Our code is implemented using MMDetection framework [9], and follows its standard setup. Specifically, we use SGD optimizer with weight decay of 0.0001 and the momentum of 0.9. $1\times$ training schedule includes 12 epochs with initial learning rate of 0.01, which is then divided by 10 at epoch 8 and epoch 11. $2\times$ training schedule includes 24 epochs with initial learning rate of 0.01, which is then divided by 10 at epoch 16 and epoch 22. Random horizontal flipping is used in data augmentation. We use batch size 16, with image size of $(800\times 1333)$ pixels for training and testing. All other hyperparameters, such as loss weights and post-processing, are identical to the PAA baseline.

In CoLAD, the roles of teacher-student are not persistent, so the terminology teacher and student networks may not exactly apply. However, for convenience, we still call the network with higher performance at initial as the teacher. If two networks are not pretrained, then the one having more parameters and layers is supposed to be the initial teacher.

We compare our proposed CoLAD with other methods on MS COCO test-dev set, especially the recent ones addressing the label assignment for single-stage detectors. Following the previous works [21, 15, 53], we train the model with $2\times$ scheduler, randomly scale the shorter size of the image into the range of 640 and 800. We use the PAA-R50 trained with $3\times$ schedule and multi-scale on COCO as the initial teacher ¹¹1PAA-R50’s pretrained weight is provided by MMDetection https://github.com/open-mmlab/mmdetection/tree/master/configs/paa. The teacher was evaluated with $43.3\rm AP$ and $43.8\rm AP$ on the minval and test-dev sets, respectively.

Inspired by [53, 10, 52], we also modify the PAA’s head architecture by connecting the classification and localization branches by an auxiliary “objectness” prediction, which is supervised implicitly. However, we use a slightly different implement with [53, 10, 52], which we call Conditional Objectness Prediction (COP). Since COP is not our main contribution, we refer the reader to Appendix A for the motivation and more ablation studies. We report the results of the PAA network amended with COP and trained with CoLAD in Tab. 7.

With R101 backbone, CoLAD achieves $46.0\rm AP$, consistently outperforming other recent methods for label assignment, such as OTA [15] ($45.3\rm AP$), IQDet [33] ($45.1\rm AP$) and PAA [21] ($44.8\rm AP$) by large margins on all evaluation metrics. With the larger backbone ResNeXt-64x4d-101, our model can further improve to $47.5\rm AP$ and surpasses all existing methods. Finally, for the ResNeXt-64x4d-101 with Deformable Convolutional Networks (DCN) [56], CoLAD and OTA both achieve the highest score $49.2\rm AP$. Note that, with this DCN backbone, the top 4 methods perform very similar ($\pm 0.2\rm AP$).

Nevertheless, it is not our intention for LAD/CoLAD to become a replacement for current SOTA methods. Instead, we find it useful to use a simple technique, with a smaller and lower performance teacher ResNet50, to successfully boost the current SOTA results. Indeed, as shown in the next section, by simply replacing ResNet50 with a stronger teacher, we can still improve the results above further.

Vision Transformer recently demonstrates its potential to replace convolution networks. Therefore, we conduct experiments to verify the generalization of LAD with Transformer backbones. To this end, we select the Swin-B Transformer [32] and use the teacher PAA-ResNeXt-64x4-101-DCN obtained in the section 4.3. Since CoLAD training with these two networks is extremely heavy, we use LAD but with the dynamic switch instead. We follow exactly the training setting in [32], such as Swin-B backbone is pre-trained on ImageNet, multi-scale training, batch-size 16, AdamW optimizer, and $3\times$ schedule.

As reported in Tab. 8, the student PAA-Swin-B once again surpasses its teacher with a large margin ($+2.8\rm AP$). Furthermore, on COCO val set, it achieves $51.4\rm AP$, approaching the heavy Cascade Mask-RCNN ($51.9\rm AP$) [32] using the same backbone, but without training segmentation. Finally, we recycle the PAA-Swin-B as a teacher to train PAA-R50/R101 with SoLAD, and compare the results to other distillation methods in Tab. 9.

Note that, it is not advisable to judge which method is better based on Tab. 9, since each serves for a particular detector, thus having different baselines. Nevertheless, our SoLAD still achieves the highest performance for the same backbone and training schedule.