NoiseTrans: Point Cloud Denoising with Transformers

There has been a long history of research on point cloud denoising methods. Point cloud denoising has become particularly important in recent years as technologies such as remote sensing, robotics, and autonomous driving continue to develop. Non-deep learning point cloud denoising methods can generally be divided into four categories: local-surface-fitting based methods [4, 1, 17], non-local means based methods [46, 34], sparsity-based methods [40, 24, 32] and graph-based methods [15, 44, 16].

The most representative local surface-based adaptation method is MLS [1], which assumes that the surface of the point cloud is smooth and fits point to plane for denoising. In addition to that, other surface fitting methods have been proposed such as jet-fitting with re-projection [4] and various forms of bilateral filters [17] have achieved remarkable results. However, this can be over-smooth or over-sharpened at high noise levels.

Non-local means-based methods are developed from the 2D image denoising, such as [46, 34], which detect the similar structure of non-local point cloud, and then remove the noise through filtering to get the point cloud without noise. These methods can retain more surface features and avoid excessive smoothing. However, a large number of parameters need to be set, which is difficult to apply to different scenarios.

Sparsity-based methods [40, 24, 32] are generally to first solve the optimization problem of sparsity constraints to reconstruct the normal and then update the coordinate position of the point according to the reconstructed normal to denoise. MRPCA [24] is a representative of this, but performance degrades at high noise levels.

The graph-based methods [15, 44, 16] are to define the point cloud as a graph, use the points in the point cloud as nodes in the graph, and then denoise it through graph filters [44, 16]. In particular, GLR [44] uses the Laplace regularisation of the graph as a filter for denoising and achieves high effectiveness. Under high noise levels, however, it leads to instability in the structure of the graph, resulting in a degradation of the denoising effect.

In recent years, most of the significant advances in visual recognition have been achieved by deep learning models, especially deep Convolution Neural Networks (CNNs).However, directly processing point clouds using off-the-shelf image-based methods is not straightforward since point clouds are essentially irregular and unordered, which are not suitable to be processed directly using convolution neural networks designed for gridded features. With the deepening of 3D object classification, especially the proposal of PointNet [27] and PointNet++ [28], the point-based deep learning methods have become possible, and the point cloud denoising method based on the deep learning methods has also gained more extensive attention.

PointCleanNet [30], which employs PCPNet [13] (a variant of PointNet [27]) as the backbone, predicted the displacement vector from noise point to object surface according to the local feature of each point. Because of the multiple iterations in the denoising process, shrinkage of the point cloud occurs. TotalDn [14] proposed the unsupervised point cloud denoising for the first time. They introduced the prior knowledge that the denser the point cloud is, the closer it is to the real surface. They propose a new loss function that uses neighbors to determine the location of ground truth points without considering the noise point themselves. This method works well for flatter surfaces with lower noise levels. Due to the similarity between point cloud and graph in 3D space, graph convolution network [26] is used for denoising, and positive results have also been achieved. Recently, DMR [22] has shifted its attention to the physical characteristics of point clouds and reconstructed the bottom manifold of point clouds through a down-up sampling network structure to the denoise point cloud. Meanwhile, Score [23] regarded the point cloud with noisy points as the distribution after convolution of the underlying manifold and noise distribution and proposed that the likelihood function of the distribution should update the position of each point by iteration of gradient rise to achieve the effect of denoising. But this approach leads to the creation of some outliers.

Transformer [36] was originally proposed as a sequence-to-sequence model for machine translation [35] and consists of an encoder and a decoder. Each encoder block consists mainly of a multi-head self-attention module and a feed-forward network(FFN). Compared to the encoder blocks, decoder blocks additionally insert cross-attention modules between the multi-head self-attention modules and the FFN [18]. With the continuous development of natural language processing(NLP) techniques, transformer-based pre-training models have been proposed one after another. Among them, pre-trained models such as the BERT [8] and GPT series [29, 2] have achieved advanced performance in various tasks. As a result, the transformer has become the architecture of choice for NLP.

Inspired by the great success of the self-attention mechanism in NLP, Transformer has also been transplanted to the computer vision field [5, 21, 20]. ViT [10] proposes an image transformer structure that takes image patches as input and achieves advanced results on image classification tasks. In addition, it achieves impressive performance on tasks such as object detection [3, 19], semantic segmentation [47] and target tracking [7]. Not only that, Transformer also provides new thinking for point-based 3D point cloud processing, e.g. [45] and [43], etc.

Unlike the denoising methods described above, our NoiseTrans is designed on the basis of the transformer encoder architecture. We incorporate the core self-attention mechanism into the point-based point cloud denoising method in order to achieve optimal performance.

