Source-linked AI summary
Rethinking Transformer-based Set Prediction for Object Detection
Zhiqing Sun, Shengcao Cao, Yiming Yang, Kris Kitani
TL;DR
DETR achieves end-to-end set prediction but requires extra-long training, prompting an investigation of its slow convergence. The paper analyzes Hungarian matching and cross-attention, then proposes TSP-FCOS and TSP-RCNN, which converge faster and achieve higher detection accuracy than DETR and other baselines.
Problem
DETR’s end-to-end set-prediction approach requires extra-long training to converge, creating a costly optimization problem.
Method
The paper analyzes Hungarian matching and cross-attention, then proposes TSP-FCOS and TSP-RCNN with encoder-only Transformers, feature pyramids, and new matching schemes.
Results
The proposed methods converge much faster than original DETR and significantly outperform DETR and other baselines in detection accuracy.
Takeaways & Limitations
The proposed solutions reduce DETR-like training time while achieving state-of-the-art detection performance.
Abstract
from arXiv · showhide
DETR is a recently proposed Transformer-based method which views object detection as a set prediction problem and achieves state-of-the-art performance but demands extra-long training time to converge. In this paper, we investigate the causes of the optimization difficulty in the training of DETR. Our examinations reveal several factors contributing to the slow convergence of DETR, primarily the issues with the Hungarian loss and the Transformer cross-attention mechanism. To overcome these issues we propose two solutions, namely, TSP-FCOS (Transformer-based Set Prediction with FCOS) and TSP-RCNN (Transformer-based Set Prediction with RCNN). Experimental results show that the proposed methods not only converge much faster than the original DETR, but also significantly outperform DETR and other baselines in terms of detection accuracy.
1. Introduction
Object detection can be formulated as set prediction, but DETR’s end-to-end Transformer design requires far longer training than Faster RCNN. This paper analyzes the causes and proposes TSP-FCOS and TSP-RCNN to accelerate convergence while improving accuracy.
- Object detection is essentially a set prediction problem because predicted-object ordering is unnecessary.
- DETR directly outputs a final prediction set without post-processing, unlike detect-and-merge detectors based on proposals or sliding windows.
- 500 epochs and at least 10 days on 8 V100 GPUs are required for DETR to converge, compared with about 30 epochs for Faster RCNN.
- The paper investigates DETR’s optimization difficulty and proposes TSP-FCOS and TSP-RCNN as solutions for faster convergence.
- The proposed methods converge much faster than original DETR and significantly outperform DETR and other baselines in detection accuracy.
2. Background
The background presents object detection as either one-stage or two-stage prediction and introduces DETR’s end-to-end set-prediction objective. DETR uses Hungarian matching within an encoder-decoder Transformer, while related work improves efficiency, accuracy, or attention without adopting the same objective.
- One-stage detectors predict directly from feature maps and sliding-window locations, whereas two-stage detectors first generate and then refine region proposals.
- DETR formulates end-to-end set prediction through bipartite matching between padded ground-truth objects and predictions.
- The Hungarian loss combines classification and bounding-box regression losses for matched ground-truth and predicted objects.
- DETR’s encoder processes flattened CNN features, while its non-autoregressive decoder uses learned object queries to predict labels and boxes.
- Decoder cross-attention lets different object queries attend to image locations, supporting extraction of object information from encoder outputs.
- Deformable DETR and UP-DETR improve DETR’s efficiency or accuracy, while this work explores simplifying the detection head with an encoder-only Transformer.
3. What Causes the Slow Convergence of DETR?
DETR’s slow convergence reflects both early bipartite-matching instability and, more importantly, slowly optimizing Transformer cross-attention. Removing cross-attention accelerates small-object learning but exposes weaknesses on large and differently scaled objects.
- Hungarian matching: Matching distillation improves DETR convergence during the first few epochs, but its effect becomes insignificant after about 15 epochs.This indicates bipartite-matching instability contributes partially, especially early in training, but is not necessarily the main cause of slow convergence.
- Cross-attention: Imprecise cross-attention can prevent decoder queries from extracting accurate image context, resulting in poor localization, especially for small objects.The paper measures cross-attention sparsity using negative entropy averaged across attention heads and source positions, excluding masked positions.
- Cross-attention: Cross-attention sparsity consistently increases and does not plateau even after 100 training epochs, especially in the first encoder-decoder cross-attention layer.The analysis therefore identifies cross-attention as more dominant for slow convergence than early-stage matching instability.
- Encoder-only DETR: Encoder-only DETR performs as well as original DETR overall while significantly outperforming it on small objects and partly on medium objects.The encoder-only design removes decoder cross-attention and directly predicts objects from encoder outputs.
- Encoder-only DETR: Encoder-only DETR underperforms on large objects, potentially because large objects provide many matchable feature points and a single feature map is less robust across scales.These interpretations motivate multi-level feature representations for handling objects of different sizes.
4. The Proposed Methods
The proposed methods accelerate DETR-style set prediction by addressing matching instability and cross-attention while using feature pyramids and encoder-only Transformer designs. TSP-FCOS uses FoI selection and multi-level prediction, whereas TSP-RCNN adds proposal-based two-stage refinement.
- TSP-FCOS and TSP-RCNN address instability in DETR’s bipartite matching and inefficiency from Transformer cross-attention.
- TSP-FCOS: TSP-FCOS combines FCOS-inspired multi-level prediction with encoder-only DETR through Feature Pyramid Network features.The FPN produces multi-level features to help detect objects at various scales.
- TSP-FCOS: A FoI classifier selects a limited portion of features before Transformer encoding, reducing the sequence processed by quadratic-cost self-attention.Top-scored Features of Interest are fed into the Transformer encoder.
- TSP-FCOS: The TSP-FCOS encoder aggregates selected features with self-attention and predicts category labels, including “no object,” and bounding boxes for each FoI.
- TSP-RCNN: TSP-RCNN replaces FoI selection with an RPN that generates RoIs containing both objectness scores and predicted bounding boxes.RoIAlign extracts RoI information from multi-level feature maps before further processing.
- TSP-RCNN: TSP-RCNN uses two-stage bounding-box refinement and Faster RCNN’s assignment rule, assigning proposals only when their IoU with a ground-truth box exceeds 0.5.This design requires more computational resources but can detect objects more accurately.
5. Experiments
Experiments evaluate TSP-FCOS and TSP-RCNN on COCO under standard and extended training settings, comparing accuracy, computation, latency, and convergence-related design choices. The proposed models outperform their detector baselines, while TSP-RCNN achieves stronger overall performance than TSP-FCOS with slightly higher inference latency.
- Dataset and evaluation: COCO evaluation uses 115k trainval35k images, 5k minival images, and Average Precision and FLOPs as primary metrics.Test results use the test-dev split and comparisons also consider computation overhead.
- Implementation: TSP-FCOS selects the top 700 scored feature positions, while TSP-RCNN selects the top 700 scored RPN features for Transformer processing.TSP-RCNN additionally applies RoI Align and a fully connected layer to extract proposal features.
- Implementation: Both TSP models use a six-layer, width-512 Transformer encoder with eight attention heads and randomly drop 70% of encoder inputs during training.The configuration is chosen for FLOPs comparability with DETR-DC5.
- Main results: TSP-FCOS and TSP-RCNN significantly outperform original FCOS and Faster RCNN, with TSP-RCNN stronger overall and on small objects but slightly worse in inference latency.The comparison is reported on the COCO 2017 validation set.
- Main results: TSP-RCNN achieves state-of-the-art results with a shorter training schedule, although TSP-RCNN+ underperforms DETR-DC5+ on large-object detection.The authors attribute the large-object gap to DETR’s encoder-decoder inductive bias and longer training schedule.
- Ablation study: Set prediction loss contributes more than positional encoding to TSP-RCNN’s improvement, while both components are important to the TSP mechanism.The ablation uses ResNet-50 and ResNet-101 backbones.
- Backbone and model comparisons: The TSP mechanism is complementary to deformable convolutions, and TSP-RCNN obtains the highest AP scores among compared detectors with ResNet-101 or ResNet-101-DCN backbones.The comparison uses COCO test-set single-model, single-scale results for the highest-AP claim.
6. Analysis of convergence
The proposed faster set-prediction training consistently accelerates convergence for TSP-FCOS and TSP-RCNN. Compared with DETR-DC5, the proposed models converge faster and achieve better detection performance.
- Convergence comparison: The faster training technique consistently accelerates convergence for both TSP-FCOS and TSP-RCNN relative to DETR’s original set-prediction training.The comparison appears in the upper plots of Figure 6.
- Convergence comparison: TSP-FCOS and TSP-RCNN converge faster than DETR-DC5 while also achieving better detection performance.This comparison is shown in the lower convergence plots of Figure 6.
7. Conclusion
The paper investigates DETR’s slow convergence and proposes TSP-FCOS and TSP-RCNN to reduce training time while improving object-detection prediction power. The methods achieve state-of-the-art detection performance, with sparse attention identified as future work.
- Conclusion: The paper uses extensive experiments to investigate the causes of DETR’s slow convergence and proposes TSP-FCOS and TSP-RCNN as solutions.
- Conclusion: The proposed methods require much less training time and achieve state-of-the-art detection performance.
- Conclusion: Future work will investigate sparse attention for directly modeling relationships among multi-level features.
A.1. Transformer and Detection Transformer
DETR uses an encoder-decoder Transformer: the encoder contextualizes CNN features, and the decoder uses learned object queries to produce detection results through attention-based processing. The appendix also describes Faster R-CNN and FCOS as the two detector families motivating the proposed designs.
- DETR architecture: DETR’s encoder processes CNN backbone features into context representations, while its non-autoregressive decoder maps learned object queries to detection results.
- Multi-head attention: DETR uses both multi-head self-attention and cross-attention; cross-attention combines decoder queries with encoder context vectors.
- Multi-head attention: Multi-head attention concatenates attention heads to jointly use information from different representation subspaces.
- Multi-head attention: Positional encodings are incorporated into the query and key projections, while the value vectors carry the attended representations.
- Feed-forward network: The position-wise feed-forward network follows multi-head attention and consists of two linear transformations with ReLU activation.
- Related detectors: Faster R-CNN uses an RPN to generate object proposals and a Fast R-CNN stage to refine them, with NMS retaining top-ranked proposals.
- Related detectors: FCOS is an anchor-free, per-pixel detector that regresses bounding boxes at feature-map locations without predefined anchors.
- Related detectors: FCOS produces classification and center-ness scores, with center-ness indicating normalized distance from a responsible location to an object’s center.
B. Detailed Experimental Settings
The experiments use CNN backbones with feature pyramids, specified augmentation and optimization schedules, and a faster set-prediction loss. Evaluation includes COCO validation and test comparisons under defined schedules and backbones.
- Implementation settings: ResNet-50 or ResNet-101 backbones feed a Feature Pyramid Network producing feature maps P3 through P7.The pyramid uses 256 channels; deformable convolutions may replace the last three ResNet stages when specified.
- Implementation settings: Training resizes images with shortest sides from 640 to 800 pixels, caps the longest side at 1333 pixels, and randomly flips images horizontally.The augmentation follows Detectron2 defaults.
- Loss and optimization: The loss combines faster set-prediction classification with focal weighting and L1 plus Generalized IoU box regression.Unlike DETR, the models omit auxiliary losses after each encoder layer.
- Loss and optimization: The default 36-epoch schedule uses AdamW for Transformer parameters and momentum SGD for the detector's other components, with staged learning-rate decay and warm-up.A 96-epoch schedule resumes from the 36-epoch checkpoint and adds random crop augmentation.
- Regression design: Encoder-only DETR uses feature-point coordinates as reference points for bounding-box regression, unlike original DETR's reference-free normalized-coordinate prediction.The supplied description contrasts the reference formulation but does not provide the complete regression equation.
- Evaluation: The study reports COCO 2017 comparisons for TSP-FCOS, various-FLOPs models, and iterative-refinement comparisons with Deformable DETR.The tables specify single-model, single-scale test results or validation results under stated schedules and backbones.
D. Comparison between TSP-RCNN and Deformable DETR with Iterative Refinement
TSP-RCNN is compared with Deformable DETR under iterative refinement and with stronger baselines to assess accuracy and computational fairness. Iterative refinement improves TSP-RCNN, but its two-iteration setting remains below Deformable DETR when both refine predictions.
- Iterative refinement: A two-cascade TSP-RCNN configuration reduces head dimensions to maintain parameter and FLOP levels comparable to the original model.The design follows iterative box refinement in a cascade style.
- Comparison with Deformable DETR: Without iterative refinement, TSP-RCNN outperforms Deformable DETR under the same training setting.With iterative refinement enabled for both, TSP-RCNN slightly underperforms Deformable DETR.
- Comparison with Deformable DETR: The authors attribute the refined-performance gap to Deformable DETR using six decoder refinement iterations versus two in their TSP-RCNN experiments.Incorporating multiple refinement iterations efficiently is left for future work.
- Baseline comparisons: The study also evaluates stronger Faster R-CNN and other baselines to test whether TSP improvements arise only from added computation and parameters.TSP-FCOS is reported against state-of-the-art one-stage detectors on COCO 2017 test results.
G. Ablation Study for the Number of Feature Positions & Proposals
The ablation varies the number of selected feature positions for TSP-FCOS and proposals for TSP-RCNN. The models retain high prediction accuracy even when using half of the selected inputs.
- Selection-count ablation: TSP-FCOS selects 700 feature positions and TSP-RCNN selects 700 RPN proposals before Transformer processing.The paper notes that these quantities are not necessarily optimal and evaluates them in an ablation study.
- Selection-count ablation: Using half of the selected feature positions or proposals still preserves high prediction accuracy.The ablation covers R-50 TSP-FCOS and R-50 TSP-RCNN on the validation set.
H. Qualitative Analysis
The qualitative analysis examines one Transformer attention head across six validation images. Attention links RoI boxes referring to the same instances, helping reduce prediction redundancy.
- Attention visualization: Dashed boxes mark the top-5 attended RoI boxes for corresponding solid boxes sharing the same color.All analyzed boxes are predicted by the RPN.
- Attention visualization: The Transformer encoder can capture RoI boxes that refer to the same instances, which helps reduce prediction redundancy.The analysis uses one specific attention head in TSP-RCNN.