TOOD: Task-aligned One-stage Object Detection

Chengjian Feng, Yujie Zhong, Yu Gao, Matthew R. Scott, Weilin Huang

Introduction

Object detection aims to localize and recognize objects of interest from natural images, and is a fundamental yet challenging task in computer vision. It is commonly formulated as a multi-task learning problem by jointly optimizing object classification and localization . The classification task is designed to learn discriminative features that focus on the key or salient part of an object, while the localization task works on precisely locating the whole object with its boundaries. Due to the divergence of learning mechanisms for classification and localization, spatial distributions of the learned features by the two tasks can be different, causing a certain level of misalignment when predictions are made by using two separate branches.

Recent one-stage object detectors attempted to predict consistent outputs of the two separate tasks, by focusing on the center of an object . They assume that an anchor (i.e., an anchor-point for an anchor-free detector, or an anchor-box for an anchor-based detector) at the center of the object is likely to give more accurate predictions for both classification and localization. For example, recent FCOS and ATSS both use a centerness branch to enhance classification scores predicted from the anchors near the center of the object, and assign larger weights to the localization loss for the corresponding anchors. Besides, FoveaBox regards the anchors inside a predefined central region of the object as positive samples. Such heuristic designs have achieved excellent results, but these methods might suffer from two limitations:

Recent one-stage detectors perform object classification and localization independently by using two separate branches in parallel (i.e., heads). Such a two-branch design might cause a lack of interaction between the two tasks, leading to an inconsistency in predictions when performing them. As shown in the ‘Result’ column in Figure 1, an ATSS detector recognizes an object of ‘Dining table’ (indicated by the anchor shown with a red patch), but localizes another object of ‘Pizza’ more accurately (red bounding box).

Most anchor-free detectors use a geometry-based assignment scheme to select anchor-points near the center of an object for both classification and localization , while anchor-based detectors often assign anchor-boxes by computing IoUs between the anchor boxes and ground truth . However, the optimal anchors for classification and localization are often inconsistent, and may vary considerably depending on the shape and characteristics of the objects. The widely used sample assignment scheme is task agnostic, and thus may be difficult to make an accurate yet consistent prediction for the two tasks, as demonstrated in ‘Score’ and ‘IoU’ distributions of ATSS in Figure 1. The ‘Result’ column also illustrates that a spatial location of the best localization anchor (green patch) can be not at the center of the object, and it is not well aligned with the best classification anchor (red patch). As a result, a precise bounding box may be suppressed by the less accurate one during Non-Maximum Suppression (NMS).

To address such limitations, we propose a Task-aligned One-stage Object Detection (TOOD) that aims to align the two tasks more accurately by designing a new head structure with an alignment-oriented learning approach:

In contrast to the conventional head in one-stage object detection where classification and localization are implemented separately by using two branches in parallel, we design a Task-aligned head (T-head) to enhance an interaction between the two tasks. This allows the two tasks to work more collaboratively, which in turn aligns their predictions more accurately. T-head is conceptually simple: it computes task-interactive features, and makes predictions via a novel Task-Aligned Predictor (TAP). Then it aligns spatial distributions of the two predictions according to the learning signals provided by a task alignment learning, as described next.

To further overcome the misalignment problem, we propose a Task Alignment Learning (TAL) to explicitly pull closer the optimal anchors for the two tasks. It is performed by designing a sample assignment scheme and a task-aligned loss. The sample assignment collects training samples (i.e., positives or negatives) by computing a degree of task-alignment at each anchor, whereas the task-aligned loss gradually unifies the best anchors for predicting both classification and localization during the training. Therefore, at inference, a bounding box with the highest classification score and jointly having the most precise localization can be preserved.

The proposed T-head and learning strategy can work collaboratively towards making predictions with high quality in both classification and localization. The main contributions of this work can be summarized as follows: (1) we design a new T-head to enhance the interaction between classification and localization while maintaining their characteristics, and further align the two tasks at the predictions; (2) we propose TAL to explicitly align the two tasks at the identified task-aligned anchors, as well as providing learning signals for the proposed predictor; (3) we conducted extensive experiments on MSCOCO , where our TOOD achieved a 51.1 AP, surpassing recent one-stage detectors such as ATSS , GFL and PAA , by a large margin. Qualitative results further validate the effectiveness of our task-alignment approaches.

