Unsupervised Learning of 3D Point Set Registration

The development of optimization algorithms to estimate rigid and non-rigid geometric transformations in an iterative routine has attracted extensive research attention in past decades. Assuming that a pair of point sets are related by a rigid transformation, the standard approach is to estimate the best translation and rotation parameters in the iterative search routine, therein aiming to minimize a distance metric between two sets of points. The iterative closest point (ICP) algorithm [14] is one successful solution for rigid registration. It initializes an estimation of a rigid function and then iteratively chooses corresponding points to refine the transformation. However, the ICP algorithm is reported to be vulnerable to the selection of corresponding points for initial transformation estimation. Go-ICP [21] was further proposed by Yang et al. to leverage the BnB scheme for searching the entire 3D motion space to solve the local initialization problem brought by ICP. Zhou et al. proposed fast global registration [22] for the registration of partially overlapping 3D surfaces. The TPS-RSM algorithm was proposed by Chui and Rangarajan [23] to estimate parameters of non-rigid transformations with a penalty on second-order derivatives. As a classical non-parametric method, coherence point drift (CPD) was proposed by Myronenko et al. [10], which successfully introduced a process of fitting the Gaussian mixture likelihood to align the source point set with the target point set. Existing classical algorithms have achieved great success on the registration task. Although the independent iterative optimization process limits the efficiency of registering a large number of pairs, inspiring us to design a learning-based system for this task.

In recent years, learning-based methods have achieved great success in many fields of computer vision [24, 25, 26, 27, 28, 29, 30, 31, 32, 33]. In particular, recent works have started a trend of directly learning geometric features from cloud points (especially 3D points), which motivates us to approach the point set registration problem using deep neural networks [20, 19, 30, 27, 28, 29, 34, 35, 36, 37]. PointNetLK [38] was proposed by Aoki et al. to leverage the newly proposed PointNet algorithm for directly extracting features from the point cloud with the classical Lucas $\&$ Kanade algorithm for the rigid registration of 3D point sets. Liu et al. proposed FlowNet3D [18] to treat 3D point cloud registration as a motion process between points. Wang et al. proposed a deep closest point [15] model, which first leverages the DGCNN structure to exact the features from point sets and then regress the desired transformation based on it. Balakrishnan et al. [19] proposed a voxelMorph CNN architecture to learn the registration field to align two volumetric medical images. For the registration of 2D images, an outstanding registration model was proposed by Rocco et al. [20]. For the learning-based registration solutions listed above, the main challenge concerns how to effectively model the “geometric relationship” between source and target objects in a learning-based approach. For example, [20] proposed a correlation tensor between the feature maps of source and target images. [19] leveraged a U-Net-based structure to concatenate features of source and target voxels. [18] [38] used a PointNet-based structure, and [15] used a DGCNN structure to learn the features from a point set for further registration decoding. In contrast, we first propose a model-free structure to skip the encoding step. Instead, we initialize an SCR feature without pre-defining a model, which is to be optimized with the weights of the network from the alignment loss back-propagation process.

Given training dataset $\mathbb{D}=\{(\mathbb{S_{i}},\mathbb{G_{j}})\text{ ,where }\mathbb{S_{i}},\mathbb{G_{j}}\subset\mathbb{R}^{N}(N=2\text{ or }N=3)\}$, the optimization task of a deep-learning-based method for registration problem can be generally defined in the following way. We assume the existence of a function $g_{\theta}(\mathbb{S_{i}},\mathbb{G_{j}})=\phi$ using a neural network structure, where $\phi$ represents the parameters of the network. The rigid point set registration is represented by a homogeneous transformation matrix, which is composed by a rotation matrix $\mathbb{R}\in SO(3)$ and a translation vector $\mathbb{t}\in\mathbb{R}^{3}$. Given a pair of input source and target point sets $(\mathbb{S_{i}},\mathbb{G_{j}})$, a trained model is able to predict the parameters $\phi$ based on the optimized weights $\mathbb{\theta^{optimal}}$ in the neural network structure. A pre-defined alignment metric between transformed source and target point sets can be defined as objective loss function to update weights $\theta$. For a given dataset $\mathbb{D}$, a stochastic gradient-descent-based algorithm can usually be utilized for the optimization of the weights $\theta$ to minimize the pre-defined loss function:

where $\mathcal{L}$ represents a similarity metric.

