MEDL-U: Uncertainty-aware 3D Automatic Annotation based on Evidential Deep Learning

Recently, the literature has witnessed a rise in 3D automatic annotation frameworks. An example is WSPCD [12] which allows learning 3D object parameters from a few weakly annotated examples. It has a two-stage architecture design: the first stage for cylindrical object proposal generation and the second stage for cuboids and confidence score prediction. A non-learning based approach that detects vehicles in point clouds without any 3D annotations is FGR [4] which also consists of two stages: coarse 3D segmentation stage and the bounding box estimation stage. More recently, Liu et al.[6] propose a Transformer-based 3D annotator called MTrans, which aims to address the prevalent sparsity problem of unstructured point clouds from LiDAR scans by generating extra 3D points through a multimodal self-attention mechanism combined with additional multi-task and self-supervision design. Different from previous appoaches, Qian et al.[7] propose a simplified end-to-end Transformer model, CAT, which captures local and global relationships through an encoder-decoder architecture. The encoder consists of intra-object encoder (local) and inter-object encoder (global) which performs self-attention along the sequence and batch dimensions. Through this, it can model the inter-object feature relation that gives additional information to hard samples. Additionally, several approaches (GAL [13], VS3D [14], WS3DPR [15]) have also been proposed. However, all these 3D automatic annotators only generate 3D pseudo labels without any estimation of the uncertainty or noise associated with them. In this study, we utilize the generatd 3D pseudo labels and the estimated uncertainties to train 3D detectors. This approach aims to mitigate the impact of generated noisy labels by reducing the influence of inaccurate supervision signals and enabling the model to effectively learn from more reliable pseudo labels.

Uncertainties in deep learning-based predictions can be categorized into two: one that is caused by inherent noise in data (aleatoric), and the other is model uncertainty due to incomplete training or model design (epistemic). As a tool for uncertainty estimation in regression tasks [8], Evidential Deep Learning (EDL) has found diverse applications in various tasks such as stereo matching [9], open set recognition [10], molecular structure prediction [16], and remote sensing [11]. In this work, we use the framework of EDL to estimate prediction uncertainties in 3D Object Detection. Utilizing these uncertainties to define pseudo label distributions, probabilistic object detectors, which typically adapt deterministic detectors’ architecture, can enable the prediction of probability distributions for object categories and bounding boxes. A framework that utilizes probabilistic detectors is GLENet [17], where the 3D detector predict distributions for the 3D box and assume that the ground truth labels follow Gaussian distribution with the uncertainties as the variance. They train the models with the KL divergence loss to supervise the predicted localization uncertainty. In this work, we follow the same approach but instead incorporate 3D pseudo label uncertainties.

Given the point cloud data, 2D image and the 2D bounding boxes of each object, the objective of this work is to generate the 3D bounding boxes annotation for each object and estimate the corresponding uncertainty for each 3D box parameter. First, the autolabeler is initially trained with a small set of ground truth 3D bounding boxes (e.g., 500 frames of data), where the input of the autolabeler are the point cloud data, 2D image and the 2D bounding boxes of each object and the output are the estimated 3D bounding boxes and corresponding uncertainty estimates. Secondly, with the trained autolabeler, we employ it to predict the 3D bounding boxes for the remaining data and do the uncertainty estimation for the predicted 3D boxes. Finally, we leverage the predicted 3D bounding boxes and estimated uncertainty to train a downstream probabilistic 3D detector on the massive weakly annotated data. Compared to fully supervised setting, our work only needs a few frames of labeled data to train the 3D autolabeler, which significantly reduces the manual annotation cost.