Related Work

OverFeat is one of the earliest CNN-based one-stage detectors. Afterward, YOLO was developed to directly predict bounding boxes and classification scores, without an additional stage to generate region proposals. SSD introduces anchors with multi-scale predictions from multi-layer convolutional features, and Focal loss was proposed to address the problem of class imbalance for one-stage detectors like RetinaNet. Keypoint-based detection methods, such as , address the detection problem by identifying and grouping multiple key points of a bounding box. Recently, FCOS and FoveaBox were developed to locate objects of interest via anchor-points and point-to-boundary distances. Most mainstream one-stage detectors are composed of two FCN-based branches for classification and localization, which may lead to the misalignment between the two tasks. In this paper, we enhance the alignment between the two tasks via a new head structure and an alignment-oriented learning approach.

Most anchor-based detectors such as , collect training samples by computing IoUs between proposals and ground truth, while an anchor-free detector regards the anchors inside the center region of an object as positive samples . Recent studies attempted to train the detectors more effectively by collecting more informative training samples using output results. For example, FSAF selects meaningful samples from feature pyramids based on the computed loss, and similarly, SAPD provides a soft-selection version of FSAF by designing a meta-selection network. FreeAnchor and MAL identify the best anchor-box by computing the losses in an effort to improve the matching between anchors and objects. PAA adaptively separates the anchors into positive and negative samples by fitting a probability distribution to the anchor scores. Mutual Guidance improves anchor assignment for one task by considering the prediction quality on the other task. Different from the positive/negative sample assignment, PISA re-weights the training samples according to the precision rank of the outputs. Noisy Anchor assigns soft labels to the training samples, and re-weights the anchor-boxes using a cleanliness score to mitigate the noise incurred by binary labels. GFL replaces the binary classification label with an IoU score to integrate the localization quality into classification. These excellent approaches inspired the current work to develop a new assignment mechanism from a task-alignment point of view.

Task-aligned One-stage Object Detection

Similar to recent one-stage detectors such as , the proposed TOOD has an overall pipeline of ‘backbone-FPN-head’. Moreover, by considering efficiency and simplicity, TOOD uses a single anchor per location (same as ATSS ), where the ‘anchor’ means an anchor point for an anchor-free detector, or an anchor box for an anchor-based detector. As discussed, existing one-stage detectors have limitations of task misalignment between classification and localization, due to the divergence of two tasks which are often implemented using two separate head branches. In this work, we propose to align the two tasks more explicitly using a designed Task-aligned head (T-head) with a new Task Alignment Learning (TAL). As illustrated in Figure 2, T-head and TAL can work collaboratively to improve the alignment of two tasks. Specifically, T-head first makes predictions for the classification and localization on the FPN features. Then TAL computes task alignment signals based on a new task alignment metric which measures the degree of alignment between the two predictions. Lastly, T-head automatically adjusts its classification probabilities and localization predictions using learning signals computed from TAL during back propagation.

1 Task-aligned Head

Our goal is to design an efficient head structure to improve the conventional design of the head in one-stage detectors (as shown in Figure 3(a)). In this work, we achieve this by considering two aspects: (1) increasing the interaction between the two tasks, and (2) enhancing the detector ability of learning the alignment. The proposed T-head is shown in Figure 3(b), where it has a simple feature extractor with two Task-Aligned Predictors (TAP).

where $conv_{k}$ and $\delta$ refer to the $k$ -th conv layer and a $relu$ function, respectively. Thus we extract rich multi-scale features from the FPN features using a single branch in the head. Then, the computed task-interactive features will be fed into two TAP for aligning classification and localization.

