Rethinking Transformer-based Set Prediction for Object Detection

In plain words

DETR found objects in pictures in one shot, with no hand-tuned anchors and no clean-up step that merges duplicate boxes. Its catch is training time: it needs about 500 epochs to converge, where a classic detector needs about 30. The authors traced the slowness to two parts: the matching step that pairs each predicted box with a real object, and the attention module that lets the decoder look up image content.

Their fix keeps the good parts and drops the slow one. They remove the decoder and let a transformer encoder work only on a short list of promising image locations. The list comes from a classic detector's head (FCOS or a region proposal network), and matching is limited to locations that already sit inside an object. Trained this way, the two models reach strong accuracy in far fewer epochs than DETR.

TL;DR

• DETR's set loss is end-to-end, but it converges slowly §1 §3.
• The cross-attention module is "mainly responsible" for slow convergence, per the authors §1 §3.2 §3.3.
• TSP-FCOS and TSP-RCNN run only a transformer encoder over 700 selected features and use a restricted bipartite matching §4 §5.2.
• Author-reported COCO val AP at ResNet-50 with 36 epochs: 43.1 and 43.8, against DETR's 42.0 at 500 epochs tab. 1.

Problem

DETR is the first fully end-to-end detector, and its output needs no non-maximum suppression §1 §2.2. Training is the problem: the authors cite about 500 epochs for DETR against about 30 for Faster R-CNN, and state that DETR's training takes at least 10 days on 8 V100 GPUs §1.

The authors ask which part is responsible. Three findings §3:

• Instability of the bipartite matching. A matching distilled from a well-trained teacher speeds up only the first few epochs; the effect is insignificant after about 15 epochs §3.1 fig. 1.
• Sparsity of the cross-attention. Its sparsity keeps increasing and has not plateaued after 100 epochs; the authors read it as the more dominant factor §3.2 fig. 2.
• Encoder-only DETR (decoder and cross-attention removed) matches the original overall AP, is much better on small objects and worse on large ones §3.3 fig. 4.

Setting: COCO 2017 with ResNet-50 or ResNet-101 backbones §5.1 app. B.

Method

Both models are encoder-only set predictors with the Hungarian set loss. They differ in how the input tokens are chosen §4 fig. 5.

Set prediction and the loss. DETR pads the ground-truth set to size N with "no object" and matches predictions to targets by a permutation that minimises total matching cost. The Hungarian loss is then summed over the matched pairs eq. 1 eq. 2 §2.2.

\mathcal{L}_{\text{Hungarian}}(\bar{y}, \hat{y}) = \sum_{i=1}^{N} \left[ \mathcal{L}_{class}^{i,\hat{\sigma}(i)} + \mathbbm{1}_{\{\bar{y}_i \neq \emptyset\}} \mathcal{L}_{box}^{i,\hat{\sigma}(i)} \right] eq. 1

\hat{\sigma} = \argmin_{\sigma \in \mathfrak{S}_N} \sum_{i=1}^{N} \mathcal{L}_{\text{match}}(\bar{y}_i, \hat{y}_{\sigma(i)}) eq. 2

TSP-FCOS §4.1 fig. 5, upper part.

1. A backbone and a Feature Pyramid Network build levels P3 to P7 §4.1 app. B.
2. Two shared heads run on every level, as in FCOS: a classification subnet (used for FoI scoring) and an auxiliary subnet. Each is four 3x3 convolutions with 256 channels and group normalization §5.2.
3. A binary FoI classifier scores every location, trained with FCOS's ground-truth assignment rule. The top 700 locations become features of interest (FoIs) §4.1 §5.2.
4. Each FoI gets a sinusoidal positional encoding of its normalised (x, y) §4.1 eq. 3.
5. A 6-layer Transformer encoder (width 512, 8 heads, FFN width 2048) runs on the FoIs §5.2 app. B. The paper states that 70% of the encoder inputs are randomly dropped during training §5.2; see "Paper vs code".
6. A shared feed-forward head predicts a class (including "no object") and a box per FoI §4.1.
7. The matching is restricted: a location may be matched to an object only if it lies inside the object's box at a suitable pyramid level, as in FCOS §4.1. The matching is then the Hungarian matching of eq. 2 over the allowed pairs.