In this paper, we build the model architecture for the 3D automatic labeler from MTrans [6]. MTrans extracts object features using a multimodal self-attention module that processes point cloud and image inputs fused with point-level embedding vectors. The extracted object features are utilized for various tasks such as foreground segmentation, point generation, and 3D box regression. However, the generated 3D box pseudo labels may contain noise. To account for these inaccuracies of the pseudo labels, we incorporate uncertainty estimation task to MTrans by considering EDL, a powerful uncertainty estimation framework, and applying it to MTrans. For a straightforward incorporation of uncertainty estimation in MTrans via EDL, we include an evidential box head to regress the parameters of the evidential distribution and the 3D bounding box. For training the model, we replace the dIoU loss with the evidential loss to supervise the model in learning the box parameters. Moreover, the evidence regularizer is added to calibrate the uncertainties. However, there are problems with this approach of directly applying EDL in MTrans, which we will discuss in the next sections.

In 3D object detection, we define a 3D bounding box by its center coordinates (x, y and z), length (l), width (w), height (h), and rotation (yaw angle denoted as $rot$). From the viewpoint of EDL, we assume each label $j\in\mathbb{J}=\{x,y,z,l,w,h,rot\}$ is drawn i.i.d. from a Gaussian distribution where the mean $\mu_{j}$ and variance $\sigma^{2}_{j}$ are unknown. EDL framework assumes that $\mu_{j}$ is drawn from a Gaussian prior and $\sigma^{2}_{j}$ is drawn from an inverse-gamma prior.

where $\gamma_{j}\in\mathbb{R}$, $\nu_{j}>0$, and $\Gamma(\cdot)$ is the gamma function with $\alpha_{j}>1$ and $\beta_{j}>1$.

Let $\Phi_{j}$ and $\Theta_{j}$ denote the set of parameters $\{\mu_{j},\sigma_{j}^{2}\}$ and $\{\gamma_{j},\nu_{j},\alpha_{j},\beta_{j}\}$, respectively. Assuming independence of the mean and variance, the posterior $p(\Phi_{j}|\Theta_{j})$ is defined to be an normal-inverse gamma (NIG) distribution, which is a Gaussian conjugate prior.

As presented in [8], the hyperparameters of the evidential distribution can be obtained by training a deep neural network (called evidential head) to output such values. For each 3D bounding box parameter, our model predicts four evidential parameters: $\gamma_{j},\nu_{j},\alpha_{j},\beta_{j}$.

Through an analytic computation of the maximum likelihood Gaussian without repeated inference for sampling [8], EDL provides a framework for the estimation of uncertainties in regression. We can calculate the prediction, aleatoric, and the epistemic uncertainties for each 3D box parameter as follows:

Maximizing the data fit: We are given the hyperparameters $\Theta_{j}$ of the evidential distribution as outputs of the proposed evidential head for $j\in\mathbb{J}=\{x,y,z,l,w,h,rot\}$. The likelihood of an observation $y_{j}$ is computed by marginalising over the likelihood parameters $\Phi_{j}$. If we choose to impose a NIG prior onto the Gaussian likelihood function results, an analytical solution is derived as follows:

where $St_{\nu}(t|r,s)$ corresponds to the evaluation of the Student’s t-distribution at the value t, with parameters for location ($r$), scale ($s$), and degrees of freedom ($\nu$).

For training the EDL framework, we define the evidential loss as the mean of the negative log likelihood (NLL) for each 3D box parameter as follows:

Uncertainty Calibration: Similar to Amini et al.[8], we can scale the total evidence with the prediction error for each 3D box parameter in the following manner:

where $\phi_{j}$ is the total evidence defined as $\phi_{j}=2\nu_{j}+\alpha_{j}$, and $\lVert\cdot\lVert$ is L1 norm.

• •

  Significant variability exists in uncertainty estimates for box regression parameters, with some resulting in excessively high values. As also discussed in [18], there exist the gradient shrinkage problem in the evidential NLL loss, where the model can decrease the loss value by increasing the uncertainty instead of getting accurate point estimates for the predicted bounding box parameters. The gradient with respect to the the predicted values can become very small by merely increasing the uncertainties.