We perform both object classification and localization on the computed task-interactive features, where the two tasks can well perceive the state of each other. However, due to the single branch design, the task-interactive features inevitably introduce a certain level of feature conflicts between two different tasks, which have also been discussed in . Intuitively, the tasks of object classification and localization have different targets, and thus focus on different types of features (e.g., different levels or receptive fields). Consequently, we propose a layer attention mechanism to encourage task decomposition by dynamically computing such task-specific features at the layer level. As shown in Figure 3(c), the task-specific features are computed separately for each task of classification or localization:

where $fc_{1}$ and $fc_{2}$ refer to two fully-connected layers. $\sigma$ is a $sigmoid$ function, and $\bm{x}^{inter}$ is obtained by applying an average pooling to $X^{inter}$ which are the concatenated features of $X^{inter}_{k}$ . Finally, the results of classification or localization are predicted from each $X^{task}$ :

where $M$ is computed from the interactive features, allowing it to learn a degree of consistency between the two tasks at each spatial location.

where an index $(i,j,c)$ denotes the $(i,j)$ -th spatial location at the $c$ -th channel in a tensor. Eq.(6) is implemented by bilinear interpolation, and its computational overhead is negligible due to the very small channel dimension of $B$ . Noteworthily, offsets are learned independently for each channel, which means each boundary of the object has its own learned offset. This allows for a more accurate prediction of the four boundaries because each of them can individually learn from the most precise anchor point near it. Therefore, our method not only aligns the two tasks, but also improves the localization accuracy by identifying a precise anchor point for each side.

The alignment maps $M$ and $O$ are learned automatically from the stack of interactive features:

where $conv_{1}$ and $conv_{3}$ are two 1 $\times$ 1 conv layers for dimension reduction. The learning of $M$ and $O$ is performed by using the proposed Task Alignment Learning (TAL) which will be described next. Notice that our T-head is an independent module and can work well without TAL. It can be readily applied to various one-stage object detectors in a plug-and-play manner to improve detection performance.

2 Task Alignment Learning

We further introduce a Task Alignment Learning (TAL) that further guides our T-head to make task-aligned predictions. TAL differs from previous methods in two aspects. First, from the task-alignment point of view, it dynamically selects high-quality anchors based on a designed metric. Second, it considers both anchor assignment and weighting simultaneously. It comprises a sample assignment strategy and new losses designed specifically for aligning the two tasks.

To cope with NMS, the anchor assignment for a training instance should satisfy the following rules: (1) a well-aligned anchor should be able to predict a high classification score with a precise localization jointly; (2) a misaligned anchor should have a low classification score and be suppressed subsequently. With the two objectives, we design a new anchor alignment metric to explicitly measure the degree of task-alignment at the anchor level. The alignment metric is integrated into the sample assignment and loss functions to dynamically refine the predictions at each anchor.

Considering that a classification score and an IoU between the predicted bounding box and the ground truth indicate the quality of the predictions by the two tasks, we measure the degree of task-alignment using a high-order combination of the classification score and the IoU. To be specific, we design the following metric to compute anchor-level alignment for each instance:

where $s$ and $u$ denote a classification score and an IoU value, respectively. $\alpha$ and $\beta$ are used to control the impact of the two tasks in the anchor alignment metric. Notably, $t$ plays a critical role in the joint optimization of the two tasks towards the goal of task-alignment. It encourages the networks to dynamically focus on high-quality (i.e., task-aligned) anchors from the perspective of joint optimization.

As discussed in , training sample assignment is crucial to the training of object detectors. To improve the alignment of two tasks, we focus on the task-aligned anchors, and adopt a simple assignment rule to select the training samples: for each instance, we select $m$ anchors having the largest $t$ values as positive samples, while using the remaining anchors as negative ones. Again, the training is performed by computing new loss functions designed specifically for aligning the tasks of classification and localization.

2.2 Task-aligned Loss

