Deep Closest Point: Learning Representations for Point Cloud Registration

Introduction

Geometric registration is a key task in many computational fields, including medical imaging, robotics, autonomous driving, and computational chemistry. In its most basic incarnation, registration involves the prediction of a rigid motion to align one shape to another, potentially obfuscated by noise and partiality.

Many modeling and computational challenges hamper the design of a stable and efficient registration method. Given exact correspondences, singular value decomposition yields the globally optimal alignment; similarly, computing matchings becomes easier given some global alignment information. Given these two observations, most algorithms alternate between these two steps to try to obtain a better result. The resultant iterative optimization algorithms, however, are prone to local optima.

In this work, we revisit ICP from a deep learning perspective, addressing key issues in each part of the ICP pipeline using modern machine learning, computer vision, and natural language processing tools. We call our resulting algorithm Deep Closest Point (DCP), a learning-based method that takes two point clouds and predicts a rigid transformation aligning them.

Our model consists of three parts: (1) We map the input point clouds to permutation/rigid-invariant embeddings that help identify matching pairs of points (we compare PointNet and DGCNN for this step); then, (2) an attention based module combining pointer network predicts a soft matching between the point clouds; and finally, (3) a differentiable singular value decomposition layer predicts the rigid transformation. We train and test our model end-to-end on ModelNet40 in various settings, showing our model is not only efficient but also outperforms ICP and its extensions, as well as the recently-proposed PointNetLK method . Our learned features generalize to unseen data, suggesting that our model is learning salient geometric features.

We identify sub-network architectures designed to address difficulties in the classical ICP pipeline.

We propose a simple architecture to predict a rigid transformation aligning two point clouds.

We evaluate efficiency and performance in several settings and provide an ablation study to support details of our construction.

We analyze whether local or global features are more useful for registration.

We release our codehttps://github.com/WangYueFt/dcp, to facilitate reproducibility and future research.

Related Work

ICP is the best-known algorithm for solving rigid registration problems; it alternates between finding point cloud correspondences and solving a least-squares problem to update the alignment. ICP variants consider issues with the basic method, like noise, partiality, and sparsity; probabilistic models also can improve resilience to uncertain data. ICP can be viewed as an optimization algorithm searching jointly for a matching and a rigid alignment; hence, propose using the Levenberg–Marquardt algorithm to optimize the objective directly, which can yield a better solution. For more information, summarize ICP and its variants developed over the last 20 years.

Learning on graphs and point sets:

A broad class of deep architectures for geometric data termed geometric deep learning includes recent methods learning on graphs and point clouds .

The graph neural network (GNN) is introduced in ; similarly, defines convolution on graphs (GCN) for molecular data. uses renormalization to adapt to graph structure and applies GCN to semi-supervised learning on graphs. MoNet learns a dynamic aggregation function based on graph structure, generalizing GNNs. Finally, graph attention networks (GATs) incorporate multi-head attention into GCNs. DGCNN can be regarded as graph neural network applied to point clouds with dynamic edges.

Sequence-to-sequence learning and pointer networks:

Many tasks in natural language processing, including machine translation, language modeling, and question answering, can be formulated as sequence-to-sequence problems (seq2seq). first uses deep neural networks (DNN) to address seq2seq problems at large scale. Seq2seq, however, often involves predicting discrete tokens corresponding to positions in the input sequence. This problem is difficult because there is an exponential number of possible matchings between input and output positions. Similar problems can be found in optimal transportation , combinatorial optimization , and graph matching . To address this issue, In our registration pipeline, we use a related method to Pointer Networks , which use attention as a pointer to select from the input sequence. In each output step, a Pointer Network predicts a distribution over positions and uses it as a “soft pointer.” The pointer module is fully differentiable, and the whole network can be trained end-to-end.

Non-local approaches:

To denoise images, non-local means leverages the simple observation that Gaussian noise can be removed by non-locally weighted averaging all pixels in an image. Recently, non-local neural networks have been proposed to capture long-range dependencies in video understanding; uses the non-local module to denoise feature maps to defend against adversarial attacks. Another instantiation of non-local neural networks, known as relational networks , has shown effectiveness in visual reasoning , meta-learning , object detection , and reinforcement learning . Its counterpart in natural language processing, attention, is arguably the most fruitful recent advance in this discipline. replaces recurrent neural networks with a model called Transformer, consisting of several stacked multi-head attention modules. Transformer-based models outperform other recurrent models by a considerable amount in natural language processing. In our work, we also use Transformer to learn contextual information of point clouds.

Problem Statement

Define centroids of $\mathcal{X}$ and $\mathcal{Y}$ as

Then the cross-covariance matrix $\boldsymbol{H}$ is given by

