End-to-End Object Detection with Adaptive Clustering Transformer

In plain words

DETR finds objects by letting every position in the image look at every other position, so the cost grows fast as the image gets larger. Many of those positions end up looking at the rest of the image in almost the same way, so most of the comparisons repeat each other.

ACT exploits that. A cheap hashing trick sorts the positions into groups whose features are alike. The expensive comparison is then done once per group, using the group's average feature, and the result is copied back to every position in the group. Because the network's shape does not change, the trained DETR weights still fit, and a short round of distillation can recover much of the lost accuracy. The trade-off is a dial: more hashing rounds keep more detail and cost more compute.

TL;DR

• ACT is a drop-in replacement for the encoder self-attention of DETR §3 §3.1.
• It groups query vectors with E2LSH hashing and computes attention only between cluster prototypes and keys §3.2 eq. 1-6.
• Author-reported on COCO 2017 with DETR-DC5: 58.2 vs 73.4 GFLOPs (backbone excluded) for 0.7 AP loss, no retraining §4.4 tab. 1.
• Multi-task distillation reduces the loss to 0.2 AP, with a few epochs of fine-tuning §3.3 tab. 1.
• The official repository has the ACT attention and a DETR evaluation path, but no distillation training code, and the code differs from the paper in several places (see Paper vs code).

Problem

DETR's encoder self-attention compares every spatial position with every other, so its cost grows quadratically with the number of positions, which is what high-resolution inputs produce §1. The paper's diagnosis rests on two observations: positions with similar features and nearby locations produce similar attention maps (attention redundancy), and encoder features become more alike with depth, with the distribution differing per input and per layer (feature diversity) §1. Efficient transformer variants reduce this cost, but the paper notes that they change the architecture and need retraining, and that DETR itself is slow to converge, citing 500 epochs or about 1920 V100 GPU hours of training §1 §6. ACT aims to cut inference cost without retraining §1.

Method

ACT keeps DETR's structure: a ResNet backbone, a 6-layer encoder, a 6-layer decoder with learned object queries, and a shared prediction head §3.1. Only the encoder self-attention is replaced §4.2.

1. Hash the queries. Each query vector is hashed with L rounds of Exact Euclidean LSH (E2LSH). Each round is a random projection with a Gaussian vector, a uniform offset, and a bucket width r. The L round values are combined into one integer code by a base-4 sum §3.2 eq. 1-2.
2. Form prototypes. Queries with the same code form one cluster. Its prototype is the mean of its queries §3.2 eq. 3. The number of prototypes C follows from the data, so it differs per image and layer §5 fig. 7.
3. Attend from prototypes. A softmax is taken over the scaled prototype-key products, giving a C x M attention matrix, where M is the number of keys §3.2 eq. 4. It multiplies the values to give one output per prototype §3.2 eq. 5.
4. Broadcast. Each query receives the output of its prototype §3.2 eq. 6. The cost is quoted as O(NLD_K + CMD_k + CMD_v) against the exact O(NMD_k + NMD_v), where C and L are much smaller than N and M §3.2.

ACT is applied with the same weights as the pretrained DETR, so no retraining is needed §3.2 §1.

Multi-task knowledge distillation (MTKD) fine-tunes a model so that the ACT path and the original transformer agree §3.3 fig. 3. The loss adds to the original DETR loss against ground truth an L2 term that pulls the ACT box predictions towards the DETR box predictions, with the DETR side detached §3.3 eq. 7. The paper distils only the box regression branch, because it is more sensitive to the approximation error than classification §3.3.

# Adaptive clustering attention for one head  [§3.2] [eq. 1-6]
# Input:  queries (N, Dk), keys (M, Dk), values (M, Dv)
# Output: attention output (N, Dv)

def act_attention(queries, keys, values, r, L):
    # Hash each query with L rounds of E2LSH  [§3.2] [eq. 1]
    codes = zeros(N)
    for i in range(L):
        # Random direction and offset for this round  [§3.2] [eq. 1]
        a = sample_normal(Dk)                       # a ~ N(0, 1)
        b = sample_uniform(0, r)                    # b ~ U(0, r)
        h = floor((queries @ a + b) / r)            # (N,)
        # Combine rounds with base B = 4  [§3.2] [eq. 2]
        codes = codes + 4 ** i * h
    # Queries with equal codes share one cluster  [§3.2]
    cluster_ids, C = group_equal(codes)
    # Prototype is the mean query of its cluster  [§3.2] [eq. 3]
    prototypes = mean_per_cluster(queries, cluster_ids)   # (C, Dk)
    # Attention only between prototypes and keys  [§3.2] [eq. 4]
    scores = prototypes @ keys.T / sqrt(Dk)     # (C, M)
    A = softmax(scores, axis=-1)                # rows sum to 1
    # Prototype outputs  [§3.2] [eq. 5]
    W = A @ values                              # (C, Dv)
    # Each query copies the output of its prototype  [§3.2] [eq. 6]
    return W[cluster_ids]                       # (N, Dv)