TSP-RCNN §4.2 fig. 5, lower part.

1. The same backbone and FPN feed a Region Proposal Network with two unshared convolutional subnets (classification and regression) on P3 to P7 and RetinaNet-style anchors §4.2 §5.2.
2. The top 700 RPN proposals become RoIs. Each RoI is pooled with RoIAlign and passed through a fully connected layer §4.2 §5.2.
3. The positional encoding of a RoI is the sinusoidal encoding of its normalised (cx, cy, w, h) §4.2.
4. The encoder and matching are as in TSP-FCOS, except that a proposal is positive only if its IoU with a ground-truth box is greater than 0.5 §4.2.

Encoder-only DETR (the diagnostic model of §3.3). It keeps DETR's encoder but predicts from each encoder output location, not from object queries. Its box regression is referenced to the location's own coordinates §3.3 app. C.

Training. The transformer is optimised with AdamW and the rest of the detector with SGD (momentum 0.9). The default schedule is 36 epochs (3x) app. B §5.2.

# TSP-FCOS training step: step-by-step summary of [§4.1] and [eq. 1]-[eq. 2]
# Input:  images (B, 3, H, W); per image: gt_boxes (M, 4), gt_classes (M,)
# Output: scalar loss

def training_step(images, gt_boxes, gt_classes):
    # Backbone and FPN give pyramid levels P3-P7 [§4.1] [app. B]
    pyramid = fpn(backbone(images))
    # Shared heads: FoI score per location and auxiliary features [§4.1]
    foi_scores, aux_feats = foi_heads(pyramid)   # (B, N), (B, N, C)
    # Keep the 700 best-scoring locations as FoIs [§5.2]
    foi_idx = top_k(foi_scores, k=700)           # (B, 700)
    feats = gather(aux_feats, foi_idx)           # (B, 700, d_model)
    # Sinusoidal encoding of normalised (x, y) of each FoI [§4.1] [eq. 3]
    pos = sine_encoding(foi_idx)                 # (B, 700, d_model)
    # Training only: randomly drop 70% of the encoder inputs [§5.2]
    feats, pos = random_drop(feats, pos, rate=0.7)
    # 6-layer encoder, width 512, 8 heads, FFN 2048 [§5.2] [app. B]
    memory = transformer_encoder(feats, pos)     # (B, k, d_model)
    # One class (incl. "no object") and one box per FoI [§4.1]
    logits = cls_head(memory)                    # (B, k, C + 1)
    boxes = box_head(memory)                     # (B, k, 4)
    # Only FCOS-style positives may match an object [§4.1]
    allowed = fcos_positive_mask(gt_boxes, gt_classes, foi_idx)
    # Min-cost one-to-one matching over allowed pairs [eq. 2]
    perm = hungarian(match_cost(logits, boxes, gt_boxes, allowed))
    # Set loss on matched pairs; unmatched slots learn "no object" [eq. 1]
    return set_loss(logits, boxes, gt_boxes, gt_classes, perm)

Positional encoding, as defined in the paper eq. 3:

PE(x)_{2i} = \sin\!\left(\frac{x}{10000^{2i/d_{\text{model}}}}\right), \qquad PE(x)_{2i+1} = \cos\!\left(\frac{x}{10000^{2i/d_{\text{model}}}}\right) eq. 3

The multi-head attention and feed-forward formulas used by DETR are written in app. A.1.

Hyperparameters

Values are as the paper states them. Where the released code sets a different value, the difference is listed in "Paper vs code".

Results

Author-reported. COCO 2017 validation set unless marked test-dev. Single model, single scale for test-dev tab. 4. The "+" marks random-crop augmentation and a longer schedule, applied to DETR and to the TSP models alike tab. 1.

Convergence: the paper states that the TSP models converge faster and reach better AP than DETR-DC5 §6 fig. 6. The plot has no numbers in the text.

Ablation of the number of FoIs and proposals (R-50, 100 to 700): TSP-RCNN goes from 40.3 to 43.8 AP; TSP-FCOS from 40.0 to 43.1 tab. 8.

