Dynamic Head: Unifying Object Detection Heads with Attentions

In plain words

An object detector finds things in a picture in two parts: a backbone that turns the picture into feature maps, and a head that reads those maps to say where each object is and what it is. Many detectors have heads built for one job, such as handling objects of very different sizes, handling odd shapes, or handling several kinds of target at once. The authors argue that all three needs can be covered by one head.

Their Dynamic Head treats the stack of feature maps as one block of numbers with three axes: which zoom level, which position, and which channel. It then applies three kinds of attention, one per axis. The first picks which zoom level to trust for each object. The second looks at the positions that matter most. The third chooses, per channel, which signal to use for classification or for locating the box. Several of these blocks are stacked, and the head is attached to an existing detector in place of its usual head. The authors report new top scores on the COCO benchmark, and the card's Results and Pitfalls sections say which numbers are comparable.

TL;DR

• Dynamic Head is a detection head made of stacked blocks. Each block applies a scale-aware attention over pyramid levels, a spatial-aware attention over positions, and a task-aware attention over channels §3.2.
• It is a plug-in: the authors attach it to ATSS and to other detectors, and report gains of 1.2 to 3.2 AP over the base detectors §4.3 tab. 3.
• Author-reported headline: 54.0 AP on COCO test-dev with a ResNeXt-64x4d-101-DCN backbone, multi-scale training and testing tab. 5. A 60.6 AP result uses a Swin-L backbone, extra pseudo-labelled data and a larger input app. tab. 7.
• The official code is an ATSS plug-in in PyTorch on detectron2, and it differs from the paper in several places; see Paper vs code.

Problem

A detection head must handle three things at once: objects of very different scales in one image §1; objects with different shapes, rotations and locations §1; and objects that are represented in different ways, such as boxes, centers or corner points, each with its own training target §1. Prior heads address one of these in various ways. The authors state that no unified head addresses all three §1.

Full self-attention over the whole level × space × channel tensor is the obvious unified answer. The authors say it is too hard to optimise and too costly to compute §1 §3.2. The paper's answer is to factorise the attention into three cheaper, sequential attentions, one per axis §3.2.

Method

Setting. Feature pyramid levels are resized to the size of the median level and stacked into a tensor F with shape L × S × C, where S = H × W §3.1. The paper defines the head as applying a general attention W(F) = π(F) · F eq. 1, and factorises π into three attentions applied in sequence eq. 2.

Scale-aware attention π_L §3.2 eq. 3. Features are averaged over S and C, which leaves one value per level. A linear map, approximated by a 1 × 1 convolution, and a hard-sigmoid σ(x) = max(0, min(1, (x + 1) / 2)) turn this into one weight per level. Each level is then scaled by its weight.

Spatial-aware attention π_S §3.2 eq. 4. To keep the sampling sparse, the paper uses deformable convolution. For each level l and each of K sampling points p_k, the head predicts a learned offset Δp_k and an importance scalar Δm_k. Features are read at p_k + Δp_k, weighted by a learned weight w_{l,k}, summed over levels and points, and divided by L. The paper says the offsets and scalars are predicted from the median-level feature §3.2. The paper does not state K.

Task-aware attention π_C §3.2 eq. 5. Each channel is switched between two linear functions of its own value, with the maximum taken: max(α¹ · F_c + β¹, α² · F_c + β²). The four coefficients come from a hyper-function θ, which the paper builds like Dynamic ReLU [dyrelu]: global average pooling over L × S, two fully connected layers, a normalisation layer, and a shifted sigmoid that maps outputs to [-1, 1] §3.2.

Composition. The three attentions are nested as in eq. 2: π_L is applied to F, π_S to the result, π_C to that, and the final output is multiplied by F again. Several DyHead blocks are stacked §3.2 fig. 2(a) tab. 2.

Plugging into detectors §3.3.

• One-stage: the backbone feeds one unified DyHead branch. Its outputs (classification, box regression, centerness or key-points) are appended at the end of the head, instead of building task-specific branches §3.3 fig. 2(b).
• Two-stage: π_L and π_S are applied to the feature pyramid before RoI pooling, and π_C replaces the fully connected layers after RoI pooling §3.3 fig. 2(c).

Pseudocode for one block, following eq. 2 to eq. 5. Shapes are for one image. The paper does not give the predictor networks for offsets and scalars, so they are shown as a helper.

# DyHead block  [§3.2] [eq. 2]
# Input:  features (L, S, C): L levels, S = H * W positions, C channels