We can use the singular value decomposition (SVD) to decompose $\boldsymbol{H}=\boldsymbol{U}\boldsymbol{S}\boldsymbol{V}^{\top}$. Then, the alignment minimizing $E(\cdot,\cdot)$ in (1) is given in closed-form by

Here, a mapping $m$ from each point in $\mathcal{X}$ to its corresponding point in $\mathcal{Y}$ is given by

Equations (5) and (6) form a classic chicken-and-egg problem. If we know the optimal rigid transformation $[\boldsymbol{R}_{\mathcal{XY}},\boldsymbol{t}_{\mathcal{XY}}]$, then the mapping $m$ can be recovered from (6); conversely, given the optimal mapping $m$, the transformation can be computed using (4).

ICP iteratively approaches a stationary point of $E$ in (5), including the mapping $m(\cdot)$ as one of the variables in the optimization problem. It alternates between two steps: finding the current optimal transformation based on a previous mapping $m^{k-1}$ and finding an optimal mapping $m^{k}$ based on the current transformation using (6), where $k$ denotes the current iteration. The algorithm terminates when a fixed point or stall criterion is reached. This procedure is easy to implement and relatively efficient, but it is extremely prone to local optima; a distant initial alignment yields a poor estimate of the mapping $m$, quickly leading to a situation where the algorithm gets stuck. Our goal is to use learned embeddings to recover a better matching $m(\cdot)$ and use that to compute a rigid transformation, which we will detail in next section.

Deep Closest Point

Having established preliminaries about the rigid alignment problem, we are now equipped to present our Deep Closest Point architecture, illustrated in Figure 2. In short, we embed point clouds into high-dimensional space using PointNet or DGCNN (§4.1), encode contextual information using an attention-based module (§4.2), and finally estimate an alignment using a differentiable SVD layer (§4.4).

The first stage of our pipeline embeds the unaligned input point clouds $\mathcal{X}$ and $\mathcal{Y}$ into a common space used to find matching pairs of points between the two clouds. The goal is to find an embedding that quotients out rigid motion while remaining sensitive to relevant features for rigid matching. We evaluate two possible choices of learnable embedding modules, PointNet and DGCNN .

Since we use per-point embeddings of the two input point clouds to generate a mapping $m$ and recover the rigid transformation, we seek a feature per point in the input point clouds rather than one feature per cloud. For this reason, in these two network architectures, we use the representations generated before the last aggregation function, notated $\mathcal{F}_{\mathcal{X}}=\{\boldsymbol{x}_{1}^{L},\boldsymbol{x}_{2}^{L},...,\boldsymbol{x}_{i}^{L},...,\boldsymbol{x}_{N}^{L}\}$ and $\mathcal{F}_{\mathcal{Y}}=\{\boldsymbol{y}_{1}^{L},\boldsymbol{y}_{2}^{L},...,\boldsymbol{y}_{i}^{L},...,\boldsymbol{y}_{N}^{L}\}$, assuming a total of $L$ layers.

While PointNet largely extracts information based on the embedding of each point in the point cloud independently, DGCNN explicitly incorporates local geometry into its representation. In particular, given a set of points $\mathcal{X}$, DGCNN constructs a $k$-NN graph $\mathcal{G}$, applies a nonlinearity to the values at edge endpoints to obtain edgewise values, and performs vertex-wise aggregation ($\max$ or $\sum$) in each layer. The forward mechanism of DGCNN is thus

where $\mathcal{N}_{i}$ denotes the neighbors of vertex $i$ in graph $\mathcal{G}$. While PointNet features do not incorporate local neighborhood information, we find empirically that DGCNN’s local features are critical for high-quality matching in subsequent steps of our pipeline (see §6.1).

2 Attention

Our transition from PointNet to DGCNN is motivated by the observation that the most useful features for rigid alignment are learned jointly from local and global information. We additionally can improve our features for matching by making them task-specific, that is, changing the features depending on the particularities of $\mathcal{X}$ and $\mathcal{Y}$ together rather than embedding $\mathcal{X}$ and $\mathcal{Y}$ independently. That is, the task of rigidly aligning, say, organic shapes might require different features than those for aligning mechanical parts with sharp edges. Inspired by the recent success of BERT , non-local neural networks , and relational networks using attention-based models, we design a module to learn co-contextual information by capturing self-attention and conditional attention.

Notice we treat $\phi$ as a residual term, providing an additive change to $\mathcal{F}_{\mathcal{X}}$ and $\mathcal{F}_{\mathcal{Y}}$ depending on the order of its inputs. The idea here is that the map $\mathcal{F}_{\mathcal{X}}\mapsto\Phi_{\mathcal{X}}$ modifies the features associated to the points in $\mathcal{X}$ in a fashion that is knowledgeable about the structure of $\mathcal{Y}$; the map $\mathcal{F}_{\mathcal{Y}}\mapsto\Phi_{\mathcal{Y}}$ serves a symmetric role. We choose $\phi$ as an asymmetric function given by a Transformer ,For details, see http://nlp.seas.harvard.edu/2018/04/03/attention.html. since the matching problem we encounter in rigid alignment is analogous to the sequence-to-sequence problem that inspired its development, other than their use of positional embeddings to describe where words are in a sentence.