Currently, deep learning denoising methods generally fit clean surfaces either through local features [30, 14, 26] or abstracting point clouds into mathematical models [22, 23]. In contrast, we focus on the global relationship of the point cloud. Previous work [46] has shown that synthetic and raw point clouds tend to exhibit self-similarity. Extracting and exploiting this similar geometry and overall structure can provide a powerful means of denoising. Meanwhile, Transformer has proven through extensive experiments that its self-attention mechanism has the ability to capture long-range dependencies between elements, whether in the field of natural language processing [36, 8, 2] or computer vision [45, 21, 10]. This is exactly what we need for denoising tasks. Moreover, its permutation invariance is naturally suitable for disordered point clouds. Further, the unique point-to-point correspondence of the transformer structure is convenient and crucial for our denoising task.

Based on the above analysis, we propose a transformer-based point cloud denoising method, NoiseTrans, which is shown in Figure 1. Our model consists of three parts: 1) Point embedding module: We characterize local point cloud feature sets at different scales as embeddings of points. We add the Local Point Attention of our design to the representation to make the model more sensitive to edge bulges. This preserves more detail and prevents the point cloud from being over-smooth.2) Transformer module: We propose Sparse Encoding in the transformer to perceive the structural relationships of the point cloud. Sparsity is added to the encoding as a reference for sensitivity to outliers. 3) Output header module: We design a tail output for the denoising task. Residual connections are added to facilitate the identification of noise.

Note that the downsampling layer, which is often used in transformer models [5, 45, 21], is discarded from our network structure. This is due to the fact that we are working to recover surfaces that have been corrupted by noise. Moreover, our network has no requirement for the number of input points thus it can be applied to a wide range of scenes. We will show more details in the following sections.

The purpose of this section is to obtain input for the Transformer module. For the transformer to work effectively on point clouds, the first step is converting the point cloud into a vector sequence. One of the most straightforward solutions is to use coordinates as a set of vector inputs. However, this approach can lead to the loss of local features. In order to enable each point to probe the geometric features of the local area in which it is located, a point embedding module has been designed. This also compensates for the deficiencies of Transformer in local feature extraction [37, 21].

In 2D images, convolution kernels of various sizes are usually designed to obtain the corresponding perceptual fields to extract features at different scales. Inspired by this experience, we choose to use the number of neighbor points to emulate the convolution kernels. As shown in Figure 2, we design three feature extraction units with different numbers of neighbors in a parallel manner to capture features at various scales.

In the feature extraction unit, we construct K nearest neighbor points as a patch for each point in Euclidean space, followed by a convolution layer to generate high-dimensional features. The central point aggregates the patch as its own local features. After the several layers of feature aggregation, the output of the three feature extraction units is concatenated to obtain a vector sequence containing the local features of the point cloud. We call the above process Point Embedding and the resulting vector can be fed directly to the transformer.

Formally, given the feature ${\bf F}^{l}_{p}=\{{\bf f}_{i}^{l}\}_{i=1}^{N}\in\mathbb{R}^{N\times\mathrm{d}^{l}}$ in the $lth$ layer, where $N$ denotes the number of points. Then the features of the $(l+1)th$ layer can be expressed as:

where ${\bf g}_{\theta}(.)$ represents a non-linear function with $\theta$ as a learnable parameter. $N(i)$ denotes the neighbors of point $i$, and $\max\left(.\right)$ is an aggregation function.

To obtain local features at different scales, we set varying values of k in the three feature extraction units. The output is then stitched together, following a weight-shared multi-layer perceptron layer, which can be defined as:

where ${\bf F}^{i}_{p1}$, ${\bf F}^{i}_{p2}$ and ${\bf F}^{i}_{p3}$ represent the outputs of the $ith$ feature extraction unit, and $h_{\varphi}(.)$ represents a weight shared multi-layer perceptron with $\varphi$ as a learnable parameter.$\left[...\right]$ represents the concatenation operation.

However, there are still problems with the above methods. A key issue in the denoising task is the balance between denoising and preserving protruding edges. This problem is unspecialised in most deep learning methods, and can lead to the network recognising the protruding edges as noise. This is especially so after several iterations, resulting in a loss of detail. Figure 3 illustrates this problem briefly.

To alleviate this situation, we introduced the Local Point Attention in feature extraction. This allows points that are on the same bump to be given more weight, thus preventing the object surface from being over-smooth.

Specifically, we compare the features of neighbor points at different scales to get the structural features where they are located. After concatenation, they are multiplied by a learnable matrix to obtain the weights. Our formula can be expressed as follows:

where $Agg(.)$ represents the aggregation function. Then equation (1) can be rewritten as:

where ${\bf a}_{i}^{l}$ denotes the $ith$ point at $lth$ layer with the surface similarity weights of its neighbors compared to itself. We will show the specific effects in the ablation experiments.

An advantage of the transformer is that each token can attention to any location information [41]. Our model is built on the standard transformer and implements the point cloud denoising task. But the transformer [36] was originally designed for sequence modeling and its self-attention mechanism only takes into account the semantic similarities between individual points. To make it better at capturing the structural relationships between points, we propose Sparse Encoding.

When describing the structure of a point cloud, the real directional orientation between points is only available in non-Euclidean space. This is difficult to compute for point clouds in Euclidean space. On the basis of ensuring the invariance of the overall model permutation, we use learnable coordinate differences as an approximation to the directions in non-Euclidean space. It is worth noting that we only make direction predictions for neighbor points. This is to allow the centroids to identify the spatial structure as well as to prevent errors that occur at longer distances.

For point cloud noise, the sparsity of the points is an important reference for judging the noise size and location information. We discard counting points as a criterion for evaluating sparsity with a more dynamic approach being chosen instead. We use the learnable inverse of the neighbor distance to evaluate the sparsity of points in range. This method not only allows the number of neighbors to be represented but also to some extent resembles the ‘degree’ in a graph.

Specifically, given a coordinate vector set of points ${\bf X}=\left\{{\bf x}_{i}\right\}_{i=1}^{N}\in\mathbb{R}^{N\times d}$,where $N$ represents the number of points in the set and d represents the dimensionality of the vector set, generally we take $d=3$.The location code we have designed can be expressed as follows:

where $M_{\nu}$ represents a weight shared multi-layer perceptron with $\nu$ as a learnable parameter and $j$ represents neighbor points for $i$. In our experiments, we set the number of neighbor points to $3$.

Our point embedding and sparse coding ultimately form the input to our transformer. The purpose of point cloud denoising is to recover the noisy points to the underlying surface of the object, a task that transfers one set to another with the same properties. At the same time, in order to reduce the number of parameters and the calculation time, we, therefore, use only the encoder part of the transformer, simplifying the decoder part. Our module has a total of six transformer layers and each layer is a Pre-LN structure [39]. Aimed at the organization of the point cloud, we redesign the Feed Forward Network. We use a weight-shared multi-layer perceptron to further extract features from the point cloud.

where $F^{l}_{o}$ represents the output of layer $l$, and $LN$ represents Layer Normalization. The multi-head attention mechanism plays a particularly crucial role in our task, describing the relationships between the points in different dimensions. Note that the order of inputs do not change in the above process. Meanwhile, the output has the same dimension and structure as the input, providing a great help to our loss function and the completeness of the point cloud. Since the calculations involved in the process are permutation invariant, a different order of inputs does not change the final output. As mentioned above, this is naturally suited to unordered point clouds.

We specifically design output header module to handle the point cloud denoising task. The module consists of a densely connected multi-layer perceptron containing four linear layers. Each linear layer has a GELU nonlinearity. The module uses residual connections to transform the task into simulated noise for better denoising. The final output of our network is the denoised coordinates of each point. The computational process can be expressed as follows:

where $X$ is the input coordinate set of the point cloud with noise and ${\bf F}_{o}$ represents the output of the transformer module. The output $Y$ represents the denoised coordinate.

Choosing an appropriate loss function is particularly critical for the denoising task as it directly affects the denoising performance. We consider the differences between synthetic datasets and raw point clouds to design a two-part loss function in supervised training.

For quantifying the distance between the denoised point cloud and the ground truth point cloud, we adopt the Chamfer Distance(CD) [11] as our loss function. This has also been shown to be effective by previous work [30, 26, 22]. We make some changes to the CD to make it easier to visualize during the training process. The exact expression is as follows:

Nevertheless, it has been shown in previous work [30, 26] that the use of a relative minimum distance loss function leads to inferior visual effects. And this phenomenon was confirmed again in our experiments. Particularly at several iterations, the point cloud can appear filamentous or clustered. To enable points to move away from each other, we choose an ‘absolute’ distance function as an additional loss function. Note that thanks to the point-to-point mapping nature of our network and the uniform distribution of points in the training set, we are able to choose this approach. Our specific loss function is as follows:

where ${\bf x}$ represents the input coordinates and ${\bf y}$ represents the output coordinates. Overall, our loss function is as follows:

The above processes are all evaluated on uniformly distributed point clouds. However, raw point clouds are generally non-uniform. Given this case, we reduce the weight of the ‘absolute’ distance function. We set $\alpha=0.9$ and $\beta=0.1$ in our experiments. By minimizing the overall loss function, we can ensure that the denoised points are restored to the underlying surface as closely as possible.