# Multi-task knowledge distillation loss  [§3.3] [eq. 7]
# Input:  ground truth Y; ACT boxes B1; DETR boxes B2; DETR predictions Y2
# Output: scalar loss

def mtkd_loss(Y, B1, B2, Y2):
    # Original DETR loss against ground truth  [§3.3] [eq. 7]
    detr_loss = detr_loss_fn(Y2, Y)
    # L2 pull of ACT boxes towards DETR boxes; DETR side is detached  [§3.3]
    kd_loss = l2_distance(B1, stop_gradient(B2))
    # KD weight is 1  [§4.2]
    return detr_loss + 1.0 * kd_loss

Hyperparameters

Training schedule for the base DETR model: not stated in the paper; the paper uses pretrained DETR-DC5 weights §4.2.

Results

Author-reported. Setup: COCO 2017 detection §4.1. The baseline is DETR-DC5 with a dilated ResNet-50 §4.2. The evaluation split is stated only for the FLOPs and the ablations (validation set) §4.1 §4.3; for tab. 1 the paper names the metric (bbox AP) but not the split.

Whole-model FLOPs: 184.1 for DETR-DC5 and 168.9 for ACT §4.4. Both add the 110.7 backbone to the increments in the table; 168.9 equals 110.7 plus the L=32 increment of 58.2.

Clustering target ablation (similar FLOPs, AP as fraction) tab. 2:

Inference time and memory, batch 1 on a GTX TITAN X tab. 3:

Other author-reported findings: under the same compute, ACT beats K-means clustering with the same number of clusters in every layer §4.4 fig. 5; the attention-map MSE falls as L grows and as r falls, and above r=6 increasing r barely changes error or FLOPs §4.3 fig. 4; clusters correspond to coherent image regions such as a cow, the sky and a field §5 fig. 6; the number of prototypes per query falls with encoder depth §5 fig. 7.

Paper vs code

Repository: https://github.com/gaopengcuhk/SMCA-DETR, commit e4df3ab2081571a016c82794a85abd2342a2fc1b. The owner is an author of the paper (Peng Gao). The ACT code is in Adaptive_Cluster_Transformer/. The repository has no LICENSE file, so its license is not stated; we cite it and do not copy code. Code was read, not run: the CUDA extension was not built and no GPU run was done.