def dyhead_block(features):
    # Scale-aware gate per level  [eq. 3]
    # Pool over S and C; a linear map and hard-sigmoid give one per level.
    level_summary = mean(features, over=(S, C))         # (L,)
    level_gate = hard_sigmoid(linear(level_summary))    # (L,)
    gated = level_gate[:, None, None] * features        # (L, S, C)

    # Sampling offsets and importance for K points  [§3.2] [eq. 4]
    # Predicted from the input; the paper does not give the predictor.
    offsets, importance = predict_sampling(gated)       # (L, S, K, 2)
    sampled = sample_at(features, positions + offsets)  # (L, S, K, C)

    # Spatial-aware aggregation over levels and points  [eq. 4]
    # level_weights is learned per (level, point); the sum is divided by L.
    weighted = (level_weights[:, None, :, None]         # (L, 1, K, 1)
                * sampled                               # (L, S, K, C)
                * importance[..., None])                # (L, S, K, 1)
    spatial = sum(weighted, over=(L, K)) / L            # (S, C)

    # Task-aware channel switch  [eq. 5]
    # Levels are already reduced above, so theta pools over S only.
    alpha1, alpha2, beta1, beta2 = theta(mean(spatial, over=S))
    tasked = maximum(alpha1 * spatial + beta1,
                     alpha2 * spatial + beta2)          # (S, C)

    # Eq. 2 multiplies the attention output by the input features again.
    return tasked * features                            # (L, S, C)

Hyperparameters

Results

Author-reported. Ablations use ResNet-50 at 1x and are evaluated on COCO val2017 §4.1 §4. Headline comparisons are evaluated on COCO test-dev through the official test server §4. The main models train on train2017 with no extra data §4. The appendix results use extra pseudo-labelled ImageNet data app..

Ablation on the attention modules (ResNet-50, 1x, val2017) tab. 1. L = scale-aware, S = spatial-aware, C = task-aware.

Number of DyHead blocks (ablation setting) tab. 2. The baseline row gives its GFLOPs; the other rows show a signed difference from the baseline, as printed.

Generalisation to other detectors (val2017) tab. 3. The paper reports gains of 1.2 to 3.2 AP §4.2.

Backbones on test-dev tab. 4. Rows are as printed; "Iteration" is the training schedule shown in the table.

State of the art on test-dev tab. 5. Top block: multi-scale training only. Bottom block: multi-scale training and multi-scale testing.

Transformer backbone, validation set tab. 6. This table uses val2017, not test-dev tab. 6 app..

Latest methods, test-dev tab. 7. "†" marks extra data.

Paper vs code

Code: github.com/microsoft/DynamicHead at commit 9c7cd661e362e9e58a43d8b18fe56afaa2aae22c, MIT licence (Microsoft Corporation). The repository is linked from the paper's abstract and introduction §1 and is cloned at raw/code/2106.08322. I read the model, the deformable-convolution wrapper, the DYReLU module, the ATSS integration, the training script and the configs. The repository has no ResNeXt-DCN, Swin-L or extra-data configs.

