OD-DETR: Online Distillation for Stabilizing Training of Detection Transformer

1 Introduction

Object detection is a fundamental task in computer vision, and has been investigated by the community for decades. CNN-based detectors can be categorized into either anchorbased [Girshick, 2015; Liu et al., 2016] or anchor-free [Tian et al., 2019; Yang et al., 2019] methods. The former builds upon the sliding anchor boxes, and can be designed into a single, two or multi-stage, while the latter only has grid point assumption and is generally of a single stage. Although, CNN-based detectors has achieved impressive performance, it needs to determine complex meta parameters like anchor shape and size, threshold for positive and negative samples, and non-maximum suppression (NMS) post-processing.

*Corresponding Author

DEtection TRansformer (DETR) greatly simplifies the cumbersome design. It employs the encoder attention to enhance image feature. At the same time, several decoder attention layers are also incorporated, which translate the initial parameters of learnable object queries into predicted boxes. DETR uses bipartite matching to setup the one-to-one relation between ground truths (GTs) and predictions from queries, therefore, one GT is assigned to a single query, and vice versa. The one-to-one matching scheme reduces the redundant predictions, and alleviates detector from NMS. However, DETR is often blamed for its unstable training and slow convergence. As is shown in Fig. 1, the GT for query is often switched during training. Many efforts intend to improve it, by either introducing local box prior [Meng et al., 2021; Zhu et al., 2021; Liu et al., 2022], more query groups [Li et al., 2022; Chen et al., 2022a; Jia et al., 2023], an initial stage [Yao et al., 2021; Zhang et al., 2023] or improved qualityaware loss functions [Liu et al., 2023; Cai et al., 2023].

This paper proposes a solution, named online distillation Predictions Teacher Choose Aux Queries Aux Queries Auxiliary Group Hungarian matching Matching Distillation Online Decoder GT

Hungarian

matching

Predictions Aux

Figure 2: The overall architecture of OD-DETR. In the Main Group, the EMA model's queries are given to the decoder to create Predictions Teacher, and they are matched with the GT set using Hungarian matching. Simultaneously, these queries are also input into the Online Decoder to produce Predictions Student, and they directly learn from Predictions Teacher through prediction distillation. The Online model's own queries are decoded into Predictions Online, which are also matched with the GT set. Then, through matching distillation, they refer to matching result of Predictions Teacher. In the Auxiliary Group, we select two updated queries with the lowest matching cost for each GT from the Predictions Teacher. These selections are added as an extra group to the Online Decoder.

(OD-DETR), in another perspective to stabilizing the training of DETR. Inspired by the works in semi-supervised classification [Sohn et al., 2020] and object detection [Liu et al., 2021], we utilize Exponential Moving Average (EMA) as the teacher and distill its knowledge into the student model in an online manner. Unlike traditional distillation, the teacher also gets improved by the accumulation of the student. Particularly, we leverage the predicted bounding boxes, Hungarian matching result and updated object queries from the teacher model, and design schemes for prediction distillation, matching distillation and building auxiliary group within the student. For matching distillation, we assign another label for each object query through Hungarian matching according to cost matrix between teacher's predictions and GT boxes, and use it to guide the online student together with its original matching results. To employ two possible matched GTs, we propose a multi-target QFL loss to accommodate two labels from different classes, while keep only one regression target to avoid ambiguity. Meanwhile, two predictions associated with the same GT are given different regression loss weights, and the one with larger matching cost is down-weighted.

To take full advantage of the teacher, its initial queries are fed into the online student, giving predictions that can be directly constrained by the corresponding output from the teacher with no need to rematch. We name this simple constraint between two bounding boxes as prediction distillation. To further strengthen the link between teacher and student, we build independent augmented groups by object queries from each decoding stage of the teacher. Here, high quality queries with minimal matching costs for each GT are selected, and each group is given to the student's decoder to make predictions, re-match with GTs and compute losses. The augmented query groups are mainly used for speeding up the training convergence, hence it is abandoned during inference. The framework of OD-DETR is given in Fig. 2.

To verify the effectiveness of OD-DETR, we perform a large amount of experiments on MS-COCO [Lin et al., 2014]. Particularly, OD-DETR is implemented to support different variations of DETR including Def-, DAB- and DINO. We find that our method is compatible with all of them and increases the performance obviously. In summary, the contributions of this paper lie in following aspects.