In this paper, we define the spatial correlation representation as the latent feature that characterizes the essence of spatial correlation between a given pair of source and target point sets. As shown in Figure 1, to compute the SCR feature, source and target point sets are usually fed to a feature in previous works (i.e. PointNet [27]) for the deep spatial feature extraction, and followed with a pre-defined correlation module (i.e. [20]). However, the design of an appropriate feature encoder for unstructured point clouds is challenging compared to the standard discrete convolutions assume the availability of a grid structured input (e.g. 2D image). Furthermore, the design of a correlation module for a pair of input spatial features has a significant impact on the transformation decoder. The limitation of the hand-crafted design of modules for the extraction of individual spatial feature and spatial correlation feature motivates us to design a model-free based SCR as described below.

To eliminate the side effects of the hand-craft design in feature encoder and correlation module, as shown in Figure 2, we define a trainable latent SCR (Spatial-Correlation Representation) feature for each pair of point sets. As shown in Figure 1, our SCR optimizer, which takes input as a randomly initialized latent vector to reason the optimal SCR for the point set registration task, replaces the previously hand-crafted feature encoder and correlation module. More specifically, as shown in Figure 2, for a pair of source and target point sets $(\mathbf{S},\mathbf{G})$, the randomly initialized latent vector $\mathbb{z}$ from Gaussian distribution as an initialized SCR. The initialized SCR is optimized during the training process together with the transformation decoder. In this way, in comparison with previous methods, we avoid the challenging problem of explicit defining the spatial correlation between two point sets but simply optimize the SCR feature from a random initialized status during the model training process. The implicit design of SCR allows Deep-3DAligner more flexibility in spatial correlation feature learning that is more adaptive for the alignment of unseen point sets.

Given the above spatial-correlation representation (SCR) feature, we then design a decoding network to regress the desired transformation parameters, as illustrated in Figure 2. More specifically, we first formulate the input by stacking the coordinates of each point $x$ in source point set $S$ with its corresponding SCR feature $\mathbb{z}$. We note this input as $[x,\mathbf{z}]$. Then we define a multi-layer perceptron (MLP) architecture for learning the parameters of the rigid transformation to transform the source point set toward the target point set. This architecture includes successive MLP layers with the ReLU activation function, $\{g_{i}\}_{i=1,2,...,s}$, such that $g_{i}:\mathbb{R}^{v_{i}}\to\mathbb{R}^{v_{i+1}}$, where $v_{i}$ and $v_{i+1}$ are the dimensions of the layer inputs and outputs respectively. Then, we use a max pool layer to exact the global feature $\mathbf{L}$, calculated as:

where the notation [*,*] represents the concatenation of vectors in the same domain.

We further decode the global feature $\mathbb{L}$ to the transformation parameters by a second network, which includes $t$ successive MLP layers with a ReLU activation function $\{\gamma_{i}\}_{i=1,2,...,t}$ such that $\gamma_{i}:\mathbb{R}^{w_{i}}\to\mathbb{R}^{w_{i+1}}$, where $w_{i}$ and $w_{i+1}$ are the dimensions of the layer inputs and outputs respectively.

Now, we can get the transformed source point set $\mathbb{S^{\prime}}$ by

where $\mathbb{T}_{\phi}$ denotes the transformation function defined by the predicted transformation parameters $\phi$. Based on the transformed source point set and the target point set, we can further define the alignment loss function in the next section.

In our unsupervised setting, we do not have the ground truth transformation for supervision and we do not assume a direct correspondence between these two point sets. Therefore, a distance metric between two point sets, instead of the point/pixel-wise loss is desired. In addition, A suitable metric should be differentiable and efficient to compute. In this paper, we adopt the Chamfer distance proposed in [39] as our loss function. The Chamfer loss is a simple and effective alignment metric defined on two non-corresponding point sets. We formulate the Chamfer loss between our transformed source point set $T_{\phi}(\mathbf{S})$ and target points set $\mathbf{G}$ as:

where $\phi$ represents all the parameters to define the transformation and $\phi$ is predicted from our model based on the optimized SRC feature and the trained decoder.

In section 3.2, we define a set of trainable latent vectors $\mathbb{z}$, one for each pair of point sets as the SCR feature. During the training process, these latent vectors are optimized along with the weights of network decoder using a stochastic gradient descent-based algorithm. For a given training dataset $\mathbb{D}$, our training process can be expressed as:

where $\mathcal{L}$ represents the pre-defined loss function.

For a given testing dataset $\mathbb{W}$, we fix the network parameters $\tilde{\theta}=\mathbb{\theta^{optimal}}$ and only optimize the SRC features:

The learned decoder network parameters $\tilde{\theta}$ here provides a prior knowledge for the optimization of SRC. After this optimization process, the desired transformation can be determined by $T_{\phi}=T_{g_{\tilde{\theta}}(\mathbb{S_{i}},\mathbf{z_{i}^{optimal}})}$ and the transformed source shape can be generated by $\mathbf{S_{i}^{\prime}}=T_{\phi}(\mathbf{S_{i}}),\forall\mathbb{S_{i}}\in\mathbb{W}$.

