[Paper Review] Bounding Box Regression with Uncertainty for Accurate Object Detection

Yihui He, Chenchen Zhu+3 more|arXiv (Cornell University)|Sep 23, 2018

Advanced Neural Network Applications63 references41 citations

TL;DR

The paper introduces KL Loss to jointly learn bounding box localization and per-coordinate uncertainty, plus a var voting scheme during NMS to improve localization while keeping computation nearly unchanged.

ABSTRACT

Large-scale object detection datasets (e.g., MS-COCO) try to define the ground truth bounding boxes as clear as possible. However, we observe that ambiguities are still introduced when labeling the bounding boxes. In this paper, we propose a novel bounding box regression loss for learning bounding box transformation and localization variance together. Our loss greatly improves the localization accuracies of various architectures with nearly no additional computation. The learned localization variance allows us to merge neighboring bounding boxes during non-maximum suppression (NMS), which further improves the localization performance. On MS-COCO, we boost the Average Precision (AP) of VGG-16 Faster R-CNN from 23.6% to 29.1%. More importantly, for ResNet-50-FPN Mask R-CNN, our method improves the AP and AP90 by 1.8% and 6.2% respectively, which significantly outperforms previous state-of-the-art bounding box refinement methods. Our code and models are available at: github.com/yihui-he/KL-Loss

Motivation & Objective

• Motivate how ground-truth bounding boxes are inherently ambiguous in large-scale datasets like MS-COCO.
• Propose a probabilistic bounding box regression loss (KL Loss) that learns localization and uncertainty.
• Enable better post-processing through variance-aware methods (var voting) during NMS.
• Demonstrate improvements across multiple detectors (VGG-16 Faster R-CNN, ResNet-50-FPN Mask R-CNN) on MS-COCO and PASCAL VOC 2007.
• Show that the learned uncertainties are interpretable and useful for downstream tasks.

Proposed method

• Model each bounding box coordinate as an independent Gaussian with mean x_e and variance σ^2.
• Ground-truth is modeled as a Dirac delta (σ -> 0).
• Define regression loss as KL divergence between P_D and P_Θ, leading to L_reg ∝ (x_g - x_e)^2/(2σ^2) + (1/2) log(σ^2).
• Parameterize σ via α = log(σ^2) to stabilize training, using L_reg ∝ e^{-α}(x_g - x_e)^2/2 + α/2.
• Adopt a robust partial loss for large errors when |x_g - x_e| > 1, similar to smooth L1.
• Predict per-coordinate σ to capture localization uncertainty.
• Introduce var voting during NMS, weighting neighboring boxes by IoU and their variances to refine the final box location.
• Provide an algorithmic variant of NMS that uses learned localization confidence to adjust coordinates.

Experimental results

Research questions

• RQ1 Can a probabilistic bounding box regression loss improve localization by accounting for annotation ambiguity?
• RQ2 Does learning per-coordinate localization uncertainty help refine detections in post-processing (NMS/soft-NMS) through variance-aware voting?
• RQ3 How does KL Loss interact with existing detection heads and post-processing techniques across datasets and backbones?
• RQ4 What is the empirical impact on AP and localization metrics (AP, AP90, etc.) when incorporating KL Loss and var voting?
• RQ5 Are the learned uncertainties interpretable and beneficial for downstream tasks requiring localization confidence?

Key findings

• KL Loss alone improves AP by 2.8% on VGG-16 Faster R-CNN (MS-COCO).
• KL Loss with var voting and soft-NMS yields the best MS-COCO results, achieving AP 29.1% on VGG-16 Faster R-CNN.
• On ResNet-50-FPN Mask R-CNN, KL Loss + var voting + soft-NMS improves AP by 1.8% and AP90 by 6.2%.
• Var voting refines localization by weighting neighboring boxes by predicted variances, improving high-IOU metrics (AP75, APM, APL).
• KL Loss consistently boosts AP across VGG-16, ResNet-50-FPN, and Mask R-CNN with minor latency increase (~2 ms on GTX 1080 Ti).
• On PASCAL VOC 2007, KL Loss improves mAP, and var voting provides further gains, particularly for higher overlap metrics like AP90.