Source-linked AI summary
Bounding Box Regression with Uncertainty for Accurate Object Detection
Yihui He, Chenchen Zhu, Jianren Wang, Marios Savvides, Xiangyu Zhang
TL;DR
Ambiguous ground-truth boxes and imperfect correspondence between classification and localization confidence challenge bounding-box regression. The paper introduces KL Loss to learn box regression and localization variance jointly, then uses the variance for NMS refinement. Across CNN-based detectors, the method improves localization accuracy, including a VGG-16 Faster R-CNN AP increase from 23.6% to 29.1% on MS-COCO and AP90 gains for Mask R-CNN.
Problem
Ambiguous ground-truth boxes make bounding-box regression difficult, and classification confidence is not always aligned with localization confidence.
Method
KL Loss models predicted boxes as Gaussian distributions and ground truth as Dirac delta functions, learning regression and localization uncertainty together for variance-based NMS voting.
Results
The approach improves localization accuracy across CNN-based detectors; on MS-COCO, VGG-16 Faster R-CNN AP increases from 23.6% to 29.1%.
Takeaways & Limitations
Learned localization variance enables variance voting to refine selected boxes during NMS, complementing classification-based scoring.
Takeaways & Limitations
The method assumes independent coordinates and uses single-variate Gaussian distributions for simplicity.
Abstract
from arXiv · showhide
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
1. Introduction
The paper identifies ambiguous ground-truth boxes and misaligned classification and localization confidence as limitations of conventional bounding-box regression. It proposes KL Loss, which learns localization uncertainty alongside box regression and improves detection accuracy across architectures.
- Motivation: Ground-truth boxes can be inherently ambiguous, especially when objects are partially occluded, making labeling and regression learning difficult.Examples include inaccurately labeled MS-COCO boxes and unclear boundaries in YouTube-BoundingBoxes.
- Motivation: Smooth L1 regression ignores ground-truth ambiguity, while high classification scores do not always indicate accurate localization.Figure 2 illustrates inaccurate candidate coordinates and an inaccurate boundary on the higher-scoring box.
- Method: KL Loss learns bounding-box regression and localization uncertainty together by comparing a predicted Gaussian distribution with a ground-truth Dirac delta distribution.The learned distribution represents uncertainty for bounding-box coordinates.
- Method: The learned variance supports variance voting during NMS, using neighboring box locations weighted by predicted variances to refine selected boxes.This uses localization confidence in addition to classification confidence during post-processing.
- Results: 6.2% AP90 improvement is reported for ResNet-50-FPN Mask R-CNN, while VGG-16 Faster R-CNN AP rises from 23.6% to 29.1% on MS-COCO.The evaluation covers multiple CNN-based detectors on PASCAL VOC 2007 and MS-COCO, with only 2ms added inference latency for VGG-16 Faster R-CNN.
2. Related Work
The related work spans two-stage detection, detection losses, NMS, and bounding-box refinement. The paper distinguishes its approach by learning separate coordinate variances and using them to refine boxes during NMS.
- Two-stage Detectors: Two-stage detectors generate many duplicate proposals, and standard NMS can discard lower-scored boxes whose locations are more accurate.Variance voting instead uses neighboring boxes based on localization confidence.
- Object Detection Loss: Prior detection losses address overlap, class imbalance, global optimization, and uncertainty in other estimation tasks.The cited examples include IoU loss, Focal Loss, policy-gradient optimization, and uncertainty modeling for depth estimation.
- Non-Maximum Suppression: Soft-NMS reduces neighboring detection scores continuously rather than eliminating all lower-scored overlapping boxes.Learning NMS is another cited strategy for improving suppression results.
- Bounding Box Refinement: Earlier box-refinement methods merge boxes, learn box relations, or guide NMS with predicted IoU.IoU-Net predicts IoU, whereas this paper learns variances separately for the four box coordinates.
- Bounding Box Refinement: Variance voting sets a selected box’s new location from neighboring boxes’ KL-Loss-learned variances and can work with soft-NMS.The paper reports compatibility with soft-NMS in its experiments.
3. Approach
The approach jointly predicts bounding-box locations and localization uncertainty, trains them with KL divergence, and uses predicted variance to refine detections during NMS.
- 3.1. Bounding Box Parameterization: Bounding-box coordinates are regressed separately using the (x1, y1, x2, y2) representation rather than (x, y, w, h).Each coordinate is optimized independently in the subsequent formulation.
- 3.1. Bounding Box Parameterization: The method models predicted coordinates as independent univariate Gaussians, where σ measures localization uncertainty and σ → 0 denotes high confidence.The ground-truth box is represented by a Dirac delta function, corresponding to vanishing variance.
- 3.1. Bounding Box Parameterization: The detector predicts bounding-box locations and standard deviations, with uncertainty incorporated into the KL Loss regression objective.The architecture adds standard-deviation estimation to the fast R-CNN head, while classification loss remains unchanged.
- 3.2. Bounding Box Regression with KL Loss: KL divergence between predicted and ground-truth distributions becomes the regression loss, while larger predicted variance can reduce loss for inaccurate estimates.The implementation predicts α = log(σ2) to avoid exploding gradients and converts α back to σ during testing.
- 3.3. Variance Voting: Variance voting replaces the selected box location with a variance-weighted combination of neighboring boxes during NMS or soft-NMS.Weights decrease for boxes with high variance or small IoU with the selected box, and the procedure modifies NMS with three lines of code.
4. Experiments
Experiments across MS-COCO and PASCAL VOC 2007 show that KL Loss and variance voting improve object localization across detector architectures, with small computational overhead and compatibility with soft-NMS.
- 4.1. Ablation Study: 2.8% AP improvement comes from training with KL Loss alone, with additional gains observed for ResNet-50 Faster R-CNN and Mask R-CNN.The reported improvements are 1.5% for ResNet-50 Faster R-CNN and 0.9% for Mask R-CNN.
- 4.1. Ablation Study: 1.4% AP improvement is obtained by variance voting with standard NMS, while applying it after soft-NMS still improves AP by 1.3%.The paper characterizes soft-NMS as scoring candidates and variance voting as refining selected boxes for localization.
- 4.1. Ablation Study: 2ms additional GPU latency is required by the improved VGG-16 Faster R-CNN pipeline.The localization-confidence prediction uses a 4096×324 fully connected layer and does not require the separate 2mlp head used by IoUNet.
- 4.2. Accurate Object Detection: 1.8% AP and 6.2% AP90 improvements are obtained for ResNet-50-FPN Mask R-CNN using variance voting with soft-NMS.Variance voting further improves AP80 and AP90 by 0.4% and 1.2%, respectively, before the combined final gains.
- 4.3. Experiments on PASCAL VOC 2007: 2.9% mAP improvement is reported for VGG-16 on PASCAL VOC 2007 when the approach is combined with soft-NMS.For VGG-CNN-M-1024, training with variance voting improves mAP by 1.6%, and variance voting adds another 0.8%.
5. Conclusion
The paper proposes uncertainty-aware bounding box regression to improve object localization. KL Loss learns localization variance, which enables var voting to refine selected boxes.
- KL Loss learns more accurate object localization by modeling uncertainty during bounding box regression.
- The network predicts localization variance for each bounding-box coordinate.
- The resulting variances enable var voting to refine selected bounding boxes through voting.