3 Pointer Generation

The most common failure mode of ICP occurs when the matching estimate $m^{k}$ is far from optimal. When this occurs, the rigid motion subsequently estimated using (6) does not significantly improve alignment, leading to a spurious local optimum. As an alternative, our learned embeddings are trained specifically to expose matching pairs of points using a simple procedure explained below. We term this step pointer generation, again inspired by terminology in the attention literature introduced in §4.2.

To avoid choosing non-differentiable hard assignments, we use a probabilistic approach that generates a (singly-stochastic) “soft map” from one point cloud into the other. That is, each $\boldsymbol{x}_{i}\in\mathcal{X}$ is assigned a probability vector over elements of $\mathcal{Y}$ given by

4 SVD Module

The final module in our architecture extracts the rigid motion from the soft matching computed in §4.3. We use the soft pointers to generate a matching averaged point in $\mathcal{Y}$ for each point in $\mathcal{X}$:

To backpropagate gradients through the networks, we need to differentiate the SVD. describes a standard means of computing this derivative; version of this calculations are included in PyTorch and TensorFlow . Note we need to solve only $3\times 3$ eigenproblems, small enough to be solved using simple algorithms or even (in principle) a closed-form formula.

5 Loss

Combined, the modules above map from a pair of point clouds $\mathcal{X}$ and $\mathcal{Y}$ to a rigid motion $[\boldsymbol{R}_{\mathcal{XY}},\boldsymbol{t}_{\mathcal{XY}}]$ that aligns them to each other. The initial feature module (§4.1) and the attention module (§4.2) are both parameterized by a set of neural network weights, which must be learned during a training phase. We employ a fairly straightforward strategy for training, measuring deviation of $[\boldsymbol{R}_{\mathcal{XY}},\boldsymbol{t}_{\mathcal{XY}}]$ from ground truth for synthetically-generated pairs of point clouds.

We use the following loss function to measure our model’s agreement to the ground-truth rigid motions:

Experiments

We compare our models to ICP, Go-ICP , Fast Global Registration (FGR) , and the recently-proposed PointNetLK deep learning method . We denote our model without attention (§4.2) as DCP-v1 and the full model with attention as DCP-v2. Go-ICP is ported from the authors’ released code. For ICP and FGR, we use the implementations in Intel Open3D . For PointNetLK, we adapt the code partially released by the authors.https://github.com/hmgoforth/PointNetLK Notice that FGR uses additional geometric features.

The architecture of DCP is shown in Figure 2. We use 5 EdgeConv (denoted as DGCNN ) layers for both DCP-v1 and DCP-v2. The numbers of filters in each layer are $$. In the Transformer layer, the number of heads in multi-head attention is 4 and the embedding dimension is 1024. We use LayerNorm without Dropout . Adam is used to optimize the network parameters, with initial learning rate 0.001. We divide the learning rate by 10 at epochs 75, 150, and 200, training for a total of 250 epochs. DCP-v1 does not use the Transformer module but rather employs identity mappings$\Phi_{\mathcal{X}}=\mathcal{F}_{\mathcal{X}}$and$\Phi_{\mathcal{Y}}=\mathcal{F}_{\mathcal{Y}}$.

We experiment on the ModelNet40 dataset, which consists of 12,311 meshed CAD models from 40 categories. Of these, we use 9,843 models for training and 2,468 models for testing. We follow the experimental settings of PointNet , uniformly sampling 1,024 points from each model’s outer surface. As in previous work, points are centered and rescaled to fit in the unit sphere, and no features other than $(x,y,z)$ coordinates appear in the input.

We measure mean squared error (MSE), root mean squared error (RMSE), and mean absolute error (MAE) between ground truth values and predicted values. Ideally, all of these error metrics should be zero if the rigid alignment is perfect. All angular measurements in our results are in units of degrees.

In our first experiment, we randomly divide all the point clouds in the ModelNet40 dataset into training and test sets, with no knowledge of the category label; different point clouds are used during training and during testing. During training, we sample a point cloud $\mathcal{X}$. Along each axis, we randomly draw a rigid transformation; the rotation along each axis is uniformly sampled in $[0,45^{\circ}]$ and translation is in $[-0.5,0.5]$. $\mathcal{X}$ and a transformation of $\mathcal{X}$ by the rigid motion are used as input to the network, which is evaluated against the known ground truth using (11).