• We propose an EMA-teacher based online distillation scheme. Particularly, matching distillation between the teacher and student is proposed. It assigns each query to an extra GT box based on the matching results of teacher's prediction. Moreover, we adapt the loss functions and propose a multi-target QFL classification loss and a cost sensitive regression loss.
• We perform prediction distillation and build augmented query groups to fully utilize of the EMA teacher. Prediction distillation is carried out by feeding the EMA query into student's decoder, and constraining the prediction by teacher's output, while the augmented groups are built by high-quality queries from different decoding stages within the teacher.
• Intensive experiments are carried out, showing that the proposed OD-DETR is able to effectively increase the performance of different DETR's variations.

2 RelatedWork

2.1 Supervised Object Detection

Modern object detection models mainly use convolutional networks and achieved great success recently. These CNNbased detectors fall into two categories: anchor-based and anchor-free. [Girshick, 2015; Ren et al., 2015; Lin et al., 2017] are some well-known anchor-based model, while [Tian et al., 2019; Yang et al., 2019; Duan et al., 2019] represent anchor-free models. Both both of them need manual components like non-maximum suppression (NMS) and heuristic label assignment rules.

DEtection TRansformer (DETR) [Carion et al., 2020] changes this scenario. It's the first end-to-end, query-based object detector without handcrafted designs such as anchors and NMS. However, it suffers from the slow training convergence. Several recent studies have focused on speeding up DETR's training process. Methods such as [Zhu et al., 2021; Gao et al., 2021; Gao et al., 2021; Sun et al., 2020; Zhao et al., 2023b; Zhao et al., 2023a; Liu et al., 2022] employ local box priors to concentrate on local features, reducing the search space. Other methods like [Chen et al., 2022a; Jia et al., 2023; Li et al., 2022; Zhang et al., 2023] speed up training by adding more query groups, providing additional positive samples for ground truth boxes.

2.2 EMA in SSOD and Distillation

Using EMA model as teacher to distill knowledge into student has been consistently applied in semi-supervised learning (SSL) based classifications, examples of which include [Zhu, 2005; Laine and Aila, 2017; Tarvainen and Valpola, 2017; Berthelot et al., 2019; Sohn et al., 2020]. A critical challenge in SSL is to fully leverage unlabeled images, which is realized by either self-training or consistency regularization. Another area where EMA-based teacher approaches is applied is in self-supervised learning, as exemplified by methods like [Grill et al., 2020; Gidaris et al., 2020; Caron et al., 2021]. The key to self-supervised learning also lies in the student model learning from the teacher model's outputs on unlabeled data.

Based on the similar idea in classification task, numerous SSOD (Semi-Supervised Object Detection) models have been developed. For instance, [Tang et al., 2021; Xu et al., 2021; Liu et al., 2021] adopt the principles of FixMatch [Sohn et al., 2020] and employs an EMA accumulated from the online student to provide pseudo GT boxes. However, for knowledge distillation (KD) task in object detection, current methods typically use a fixed, pre-trained model as the teacher to output labels, which are then provided to the online student for learning, as can be seen in [Chen et al., 2017; Li et al., 2017; Chang et al., 2023; Chen et al., 2022b; Wang et al., 2022]. Different from works in SSOD or KD, this paper focuses on online supervised learning for detection, and we utilize the EMA model as the teacher to provide guidance in an online manner, without requiring a fixed teacher model optimized in advance.

3 Method

3.1 Preliminaries on DETR and Its Adaptation

DETR is composed of a backbone, an encoder of several self-attention layers, a set of learnable object queries, and a decoder followed by a detection head to translate updated queries into boxes with class predictions. We define Q = {qi|qi ∈ Rc } as the query set. Each qi is a learnable parameters at the beginning stage. To make the model aware of the positions of bounding boxes, the enhanced versions of DETR [Liu et al., 2022; Zhu et al., 2021] explicitly encode the bounding box (x, y,w, h) or only box center (x, y) into positional embedding, specifying a set P = {pi|pi ∈ Rc } corresponding to Q. P, Q and image feature F are given to the decoder Dec, resulting in an updated query set Q and box set B for the next stage.

In Eq. (1), Dec indicates the decoder parameterized by θ, and the superscript denotes the decoder stage. The predicted box set B = {bi|bi = (x, y, w, h, c)} not only has box coordinate but also a class score vector c. Apart from the detection head, the DETR's decoder consists of self attention layer computed by interactions among query elements qi in Q, cross attention layer interacting between Q and F, and feed-forward network between them. In Def-DETR [Zhu et al., 2021], the cross attention can be improved by only sampling features around the reference point and giving a weighted average, and this local prior speeds up its convergence.