To explicitly increase classification scores for the aligned anchors, and at the same time, reduce the scores of the misaligned ones (i.e., having a small $t$ ), we use $t$ to replace the binary label of a positive anchor during training. However, we found that the network cannot converge when the labels (i.e., $t$ ) of the positive anchors become small with the increase of $\alpha$ and $\beta$ . Therefore, we use a normalized $t$ , namely $\hat{t}$ , to replace the binary label of the positive anchor, where $\hat{t}$ is normalized by the following two properties: (1) to ensure effective learning of hard instances (which usually have a small $t$ for all corresponding positive anchors); (2) to preserve the rank between instances based on the precision of the predicted bounding boxes. Thus, we adopt a simple instance-level normalization to adjust the scale of $\hat{t}$ : the maximum of $\hat{t}$ is equal to the largest IoU value ( $u$ ) within each instance. Then $Binary$ $Cross$ $Entropy$ ( $BCE$ ) computed on the positive anchors for the classification task can be rewritten as,

where $i$ denotes the $i$ -th anchor from the $N_{pos}$ positive anchors corresponding to one instance. Following , we employ a focal loss for classification to mitigate the imbalance between the negative and positive samples during training. The focal loss computed on the positive anchors can be reformulated by Eq.(10), and the final loss function for the classification task is defined as follows:

where $j$ denotes the $j$ -th anchor from the $N_{neg}$ negative anchors, and $\gamma$ is the focusing parameter .

A bounding box predicted by a well-aligned anchor (i.e., having a large $t$ ) usually has both a large classification score with a precise localization, and such a bounding box is more likely to be preserved during NMS. In addition, $t$ can be applied for selecting high-quality bounding boxes by weighting the loss more carefully to improve the training. As discussed in , learning from high-quality bounding boxes is beneficial to the performance of a model, while the low-quality ones often have a negative impact on the training by producing a large amount of less informative and redundant signals to update the model. In our case, we apply the $t$ value for measuring the quality of a bounding box. Thus, we improve the task alignment and regression precision by focusing on the well-aligned anchors (with a large $t$ ), while reducing the impact of the misaligned anchors (with a small $t$ ) during bounding box regression. Similar to the classification objective, we re-weight the loss of bounding box regression computed for each anchor based on $\hat{t}$ , and a $GIoU$ loss ( $L_{GIoU}$ ) can be reformulated as follows:

where $b$ and $\bar{b}$ denote the predicted bounding boxes and the corresponding ground-truth boxes. The total training loss for TAL is the sum of $L_{cls}$ and $L_{reg}$ .

Experiments and Results

All experiments are implemented on the large-scale detection benchmark MS-COCO $2017$ . Following the standard practice , we use the $trainval135k$ set (115 $K$ images) for training and $minival$ set (5 $K$ images) as validation for our ablation study. We report our main results on the $test$ - $dev$ set for comparison with the state-of-the-art detectors. The performance is measured by COCO Average Precision (AP) .

As with most one-stage detectors , we use the detection pipeline of ‘backbone-FPN-head’, with different backbones including ResNet-50, ResNet-101 and ResNeXt-101-64 $\times$ 4d pre-trained on ImageNet . Similar to ATSS , TOOD tiles one anchor per location. Unless specified, we report experimental results of an anchor-free TOOD (an anchor-based TOOD can achieve a similar performance as shown in Table 3). The number of interactive layers $N$ is set as 6 to make T-head have a similar number of parameters as the conventional parallel head, and the focusing parameter $\gamma$ is set to 2 as used in . More implementation and training details are presented in Supplementary Material (SM).

1 Ablation Study

For an ablation study, we use the ResNet-50 backbone and train the model for 12 epochs unless specified. The performances are reported on COCO $minival$ set.

We compare our T-head with the conventional parallel head in Table 1. It can be adopted in various one-stage detectors in a plug-and-play manner, and consistently outperforms the conventional head by 0.7 to 1.9 AP, with fewer parameters and FLOPs. This validates the effectiveness of our design, and demonstrates that T-head can work more efficiently with higher performance, by introducing task interaction and prediction alignment.

To demonstrate the effectiveness of TAL, we compare TAL with other learning methods using different sample assignment methods, as shown in Table 2. Training sample assignment can be divided into the fixed assignment and adaptive assignment according to whether it is a learning-based method. Different from the existing assignment methods, TAL adaptively assigns both positive and negative anchors, and at the same time, computes the weights of positive anchors more carefully, resulting in higher performance. To compare with PAA (+IoU pred.) which has an additional prediction structure, we integrate TAP into TAL, resulting in a higher AP of 42.5. More discussions on the differences between TAL and previous methods are presented in SM.

