Localization Distillation for Dense Object Detection

In plain words

An object detector finds objects in a picture by guessing where each object's edges are. A small "student" detector can be trained to copy a large "teacher" detector, which is called knowledge distillation. Earlier work on detectors mostly had the student copy the teacher's internal feature maps. Those maps mix up what an object is with where its edges are, so it is unclear what the student actually learns.

This paper has the student copy the teacher's guesses about the edges instead. Each edge is treated as a spread of probabilities over many possible positions. A flat spread means the detector is unsure where that edge is. The student is trained to match the teacher's spread, which is a simple way to pass on the teacher's knowledge about localization. Classification is taught separately, in the usual way.

The authors also apply the edge lesson to a band of positions around each object, which they call the valuable localization region. The paper reports that this lifts a common single-scale detector from 40.1 to 42.1 AP on the COCO benchmark, with no change in its inference speed.

TL;DR

Localization distillation (LD) trains a dense-detector student to match the teacher's per-edge box distributions with a temperature-softened KL loss, while classification uses ordinary logit distillation. A valuable localization region (VLR) adds the LD loss to anchors just outside the positives. Author-reported gain: GFocal ResNet-50 at 1x goes from 40.1 to 42.1 AP on COCO val2017. The official code is an MMDetection fork under Apache-2.0. It differs from the paper in several places (see Paper vs code), so check those before implementing.

Problem

Localization is a core part of object detection. When an object's edges are not sharp, its box is hard to place, and small models suffer most §1 fig. 1. Knowledge distillation is a standard way to shrink a model, but for detection the classic logit-matching step was judged inefficient at transferring localization information. So most detection-distillation methods instead imitate deep features in chosen regions §1 §2. Because semantic and localization knowledge are mixed on those feature maps, the paper argues it is hard to tell which regions and which knowledge help §1.

The method assumes a teacher and a student from the same dense-detector family whose box-regression head can be switched to a discrete probability distribution per edge §3.1 §3.2.

Method

The method splits distillation in two. Classification is distilled from the classification-head logits with standard logit KD §3.2. Localization is distilled from the localization-head logits with LD. Both act on logits, not on feature maps §3. The regions where each term applies are chosen separately (divide and conquer, §3).

1. Box as a distribution. Each edge (t, b, l, r) is written as the expectation of a coordinate x under a probability Pr(x) over its range §3.1 eq. 1. The range is quantised into n uniform bins, and a softmax over n logits gives Pr §3.1. The paper says the (t, b, l, r) form and the (x, y, w, h) form perform the same, and it converts to (t, b, l, r) §3.2. The number n is not stated in the paper.

   \hat{e}=\int_{e_{\min}}^{e_{\max}} x \Pr(x)\,dx eq. 1

2. Temperature. For each edge, the student logits z_S and the teacher logits z_T are softened as S(z, tau) = softmax(z / tau) §3.2. At tau = 1 this is ordinary softmax. tau tending to 0 gives a Dirac delta, and tau tending to infinity gives a uniform distribution. The paper uses tau > 1 to keep more information in the distribution §3.2.

3. LD loss per edge. The KL divergence between the two softened distributions is the LD loss for one edge §3.2 eq. 2 eq. 3. The written form puts the student first.

   \mathcal{L}_\text{LD}^e = \mathcal{L}_\text{KL}(\bm{p}_S^\tau,\bm{p}_T^\tau) = \mathcal{L}_\text{KL}\left(\mathcal{S}(\bm{z}_S, \tau), \mathcal{S}(\bm{z}_T, \tau)\right) eq. 2 eq. 3

   The LD loss for a box sums the four edges §3.2 eq. 4.

4. Main distillation region. This is the set of locations that are positive under the detector's label assignment §3.3.

5. Valuable localization region (VLR). For each FPN level, compute the DIoU between every anchor and every ground-truth box. Set the lower bound to alpha_vl = gamma * alpha_pos, where alpha_pos is the positive IoU threshold of label assignment. The VLR is the set of anchors whose DIoU lies in [alpha_vl, alpha_pos], as a 0/1 mask §3.3 alg. 1. gamma is the only parameter of the VLR. gamma = 0 gives the full window, and gamma tending to 1 empties it §3.3. DIoU is used because it favours locations near an object's centre §3.3. For anchor-free detectors the preset grid anchors are used, so the regression form stays anchor-free. For anchor-based detectors the anchors are unfolded before the DIoU matrix is computed §3.3.