Table 1 evaluates performance of our method and its peers in this experiment (vanilla ICP nearly fails). DCP-v1 already outperforms other methods under all the performance metrics, and DCP-v2 exhibits even stronger performance. Figure 4 shows results of DCP-v2 on some objects.

2 ModelNet40: Category Split

To test the generalizability of different models, we split ModelNet40 evenly by category into training and testing sets. We train DCP and PointNetLK on the first 20 categories, then test them on the held-out categories. ICP, Go-ICP and FGR are also tested on the held-out categories. As shown in Table 2, on unseen categories, FGR behaves more strongly than other methods. DCP-v1 has much worse performance than DCP-v2, supporting our use of the attention module. Although the learned representations are task-dependent, DCP-v2 exhibits smaller error than others except FGR, including the learning-based method PointNetLK.

3 ModelNet40: Resilience to Noise

We also experiment with adding noise to each point of the input point clouds. We sample noise independently from $\mathcal{N}(0,0.01)$, clip the noise to $[-0.05,0.05]$, and add it to $\mathcal{X}$ during testing. In this experiment, we use the model from §5.1 trained on noise-free data from all of ModelNet40.

Table 3 shows the results of this experiment. ICP typically converges to a far-away fixed point, and FGR is sensitive to noise. Go-ICP, PointNetLK and DCP, however, remain robust to noise.

4 DCP Followed By ICP

Since our experiments involve point clouds whose initial poses are far from aligned, ICP fails nearly every experiment we have presented so far. In large part, this failure is due to the lack of a good initial guess. As an alternative, we can use ICP as a local algorithm by initializing ICP with a rigid transformation output from our DCP model. Figure 3 shows an example of this two-step procedure; while ICP fails at the global alignment task, with better initialization provided by DCP, it converges to the global optimum. In some sense, this experiment shows how ICP can be an effective way to “polish” the alignment generated by DCP.

5 Efficiency

We profile the inference time of different methods on a desktop computer with an Intel I7-7700 CPU, an Nvidia GTX 1070 GPU, and 32G memory. Computational time is measured in seconds and is computed by averaging 100 results. As shown in Table 4, DCP-v1 is the fastest method among our points of comparison, and DCP-v2 is only slower than vanilla ICP.

Ablation Study

We conduct several ablation experiments in this section, dissecting DCP and replacing each part with an alternative to understand the value of our construction. All experiments are done in the same setting as the experiments in §5.1.

We first try to answer whether the localized features gathered by DGCNN provide value over the coarser features that can be measured using the simpler PointNet model. As discussed in , PointNet learns a global descriptor of the whole shape while DGCNN learns local geometric features via constructing the $k$-NN graph. We replace the DGCNN with PointNet (denoted as PN) and conduct the experiments in §5.1 on ModelNet40 , using DCP-v1 and DCP-v2. Table 5. Models perform consistently better with DGCNN that their counterparts with PointNet.

2 MLP or SVD?

While MLP is in principle a universal approximator, our SVD layer is designed to compute a rigid motion specifically. In this experiment, we examine whether an MLP or a custom-designed layer is better for registration. We compare MLP and SVD with both DCP-v1 and DCP-v2 on ModelNet40. Table 6 shows both DCP-v1 and DCP-v2 perform better with SVD layer than MLP. This supports our motivation to compute rigid transformation using SVD.

3 Embedding Dimension

Conclusion

In some sense, the key observation in our Deep Closest Point technique is that learned features greatly facilitate rigid alignment algorithms; by incorporating DGCNN and an attention module, our model reliably extracts the correspondences needed to find rigid motions aligning two input point clouds. Our end-to-end trainable model is reliable enough to extract a high-quality alignment in a single pass, which can be improved by iteration or “polishing” via classical ICP.

DCP is immediately applicable to rigid alignment problems as a drop-in replacement for ICP with improved behavior. Beyond its direct usage, our experiments suggest several avenues for future inquiry. One straightforward extension is to see if our learned embeddings transfer to other tasks like classification and segmentation. We could also train DCP to be applied iteratively (or recursively) to refine the alignment, rather than attempting to align in a single pass; insight from reinforcement learning could help refine approaches in this direction, using mean squared error as reward to learn a policy that controls when to stop iterating. Finally, we can incorporate our method into larger pipelines to enable high-accuracy Simultaneous Localization and Mapping (SLAM) or Structure from Motion (SFM).

Acknowledgement

The authors acknowledge the generous support of Army Research Office grant W911NF-12-R-0011, of National Science Foundation grant IIS-1838071, from an Amazon Research Award, from the MIT-IBM Watson AI Laboratory, from the Toyota-CSAIL Joint Research Center, and from the Skoltech-MIT Next Generation Program. Yue Wang wants to thank David Palmer for helpful discussion.

References