• •

  Since the NLL-based evidential loss is not sufficient to optimize the accuracy of the prediction, the generated pseudo labels exhibit lower quality than MTrans. Empirical results also demonstrate that IoU-based loss is better than evidential loss to learn the 3D box parameters as the latter treats each 3D box parameter independently.

Intuitively, the loss information obtained from several multi-task loss functions during training corresponding to the same object can be utilized to help the model understand the evidence in support to the prediction. In line with the works on uncertainty estimation [19, 9], we introduce regularized multi-task loss functions based on the form of NLL minimization of aleatoric uncertainty estimation and intuitively aligned with the learned loss attenuation. The main insight behind these loss functions is that the model’s predicted evidence is inversely related to the losses for multiple tasks, effectively serving as an evidence regularizer. Let $\mathcal{L}^{{}^{\prime}}_{t}$ be the loss function corresponding to the task t where $t\in\{$seg, depth, conf, dir$\}$. The proposed evidence-aware multi-task loss is

where $\alpha=|\mathbb{J}|^{-1}\sum_{j\in\mathbb{J}}\alpha_{j}$ and $\nu=|\mathbb{J}|^{-1}\sum_{j\in\mathbb{J}}\nu_{j}$. In Figures C.1 and C.2, we demonstrate the effect of including the evidence-aware multi-task loss in the magnitude of the epistemic uncertainties.

While the NLL-based evidential loss can enable the model to learn the 3D box and evidential parameters, it is not sufficient to regress 3D box variables with a quality comparable to those produced by existing autolabelers. Hence, we propose to include an IoU-based loss inspired by the DIoU loss [20] similar in form to 6. This makes $\mathcal{R}$ (penalty term for the prediction and ground truth) and IoU-related term uncertainty-aware.

Incorporating this new evidence-aware IoU loss improves the model’s ability to handle uncertainty while generating high quality 3D box that is comparable to other 3D autolabelers. In Figure D.3, we demonstrate the effect of including the uncertainty-aware IoU loss in the model performance on the validation set. Note that the original evidence regularizer [8] could also train $\gamma_{j}$. Our empirical findings suggest that eliminating this regularizer is necessary as components of the proposed multi-task losses and the uncertainty-aware IoU loss already serve for regularization purposes.

In summary, the final overall loss function $\mathcal{L}$ is computed as a weighted combination of the evidential loss, the uncertainty-aware IoU loss, and the multi-task losses with evidence regularizers:

where $\eta_{seg},\eta_{depth},\eta_{conf},\eta_{dir},\eta_{box}$ and $\eta_{evi}$ are hyperparameters. Please refer to Figure 2 for the overall workflow of the model.

Prior to utilizing the pseudo labels and uncertainties as supervision signals in the training process of existing 3D detectors, it is necessary to apply a post-processing step to address the variability in the magnitudes of uncertainties associated with each 3D box parameter and to make the uncertainties more appropriate for downstream task. We propose a post-processing procedure that rescales the uncertainties but ensures that the resultant uncertainty values maintain its Spearman’s rank correlation with two crucial metrics: the L2 norm of the residuals and the 3D IoU between the predicted box and the ground truth box. Initially, the predicted epistemic uncertainty pertaining to each 3D box parameter $j$, where $j\in\mathbb{J}=\{x,y,z,l,w,h,rot\}$, undergoes a transformation which constrains these uncertainties within the range of 0 to 1 through the application of a simple monotonic function such as min-max scaling. Moreover, we pass uncertainty estimate for each box parameter $x_{j}$ to a function $f$, formulated as $f(x_{j})=x_{j}^{1/\kappa_{j}}$, where we carefully select values of $\kappa_{j}\in[0,10]$ to minimize the NLL of the uncertainties for each 3D box parameter $j$ with respect to the same limited training data used to train the 3D autolabeler. The assumption is that lower NLL means better uncertainty estimates, consequently improving supervision for the downstream task. Lastly, we generate multiple variations of uncertainty by passing them to the function $g(x_{j})=x_{j}^{1/\epsilon_{j}}$, with $\epsilon_{j}$ acting as a downstream training hyperparameter. The underlying rationale for introducing $\epsilon_{j}$ is rooted in the observation that the $\kappa_{j}$ parameter only effectively minimizes NLL over the limited training data and there must exist an additional parameter which may prove essential for appropriately adjusting the uncertainties to achieve lower NLL over the entire training dataset.

