[Paper Review] Probabilistic two-stage detection

Xingyi Zhou, Vladlen Koltun+1 more|arXiv (Cornell University)|Mar 12, 2021

Advanced Neural Network Applications59 references161 citations

TL;DR

The paper presents a probabilistic interpretation of two-stage object detection and builds probabilistic two-stage detectors by using a strong one-stage first stage to predict object likelihood, achieving higher accuracy and speed than prior two-stage and one-stage methods.

ABSTRACT

We develop a probabilistic interpretation of two-stage object detection. We show that this probabilistic interpretation motivates a number of common empirical training practices. It also suggests changes to two-stage detection pipelines. Specifically, the first stage should infer proper object-vs-background likelihoods, which should then inform the overall score of the detector. A standard region proposal network (RPN) cannot infer this likelihood sufficiently well, but many one-stage detectors can. We show how to build a probabilistic two-stage detector from any state-of-the-art one-stage detector. The resulting detectors are faster and more accurate than both their one- and two-stage precursors. Our detector achieves 56.4 mAP on COCO test-dev with single-scale testing, outperforming all published results. Using a lightweight backbone, our detector achieves 49.2 mAP on COCO at 33 fps on a Titan Xp, outperforming the popular YOLOv4 model.

Motivation & Objective

• Motivate a probabilistic interpretation of two-stage object detection to unify training signals across stages.
• Develop training objectives that optimize a joint probabilistic objective via tractable lower bounds.
• Design a first stage that provides calibrated object likelihoods while leveraging efficient one-stage detectors.
• Integrate the first-stage likelihood with a second-stage classifier to form a probabilistic two-stage detector.
• Demonstrate the approach across COCO, LVIS, and Objects365 with state-of-the-art results and real-time variants.

Proposed method

• Factorize the per-object class distribution into P(O) (object likelihood, first stage) and P(C|O) (conditional classification, second stage).
• Train using maximum likelihood with lower-bound relaxations to jointly optimize the two stages (Eq. 2, Eq. 3, Eq. 4).
• Use strong one-stage detectors as the first stage to predict calibrated object likelihoods (examples: RetinaNet-style, CenterNet with FPN, ATSS/GFL baselines).
• Combine first-stage P(O) with second-stage P(C|O) by multiplying the scores to form the final detection score.
• Experiment with FasterRCNN and CascadeRCNN as the second stage, built on diverse backbones (ResNet, ResNeXt, BiFPN, etc.).
• Tune hyperparameters (IoU thresholds, number of proposals, NMS) to suit probabilistic two-stage inference and maintain speed advantages.

Experimental results

Research questions

• RQ1 Can a probabilistic interpretation unify and improve training of two-stage detectors by making the first stage predict calibrated object likelihoods?
• RQ2 Does replacing the traditional RPN with a strong one-stage detector as the first stage improve both accuracy and speed when forming a probabilistic two-stage detector?
• RQ3 How does the probabilistic coupling of first- and second-stage scores affect final detection performance on COCO, LVIS, and Objects365?
• RQ4 What is the trade-off between number of proposals and accuracy in probabilistic two-stage detectors compared to conventional two-stage detectors?

Key findings

• A probabilistic two-stage detector outperforms both its one-stage and traditional two-stage precursors on COCO test-dev with single-scale testing (56.4 mAP).
• A CenterNet2-based configuration with strong backbones achieves 56.4 mAP, surpassing published results on COCO test-dev.
• Real-time CenterNet2 variants reach 49.2 mAP at 33 fps on Titan Xp, outperforming YOLOv4 on the same hardware.
• Using fewer, higher-quality proposals (256) with a calibrated first stage yields faster and more accurate detection than 1K proposals in traditional RPN-based two-stage detectors.
• Across LVIS and Objects365, CenterNet2 improves over CascadeRCNN baselines by 2.7–2.9 mAP, demonstrating robustness to large vocabulary datasets.