Deep-3DAligner: Unsupervised 3D Point Set Registration Network With Optimizable Latent Vector

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. The iterative closest point (ICP) algorithm [15] 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. Go-ICP [16] 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 [17] for the registration of partially overlapping 3D surfaces. The TPS-RSM algorithm was proposed by Chui and Rangarajan [18] to estimate parameters of non-rigid transformations with a penalty on second-order derivatives. Coherence point drift (CPD) was further proposed by Myronenko et al. [6] for non-rigid point set registration. Although the independent iterative optimization process limits the efficiency of registering a large number of pairs, inspiring us to leverage its advantage and equip it with a learning-based system for this task.

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 [14, 19, 20]. PointNetLK [21] 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 [22] to treat 3D point cloud registration as a motion process between points. Wang et al. proposed a deep closest point [11] 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. [14] proposed a voxelMorph CNN architecture to learn the registration field to align two volumetric medical images. 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. PR-Net [13] as the first work proposed a method to detect corresponding points from partial shapes and then solve the desired geometric transformation based on it. In comparison, our method detects the overlapping subsets based on our proposed adaptive Chamfer loss.

Given training dataset $\mathbf{D}=\{(S_{i},G_{i})\text{ ,where }S_{i}\in\mathbf{S},G_{i}\in\mathbf{G}\}$. $\mathbf{S}$ is the source point set and $\mathbf{G}$ is the target point set, and we have $\mathbf{S},\mathbf{G}\subset\mathbb{R}^{N}(N=2\text{ or }N=3)$. Assuming the existence of a function $g_{\theta}(S_{i},G_{i})=\phi$ and $g$ has parameter set $\theta$ in a neural network structure. When considering only rigid point set registration, the output $\phi$ usually contains a homogeneous transformation matrix including a rotation matrix and a translation vector. A model with optimized weights $\theta^{optimal}$ can generate the desired rotation and translation parameters $\phi$ in $g$ to further align the source and target point sets. The objective loss function $\mathcal{L}$ is usually a pre-defined similarity metric for evaluation of the alignment quality between transformed source and target point sets. Based on a given dataset $\mathbf{D}$, a stochastic gradient-descent-based algorithm can be used to update the weight parameters $\theta$ and to minimize the pre-defined loss function:

$$\begin{split}\theta^{optimal}=\operatorname*{argmin}_{\theta}[\mathbb{E}_{(S_{i},G_{i})\sim\mathbf{D}}[\mathcal{L}(S_{i},G_{i},g_{\theta}(S_{i},G_{i}))]]\end{split}$$

(1)

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 for the deep spatial feature extraction, and followed with a pre-defined correlation module. 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). 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 and to better equip the optimization process within the system, as shown in Figure 2, we define a trainable latent SCR (Spatial-Correlation Representation) feature for each pair of point sets. The design of SCR makes it possible to both leverage the common “knowledge” of registration from the dataset via the shared generator and individual adjustment for each input pair through the optimization of their SCR. After optimization, SCR should contain the spatial correlation information for each input pair. As shown in Figure 2, for a pair of source and target point sets $S_{i}$ and $G_{i}$, the randomly initialized latent vector $z_{i}\sim\mathcal{N}(0,0.01)$ from Gaussian distribution as an initialized SCR. The initialized SCR is optimized during the training process together with the transformation decoder. 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 and partial point sets as well.

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, $\forall x\in S_{i}$, we stack the coordinates of $x$ with $z_{i}$ and we note it as $[x,z_{i}]$. We leverage a multi-layer perceptron (MLP) structure to regress the rotation and translation parameters in the desired transformation. For each layer of the MLP, we define $\{g_{i}\}_{i=1,2,...,s}$, such that $g_{i}:\mathbb{R}^{v_{i}}\to\mathbb{R}^{v_{i+1}}$, where $v_{i}$ is the dimension of input layer and $v_{i+1}$ is the dimension of output layer. For each MLP layer, we use the ReLU activation function. For the output of last layer, we leverage a max pool function to accumulate the point features into a latent vector $L_{i}$, calculated as:

$$\begin{split}L_{i}=Maxpool\{g_{s}g_{s-1}...g_{1}([x_{j},z_{i}])\}_{x_{j}\in S_{i}}\end{split}$$

(2)

Taking the latent vector $L_{i}$ as the input, we have a further $t$ successive MLP layers with a ReLU activation function to regress the parameters $\phi_{i}$ of the desired transformation. We define function $\{\gamma_{i}\}_{i=1,2,...,t}$, such that $\gamma_{i}:\mathbb{R}^{w_{i}}\to\mathbb{R}^{w_{i+1}}$, where $w_{i}$ is the dimension of input layer and $w_{i+1}$ is the dimension of output layer.

$$\begin{split}\phi_{i}=\gamma_{t}\gamma_{t-1}...\gamma_{1}(L_{i})\end{split}$$

(3)

We further compute the transformed source point set, defined as $S_{i}^{\prime}$,