A probabilistic object detector enables the inclusion of pseudo label uncertainties during the training phase, where these uncertainties can be interpreted as factors for reweighting losses. Our work follows [21, 17] in transforming a 3D detector from deterministic to probabilistic, where the detection head is enforced to estimate a Gaussian probability distribution over bounding boxes. Let $\theta$ indicate the learnable network weights of a typical detector, $\hat{y}$ indicate the predicted box parameters, and $\hat{\sigma}^{2}$ the predicted localization variance. Moreover, the pseudo ground truth bounding boxes are also assumed to be Gaussian distributed having a variance $\sigma^{2}$ where $\sigma^{2}$ are estimated by MEDL-U. Let the pseudo ground truth bounding box be denoted by $y_{g}$ and D refer to the pseudo label distribution. Hence, the generated pseudo label uncertainty can be incorporated in the KL Divergence loss between the distribution of prediction and pseudo ground truth in the detection head:

Similar to [17, 22], we also employ 3D Variance Voting which uses the predicted variance $\hat{\sigma}^{2}$ to combine nearby bounding boxes for better 3D localization prediction.

The KITTI Object Detection dataset [23], renowned for 3D detection in autonomous driving, is employed in this study. The dataset has a total of 7481 frames with labels in 3D. Following the official procedure, the dataset is divided into training and validation sets, consisting of 3,712 and 3,769 frames, respectively. Similar to previous works [4, 5, 6, 7], we concentrate on the Car class and exclude objects with fewer than 5 foreground LiDAR points.

Our method is implemented in PyTorch [24]. Similar to the original MTrans [6], MEDL-U architecture incorporates four multimodal self-attention layers, each having a hidden size of 768 and 12 attention heads. Unless otherwise stated, training the autolabeler requires 500 annotated frames only. We employed a dropout rate of 0.4 and utilized the Adam optimizer with a learning rate of $0.60e-04$. Autolabeler training is conducted for 300 epochs, with a batch size of 5. Training of probabilistic 3D detectors is conducted for 80 epochs. Note that only epistemic uncertainties from MEDL-U are utilized. Unless specified differently, hyperparameter tuning on the KITTI validation set suggests using $\epsilon=1$ for PointPillars and CIA-SSD, and $\epsilon=5$ for other detectors. All trainings are executed on NVIDIA RTX 2080Ti GPU.

MEDL-U evidential regression head consists of four output units for each of the seven box attributes. The input to this head are the transformed element representations extracted from the self-attention layer. The evidential regression head comprises a sequence of linear layers, followed by LayerNorm, a Dropout layer, and a ReLU activation function. To ensure that certain values are positive, a Softplus activation is applied to $\nu$, $\alpha$, and $\beta$, where $\alpha$ is then incremented by 1 to ensure $\alpha>1$. For $\gamma$, a linear activation is used. Overall, MEDL-U has over 23 million trainable parameters. It is trained to learn seven evidential distributions simultaneously for each of the 3D box parameters, along with the other tasks of segmentation, point generation, direction, and confidence prediction. In the next sections, MEDL refers to utilizing the pseudo labels only while MEDL-U refers to utilizing both pseudo labels and the uncertainties in the downstream training of 3D detectors.

We evaluate several one-stage and two-stage 3D detectors on the KITTI val set when trained using outputs from various annotators. As seen in Table I, detectors trained on MEDL-U outputs outperform vanilla deterministic 3D detectors trained on MTrans and MEDL pseudo labels, demonstrating the effectiveness of utilizing not only the pseudo labels but also the uncertainty estimates for each 3D box parameter.