• Hash offset distribution. Paper: b ~ U(0, r) §3.2 eq. 1. Code draws the offset from U(0, 1) before dividing by r ada_clustering_attention.py. The effective offset is therefore r times narrower than the paper states.
• Bucket width. Paper: r is a hyperparameter swept in the ablation and fixed at 8 §4.2 §4.3. Code: r=8 is a default argument and is never passed from the command line ada_clustering_attention.py ada_clustering_attention.py.
• Hash rounds L do not reach the hash. Paper: L rounds set the FLOPs and the accuracy §3.2 §4.3. Code: the hashing always uses 32 random projections, because _compute_hashes is called with its defaults ada_clustering_attention.py. The --q_hashes value is passed on as n_hashes= to the clustering call ada_clustering_attention.py, which the Python wrapper defines with only hashes as argument [code: Adaptive_Cluster_Transformer/ACT/extensions/init.py#L7-L14]. As read, the group_Q path raises a TypeError; this includes the README evaluation command, which passes --group_Q README.md. Not run; this must be checked on a GPU build.
• Which vectors are clustered. Paper: queries Q in R^{N x D_k} §3.2. Code: the attention receives per-head tensors after the query projection and the 1/sqrt(head_dim) scaling, so each head is clustered separately on the scaled head-dimension queries attention_layer.py attention_layer.py ada_clustering_attention.py.
• Prototype count. Paper: C per layer, varying per image §3.2 §5. Code: the output has as many clusters as the largest cluster count over the whole batch and all heads; smaller ones are zero-padded [code: Adaptive_Cluster_Transformer/ACT/extensions/init.py#L13-L14]. The FLOP saving therefore depends on the batch.
• Base-4 code packing. Paper: sum_i 4^i h_i eq. 2. Code packs each hash into 2 bits by shifting and summing ada_clustering_attention.py. This is a unique code only if each h_i (after the minimum shift) lies in 0..3. The code does not check this, and the sum can overflow 64 bits at 32 rounds. Two different queries can therefore share a code and be clustered together.
• Resampled projections. Paper: a and b are random variables §3.2; whether they are fixed per model is not stated. Code: new Gaussian projections and offsets are drawn in every forward call, including evaluation ada_clustering_attention.py. Results depend on the random state.
• Key clustering. Paper: key clustering is an ablation with its own L §4.3 tab. 2. Code: when keys are clustered, the softmax is weighted by the key counts and renormalised ada_clustering_attention.py ada_clustering_attention.py. Equation 4 has no such weight.
• Default settings in the code. group_Q is on by default in the encoder transformer.py transformer.py. The DETR training defaults (learning rate 1e-4, weight decay 1e-4, 300 epochs, StepLR at 200) are in the shared training script main.py main.py. They are not the MTKD recipe in §4.2.
• MTKD is not in the code. The repository contains no distillation loss, no KD weight and no MTKD schedule. A search for distillation and detach terms in the Python files found only unrelated uses. Equation 7 and the schedule in §4.2 are paper-only.

The reported AP, FLOP and timing numbers were not checked against the code: the repository ships no logs or result files, and the FLOP counter is not in the repository.

Pitfalls

• Notation. The paper says V in R^{C x D_v} §3.2, but A is C x M and V must have M rows (one per key). Use M x D_v.
• Abstract complexity. The abstract's O(NK) leaves out the hashing term O(NLD_K) and the prototype-to-key cost O(CMD_k + CMD_v) §3.2. Hashing is not free, and it grows with L.
• Accuracy at low FLOPs. Without MTKD, ACT with L=16 loses 6.2 AP against DETR-DC5 (43.3 to 37.1) and loses 8.6 points on AP_S (22.5 to 13.9) tab. 1. Small objects bear most of the loss, which the paper notes §4.4.
• Text vs table. The text says L=16 with MTKD improves AP "by 4.3%" §4.4. Table 1 gives 40.6 minus 37.1, which is 3.5 points. The text also says ACT's AP_L is "0.3% higher" than DETR's §4.4; Table 1 shows equality without MTKD (61.1 vs 61.1) and +0.3 only with MTKD (61.4 vs 61.1).
• "Much stronger" for medium objects. The text says ACT is "much stronger" than Faster RCNN-DC5 on medium objects §4.4. Table 1 gives 46.8 vs 45.9 for ACT (L=32), a difference of 0.9 points.
• Two FLOP framings. The introduction quotes 73.4 to 58.2 GFLOPs excluding the backbone §1. Section 4.4 quotes 184.1 to 168.9 including it §4.4. Compare like with like.
• Table 2 units. AP in tab. 2 is a fraction (0.413), not a percentage as in tab. 1.
• Batch-level prototype count. The number of clusters is set by the largest count in the batch [code: Adaptive_Cluster_Transformer/ACT/extensions/init.py#L13-L14]. Measured FLOP savings at larger batches may be smaller than per-image counts suggest. The paper does not say which batch size its FLOP numbers use.
• Random projections. The hash projections change at every call ada_clustering_attention.py. Reproducing a number requires fixing the torch seed; the paper does not discuss the variance.
• Training with MTKD. The paper does not state which weights are trainable during MTKD, whether ACT's predicted boxes get any ground-truth loss (eq. 7 has none), whether the detached branch shares weights with the ACT path, or the batch size and weight decay §3.3 §4.2.
• Hyperparameter sweep. The choice of r=8 used 1000 training images for the attention-map MSE §4.2 §4.3. The choice of L is then a FLOPs knob, not an accuracy-tuned value §4.3.
• Timing. tab. 3 is for a single GPU, batch 1 and one model family. The paper calls the speed-up "consistent" with FLOPs §4.4; it does not report the clustering overhead on its own.

Lineage

• Builds on: DETR, arXiv:2005.12872 §1 §3.1; ACT keeps its architecture and weights and changes only the encoder self-attention §4.2.
• Uses: E2LSH, Datar et al. 2004 §3.2 eq. 1. Reformer also uses LSH for attention §2.
• Compared against: K-means clustering with a fixed number of clusters, as used by Vyas et al. §4.4; Faster RCNN-DC5 and DETR-DC5 tab. 1; Deformable DETR and SMCA are cited as training-side improvements but are not compared in the reported tables §2 §4.4.
• Superseded by: not stated.