6. Total loss. The student loss sums the classification, GIoU regression and DFL terms as usual, plus four distillation terms: LD and KD, each on the main region and on the VLR §3.4 eq. 5. Each term has an indicator for its region. All distillation losses of one type share one weight, so LD follows the box-regression weight and KD follows the classification weight §3.4. The paper gives no values for the lambda weights. DFL can be switched off, because the LD loss gives enough guidance §3.4. Each of the four distillation types can be switched on or off §3.4.

   \mathcal{L}=\lambda_0\mathcal{L}_\text{cls}+\lambda_1\mathcal{L}_\text{reg} +\lambda_2\mathcal{L}_\text{DFL} +\lambda_3\mathbb{I}_\text{Main}\mathcal{L}_\text{LD} +\lambda_4\mathbb{I}_\text{VL}\mathcal{L}_\text{LD} +\lambda_5\mathbb{I}_\text{Main}\mathcal{L}_\text{KD} +\lambda_6\mathbb{I}_\text{VL}\mathcal{L}_\text{KD} eq. 5

7. Main configuration. The paper's main model uses Main KD, Main LD and VLR LD, and leaves out VLR KD §4.2 tab. 2. The paper calls this the separate distillation region manner.

Algorithm 1 gives the VLR for one level alg. 1:

# Valuable localization region, one FPN level  [alg. 1]
# Input:  anchors (I, 4); gt_boxes (J, 4); a_pos: IoU threshold; gamma
# Output: vlr (I, J), True = in region, False = ignored

def valuable_localization_region(anchors, gt_boxes, a_pos, gamma):
    # DIoU between every anchor and every ground-truth box  [alg. 1]
    diou = diou_matrix(anchors, gt_boxes)            # (I, J)
    # Lower bound of the window: gamma times the positive threshold  [§3.3]
    a_vl = gamma * a_pos
    # Anchors whose DIoU lies inside [a_vl, a_pos] form the region  [alg. 1]
    vlr = (diou >= a_vl) & (diou <= a_pos)
    return vlr

The LD loss for one image follows §3.2 eq. 2 eq. 3 eq. 4:

# Localization distillation for one image  [§3.2] [eq. 2-4]
# Input:  z_s, z_t: student and teacher edge logits, (N, 4, n)
#         main: bool (N,), positive locations
#         vlr_weight: float (N,), 1 for region anchors  [§3.3]
# Output: scalar LD loss, before the lambda weights of [eq. 5]

def ld_loss(z_s, z_t, main, vlr_weight, tau):
    # Soften each edge's logits with temperature tau  [§3.2]
    p_s = softmax(z_s / tau, axis=-1)
    p_t = softmax(z_t / tau, axis=-1)
    # Per-edge KL in the written order, student first  [eq. 2] [eq. 3]
    kl = sum(p_s * (log(p_s) - log(p_t)), axis=-1)   # (N, 4)
    # Main term: sum over the four edges of positive locations  [eq. 4]
    loss_main = sum(kl[main])
    # VLR term: same KL on region anchors  [§3.3]
    loss_vlr = sum(vlr_weight[:, None] * kl)
    return loss_main + loss_vlr

Hyperparameters

Values the paper gives are cited to the paper. Values taken from the released code are labelled as such. Where the paper and code differ, Paper vs code gives the details.

Results

Author-reported. Setup: COCO train2017 (118K images) for training, val2017 (5K) for ablations, test-dev 2019 for some SOTA rows, with a ResNet-101 teacher and a student from the GFocal (GFL) family §4.1 §4.2 tab. 1. Hardware: 8 GPUs, 2 images each app. A1. The baseline is the same detector retrained with the same settings §4.1.

Headline and ablations (R101 teacher to R50 student, 1x, val2017):

Logit mimicking versus feature imitation (same setting, val2017):

Sweeps (main region only unless stated):

Other students and detectors (teacher R101, val2017, 1x):

State of the art (author-reported, LD applied on top of GFocalV2):

Table 6 also lists nine other detectors (RetinaNet, FCOS, SAPD, ATSS, BorderDet, AutoAssign, PAA, OTA, GFocal) at several settings; they are not repeated here.

