Uncertainty-Aware Gradient Stabilization for Small Object Detection

TL;DR

This paper introduces Uncertainty-Aware Gradient Stabilization (UGS), a novel framework that stabilizes gradients in small object detection by reformulating localization as classification, leading to significant performance improvements across multiple benchmarks.

Contribution

UGS reformulates small object localization as classification, stabilizes gradients, and incorporates uncertainty minimization and refinement modules, advancing small object detection performance.

Findings

UGS improves small object detection accuracy across benchmarks.

UGS surpasses state-of-the-art results on VisDrone with 2.6 AP gain.

UGS enhances various detector architectures, including DINO-5scale.

Abstract

Despite advances in generic object detection, there remains a performance gap in detecting small objects compared to normal-scale objects. We reveal that conventional object localization methods suffer from gradient instability in small objects due to sharper loss curvature, leading to a convergence challenge. To address the issue, we propose Uncertainty-Aware Gradient Stabilization (UGS), a framework that reformulates object localization as a classification task to stabilize gradients. UGS quantizes continuous labels into interval non-uniform discrete representations. Under a classification-based objective, the localization branch generates bounded and confidence-driven gradients, mitigating instability. Furthermore, UGS integrates an uncertainty minimization (UM) loss that reduces prediction variance and an uncertainty-guided refinement (UR) module that identifies and refines high-uncertainty regions via perturbations. Evaluated on four benchmarks, UGS consistently improves anchor-based, anchor-free, and leading small object detectors. Especially, UGS enhances DINO-5scale by 2.6 AP on VisDrone, surpassing previous state-of-the-art results.