Technical Summary: From Keypoints to Predictive Distributions

Problem Statement

YOLO-Pose models are widely deployed for real-time keypoint localization due to their efficiency and single-pass inference capabilities. However, they output a single deterministic location for each keypoint and an instance-level confidence score that does not quantify keypoint-specific spatial uncertainty. In safety-critical downstream applications (e.g., sensor fusion, tracking), the inability to assess the reliability of individual keypoints limits the system's robustness.

Existing uncertainty quantification methods within the YOLO family primarily focus on object detection (bounding boxes) or require joint retraining of the model to predict uncertainty alongside localization. This joint-learning paradigm is often undesirable for deployed models where retraining is constrained by certification requirements or where changing the loss function might degrade the established mean prediction performance. Furthermore, standard losses like Object Keypoint Similarity (OKS) induce heavy-tailed residual distributions that are poorly captured by standard Gaussian assumptions, leading to miscalibrated uncertainty estimates.

Methodology

The authors propose a lightweight post-hoc probabilistic extension that augments a pre-trained, frozen YOLO-Pose model without altering its original keypoint predictions. The approach consists of three main components:

1. Architecture Extension

The method introduces additional probabilistic heads that mirror the original pose heads. These heads utilize the shared multi-resolution neck features to predict a heteroscedastic, input-dependent $2 \times 2$ dispersion matrix ($\Sigma$) for each keypoint. The matrix is constrained to be positive definite via Cholesky decomposition. The mean of the predictive distribution is fixed at the original YOLO prediction ($\hat{y}$), while the dispersion matrix captures the spatial uncertainty.

2. Training with Importance-Weighted NLL

To learn the dispersion structure, the authors employ a Gaussian Negative Log-Likelihood (GNLL) as a surrogate for the OKS loss. To address the mismatch between training-time label assignment and inference-time post-processing (e.g., confidence thresholding, NMS), they introduce an importance-weighted training objective:

• Retention Weight ($w_{ret}$): Based on the instance-level confidence score ($c_{det}^\alpha$), approximating the likelihood of a prediction being retained at inference.
• Scale Weight ($w_{scale}$): Based on the inverse object area ($1/s^2$), counteracting the tendency of larger objects to dominate the optimization signal and ensuring balanced calibration across object scales.
  The objective minimizes the weighted GNLL over matched prediction-ground-truth pairs.

3. Calibration Schemes

Recognizing that the OKS loss induces heavy-tailed residuals (due to its robust, outlier-suppressing nature), the authors propose two complementary calibration strategies applied to a held-out calibration set:

• Gaussian Calibration: Applies a temperature scaling ($\tau^2$) to the learned dispersion matrix. This yields a predictive distribution $\mathcal{N}(\hat{y}, \tau^2\Sigma)$, ensuring compatibility with downstream covariance-based methods (e.g., Kalman filters).
• Student-t Calibration: Applies both a temperature scaling ($\tau^2$) and a degrees-of-freedom parameter ($\nu$) to the dispersion matrix, treating it as a scale matrix for a Student-t distribution $T_\nu(\hat{y}, \tau^2\Sigma)$. This better captures the empirical heavy-tailed nature of the residuals while maintaining a principled covariance interpretation for $\nu > 2$.

4. Uncertainty-Based Pruning

The calibrated covariance is converted into a scalar localization confidence score (e.g., using the trace of the covariance). Keypoints with scores below a threshold (selected via Youden's J statistic on a validation set) are pruned before downstream processing, removing unreliable predictions while preserving the mean estimates of retained keypoints.

Evaluation Protocol

The paper introduces a comprehensive evaluation framework for heteroscedastic keypoint regression:

• Average Keypoint Precision (AKP): A COCO-style metric extended to the keypoint level. It evaluates whether the uncertainty scores effectively rank keypoints by correctness, enabling the selection or pruning of unreliable keypoints.
• Distributional Calibration Diagnostics:
  • Coverage: Evaluates if nominal probability regions (e.g., 95% ellipses) match empirical frequencies (Average Coverage Error, ACE).
  • Variance Calibration: Compares predicted variance magnitudes against empirical squared residuals across uncertainty levels (Expected Normalized Calibration Error, ENCE).
  • Q-Q Plots: Assesses the shape and tail behavior of the predictive distribution against the reference law.

Key Results

Experiments on the COCO dataset (using YOLO-NAS-Pose Nano and Large) demonstrate:

• Reliability Ranking: The learned uncertainty estimates provide effective keypoint-level reliability rankings, achieving high AKP scores significantly above random baselines.
• Calibration Performance:
  • Gaussian Calibration improves variance calibration (lower ENCE) but struggles with coverage (higher ACE) due to tail misspecification.
  • Student-t Calibration achieves the best overall performance, significantly reducing ACE (near-perfect coverage) while maintaining or improving ENCE and NLL. This confirms that the empirical residual distribution is better modeled by a heavy-tailed Student-t distribution.
• Pruning Efficacy: Uncertainty-based pruning effectively removes unreliable keypoints (particularly for ambiguous limbs like knees and ankles), improving the calibration metrics of the retained set.
• Ablation Studies: Importance weighting (specifically retention and scale weights) is critical for balanced calibration across object sizes. Full heteroscedastic covariance matrices outperform isotropic or diagonal approximations.

Application Demonstration: Vision-Based Aircraft Landing

The authors demonstrate the utility of their method in a vision-based landing (VBL) scenario using a large-scale internal dataset.

• Workflow: Runway keypoints are detected, and their calibrated 2D covariances are propagated through a covariance-weighted Perspective-n-Point (PnP) algorithm to estimate the aircraft's 3D position and covariance.
• Sensor Fusion: The resulting 3D position estimate and covariance are integrated into an Extended Kalman Filter (EKF) alongside other navigation sensors (GNSS, ILS, IRS).
• Findings: The Gaussian-calibrated model provides well-calibrated image-plane covariances that support uncertainty-aware position estimation. The method maintains the low-latency requirements of VBL.
• Limitation: In foggy conditions (out-of-distribution), the model fails to capture the increased epistemic uncertainty, highlighting that the current post-hoc approach models aleatoric uncertainty within the training distribution but does not inherently detect distribution shifts.

Significance and Claims

The paper claims to be the first extension of YOLO-Pose models that provides calibrated predictive distributions over keypoint localization errors without retraining the underlying model or modifying its original predictions.

Key contributions include:

1. Post-Hoc Flexibility: A method that decouples uncertainty estimation from point prediction, allowing for the calibration of existing, certified models.
2. Distributional Fidelity: The demonstration that Student-t calibration is necessary to capture the heavy-tailed residuals induced by robust OKS losses, outperforming standard Gaussian assumptions.
3. Evaluation Standard: The introduction of AKP and a suite of calibration diagnostics tailored for bivariate, heteroscedastic keypoint regression, addressing a gap in the literature where uncertainty is often evaluated only qualitatively or via scalar confidence.
4. Practical Utility: Validation in a safety-critical domain (aviation), showing that calibrated covariances are functional requirements for downstream sensor fusion and that uncertainty-based pruning can improve system robustness.

The authors maintain a modest stance regarding epistemic uncertainty, acknowledging that their current approach relies on in-distribution residuals and that handling out-of-distribution inputs (e.g., severe weather) requires future work in epistemic uncertainty modeling.