Appendix results:

• Individual LD (R18 student, main region, no GT box losses): reg. + DFL alone 35.8 AP / 38.2 AP75; LD alone 36.4 / 39.3; reg. + DFL + LD 36.5 / 39.3 app. A2 tab. A1 in the paper.
• Self-LD (teacher = student, main region, applied once): R18 35.8 to 36.1 AP; R50 40.1 to 40.6 AP; ResNeXt-101-32x4d-DCN 46.9 to 47.5 AP app. A2 tab. A2.
• Cost (1333x800, RTX 2080 Ti, FPS averaged over 3 runs) app. A2 tab. A3: RetinaNet 37.74M to 39.07M parameters, 239.32G to 267.64G FLOPs, 21.0 to 19.6 FPS, AP 36.9 to 39.0. FCOS 32.02M to 32.17M, 200.50G to 203.60G, 22.3 to 22.4 FPS, AP 38.6 to 40.6. ATSS 32.07M to 32.22M, 205.21G to 208.36G, 21.9 to 21.9 FPS, AP 39.2 to 41.6. GFocal 32.22M to 32.22M, 208.31G to 208.31G, 21.9 to 21.9 FPS, AP 40.1 to 42.1.
• PASCAL VOC 07+12 train, VOC 2007 test, R18 student, main region only: baseline 51.8 AP; with a ResNet-34 teacher 52.9; with ResNet-50 52.8; with ResNet-101 53.0; with ResNet-101-DCN 52.8 app. A2 tab. A4.

Paper vs code

The official repository is https://github.com/HikariTJU/LD, pinned at commit 4287990ca27f51f6f15dc66d28bb420a053c307e. It is a fork of MMDetection 2.x and carries MMDetection's Apache-2.0 license file LICENSE. The LD code lives in mmdet/models/dense_heads/ld_head.py, mmdet/models/losses/kd_loss.py and the ATSS assigner, and the configurations are in configs/ld/. Each item below was checked against the code. Discrepancies were found, as listed below.

1. VLR weight is IoU, not a 0/1 mask. The paper's algorithm returns a 0/1 mask alg. 1. The code weights the VLR LD loss of each anchor by the anchor's IoU, the value stored by the assigner ld_head.py atss_assigner.py.
2. The VLR window is adaptive, and gamma is hard-coded. The code does not use a fixed alpha_pos. For each GT it takes the mean plus standard deviation of the IoU of the 9 nearest anchors on each FPN level, pooled over levels atss_assigner.py atss_assigner.py. It then keeps anchors whose DIoU lies in [0.25 x threshold, threshold) atss_assigner.py. The factor 0.25 is a literal in the code, so the gamma sweep in tab. 1 cannot be rerun from the released code without editing it.
3. Main positives are not removed from the VLR. The code applies the VLR LD loss to every anchor with a nonzero region value, whether or not it is also positive ld_head.py ld_head.py. So an anchor can receive both LD terms. Eq. 5 with its indicators allows this, but the text does not say so.
4. KL direction and scale. Eq. 2 writes KL(p_S, p_T) with the student first. The code uses the teacher as the target of F.kl_div, so it computes KL(p_T || p_S) kd_loss.py kd_loss.py. The code also averages over the bins (.mean(1)) and multiplies by T^2 kd_loss.py. The paper has neither factor.
5. Student-confidence weight on LD. The code weights the LD and DFL terms of each positive by the student's detached maximum class sigmoid score ld_head.py ld_head.py. The paper does not mention this.
6. Fixed normalisers. LD on the main region is divided by a constant 4.0 per level, and VLR LD by a constant 16.0 per level, regardless of how many anchors are in the region ld_head.py ld_head.py. KD is divided by the number of positives on the level ld_head.py. GIoU and DFL are divided by the sum of the positive weights over all levels ld_head.py.
7. VLR KD is disabled. The VLR KD term is computed but multiplied by 0 ld_head.py. This matches the paper's main model §4.2, but the VLR KD ablation cannot be reproduced without edits.
8. Feature imitation is on in the R50 config. The released R50 config ld_r50_gflv1_r101_fpn_coco_1x.py sets an IMLoss weight of 2.0 and imitation_method='finegrained' ld_r50_gflv1_r101_fpn_coco_1x.py ld_r50_gflv1_r101_fpn_coco_1x.py. The fine-grained branch applies MSE on FPN features at anchors whose IoU with a GT exceeds half that GT's maximum IoU ld_head.py ld_head.py. The paper's 42.1 result ("Ours") excludes feature imitation; "Ours + Fine-Grained" is 42.4 tab. 3. The repository does not say which setting produced 42.1.
9. Learning rate. The paper states 0.01 for 8 GPUs with 2 images each app. A1. The released R50 config uses 0.0025 with the same batch size ld_r50_gflv1_r101_fpn_coco_1x.py. The README formula lr = (samples_per_gpu * num_gpu) / 16 * 0.01 gives 0.01 for this setting README.md. The ATSS R50 config uses 0.01 ld_r50_atss_r101_1x.py.
10. Anchor-free versus ATSS. The paper's experiments use "the FCOS-style anchor-free head" §4.1. The released GFL configs use one square anchor per location and the ATSS label assigner with top-k 9 ld_r18_gflv1_r101_fpn_coco_1x.py ld_r18_gflv1_r101_fpn_coco_1x.py.
11. Classification KD values. The paper gives no temperature or weight for classification KD. The code uses T = 2 and weight 10 ld_r50_gflv1_r101_fpn_coco_1x.py.
12. SOTA configs. The GFocalV2-based config in the repository is 1x only and enables gradient-imitation (gibox) with weight 2 ld_r50_gflv2_r101_fpn_1x.py. The 2x and multi-scale configs in configs/im/ and configs/imv2/ also set an IMLoss term (weights 2, 2 and 1.0) im_gflv2_r2n101_dcn_fpn_2x.py im_gflv2_x101-32x4dr2n101_dcn_2x.py. I did not trace whether those configs reproduce the Table 6 numbers.
13. Other details the paper does not give, which the code does. The teacher runs in eval mode under no_grad on every batch kd_one_stage.py kd_one_stage.py. The number of bins is 17 (reg_max 16) ld_r18_gflv1_r101_fpn_coco_1x.py.