We evaluate the performance of the complete TOOD (T-head + TAL). As shown in Table 3, the anchor-free TOOD and anchor-based TOOD can achieve similar performance, i.e., 42.5 AP and 42.4 AP. Compared with ATSS, TOOD improves the performance of $\sim$ 3.2 AP. To be more specific, the improvements on AP75 are significant, which yields $\sim$ 3.8 points higher AP in TOOD. This validates that aligning the two tasks can improve the detection performance. Notably, TOOD brings a higher improvement (+3.3 AP) than the sum of the individual improvements by T-head + ATSS (+1.9 AP) and Parallel head + TAL (+1.1 AP), as shown in Table 6. It suggests that T-head and TAL can compensate strongly to each other.

We first investigate the performance using different values of $\alpha$ and $\beta$ for TAL, which control the influence of classification confidence and localization precision on anchor alignment metric via $t$ . Through a coarse search shown in Table 4, we adopt $\alpha=1$ and $\beta=6$ for our TAL. We then conduct several experiments to study the robustness of the hyper-parameter $m$ , which is used to select positive anchors. We use different values of $m$ in , and achieve results in a range of 42.0 $\sim$ 42.5 AP, which suggests the performance is insensitive to $m$ . Thus, we adopt $m=13$ in all our experiments.

2 Comparison with the State-of-the-Art

We compare our TOOD with other one-stage detectors on the COCO $test$ - $dev$ in Table 5. The models are trained with scale jitter (480-800) and for 2 $\times$ learning schedule (24 epochs) as the most current method . For a fair comparison, we report results of single model and single testing scale. With ResNet-101 and ResNeXt-101-64 $\times$ 4d, TOOD achieves 46.7 AP and 48.3 AP, outperforming the most current one-stage detectors such as ATSS (by $\sim$ 3 AP) and GFL (by $\sim$ 2 AP). Furthermore, with ResNet-101-DCN and ResNeXt-101-64 $\times$ 4d-DCN, TOOD brings a larger improvement, comparing to other detectors. For example, it obtains an improvement of 2.8 AP (48.3 $\rightarrow$ 51.1 AP) while ATSS has a 2.1 AP (45.6 $\rightarrow$ 47.7 AP) improvement. This validates that TOOD can cooperate with Deformable Convolutional Networks (DCN) more efficiently, by adaptively adjusting the spatial distribution of the learned features for task-alignment. Note that in TOOD, DCN is applied to the first two layers in the head tower. As shown in Table 5, TOOD achieves a new state-of-the-art result with 51.1 AP in one-stage object detection.

3 Quantitative Analysis for Task-alignment

We quantitatively analyze the effect of the proposed methods on the alignment of two tasks. Without NMS, we calculate a Pearson Correlation Coefficient (PCC) between the rankings of classification and localization by selecting top-50 confident predictions for each instance, and a mean IoU of the top-10 confident predictions, averaged over instances. As shown in Table 6, the mean PCC and IoU are improved by using T-head and TAL. Meanwhile, with NMS, the number of the correct boxes (IoU $>=$ 0.5) increases while those of the redundant (IoU $>=$ 0.5) and error boxes (0.1 $<$ IoU $<$ 0.5) decrease substantially when applying T-head and TAL. The statistics suggest that TOOD is more compatible with NMS, by preserving more correct boxes, and suppressing the redundant/error boxes significantly. At last, detection performance is improved by 3.3 AP in total. Several detection examples are illustrated in Figure 4.

Conclusion

In this work, we illustrate the misalignment between classification and localization in the existing one-stage detectors, and propose TOOD to align the two tasks. In particular, we design a task-aligned head to enhance the interaction of two tasks, and then improve its ability of learning the alignment. Furthermore, a new task-aligned learning strategy is developed by introducing a sample assignment scheme and new loss functions, both of which are computed via an anchor alignment metric. With these improvements, TOOD achieved a 51.1 AP on MS-COCO, surpassing the state-of-the-art one-stage detectors by a large margin.