The updated Q are given to the detection head for classification and bounding box regression. But before loss computation, the predictions from queries are assigned to GTs according to matching costs. DETR adopts the Hungarian algorithm to set up the one-to-one relation between query set Q and all GTs, which is also the key reason that DETR can be free from NMS. However, this dynamic one-to-one matching strategy also leads to training instability. The work [Li et al., 2022] finds that there are some queries assigned with different GTs between two training epochs, as is shown in the blue curve in Fig. 1, and it causes the slow convergence of DETR.

The training target of original DETR includes a focal loss Lcls [Lin et al., 2017] for classification, and L1 and GIoU loss LGIoU for regression. Some works show that quality metric can improve performance. Notably, with quality focal loss LQFL [Li et al., 2020] defined in Eq. (2), classification and regression tasks are bounded together. Here t is the IoU target calculated between predicted box b and its matched GT box. s is a prediction score from sigmoid function.

Our proposed OD-DETR is built on enhanced DETR with QFL classification loss. Next, we introduce three key components: matching distillation, prediction distillation, and auxiliary group, as is shown in Fig.2.

3.2 Matching Distillation

Inspired by the success of the EMA model, we take advantage of it to stabilize DETR training. We first validate its potential application in Def-DETR. Specifically, we train it in 12 epochs and make an EMA at the same time. We compare it with the online model based on the instability metrics defined in [Li et al., 2022]. As is shown in Fig. 1, the behaviour of EMA is more stable than the online model, showing that it indeed prevents the label switch between two epochs, therefore giving a stable matching result. Considering that the EMA model has a more stable matching result, we intend to distill it into the online student.

Fig.3 presents an intuitive illustration of matching distillation. Matching result from the teacher and the student are combined into classification and regression losses for training the online model. Particularly, given teacher's query, PE and image feature Q′ , P ′ and F ′ , the decoder of the teacher, parameterized by θ′ , can output predicted boxes B′ , as in Eq. (1). Hungarian matching result between B′ and the GT set is used as a reference for constraining the prediction B from student model. Note that B also have a matching result, which means that qi and its EMA version q′i may have different GT to be matched. We now illustrate matching distillation for classification and regression in the following two sections. Matching distillation for classification. Since the online query qi are tightly connected with q′i , we design a multitarget QFL LMQF for class prediction in Eq. (3), which utilizes the matched target label of q′i , if it is different from the original target. Here s′ is another predicted class score, and t′ is its corresponding IoU target computed based on the teacher's matching result.

Particularly, if two matched GT are from the same class, LMQF becomes the same with LQFL defined in Eq. (2), computed by its original IoU target t. Otherwise, if two GTs are in different classes, two IoU targets t and t′ impose constraints on s and s′ , respectively. Fig.4 shows how the multi-target label is set for a query's prediction that matches two different categories of GTs. In dense scenes, one predicted box often contains multiple objects of different categories. Therefore, the single-category one-hot label is not appropriate. Our LMQF dynamically sets targets for different categories based on the IoU between the predicted box and two matched GTs, providing richer semantic information for training. It also increases training stability by preventing sudden label changes when online matching result changes.

Matching distillation for regression. For bounding box regression, the matching result from the teacher are also referred to. But we avoid the ambiguity due to the case in which the two different GT boxes are matched with one query, and only use the online matched result as the target to compute L1 and LGIoU . Moreover, since one GT may still be matched by two queries, we simply down-weight the regression loss for the query with a bigger matching cost. The regression loss Lr is defined in Eq. (4). Here b and bgt are the predicted and matched GT boxes. wd is a hyper-parameter with value of 0.51 following the idea in [Cai et al., 2023].

3.3 Prediction Distillation

Besides the matching distillation, the output from the teacher model B′ can also be exploited for training, and we name it prediction distillation. However, since the online predictions are already constrained by GT boxes, it is ambiguous to require them to approach another set of targets. To take full advantage of B′ , we feed the EMA queries Q′ together with the PE P and image feature F into the student decoder, and obtain output ˆ B as is in Eq. (1). Note that B′ is different from ˆ B. The former are totally from the teacher and work as a target set, while the latter are predictions which need to be constrained. Moreover, B′ and ˆ B are explicitly associated since they all begin from the same Q′ . Therefore, prediction distillation loss Lpd can be directly computed between them.