As shown in Table II, evaluating the probabilistic PointRCNN trained with MEDL-U ouputs on the KITTI val set yields superior performance relative to all existing and current 3D autolabeling methods in terms of $AP_{3D}$.

In Table III, probabilistic PointRCNN trained with outputs of MEDL-U on the entire KITTI training and val sets results in better performance on the KITTI official test set compared to PointRCNN trained with vanilla MTrans. Moreover, PointRCNN trained with MEDL-U outputs yields superior performance in terms of $AP_{3D}$ and $AP_{BEV}$ for both Easy and Moderate levels relative to all existing 3D automatic labeling method. MEDL-U does not outperform CAT across all difficulty levels, which is understandable considering that MEDL-U is built upon MTrans, chosen for its open-source availability. Moreover, CAT is trained for 1000 epochs and a batch size of 24, which is different from the training setting for MTrans and MEDL-U. We argue that the enhancements seen in MEDL-U over MTrans can also be applied to CAT.

We also show evaluation performance on the KITTI val and test set using PointPillars when MTrans and MEDL-U are trained with 500 and 125 annotated frames. MEDL-U significantly improves the baseline as shown in Table IV.

Using 3D box parameter uncertainties generated by MEDL-U and other popular uncertainty estimation methods, we evaluate probabilistic version of PointRCNN on the KITTI val set. Three baseline methods were implemented: (1) A Monte Carlo dropout (MC Dropout) system with a dropout rate of 0.2 and was forwarded 5 times during inference. (2) A Deep Ensemble of 5 systems trained with different random seeds. (3) Confidence score predicted by vanilla MTrans was used as proxy to uncertainty. As shown in Table V, PointRCNN trained with MEDL-U ($\epsilon=5$) yields the overall best result in terms of $AP_{3D}\>R40$. Noticeably, using other uncertainty estimation methods to generate uncertainties also effectively increase $AP_{3D}\>R40$ of the base method, although MC Dropout and Deep Ensemble come at the cost of huge additional computational overhead. While the Deep Ensemble approach can improve the base result, the original pseudo label quality is relatively poor.

Table VI shows that the uncertainties generated by MEDL-U achieve the overall best result in terms of NLL and Spearman’s rank correlation coefficient. Specifically, it achieves the highest uncertainty Spearman’s rank correlation with residuals for the $x$ and $y$ coordinate of the center, height, and rotation, and consistently among the two lowest calculated NLL for all the 3D box parameters. Details on ablation results and results on varying $\epsilon$ for the downstream task can be seen in Tables VII and VIII, respectively.

The proposed framework MEDL-U is built from MTrans [6]. In this section, we briefly describe the architecture of MTrans. MTrans generates object features from the multimodal image and point cloud inputs. Then these features are encoded to generate embedding vectors using sinusoidal position embedding for 2D pixels, MLP for 3D coordinates, and 3-layer CNN for the image patch features; these features are fused and passed to a fully-connected layer to generate element representations. Then a multimodal self-attention mechanism manipulates and transforms these features to generate object-level representations. For the multi-task supervision, multiple prediction heads take the transformed representation as input. In the self-attention mechanism, MTrans enriches the sparse point cloud by defining a training mechanism for context points with known 3D coordinates, padding points which are randomly dropped points, and target points which have a corresponding uniformly distributed image pixels. The various heads are the 2-layer MLP segmentation head, 2-layer MLP point generation head, box regression head, binary direction head (front or back), and the confidence head. Note that MTrans only included the confidence head for the purpose of converting the autolabeler to a 3D object detector version. But in our framework, we included the confidence head so that we may use the confidence estimates as a proxy to uncertainty, which can be compared with our proposed uncertainty estimation approach.

