Uncertainty-Aware Knowledge Distillation for Compact and Efficient 6DoF Pose Estimation

TL;DR

This work tackles efficient 6DoF pose estimation by introducing an uncertainty-aware end-to-end knowledge distillation framework. It couples Uncertainty-Aware KD at the prediction level with a Prediction-related Feature KD to transfer knowledge from large teacher models to compact students, leveraging deep ensembles to quantify keypoint uncertainties and an optimal-transport plan to align predictions and features. The approach yields state-of-the-art results on LINEMOD and SPEED+ while substantially reducing model size and computations. The findings demonstrate that incorporating uncertainty into KD improves both accuracy and robustness across diverse 6DoF pose estimation scenarios.

Abstract

Compact and efficient 6DoF object pose estimation is crucial in applications such as robotics, augmented reality, and space autonomous navigation systems, where lightweight models are critical for real-time accurate performance. This paper introduces a novel uncertainty-aware end-to-end Knowledge Distillation (KD) framework focused on keypoint-based 6DoF pose estimation. Keypoints predicted by a large teacher model exhibit varying levels of uncertainty that can be exploited within the distillation process to enhance the accuracy of the student model while ensuring its compactness. To this end, we propose a distillation strategy that aligns the student and teacher predictions by adjusting the knowledge transfer based on the uncertainty associated with each teacher keypoint prediction. Additionally, the proposed KD leverages this uncertainty-aware alignment of keypoints to transfer the knowledge at key locations of their respective feature maps. Experiments on the widely-used LINEMOD benchmark demonstrate the effectiveness of our method, achieving superior 6DoF object pose estimation with lightweight models compared to state-of-the-art approaches. Further validation on the SPEED+ dataset for spacecraft pose estimation highlights the robustness of our approach under diverse 6DoF pose estimation scenarios.

Figures (3)

• Figure 1: Uncertainty in Teacher Predictions. (a) Predictions from multiple teacher models show varied keypoint locations, represented by individual teacher predictions. (b) Keypoint predictions with high uncertainty values are visualized through clustered uncertainty markers.
• Figure 2: Overview of the Proposed Knowledge Distillation (KD) Framework (Best viewed in color). (A) General overview: Keypoints predicted by the ensemble of teachers are used to estimate uncertainties and are subsequently processed by UAKD, while averaged feature maps are directed to PFKD, guided by the OT plan $\boldsymbol{\pi}$. (B) UAKD Module: Keypoints, along with corresponding confidence scores $\boldsymbol{\alpha^{T,c}}$ and $\boldsymbol{\alpha^{S,c}}$, are aligned using an unbalanced Sinkhorn algorithm sejourne2023unbalanced, where $\bigotimes$ represent the tensor product. (C) PFKD Module: Predicted keypoints are mapped back to their respective regions in the feature maps and are aligned consistently according to the OT plan $\boldsymbol{\pi}$.
• Figure 3: Impact of $\lambda$ on the ADD-0.1d Metric using WDRnet on LINEMOD.