We test the performance of our model for 3D point set registration on the ModelNet40 dataset. This dataset contains 12311 pre-processed CAD models from 40 categories. For each 3D point object, we uniformly sample 1024 points from its surface. Following the settings of previous work, points are centered and re-scaled to fit in the unit sphere. To demonstrate the robustness of our model in the presence of various noise types, we add noise to the target point sets for model evaluation. To prepare the position drift (P.D.) noise, a zero-mean Gaussian is applied to each point in the target point set. The level of P.D. noise is defined as the standard deviation of Gaussian Noise. To prepare the data incompleteness (D.I.) noise, we randomly remove a certain amount of points from the entire point set. The level of D.I. noise is defined as the ratio of the eliminated points and the entire set. To prepare the data outlier (D.O.) noise, we randomly add a certain amount of points generated by a zero-mean Gaussian to the point set. The level of D.O. noise is defined as the ratio of the added points to the entire point set.

We train our network using batch data from the training data set $\{(\mathbf{S_{i}},\mathbf{G_{i}})|\mathbf{S_{i}},\mathbf{G_{i}}\in\mathbf{D}\}_{i=1,2,...,b}$. We set the batch size $b$ to 128. The latent vectors are initialized from a Gaussian distribution $\mathcal{N}(0,1)$ with a dimension of 2048. For the decoding network, the first part includes 2 MLP layers with dimensions (256,128) and a max pool layer. Then, we use 3 additional MLPs with dimensions of (128, 64, 3) for decoding the rotation matrix and with dimensions of (128, 64, 3) for decoding the translation matrix. We use the leaky-ReLU [40] activation function and implement batch normalization [41] for every layer except the output layer. Our model is optimized with Adam optimizer. The learning rate is set as 0.001 with exponential decay of 0.995 at each epoch. For the outlier and missing point cases, we clip the Chamfer distance by a fixed value of 0.1.

We use the mean squared error (MSE), root mean squared error (RMSE), and mean absolute error (MAE) to measure the performance of our model and all comparing methods. Lower values indicate better alignment performance. All angular measurements in our results are in units of degrees. The ground-truth labels are only used for the performance evaluation and are not used during the training/testing process.

For performance evaluation, we compare our method with both supervised and unsupervised methods. The current state-of-the-art results are achieved by DCP, which is a supervised approach including four different versions. The version DCPv1+MLP uses deep neural networks (DNNs) to model the transformation. The version DCPv2+MLP improves its performance by integrating a supervised attention mechanism. DCPv1+SVD further improves its performance by integrating an additional SVD-based fine-tuning process, and DCPv2+SVD integrates both attention and SVD to further boost the performance. Since our Deep-3DAligner is an unsupervised approach and only use DNNs as auxiliary function to model the transformation, we will use DCPv1+MLP as the baseline model for compassion.

In this experiment, we follow previous works, DCP [15] and PointNetLK [38], to test our model for 3D point set registration on unseen point sets.

In this experiment, we follow previous works, DCP and PointNetLK, to test our model for 3D point set registration on objects from unseen categories.

In this experiment, we further verify our model’s performance for 3D rigid point set registration in the presence of P.D. noise.

In this experiment, we further verify our model’s performance for 3D rigid point set registration in the presence of D.I. noise.

In this experiment, we further verify our model’s performance for 3D rigid point set registration in the presence of D.O. noise.

In this experiment, we conduct further experiments to verify the design of our spatial-correlation representation (SCR) and decoder network.

Given source and target point sets with a size of $N=2048$, we can first define a relative position tensor $N\times 3$ which characterizes the relative positions between two point sets and the alignment is measured using Chamfer loss. Our approach essentially leverages a multi-layer neural network to learn the nonlinear mapping function that maps the relative position tensor $N\times 3$ to geometric transformation representation (R,t) (dimension of 6). Technically, it is possible to directly optimize the chamfer distance over the parameters (R,t) by setting up one single-layer network (i.e. the geometric transformation representation layer). However, practically it is not feasible to just train a single-layer neural network that is capable of mapping the high-dimensional relative position tensor (e.g. dimension of $3,000$ for point sets of $1,000$) to low-dimensional geometric transformation representation (R,t) (dimension of 6). Therefore, our method uses a multi-layer neural network to model this non-linear mapping/dimension reduction problem. In addition, to better formulate the concept of relative position tensor, we propose to design spatial correlation representation (SCR) features. During training, we jointly optimize the SCR and decoder to complete the non-linear mapping process. In the testing phase, we fix the trained decoder and only optimize the SCR for a new given pair of source and target point sets.