To supervise the training of these multiple heads, sum of cross-entropy and dice loss is used for the binary segmentation task loss $\mathcal{L}_{seg}$, smooth $L_{1}$ loss for the point generation task loss $\mathcal{L}_{depth}$ with a self-supervised mask-and-predict strategy, dIoU loss for the box loss, cross-entropy loss $\mathcal{L}_{dir}$ for the binary direction classification, and smooth $L_{1}$ loss $\mathcal{L}_{conf}$ for the confidence prediction. We advise the reader to read [6] for a more complete information about MTrans.

Similar to the configuration of MTrans, we apply the standard image augmentation methods of auto contrast, random sharpness, and color jittering as well as point cloud augmentation methods of random translation, scaling and mirroring. For each object, the cloud size is set to 512 for context plus padding points, and there are 784 image-sampled target points. For the training of the autolabeler, we set the following loss weights for 8: $\eta_{seg}=1$, $\eta_{depth}=1$, $\eta_{conf}=1$, $\eta_{dir}=1$, $\eta_{evi}=2$, and $\eta_{IoU}=5$.

Moreover, in MTrans, the ranges of the predicted and ground truth yaw are limited to $[-\pi/2,\pi/2]$ and $[-\pi$ to $\pi]$, respectively. Hence, if we were to employ the evidential loss to supervise the predicted yaw, this would lead to inconsistencies, since a component of the evidential loss involves a direct calculation of the difference between the predicted yaw and the ground truth yaw. As a result, the constraints imposed on the predicted yaw in vanilla MTrans may not be compatible with the requirements of the evidential loss. Different from MTrans, we wrap the ground truth yaw Euler angles to the range of [-$\pi$/2, $\pi$/2] to make it appropriate for the evidential loss corresponding to rotation. Moreover, the presence of the direction loss is sufficient for accurate yaw prediction. By wrapping the yaw Euler angle to the range [-$\pi$/2, $\pi$/2], we ensure that the ground truth yaw angle stays within a compact interval, which makes it appropriate for the evidential loss corresponding to rotation.

Figure C.1 visualizes the excessively high epistemic uncertainties generated for the 3D box parameters during training and the huge variations in the magnitudes of the uncertainty for each 3D box parameter when only the evidential loss with uncertainty calibration term is used to supervise the model training. Figures C.2 shows that including the evidence-aware multi-task loss during training mitigates the problem of generating excessively high epistemic uncertainties and reduces the differences of the magnitude of the epistemic uncertainties for each 3D box parameter. Moreover, the trend of the mean epistemic uncertainties during training is generally downward, which is indicative of the model being less uncertain as the training progresses.

In Figure D.3, we can see an improved IoU of the predictions with the ground truth during training on the validation set because of the inclusion of the uncertainty-aware IoU loss.

In this section, we demonstrate that MEDL-U can generate reasonable uncertainty estimates (both aleatoric and epistemic uncertainties) for each of the 3D box parameter $j\in\mathbb{J}=\{x,\>y,\>z,\>l,\>w,\>h,rot\}$, which could give an intuition as to why using the information on pseudo label uncertainty improves the performance of the downstream 3D detection task.

In this section, we visualize the annotation results of MEDL-U and detection results of the 3D detector PointPillars trained with the outputs of MTrans and MEDL-U. In Figures F.8 and F.9, we provide visual comparison of the detection results of PointPillars trained with MTrans pseudo labels and MEDL-U pseudo labels and uncertainties on the KITTI tracking data (test set) when the 3D autolabelers are trained with 500 frames of annotated data. Visual comparisons are also shown in Figures F.10 and F.11 for the case when the autolabelers are trained with 125 frames of annotated data. It can be seen that PointPillars trained with MEDL-U obtains better detection results. In Figures F.6 and F.7, the annotation results of MEDL-U on the KITTI tracking data (training set) are visualized. Visualizing in BEV, we observe that the proposed framework generates reasonable uncertainty estimates associated with each object without compromising the quality of the pseudo labels. We can see that in most cases, MEDL-U generates larger uncertainty estimates for faraway objects with sparser points.