$$\begin{split}S_{i}^{\prime}=\mathbf{\mathbf{T}_{\phi_{i}}}(S_{i})\end{split}$$

(4)

where $\mathbf{T}_{\phi_{i}}$ denotes the desired geometric transformation with parameters $\phi_{i}$. Based on the transformed source point set $S_{i}^{\prime}$ and the target point set $G_{i}$, we further introduce the loss function.

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 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 $S_{i}^{\prime}=T_{\phi}(S_{i})$ and target points set $G_{i}$ as:

$$\begin{split}L_{\text{Chamfer}}(S_{i}^{\prime},G_{i})&=\sum_{x\in S_{i}^{\prime}}\min_{y\in G_{i}}||x-y||^{2}_{2}\\ &+\sum_{y\in G_{i}}\min_{x\in S_{i}^{\prime}}||x-y||^{2}_{2}\end{split}$$

(5)

For aligning partial-shapes, we further propose an adaptive Chamfer loss. For a time period $t$ during the optimization process, we assume a pre-defined distance threshold $\sigma_{t}$. For the transformed source point set $S_{i}^{\prime}$, we define the overlapping subset of it with the target point set $G_{i}$ as ${S^{\prime}_{i}}^{(t)}\subset{S^{\prime}_{i}}^{(t-1)}\subset...\subset{S^{\prime}_{i}}^{(0)}=S_{i}^{\prime}$, such that $\forall x\in{S^{\prime}_{i}}^{(t-1)},\text{if \ }\min\{||x-y||^{2}_{2}\}_{y\in G_{i}^{(t-1)}}<\sigma_{t}$, then $x\in{S^{\prime}_{i}}^{(t)}$. Similarly, for the target point set $G_{i}$, we define the overlapping subset between it with $S_{i}^{\prime}$ as ${G_{i}}^{(t)}\subset{G_{i}}^{(t-1)}\subset...\subset{G_{i}}^{(0)}=G_{i}$, such that $\forall x\in{G_{i}}^{(t-1)},\text{if \ }\min\{||x-y||^{2}_{2}\}_{y\in{S^{\prime}_{i}}^{(t-1)}}<\sigma_{t}$, then $x\in{G_{i}}^{(t)}$. Therefore, based on the overlapping subsets ${S_{i}^{\prime}}^{(t)}$ and ${G_{i}}^{(t)}$ for time period $t$, we define the adaptive Chamfer loss for $S_{i}^{\prime}$ and $G_{i}$ as:

$$\begin{split}L^{t}_{\text{Adaptive-Chamfer}}(S_{i}^{\prime},G_{i})&=\sum_{x\in S_{i}^{\prime(t)}}\min_{y\in G_{i}^{(t)}}||x-y||^{2}_{2}\\ &+\sum_{y\in G_{i}^{(t)}}\min_{x\in S_{i}^{\prime(t)}}||x-y||^{2}_{2}\end{split}$$

(6)

In section Spatial-Correlation Representation, we define a set of trainable latent vectors $\mathbf{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 $\mathbf{D}$, our training process can be expressed as:

$$\begin{split}\mathbf{\theta^{optimal},z^{optimal}}=\operatorname*{argmin}_{\theta,\mathbf{z}}[\mathbb{E}_{(S_{i},G_{i})\sim\mathbf{D}}[\mathcal{L}(S_{i},G_{i},g_{\theta}(S_{i},z_{i}))]],\end{split}$$

(7)

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

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

$$\begin{split}\mathbf{z^{optimal}}=\operatorname*{argmin}_{\mathbf{z}}[\mathbb{E}_{(S_{j},G_{j})\sim\mathbf{W}}[\mathcal{L}(S_{j},G_{j},g_{\tilde{\theta}}(S_{j},z_{j}))]].\end{split}$$

(8)

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_{i}}=T_{g_{\tilde{\theta}}(S_{i},z_{i}^{optimal})}$ and the transformed source shape can be generated by $S_{i}^{\prime}=T_{\phi_{i}}(S_{i}),\forall S_{i}\in\mathbf{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. For each source shape $S_{i}$ we generate the transformed shapes $G_{i}$ by applying a rigid transformation defined by the rotation matrix which is characterized by 3 rotation angles along the x-y-z-axis, where each value is uniformly sampled from $[0,45]$ unit degree, and the translation which is uniformly sampled from $[-0.5,0.5]$. At last, we simulate partial point sets by randomly select a point in unit space and keep its 768 nearest neighbors for source and target shapes.

We train our network using batch data from the training data set $\{(S_{i},G_{i})|S_{i},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,0.01)$ with a dimension of 256. For the Deep-3Daligner 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 [23] activation function and implement batch normalization [24] for every layer except the output layer. For adaptive Chamfer, $\sigma_{t}$ decreases from 10 to 0.01 in 100 epochs. The learning rate is set as 0.001 with exponential decay of 0.995 at each epoch. 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.

In this experiment, we follow previous works to test our model for 3D point set registration on unseen point sets in Test 1, and on unseen categories in Test 2.

In this experiment, we further verify the performance of our model for registering partial shapes.