• Scale-aware attention. The paper pools over S and C and maps the L values through a linear function f eq. 3 §3.2. The code pools each candidate feature spatially only, projects it with a 1 x 1 convolution to one scalar, and applies a ReLU. It does this only for the level's own feature and its two neighbours, not for all L levels dyhead.py dyhead.py.
• Hard-sigmoid. The paper uses σ(x) = max(0, min(1, (x + 1) / 2)) §3.2. The code uses ReLU6(x + 3) / 6 dyrelu.py. The two differ for most inputs (they agree only at x = 0), so an implementer should follow the code if matching the released model.
• Level averaging in π_S. The paper divides by L eq. 4. The code averages the 2 or 3 neighbour candidates dyhead.py.
• Spatial sampling. The paper says the offsets and scalars come from the median-level feature §3.2. The code predicts a 27-channel map per level: 18 offset channels and 9 mask channels, with the mask passed through a sigmoid. It uses a 3 x 3 modulated deformable kernel, so K = 9 [code: dyhead/dyhead.py#L40] dyhead.py. The deformable operator itself takes the mask as input and does not apply the sigmoid deform.py.
• Resampling to the median level. The paper resizes all levels to the median level §3.1. The code does not. It samples from the previous level with a stride-2 convolution and from the next level with a stride-1 convolution followed by bilinear upsampling dyhead.py dyhead.py.
• Normalisation not in the paper. Each deformable convolution is followed by GroupNorm with 16 groups dyhead.py.
• Task-aware θ. The paper pools over L × S and describes two fully connected layers, a normalisation layer and a shifted sigmoid §3.2. The code pools each level over H × W only, and the module is applied to each level's output dyrelu.py dyhead.py. Its hyper-function is Linear, ReLU, Linear, hard-sigmoid, with no normalisation layer dyrelu.py. The slope coefficients are mapped to [-1, 1] as (h - 0.5) x 2 + init, and the offsets are shifted by -0.5 + init dyrelu.py dyrelu.py. The initial values are a = [1, 0] and b = [0, 0] dyrelu.py. The paper does not state these.
• Max over two lines. This matches eq. 5 dyrelu.py.
• Learning rate. The paper states 0.02 §4.1. The code's R50 1x ATSS config uses 0.01 with 16 images per batch [code: configs/dyhead_r50_atss_fpn_1x.yaml#L24-L25].
• Schedule. The paper's LR steps at 67% and 89% §4.1 match the code's R50 1x steps at 60000 and 80000 of 90000 iterations [code: configs/dyhead_r50_atss_fpn_1x.yaml#L26-L27].
• ATSS tower. The configs set the ATSS tower's own convolutions to 0 dyhead_r50_atss_fpn_1x.yaml, so the DyHead tower is the only tower. The ATSS default in extra/config.py is 4 [code: extra/config.py#L25].
• Swin. The appendix describes Swin-L, extra data and maximum side 2000 app.. The only Swin config in the repo is Swin-T, 2x, with AdamW, learning rate 1e-4, weight decay 0.05 and training minimum side 640 to 800 [code: configs/dyhead_swint_atss_fpn_2x_ms.yaml#L25-L33].
• Momentum and weight decay. The paper gives momentum 0.9 and weight decay 1e-4 §4.1. The code reads momentum from the solver config [code: train_net.py#L149], but its R50 config does not set momentum or weight decay, so the repo does not confirm the paper's values.

Pitfalls

• Equation 2 is ambiguous. It nests π_S inside π_L(F) · F and multiplies the result by F again. The paper does not say whether the sampled features in eq. 4 come from the level-gated tensor or the raw input. The code samples from the raw input dyhead.py. The pseudocode above follows the printed equations.
• Computational cost. The abstract says "without any computational overhead". Table 2 shows +21.50 GFLOPs at 6 blocks, which is about 8.5% of the baseline's 254.39 tab. 2. The text calls this negligible §4.2.
• Depth claim. The text says the head benefits from more blocks "until 8" §4.2. Table 2 peaks at 6 blocks (42.6 AP), with 8 blocks at 42.5 tab. 2.
• "More than 2.0 AP" over transformer detectors. Against RelationNet++ (50.3) the gain is exactly 2.0. Against Deformable DETR (50.1) it is 2.2 tab. 5 §4.3.
• "1/20 training time". The comparison uses iteration counts marked with "~" in Table 5 tab. 5 §4.3. The ratio is an estimate.
• Comparisons are not like for like. The 54.0 result uses test-time augmentation and a DCN backbone §4.3 tab. 5. The 60.6 result uses Swin-L, extra pseudo-labelled ImageNet data, a larger input and an unstated training recipe app. tab. 7. The 58.7 result without extra data is the fairer comparison with HTC++ tab. 7.
• Naming. The paper uses "ResNeXt-64x4d-101" in some tables and "ResNeXt-101-DCN" in the abstract tab. 4 tab. 5 §1.
• Unstated values. The paper does not state batch size, warm-up, epoch count, K, L, the DYReLU reduction, the ATSS top-k or the Swin optimiser. Code fills some of these; see Paper vs code.
• CUDA only. The deformable operator raises NotImplementedError for non-CUDA tensors deform.py. The repo builds its CUDA extension through setup.py setup.py.
• Ablations are not headline results. Ablations use ResNet-50 at 1x on val2017 §4.2. The headline numbers use other backbones, schedules and test-dev, so ablation gains do not carry over directly.
• No variance reported. The paper gives no seeds, repeat runs or confidence intervals for any table.

Lineage

• Builds on: feature pyramid networks, for the pyramid input §2 §3.1; ATSS, whose training framework the experiments use §4.1; Dynamic ReLU, whose hyper-function the task-aware attention follows §3.2; Deformable ConvNets v1 and v2, for the sparse sampling and modulation in π_S §3.2 §3.4.
• Relates to: non-local networks, Transformer-style detection heads (DETR, Deformable DETR, RelationNet++) §3.4.
• Compared against: Mask R-CNN, Cascade R-CNN, FCOS, RepPoints v2, BorderDet, SEPC, EfficientDet, SpineNet, CenterNet2, CopyPaste and HTC++ §4.3 tab. 4 tab. 5 tab. 7. Swin Transformer is the backbone for the Swin-T and Swin-L results app. tab. 6 tab. 7 (card: /papers/2103.14030).