Here a naive way is to adopt LQFL defined in Eq. (2) but replace the IoU target t by the class score c′ from the teacher, which is a strategy for distillation under a fixed teacher [Chen et al., 2022b]. However, the predictions from the EMA teacher are not very accurate in our case, especially during the early training stages. Applying it in the naive way can introduce many errors from the teacher itself, which lead to worse results. We adapt it so it becomes suitable for online distillation. First, since B′ have a matching result with GT boxes, an implicit class label index cg for each box in B′ can be inferred. Then we modify the predicted score vector c′ at the corresponding entry according to Eq. (5).

Here IoU ′ is computed between bounding box in B′ and its matched GT. α and β are two hyper-parameters. Following [Feng et al., 2021], we set α = 0.25 and β = 0.75. After update on c′ , we then replace t in Eq. (2) by it, therefore, giving Lcls distill for prediction distillation. For regression task, c′ [cg] is also used as a weight. Using c′ [cg] in Lrdistill and Lcls distill is named as TOOD-weight, which is first proposed in [Feng et al., 2021] for one-stage detector. Note that c′ [cg] is computed from the teacher, which gives the key difference with [Feng et al., 2021]. Consequently, Lrdistill = c′ [cg](L1 +LGIoU). Lcls distill and Lrdistill are combined into Lpd in Eq. (6).

Here 1 denotes an indicator function, and it only equals to one when the predictions from the teacher have a larger IoU with GT box than the student, otherwise, it returns 0, meaning that Lrdistill is not considered in this case. This approach is named as Listen2stu.

3.4 Auxiliary Group

To better utilize the EMA teacher and enhance training stability, we choose some updated queries ˜ q and corresponding predicted boxes ˜b from (Q′ )t and (B′ )t at t-th stage of the teacher's decoder. Then we use them as independent initial queries and anchors for PE. They are fed into the first decoding stage of online model, providing more positive examples for learning. To reduce the computational load of the auxiliary group, we select only the top two queries with the lowest matching cost to each GT from each decoder stage. This method ensures that the selected queries include both positive examples and challenging negative ones. Note that queries from the same decoding stage forms an independent group. Predictions from each group is matched with GT set in one-to-one manner.

Matching distillation is also used in each auxiliary group, just as we do in main group. This method combines the original matching from the teacher with the new ones, which enhances the training stability for auxiliary group. So the loss in auxiliary group can be denoted by Laux = ˜LMQF + ˜Lr. In summary, the total training loss is shown in Eq. (7).

4 Experiments

4.1 Settings

Datasets. We conduct all our experiments on MS-COCO [Lin et al., 2014] 2017 dataset and evaluate the performance of our models on validation dataset by using mean average precision (mAP) metric. The COCO dataset contains 117K training images and 5K validation images.

Implementation details. We use Def-DETR (with iterative bounding box refinement), DAB-Def-DETR, and DINO as our baseline methods. Our experiments are conducted over 12 (1x) and 24 (2x) epochs on 8 GPUs. Learning rate settings for OD-DETR are identical to those of Def-DETR, with a learning rate of 2 × 10-5 for the backbone and 2 × 10-4 for the Transformer encoder-decoder framework, coupled with a weight decay of 2 × 10-5 . The learning rates and batch sizes for OD-DAB-DETR and OD-DINO follow their respective baselines. We set the EMA decay value at 0.9996.

4.2 Main Results

As shown in Table 1, we firstly use Def-DETR [Zhu et al., 2021] as baseline, and compare our OD-DETR on COCO val2017 dataset with other competitive DETR variants. Our OD-DETR achieved a significant improvement, reaching 47.7 AP. It notably surpassed the baseline (Def-DETR iterative bounding box refinement) trained for 50 epochs, with a substantial increase of 2.3 in AP after just 2x training.

The OD-DAB-DETR model, improving upon the DABDef-DETR [Liu et al., 2022], achieved 50.2 AP. It exceeds the baseline by 1.6 AP after just 2x training and outperforms DN-Def-DETR [Li et al., 2022] with 50 epochs of training.

Under both 1x and 2x schedulers, our OD-DINO consistently outperforms the baseline DINO [Zhang et al., 2023], achieving 50.4 AP with 1x and 51.8 AP with 2x scheduling, marking a 1.4 increase in AP. Our OD-DINO series outperform other DINO-based approaches, including Stable-DINO [Liu et al., 2023], Align-DETR [Cai et al., 2023] and so on.

4.3 Ablation Study

We conduct a series of ablation studies to assess the effectiveness of the components. All experiments in this section are based on the Def-DETR model with an R50 backbone and follow a standard 2x training schedule.

