{'fraction_non_alphanumeric': 0.03284324961524608, 'fraction_numerical': 0.013163495857755657, 'mean_word_length': 4.90715667311412, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 1, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Knowledge Distillation for 6D Pose Estimation by Keypoint Distribution Alignment

Knowledge distillation facilitates the training of a compact student network by using a deep teacher one. While this has achieved great success in many tasks, it remains completely unstudied for image-based 6D object pose estimation. In this work, we introduce the ﬁrst knowledge distillation method for 6D pose estimation. Speciﬁcally, we follow a standard approach to 6D pose estimation, consisting of predicting the 2D image locations of object keypoints. In this context, we observe the compact student network to struggle predicting precise 2D keypoint locations. Therefore, to address this, instead of training the student with keypoint-to-keypoint supervision, we introduce a strategy based the optimal transport theory that distills the teacher’s keypoint distribution into the student network, facilitating its training. Our experiments on several benchmarks show that our distillation method yields state-of-the-art results with different compact student models.


Introduction
Estimating the 3D position and 3D orientation, a.k.a. 6D pose, of an object relative to the camera from a single image has a longstanding history in computer vision, with many real-world applications, such as robotics, autonomous navigation, and virtual and augmented reality. While modern methods [42,30,22,17,19,25,40,6] now all rely on deep networks to address this task, the most effective ones draw their inspiration from the traditional approach, which consists of establishing correspondences between the object's 3D model and the input image and compute the 6D pose from these correspondences using a Perspective-n-Point (PnP) algorithm. In particular, the state-of-the-art methods, GDR-Net [40] and SO-Pose [6], achieve this by combining an object detection module, a large encoder-decoder to predict multiple intermediate dense representations, including 2D-3D correspondences, and a learnable PnP module to output the final pose. Unfortunately, the intermediate dense feature maps they extract incur many parameters, thus yielding huge models that are impractical for deployment on embedded platforms and edge devices. In principle, one could reduce the models' size by employing smaller backbones. However, as shown in Figure 1 for SO-Pose, the resulting accuracy drops quickly.
By contrast, we have observed keypoint-based methods [17,30,18,19] to yield better robustness to the use of lightweight backbones. This is illustrated in Figure 1 for WDRNet [19] whose performance undergoes a much smaller drop than that of SO-Pose as the number of parameter decreases. We believe this to be due to the fact that these methods aggregate multiple votes, originating from the individual feature-map locations containing the object, for the 2D image positions of a sparse set of 3D keypoints, as shown in Figure 2 where the keypoints are the 8 corners of the 3D object bounding box. Nevertheless, the resulting compact architectures still incur a significant drop in pose accuracy.  We show the SOTA dense prediction-based SO-Pose (orange diamond) and the keypoint-based WDRNet+ 1 (blue circles). Although SO-Pose outperforms WDRNet+ with a large backbone, its performance drops more rapidly than that of WDR-Net+ as the number of parameters decreases. Our KD method (stars) further boosts the performance of WDRNet+ with lightweight backbones.  The large backbone of the teacher allows it to produce accurate keypoints, indicated by tight clusters. By contrast, because of its more compact backbone, the student struggles to predict accurate keypoints when trained with keypoint-to-keypoint supervision.
We therefore propose to align the student's and teacher's keypoint distributions.
In this paper, we address this by introducing a knowledge distillation strategy for such keypoint-based 6D pose estimation networks.
Knowledge distillation aims to transfer some information from a deep teacher network to a compact student one. The research on this topic has tackled diverse tasks, such as image classification [15,43,34], object detection [44,9,8] and semantic segmentation [27,12]. While some techniques, such as feature distillation [34,44,43,13], can in principle generalize to other tasks, no prior work has studied knowledge distillation in the context of 6D pose estimation.
In this paper, we introduce a knowledge distillation method for 6D pose estimation motivated by the following observation. As shown in Figure 2, predicting accurate keypoint locations with keypointto-keypoint supervision is much harder for a compact student network than for a deep teacher one. We therefore argue that knowledge distillation for 6D pose estimation should not be performed by matching the individual student and teacher keypoints but instead by encouraging the student's keypoint distribution to become similar to the teacher one, which leaves more flexibility to the student and thus facilitates its training. To achieve this, we follow an Optimal Transport (OT) formalism [39], which lets us measure the distance between the two keypoint sets. We express this as a loss function that can be minimized using a weight-based variant of Sinkhorn's algorithm [5], which further allows us to exploit predicted object segmentation scores in the distillation process. Our strategy is invariant to the order and the number of predicted keypoints within a cluster, making it applicable to unbalanced teacher and student predictions that are not in one-to-one correspondence.
We validate the effectiveness of our approach by conducting extensive experiments on the popular LINEMOD [14] and Occluded-LINEMOD [2] datasets. Our keypoint distribution alignment strategy consistently outperforms both a keypoint-to-keypoint distillation baseline and the state-of-the-art feature distillation method [44] using diverse lightweight backbones and architecture variations. Interestingly, our approach is orthogonal to feature distillation, and we show that combining it with the state-of-the-art approach of [44] further boosts the performance of student network.
Our main contributions can be summarized as follows. (i) We investigate for the first time knowledge distillation in the context of 6D pose estimation. (ii) We introduce an approach that aligns the teacher and student keypoint distributions together with their predicted object segmentation scores.
(iii) Our approach can be used in conjunction with feature distillation to further boost the student's performance. We will make our code publicly available.
2 Related Work 6D pose estimation. With the great development and success of deep learning in computer vision [23,10,26,11,33,28], many works have explored its use for 6D pose estimation. The first attempts [42,22,21] aimed to directly regress the 6D pose from the input RGB image. However, the representation gap between the 2D image and 3D rotation and translation made this task difficult, resulting in limited success. Therefore, most methods currently predict quantities that are closer to the input image space. In particular, the state-of-the-art approaches [25,40,6] output dense correspondences between the input image and the object 3D model, typically by predicting a 3D coordinate at every input location containing an object of interest. A consequence of the resulting accuracy boost, however, is a significant increase in the number of learnable parameters, arising from the use of larger encoder-decoder architectures to obtain dense intermediate representations in the two best-performing frameworks, namely GDRNet [40] and SO-Pose [6]. Although we have tried to reduce the size of these models, as shown in Figure 1 for SO-Pose, we have observed a rapid performance drop, even when retaining a relatively large number of parameters. By contrast, we have found that networks predicting a sparser set of 2D-to-3D correspondences, such as WDRNet+ [19] in Figure 1, offered a better robustness to the use of lightweight backbones. In essence, these methods jointly segment the object by classifying learned local features and predict either the image locations [17][18][19] or the 2D displacements from the cells' center [30] or 3D object keypoints, typically taken as the corners of the object bounding box. Nevertheless, the original backbones used by these methods remain cumbersome, and reducing them yields a performance drop. Here, we address this by introducing a knowledge distillation for keypoint-based 6D pose estimation.
Knowledge distillation. Knowledge distillation has been proven effective to transfer information from a deep teacher to a shallow student in several tasks. This trend was initiated in the context of image classification, where Hinton et al. [15] guide the student's output using the teacher's class probability distributions, and Romero et al. [34], Zagoruyko and Komodakis [43], Tian et al. [38] encourage the student's intermediate feature representations to mimic the teacher's ones. Recently, many works have investigated knowledge distillation for other tasks, evidencing the benefits of extracting task-driven knowledge. For example, in object detection, Zhang and Ma [44] adapt the feature distillation strategy of [34] to object detectors; Wang et al. [41] restrict the teacher-student feature imitation to regions around the positive anchors; Guo et al. [8] decouple the intermediate features and the classification predictions of the positive and negative regions; Guo et al. [9] distill detection-related knowledge from a classification teacher to a detection student. In semantic segmentation, Liu et al. [27] construct pairwise and holistic segmentation-structured knowledge to transfer. All of these works evidence that task-driven knowledge distillation boosts the performance of compact student models. Note that the feature distillation strategy of Zhang and Ma [44], although developed for object detection, is general and can be applied to 6D pose estimation. However, such a general distillation method does not leverage the specific properties of 6D pose estimation. Here, we introduce an approach that does so, and show that it outperforms the general FKD method [44] and can be combined with it to obtain a further performance boost.
Optimal transport (OT) has received a growing attention both from a theoretical perspective [39,5,35] and for specific tasks, including shape matching [37], generative modeling [1], domain adaptation [4], and model fusion [36]. In particular, OT has the advantage of providing a sound and theoretically grounded way of comparing multivariate probability distributions without the need for approximating them with parametric models. Furthermore, by considering the geometry of the underlying space through a cost metric, it can encode useful information about the nature of the problem. Our work constitutes the first attempt at using OT to align student and teacher keypoint distributions for knowledge distillation in 6D pose estimation.

Methodology
Let us now introduce our method to knowledge distillation for 6D pose estimation. As discussed above, we follow a keypoint-based approach [17,30,18,19] to 6D pose estimation. Given a single PnP Solver

< l a t e x i t s h a 1 _ b a s e 6 4 = " m z x t 4 + f x t 4 A o P D T u b J b a I v J g + F I = " > A A
< l a t e x i t s h a 1 _ b a s e 6 4 = "  Figure 3: Overview of our method (better viewed in color). The teacher and student follow a segmentation-driven approach to 6D pose estimation, which, given an RGB input image, outputs both a segmentation score map by classifying the individual cells in the feature map and 2D keypoint locations, with each cell voting for the locations of the 8 corners of the 3D object bounding box. The keypoint locations with the bounding box corners form correspondences, which are passed to a RANSAC-based PnP solver [24] or a simple PnP network [18] to obtain the final 3D translation and 3D rotation. Instead of performing keypoint-to-keypoint distillation, we propose an optimal transport-based strategy that lets us jointly distill the teacher's keypoint distribution and segmentation score map into the student.
RGB image captured by a calibrated camera, this approach aims to predict the 2D locations of 3D object keypoints in the image plane. This is illustrated in Figure 3 for keypoints taken as the 8 corners of the 3D object bounding box. More precisely, the deep network classifies each local cell, i.e., anchor, in the feature map to obtain a rough object segmentation, and predict the 2D keypoint locations for each active anchor, i.e., each anchor classified as belonging to the object. As such, each bounding box corner corresponds to a cluster of 2D locations in Figure 2. These locations then form 2D-to-3D correspondences that can act as input to a RANSAC-based PnP solver, e.g., [24] or to a shallow PnP network [18] to estimate the 6D pose.
In essence, the key to the success of such a correspondence-based approach is the prediction of accurate 2D keypoint locations. However, as shown in Figure 2, the predictions of a shallow student network differ significantly from those of a deep teacher one; they are less concentrated around the true keypoint locations, and thus yield less accurate 6D poses. Below, we first present a naive distillation strategy to addressing this, and then introduce our approach.

Naive Keypoint-to-keypoint Distillation
The most straightforward way of performing knowledge distillation is to encourage the student's predictions to match those of the teacher. In our context, one could therefore think of minimizing the distance between the locations predicted by the teacher and those predicted by the student. Formally, with N anchors in the feature map, this can be expressed as where P s ik , resp. P t ik , represent the student's, resp. teacher's, 2D prediction for keypoint k in anchor i, λ i ∈ {0, 1} indicates whether the anchor is active or not, and p ∈ {1, 2}.
One drawback of this strategy comes from the fact that the teacher and student network may disagree on which anchors are active and which are not, as illustrated in Figure 3. In other words, the λ i s may be taken from either the student, the teacher, or the intersection of their segmentation results. However, in any case, the distillation may be suboptimal, as some student's predictions would be either unsupervised by the teacher or supervised by potentially unconfident teacher predictions. Furthermore, and as argued above, a compact student tends to struggle when trained with keypoint-tokeypoint supervision, and such a naive KD formulation still follows this approach. Therefore, and as will be shown in our experiments, this naive strategy does not outperform the direct student training. Below, we therefore introduce a better-suited approach to keypoint-based knowledge distillation.

Keypoint Distribution Alignment
As discussed above and illustrated in Figure 3, the number of active student anchors N s may differ from that of active teacher anchors N t , making a direct match between the individual teacher and student predictions ill-suited. To address this, and account for the observation that keypoint-tokeypoint supervision is ill-suited to train the student, we propose to align the distributions of the teacher and student keypoint predictions. Specifically, we achieve this using optimal transport, which can handle the case where N s = N t .
Formally, to allow the number of student and teacher anchors to differ, we leverage Kantorovich's relaxation [20] of the transportation problem. In our context, assuming that all the keypoints have the same probability mass, i.e., 1 N t for the teacher predictions and 1 N s for the student ones, we derive a distillation loss based on Kantorovich's optimal transport problem as Note that we consider separate costs for the individual keypoints, thus preventing a 2D location corresponding to one particular corner to be assigned to a different corner. In our experiments, we found p = 2 to be more effective than p = 1 and thus the 2 norm below.
The above formulation treats all keypoint predictions equally. However, different predictions coming from different anchors might not have the same degree of confidence. In particular, this can be reflected by how confident the network is that a particular anchor contains the object of interest, or, in other words, by the segmentation score predicted by the network. Let α s i denote such a score for anchor i in the student network, and α t j a similar score for anchor j in the teacher network. 2 We then re-write our distillation loss as In essence, because this loss involves both the 2D keypoint locations and the anchor-wise classification scores, it distills jointly the correspondences and the segmentation results from the teacher to the student.
To solve this optimal transport problem, we rely on Sinkhorn's algorithm [5], which introduces a soft versions of the contraints via Kullback-Leibler divergence regularizers. This then yields the final distillation loss where α s and α t concatenate the classification score values for the student and the teacher, respectively. This formulation was shown to be amenable to fast parallel optimization on GPU platforms, and thus well-suited for deep learning [5,7]. In our experiments, we normalize the predicted 2D keypoints by the image size to the [0, 1] 2 space and set ε to 0.001, and ρ to 0.5 to handle outliers.

Network Architecture
Our approach can be applied to any network that predicts the 2D locations of 3D object keypoints. Without loss of generality, we use WDRNet [19] as our base network, which is a typical keypointbased framework. WDRNet employs a feature pyramid strategy to predict the 2D keypoint locations at multiple stages of its decoder network. These multi-stage predictions are then fused by an ensembleaware sampling strategy, ultimately still resulting in 8 clusters of 2D locations, i.e., one cluster per 3D bounding box corner. To make our baseline consistent with the state-of-the-art methods [25,40,6], we also use a detection pre-processing step that provides an image patch as input to WDRNet. We refer to this as WDRNet+. As will be shown in our results, the pre-processing detection step allows WDRNet to outperform the state-of-the-art SO-Pose [6] with lightweight backbones. Nevertheless, as also evidenced by our experiments, the success of our distillation strategy does not depend on it.
In our experiments, the teacher and student networks follow the same general architecture, only differing in their backbones. Note that different backbones may also yield different number of stages in the feature pyramid, but our distribution matching approach to knowledge distillation is robust to such differences.
To train our WDRNet+ networks, we rely on the focal loss for the segmentation branch and the 3D regression loss proposed in [19] for the keypoint prediction one. When performing distillation to a student network, we complement these loss terms with our distillation loss of Eq. 4. To implement it, we rely on the GeomLoss library [7].

Experiments
In this section, we first discuss our experimental settings, and then demonstrate the effectiveness and generalization ability of our approach on two widely-adopted datasets, LINEMOD and Occluded-LINEMOD. Finally, we analyze different aspects of our method and evaluate it on variations of our architecture.

Experimental Settings
Datasets. We conduct experiments on the standard LINEMOD [14] and Occluded-LINEMOD [2] 6D pose estimation benchmarks. The LINEMOD dataset contains around 16000 RGB images depicting 13 objects, with a single object per image. Following [3], we split the data into a training set containing around 200 images per object and a test set containing around 1000 images per object. The Occluded-LINEMOD dataset was introduced as a more challenging version of LINEMOD, where multiple objects heavily occlude each other in each RGB image. It contains 1214 testing images. For training, following standard practice, we use the real images from LINEMOD together with the synthetic ones provided with the dataset and generated using physically-based rendering [16].
Teacher & student architectures. For our teacher models, we use DarkNet53 [31] as backbone, as in the original WDRNet [19]. For the compact students, we experiment with different lightweight backbones, including DarkNet-tiny [32] and a further reduced model, DarkNet-tiny-H, containing half of the channels of DarkNet-tiny in each layer.
Baselines. We compare our method to the direct training of the student without any distillation (Student), the naive knowledge distillation strategy introduced in Section 3.1 (Naive-KD), and the state-of-the-art feature distillation method (FKD) [44], which, although only demonstrated for object detection, is applicable to 6D pose estimation. For these baselines, we report the results obtained with the best hyper-parameter values. Specifically, for FKD, the best distillation loss weight on both datasets was 0.01; for Naive-KD, the best weight was 0.1, and the best norm was p = 1 for DarkNet-tiny and p = 2 for DarkNet-tiny-H, respectively. For our method, the distillation loss was set to 5 for LINEMOD and to 0.1 for Occluded-LINEMOD. We provide the results of the hyper-parameter search in the supplementary material.
Evaluation metric. We report our results using the standard ADD-0.1d metric. It is computed as the percentage of images for which the average 3D point-to-point distance between the object model in the ground-truth pose and in the predicted one is less than 10% of the object diameter. For symmetric objects, the point-to-point distances are computed between the nearest points. Note that, on LINEMOD, we report the results obtained using the ground-truth 2D bounding boxes to remove  the effects of the pretrained detectors. On Occluded-LINEMOD, we report the results obtained with the same detector as in [40,6] to evidence the effectiveness of our knowledge distillation method.

Comparison with the State of the Art
Results on LINEMOD. We report the results of our method and the baselines for all classes of the LINEMOD dataset in Table 1 and Table 2 for DarkNet-tiny and DarkNet-tiny-H, respectively. While Naive-KD slightly improves direct student training with the DarkNet-tiny backbone, it degrades the performance with DarkNet-tiny-H. This matches our analysis in Section 3; the fact that the student's and teacher's active anchors differ make keypoint-to-keypoint distillation ill-suited.
Both FKD and our approach boost the student's results, with a slight advantage for our approach.
In particular the accuracy improvement is larger, i.e., 2.3 points, for the smaller DarkNet-tiny-H backbone, for which the gap between the student and the teacher performance is also bigger. Note that the improvement of our approach over the student is consistent across the 13 objects. Interestingly, the types of distillation performed by FKD and by our approach are orthogonal; FKD distills the features while we distill the predictions. As such, the two methods can be used together. As can be seen from the table, this further boosts the results, reaching a margin over the student of 1.7 points and 2.9 points with DarkNet-tiny and DarkNet-tiny-H, respectively, and thus constitutes the state of the art on the LINEMOD dataset for such compact architectures.
Results on Occluded-LINEMOD. We now show the generality of our approach by evaluating it on the more challenging Occluded-LINEMOD. Here, we use only FKD [44] as baseline and drop Naive-KD due to its inferior performance shown before. The results are provided in Table 3. Our keypoint-based knowledge distillation method yields results on par with the feature-based FKD on average. Note, however that FKD requires designing additional adaptive layers to match the misaligned feature maps, while our method does not incur additional parameters. More importantly, jointly using our method with FKD achieves the state-of-the-art results with 4.0 points improvements over the baseline student model. For some classes, such as can, eggbox and holepuncher, the improvements surpasses 5 points. With vs without segmentation scores. We compare the results of our approach without and with the use of the segmentation scores in the optimal transport formulation, i.e., Eq. 3 vs Eq. 2.

Additional Analyses
The comparison in Table 4 shows the benefits of jointly distilling the predicted 2D locations and the segmentation scores. In the table, we also report the results of the baseline SO-Pose [6] and WDRNet+. The numbers show that, despite its larger number of weights, SO-Pose yields worse results than our student baseline for a comparable backbone. Moreover, its performance drops significantly from 85.4 to 66.7 as the number of parameters decreases, while that of WDRNet+ degrades more gracefully.
Without detection pre-processing. Note that we incorporated the pre-processing detection step in WDRNet only to show that keypoint-based methods could match the performance of the state-of-the-art dense prediction ones, such as SO-Pose, which also estimate the pose based on detected image patches. In other words, the success of our knowledge distillation strategy does not depend on the detection. To demonstrate this, in the left portion of Table 5, we report the results of our approach applied to the original WDRNet with a DarkNet-tiny backbone. As a matter of fact, the gap between direct student training and our approach is even larger (1.2 vs 2.1), showing the benefits of our approach on weaker networks. With a simple PnP network. In the right portion of Table 5, we compare the results of our approach with those of the baselines on an architecture obtained by incorporating a simple PnP network at the end of WDRNet, following the strategy of [18]. With such an architecture, the 2D keypoint locations only represent an intermediate output of the network, with the PnP module directly predicting the final 3D translation and 3D rotation from them. As can be seen from these results, our distillation strategy still effectively boosts the performance of the student with this modified architecture, further showing the generality of our approach, which can distill keypoint-based knowledge both for PnP solvers and PnP networks.

Conclusion
We have introduced the first approach to knowledge distillation for 6D pose estimation. Our method is driven by matching the distributions of predicted 2D keypoint locations from a deep teacher network to a compact student one. We have formulated this as an optimal transport problem that lets us jointly distill the predicted 2D locations and classification scores that segment the object in the image. Our experiments have demonstrated the effectiveness of our approach and its benefits over a naive point-to-point distillation strategy. Furthermore, our formalism is complementary to feature distillation strategies and can further boost its performance. In essence, our work confirms the importance of developing task-driven knowledge distillation methods, and we hope that it will motivate others to pursue research in this direction, may it be for 6D pose estimation or for other tasks.
Limitations and societal impact. Our results show that some object classes benefit more from our KD approach than others. In the future, we will aim to analyze this in more detail, with a view to designing a class-adaptive version of our method. Furthermore, our current framework is designed for keypoint-based 6D pose estimation methods. Although these methods are successful, we intend to study if our distribution-based approach can be extended to frameworks that perform dense prediction for 6D pose estimation.
By achieving state-of-the-art 6D pose estimation results with compact networks, our approach opens the door to deployment on mobile applications, including safety-critical ones, such as pose estimation for autonomous navigation or space debris capture. Nevertheless, our approach does not offer robustness guarantees, in particular to adversarial attacks. As such, our approach could potentially be diverted for harmful objectives. weight of 0.01 also achieves the best results. Note that we did not test a weight of 0.1 on all classes because of the worse results it gave on Ape and Duck.

A.3 Hyper-parameters for our Approach
In this section, we include the hyper-parameter search for our proposed keypoint distribution alignment distillation method, including the norm ps and the weight of our distillation loss. As for the baseline, we focused this search on DarkNet-tiny-H for the Ape and Duck classes. As shown in Table A5, p = 2 yields much better results than p = 1, and we therefore use p = 2 in the main paper. As for the loss weight, on the LINEMOD dataset, 5 yields the best results, which we use to report the results on the 13 classes in the main paper. For Occluded-LINEMOD, as shown in Table A6, we obtain the best results with a weight of 0.1. Note that our preliminary experiments with a weight of 1 showed worse performance, and we thus did not compute full results with weights larger than 0.1.

A.4 Qualitative Analysis
In this section, we provide additional comparisons of the predicted 2D keypoints distributions obtained with the baseline student model and with our distilled model on several examples from Occluded-LINEMOD. As shown in Figure A1, the predicted 2D keypoints clusters from our distilled models are closer to the ground-truth object corners than those of the baseline model. Furthermore,

Student
Ours Teacher Figure A1: Qualitative Analysis (better viewed in color). In (a), we compare the 2D keypoints predictions from our distilled model (3rd column with orange dots) and the baseline student model (2nd column with blue dots). With our proposed keypoint distribution alignment distillation method, the model predicts tighter keypoint clusters, closer to the ground-truth corners (pink cross) than the baseline model. Furthermore, our distilled model is able to mimic the teacher's keypoint distributions (1st column with orange dots). Light-green boxes highlight some keypoints clusters, which are also zoomed in on the side of the image. In (b), we show the estimated poses. The blue boxes are the predicted 3D bounding boxes while the gray ones are the ground-truth bounding boxes. Our distilled model generates better pose estimates than the student model.