Pitfalls

• Loss scale is not in the paper. The paper gives no lambda values and no T^2 or bin-average factor. With the code's values, the LD term is 0.25 x T^2 (= 25) times a mean over 17 bins, which is about 1.47 times the plain sum over bins. Check the scale against the code, not against eq. 4.
• KD on the main region only. In the code the classification KD term only covers positives ld_head.py.
• Speed claim depends on the head. The paper's inference-speed claim holds for GFocal in its table, where FPS stays at 21.9 app. A2. For RetinaNet, FPS falls from 21.0 to 19.6 because it has 9 anchors per location app. A2. The output of the localization head grows from 4 to 4 x (n) channels, which is 68 with n = 17 ld_r18_gflv1_r101_fpn_coco_1x.py.
• Training cost. Every training step runs a teacher forward pass on the same image kd_one_stage.py. The paper reports no training-time or memory figures.
• Batch and learning rate. Follow the paper's stated setting (2 images per GPU on 8 GPUs, lr 0.01). The R50 config's 0.0025 disagrees with it (Paper vs code, item 9).
• GI path needs CUDA. The gradient-imitation branch hard-codes .cuda() ld_head.py. Only the GFocalV2 configs use it.
• Environment. The README recommends PyTorch 1.7 and CUDA 11 README.md. The code depends on MMDetection 2.x APIs, so an unmodified newer MMDetection will not run it.
• Teacher checkpoints are external. Configs point to MMDetection checkpoint URLs ld_r18_gflv1_r101_fpn_coco_1x.py. The teacher AP (44.7) is given only in the paper §4.2.

Lineage

• Builds on: GFocal, whose probability-distribution box representation and DFL term this paper reuses as its baseline; logit mimicking (Hinton et al.) for the classification term; FCOS-style and ATSS-style detectors for the head and assigner. These are cited by name in the paper and are not compiled here.
• Compared against: feature-imitation methods (FitNets, Fine-Grained, DeFeat, GI imitation) tab. 3; the teacher-bounded regression (TBR) loss §4.2 tab. 1; GFocalV2 and other single-stage detectors in tab. 6.
• Extended by: the README names rotated-object ports (Rotated-LD in mmRotate and JDet). Those are not part of this card's evidence.