Ablation Study on Components

In this section, we perform a set of ablation studies on key components: matching distillation (MD), prediction distillation (PD), and auxiliary group (AG), to assess their individual effectiveness. The findings are detailed in Table 2. These studies reveal that each component plays a significant role in enhancing performance. Specifically, MD contributes an increase of 1.4 AP, PD adds 0.5 AP, and AG contributes 0.4 AP. Cumulatively, these components lead to an overall improvement of 2.3 AP over the baseline model.

Fig.5 shows that our methods improves the consistency of matching GTs between the online model's queries and the EMA model's queries. As the online model's matching improves with training, it also boosts the EMA model's stability. This two-way improvement creates a strong combined effect.

Comparisons of Different Matching Distillation Methods

Experiments in Table 3 focuses on different matching distillation methods for only classification. The first row shows that employing only QFL [Li et al., 2020] without MD results in a significantly lower performance, achieving 45.8 AP, compared to our MD method combined with LMQF (QFL MD), which reaches 46.4 AP. This difference highlights the significant contribution of MD beyond the impact of QFL alone. Conditional MD, introduced in the StageInteractor [Teng et al., 2023], is a technique that allows a query to match with multiple ground truths. This method involves verifying if the additional query-GT matches have an IoU above 0.5, discarding those that don't meet this criterion. The performance of it is 45.9 AP. Integrating LMQF with Conditional MD under the label QFL + Conditional MD doesn't enhance performance. The performance is 46.3 AP. This observation implies that LMQF and MD are inherently synergistic. LMQF eliminates the necessity of manually setting thresholds like in Conditional MD for Hungarian matching.

Analyzing Down-Weight Impact in Matching Distillation

In Method 1 from Table 4, MD is only used for classification, keeping a one-to-one match for regression, resulting in a lower AP of 46.4 (QFL MD in Table 3). This implies that MD is beneficial for regression. Method 3, which applies downweight (wd) only for regression, achieves the highest AP of 46.8. This confirms the effectiveness of down-weighting regression loss. It also demonstrates QFL MD's effectiveness in handling many-to-many matching in classification. It eliminates the need to adjust loss for GTs that match multiple queries.

Comparisons of Various Prediction Distillation Methods

In Table 5, w/o PD means using MD alone without PD, and it's also referred to Method 1 in Table 2. Naive distillation applies distillation evenly across all queries, but this method only achieves 44.2 AP, lower than w/o PD (46.8 AP). This suggests that simply adding online distillation doesn't work. Using TOOD-weight and Listen2stu methods led to a notable increase of 2.3 AP, reaching 46.6 AP. As Fig.6 shows, PD notably enhances performance before learning rate decay, but tends to have a negative impact after the learning rate decay. Thus, we adopt a Stop Update strategy, halting updates to the EMA teacher's weights after learning rate decay. This method eventually achieves 47.3 AP. We also test Queryprior Assignment Distillation from DETRDistill [Chang et al., 2023]. It feeds the teacher's query into the student's decoder as an extra group to learn GT. The outcome, at 46.9 AP, is less effective than our PD approach.

Experiments on Applying MD to the Auxiliary Group

Table 6 shows results for different matching methods in auxiliary group. Using either the EMA model's original matching or re-matching gives 47.0 AP and 46.7 AP, respectively. The highest AP, 47.2, comes from using MD for auxiliary group (Method 3 in Table 2), just like in the main group. This proves MD works well in both main and auxiliary group.

Comparative Results of EMA Models

Table 7 compares the EMA model performances of DINO4scale and our OD-DINO-4scale, both trained using a 2x setting. Our OD-DETR-EMA outperforms DINO-EMA by 1.3 AP. DINO's EMA version achieves 50.7 AP, which is lower than the the online version of OD-DINO (51.8 AP). This indicates that our methods not only improves the performance of the online student but also enhances the EMA teacher itself. A better teacher then further improves the student, creating a beneficial cycle.

5 Conclusion

This paper proposes an OD-DETR, which is an online distillation method to stabilize the training of DETR. We find that the EMA model accumulated during training provides not only high-quality predicted boxes, but also the query-GT matching result and extra query groups. Under the help of EMA model, we improve the online training through prediction distillation, matching distillation and auxiliary query group. We show that the proposed OD-DETR can steadily increase the performance across different DETR variations.

Acknowledgments

This work is supported by the Science and Technology Commission of Shanghai Municipality under Grant No. 22511105800, 19511120800 and 22DZ2229004.