Neighborhood Spatial Aggregation MC Dropout for Efficient Uncertainty-aware Semantic Segmentation in Point Clouds

Uncertainty-aware semantic segmentation of the point clouds includes predictive uncertainty estimation and uncertainty-guided model optimization. One key challenge in the task is the efficiency of pointwise predictive distribution establishment. The widely used Monte Carlo (MC) dropout establishes the distribution by computing the standard deviation of samples using multiple stochastic forward propagations, which is time-consuming for tasks based on point clouds containing massive points. Hence, a framework embedded with neighborhood spatial aggregation (NSA)-MC dropout, a variant of MC dropout, is proposed to establish distributions in just one forward pass. Specifically, our method uses the one-time stochastic inference of a point with neighbors to approximate the point’s repeated stochastic inferences, outputting pointwise distribution via the prediction variance of neighbors. Based on this, uncertainties acquire from the predictive distribution. The aleatoric uncertainty is integrated into the loss function to suppress noise, preventing the model from overfitting. Besides, the predictive uncertainty quantifies the prediction confidence. Experiments show that our plug-and-play NSA-MC dropout significantly improves the segmentation performance of backbones, ranging from multilayer perceptron (MLP)-, convolution-, and attention-based to transformer-based networks, without introducing additional parameters. Besides, it is several times faster than MC dropout to quantify the results’ credibility, and the inference time does not establish a coupling relation with the sampling times.