Paper vs code

Checked the official repository linked in the paper and its abstract: github.com/Edward-Sun/TSP-Detection at commit 1b8a50388c84dd2644b6feba582cd934ade852e5. Code read: configs/, fcos/, rcnn/, train_net_fcos.py, train_net_rcnn.py, tsp_fcos.sh, tsp_rcnn.sh. The released run scripts are the reference for the TSP models; the defaults in */config.py are used only where a script does not set the key.

• Matching is not the Hungarian algorithm in the released code. The paper defines the matching as a permutation minimising total cost eq. 2. Both scripts set USE_OBJ_LOSS True tsp_fcos.sh [code: tsp_rcnn.sh#L26]. In that path, the code builds a per-ground-truth cost row and takes, for each ground-truth object, the proposal with the lowest cost [code: fcos/my_fcos_outputs.py#L538-L543] my_fast_rcnn_output.py. Nothing enforces one proposal per object. The Hungarian path with matched_indices = self.matcher(...) runs only when USE_DETR_LOSS is set my_fcos_outputs.py, which the scripts do not set. Implement the paper's matching or the released code's variant, not both.
• Matching cost weights are not in the paper. The released FCOS cost is minus the predicted probability of the true class, multiplied by the centerness target, plus GIOU_WEIGHT times the GIoU loss [code: fcos/my_fcos_outputs.py#L528-L533]. The script sets GIOU_WEIGHT 2.0 [code: tsp_fcos.sh#L29]; the default is 4.0 config.py. The RCNN cost is class minus probability plus giou_weight times GIoU loss [code: rcnn/my_fast_rcnn_output.py#L567]. The paper gives no matching-cost weights eq. 2.
• Random drop: keep, not drop. The paper says 70% of encoder inputs are dropped §5.2. The code samples a factor uniformly in [lower, upper] = [0.7, 1.0] and keeps int(700 * factor) of the 700 locations [code: fcos/my_fcos.py#L347-L353] tsp_fcos.sh. The RCNN head does the same to the sampled RoIs rcnn_heads.py [code: tsp_rcnn.sh#L24-L25]. On average that keeps about 85%, not 30%.
• Gradient clip is 0.01 in code, 0.1 in the paper. The paper states a maximal L2 gradient norm of 0.1 app. B. Both configs set CLIP_VALUE 0.01 Base-FCOS-TSP.yaml Base-RCNN-TSP.yaml. The clip applies to parameters whose module name contains "transformer" train_net_fcos.py train_net_fcos.py.
• SGD learning rate for TSP-RCNN is 0.02, not 1e-2. The paper gives 1e-2 for SGD app. B. tsp_rcnn.sh sets BASE_LR 0.02 tsp_rcnn.sh. tsp_fcos.sh uses 0.01 tsp_fcos.sh, which matches. The transformer learning rate is 1e-4 in both scripts (0.01 x 0.01 and 0.02 x 0.005) [code: tsp_fcos.sh#L25-L26] tsp_rcnn.sh, matching app. B.
• Regression loss. The paper says L1 plus GIoU app. B. TSP-FCOS uses GIoU alone, weighted by centerness targets my_fcos_outputs.py my_fcos_outputs.py. TSP-RCNN uses smooth L1 plus 2.0 times GIoU Base-RCNN-TSP.yaml config.py my_fast_rcnn_output.py.
• Inference uses NMS for TSP-RCNN. The paper frames set prediction as removing NMS §2.2. The TSP-RCNN script sets the test NMS threshold to 0.7 tsp_rcnn.sh. Inference calls the NMS path unless DETR_EVAL_PROTOCOL is set my_fast_rcnn_output.py; the default is False config.py and no script sets it. The FCOS inference path was not checked.
• FFN width is 2048 in the scripts, 1024 in the defaults. The paper says 2048 §5.2. Both scripts override the default of 1024 [code: rcnn/config.py#L29] to 2048 tsp_fcos.sh tsp_rcnn.sh.
• Dropout is 0.1 in code, not stated in the paper. [code: rcnn/config.py#L30].
• FCOS-side options the paper does not mention. The scripts turn on ONLY_REWEIGHT_FG tsp_fcos.sh, PREDICT_WITHOUT_CTR [code: tsp_fcos.sh#L30] and CLASS_DENORM_TYPE mixed_2x tsp_fcos.sh. The defaults EOS_COEF 0.1 and CENTER_SAMPLE True with POS_RADIUS 1.5 are in config.py and config.py. Center sampling is the FCOS assignment rule. The paper does not describe it.
• RPN anchors are given only in code. The paper says RetinaNet-style anchors §5.2 with no sizes. The script sets three sizes per level, roughly 32 times 2 to the power of 0, 1/3 and 2/3, with ratios 0.5, 1 and 2 [code: tsp_rcnn.sh#L16] Base-RCNN-TSP.yaml.
• Not in the repository. The 96-epoch "+" schedule app. B has no script here; the two scripts run MAX_ITER 270000 (36 epochs). The 1000-iteration warm-up is not set in the configs.
• Code license. Apache-2.0, from the repository's LICENSE file (verbatim Apache License, Version 2.0 text).

No discrepancies were found in the architecture: 6 encoder layers, 512 width, 8 heads, 700 FoIs or proposals, the sinusoidal encoding and the transformer learning rate match between paper and scripts.

Pitfalls

• Matching is where paper and code differ most. Implementing the Hungarian matching of eq. 2 does not reproduce the released code, which takes a per-object minimum (see "Paper vs code"). Pick one and say which.
• The matching cost is undefined in the paper. Any weights used must be chosen by the implementer; the released scripts use GIoU weight 2.0 and a centerness-weighted class term tsp_fcos.sh.
• "Drop 70%" means keep 70% to 100% in code. The name of the config key is LOWER_BOUND, not a drop rate tsp_fcos.sh.
• Gradient clipping is 0.01 in code, 0.1 in the paper. The paper's 0.1 is the stated value app. B; the released configs use 0.01. Cite which one you implement.
• Convergence claims rest on epoch counts that hide augmentation. The "+" rows add random crop and a longer schedule tab. 1. The 36-epoch TSP rows and the 500-epoch DETR+ row are not the same training recipe.
• Iterative refinement table does not match the text. tab. 7 lists 45.4 for both Deformable DETR and TSP-RCNN with refinement, while app. D says TSP-RCNN "slightly underperforms". Use the table as the source.
• Large objects are weaker. Encoder-only DETR is worse on large objects §3.3 fig. 4; TSP-RCNN+ is 58.0 AP-L against DETR-DC5+ 61.1 tab. 1. The paper attributes this to DETR's encoder-decoder inductive bias and its longer schedule §5.3.
• Compute. FLOPs are measured with Detectron2; FPS on a single RTX 2080 Ti tab. 1. TSP-RCNN is slower than TSP-FCOS (11 vs 15 FPS at R-50) for higher accuracy. These are the authors' numbers; no other hardware is reported.
• Attention cost drives the FoI design. Self-attention cost is quadratic in the sequence length §4.1, which is why FoI selection exists. Using 300 instead of 700 costs 0.6 AP for TSP-FCOS (42.5 vs 43.1) and 0.5 for TSP-RCNN (43.3 vs 43.8) tab. 8 app. G.
• Dynamic sample size is optional. The code has a random sample-size option (RANDOM_SAMPLE_SIZE, default off) my_fcos.py. The paper does not mention it, and the scripts do not enable it.

Lineage

• Builds on: DETR, whose set loss and Hungarian matching this paper keeps. Also FCOS and FPN, which supply the FoI assignment rule and the pyramid; Faster R-CNN, which supplies the RPN for TSP-RCNN. Those are named, not linked: no wiki card exists for them. RoIAlign comes from Mask R-CNN §4.2.
• Compared against: DETR and DETR-DC5 tab. 1; Deformable DETR (arXiv:2010.04159, named) tab. 7; UP-DETR tab. 1; FCOS, RetinaNet, ATSS, TridentNet and Dynamic RCNN tab. 4.
• Superseded by: not stated in this wiki.