# Learning 3D-3D Correspondences for One-shot Partial-to-partial Registration

Zheng Dang,
IAIR, Xi’an Jiaotong University, China
<EMAIL_ADDRESS>
&Fei Wang,
IAIR, Xi’an Jiaotong University, China
<EMAIL_ADDRESS>
&Mathieu Salzmann,
EPFL-CVLab & ClearSpace, Switzerland,
<EMAIL_ADDRESS>
###### Abstract

While 3D-3D registration is traditionally tacked by optimization-based
methods, recent work has shown that learning-based techniques could achieve
faster and more robust results. In this context, however, only PRNet can
handle the partial-to-partial registration scenario. Unfortunately, this is
achieved at the cost of relying on an iterative procedure, with a complex
network architecture. Here, we show that learning-based partial-to-partial
registration can be achieved in a one-shot manner, jointly reducing network
complexity and increasing registration accuracy. To this end, we propose an
Optimal Transport layer able to account for occluded points thanks to the use
of outlier bins. The resulting OPRNet framework outperforms the state of the
art on standard benchmarks, demonstrating better robustness and generalization
ability than existing techniques.

##  1 Introduction

Registration aims to determine the rigid transformation, i.e., 3D rotation and
3D translation, between two 3D point sets. The traditional approach to this
problem is the Iterative Closest Point (ICP) algorithm [1]. In its vanilla
version, this iterative algorithm easily gets trapped in poor local optima.
While globally-optimal solutions [2, 3] have been proposed to remedy this,
they come at a high computational cost and lack robustness to noise, thus
greatly reducing their practical applicability.

Recently, learning-based methods [4, 5] were shown to outperform the previous
traditional, optimization-based techniques in terms of both speed and
robustness to noise. In particular, DCP [5] provides a one-shot strategy that
prevents the need for an expensive iterative process. Unfortunately, these
methods cannot handle the partial-to-partial registration scenario, thus only
working under the unrealistic assumption that both point sets are fully
observed, even at test time. This was addressed by PRNet [6], but by falling
back to an iterative keypoint-matching procedure, relying on a reinforcement
learning-inspired actor-critic strategy. This drastically increased the
complexity of the resulting network, requiring four 32GB GPUs for training,
instead of just one for DCP.

In this paper, we introduce a more direct, one-shot learning-based solution to
the partial-to-partial registration problem. To this end, we revisit the DCP
framework so as to remove its assumption that each 3D point in the target
(observed) point set can always be matched to a source (model) point.
Specifically, we replace the softmax operator used in DCP, which not only
assumes that each target point corresponds to a source point but also
incorrectly allows for one-to-many correspondences, with an Optimal Transport
(OT) layer. This layer leverages a measure of similarity between the target
and source points to compute a transportation cost, which we then use to
obtain a matching with Sinkhorn’s algorithm. To model _partial_
correspondences, we introduce outlier bins, allowing us to account for
occluded points. We further show that imposing direct supervision on the
output of the OT layer, instead of computing a loss function on the rigid
transformation obtained via SVD, as done in DCP, improves the robustness of
the resulting model.

Our contributions can be summarized as follows:

  * •

We introduce the first one-shot learning-based approach for partial-to-partial
registration, which we dub OPRNet, for One-shot Partial-to-partial
Registration Network.

  * •

We develop an Optimal Transport layer able to handle partial matches thanks to
outlier bins.

  * •

We show that better robustness can be achieved by direct supervision of the OT
layer output than by supervising the rigid transformation obtained after SVD.

Our experiments show that, despite being a one-shot algorithm, our approach
outperforms the state-of-the-art iterative PRNet on the standard ModelNet40
benchmark, including in challenging scenarios with unseen point-clouds and
unseen categories at test time. Furthermore, our approach can handle high
missing-point rates, such as those observed when acquiring point sets with
depth sensors, and generalizes well to real data without any re-training or
fine-tuning. Finally, our network retains the relatively compact size of DCP,
requiring training resources around 6 times smaller than PRNet. We will
release our code upon acceptance to facilitate reproducibility and future
research.

##  2 Related Work

Traditional point cloud registration. ICP is the best-known algorithm for
solving the point cloud registration problem. It comprises two steps: One
whose goal is to find the closest target point for each source point to
generate 3D-3D correspondences, and the other that computes the rigid
transformation from these correspondences by solving a least-square problem.
These two steps are repeated until a termination condition is satisfied.
Several variants, such as Generalized-ICP [7] and Sparse ICP [8], have been
proposed to improve robustness to noise and mismatches, and we refer the
reader to [9, 10] for a complete review of ICP-based strategies. The main
drawback of these methods is their requirement for a reasonable initialization
to converge to a good solution. Only relatively recently has this weakness
been addressed by the globally-optimal registration method Go-ICP [2]. In
essence, this approach follows a branch-and-bound strategy to search the
entire 3D motion space $SE(3)$. Motivated by this, other approaches to finding
a global solution have been proposed, via, e.g., Riemannian optimization [11],
convex relaxation [12], and mixed-integer programming [13]. While globally
optimal, these methods all come at a much higher computational cost than
vanilla ICP. This was, to some degree, addressed by the Fast Global
Registration (FGR) algorithm [3], which leverages a local refinement strategy
to speed up computation. While effective, FGR still suffers from the presence
of noise and outliers in the point sets, particularly because, as vanilla ICP,
it simply relies on 3D point-to-point distance to establish correspondences.
In principle, this can be addressed by designing point descriptors that can be
more robustly matched. Over the years, several works have tackled this task,
in both a non learning-based [14, 15, 16] and learning-based [17, 18] fashion.
Nowadays, however, these approaches are outperformed by end-to-end learning
frameworks, which directly take the point sets as input, as discussed below.

End-to-end learning with point sets. A key requirement to enable end-to-end
learning-based registration was the design of deep networks acting on
unstructured sets. Deep sets [19] and PointNet [20] constitute the pioneering
works in this direction. They use shared multilayer perceptrons to extract
high-dimensional features from the input point coordinates, and exploit a
symmetric function to aggregate these features. This idea was then extended in
PointNet++ [21] via a modified sampling strategy to robustify the network to
point clouds of varying density, in DGCNN [22] by building a graph over the
point cloud, in PointCNN [23] by learning a transformation of the data so as
to be able to process it with standard convolutional layers, and in PCNN [24]
via an additional extension operator before applying the convolutions. While
the above-mentioned works focused on other tasks than us, such as point cloud
classification or segmentation, end-to-end learning for registration has
recently attracted a growing attention. In particular, PointNetLK [4] combines
the PointNet backbone with the traditional, iterative Lucas-Kanade (LK)
algorithm [25] so as to form an end-to-end registration network; DCP [5]
exploits DGCNN backbones followed by Transformers [26] to establish 3D-3D
correspondences, which are then passed through an SVD layer to obtain the
final rigid transformation. While effective, PointNetLK and DCP cannot tackle
the partial-to-partial registration scenario. That is, they assume that both
point sets are fully observed, during both training and test time. This was
addressed by PRNet [6] via a deep network designed to extract keypoints from
each input set and match these keypoints. This network is then applied in an
iterative manner, so as to increasingly refine the resulting transformation.
Unfortunately, this iterative procedure, and the corresponding actor-critic
matching strategy, significantly increase the network resource requirements
compared to the one-shot DCP. Here, we show that accurate, one-shot partial-
to-partial registration can be achieved with a network requiring resources
similar to DCP.

##  3 Methodology

###  3.1 Problem Formulation

Let $\mathcal{X}\in\mathbb{R}^{M\times 3}$ and
$\mathcal{Y}\in\mathbb{R}^{N\times 3}$ be two sets of 3D points sampled from
the same object surface. We typically refer to $\mathcal{X}$ as the source
point set and to $\mathcal{Y}$ as the target point set. Registration then aims
to find a rigid transformation $\mathcal{T}$ that aligns $\mathcal{X}$ to
$\mathcal{Y}$. In this work, we tackle the partial-to-partial scenario, in
which both sets of samples can vary significantly, but nonetheless have _some_
overlap.

As discussed in [1], given correspondences between the two point sets, the
rigid transformation $\mathcal{T}$ can be obtained in closed form by solving a
Procrustes problem. Here, we therefore develop a one-shot deep network that
produces such correspondences. To this end, as illustrated by Fig. 1, we make
use of a feature embedding backbone that extracts a representation for each
point in each set, and compute an $M\times N$ score map $\mathcal{S}$ from
these features. We then formulate the matching problem as an optimal transport
[27] one, for which we obtain an approximate solution via a log-domain
Sinkhorn algorithm [28] incorporated in the network. Below, we discuss these
steps in more detail.

###  3.2 Feature Embedding Backbone

Figure 1:  Network architecture of OPRNet.

To process the individual point sets, we make use of the same DGCNN [22] \+
Transformer [26] combination as in DCP(v2) [5]. Specifically, DGCNN takes a
point set as input, from which it constructs a k-NN graph. It then extracts
point-wise features via standard convolutions on this graph, encoding diverse
levels of context by max-pooling the local features and concatenating the
resulting representations to the point-wise ones. Let $\theta^{x}$, resp.
$\theta^{y}$, be the final feature matrix, i.e., one $P$-dimensional feature
vector per 3D point, for $\mathcal{X}$, resp. $\mathcal{Y}$. The transformer
then learns a function $\phi:\mathbb{R}^{M\times P}\times\mathbb{R}^{N\times
P}\rightarrow\mathbb{R}^{M\times P}$, that combines the information of the two
point sets. Ultimately, this backbone produces descriptor matrices
$\mathbf{f}^{x}$, resp. $\mathbf{f}^{y}$, for $\mathcal{X}$, resp.
$\mathcal{Y}$, as

$\mathbf{f}^{x}=\theta^{x}+\phi(\theta^{x},\theta^{y}),\;\;\;\mathbf{f}^{y}=\theta^{y}+\phi(\theta^{y},\theta^{x})\;.$
(1)

Given these matrices, we then form a score map
$\mathcal{S}\in\mathbb{R}^{M\times N}$ by computing the similarity between
each source-target pair of descriptors. That is, we compute the $(i,j)$-th
element of $\mathcal{S}$ as

$\mathcal{S}_{i,j}=<\mathbf{f}^{x}_{i},\mathbf{f}^{y}_{j}>,\;\;\forall(i,j)\in[1,M]\times[1,N]\;,$
(2)

where $<\cdot,\cdot>$ is the inner product, and
$\mathbf{f}^{x}_{i},\mathbf{f}^{y}_{j}\in\mathbb{R}^{P}$. In DCP [5], this
score map is passed through a row-wise softmax so as to obtain
correspondences. These correspondences are then processed within the network
via an SVD layer to solve the Procrustes problem, and the resulting rigid
transformation is compared to the ground-truth one with a mean squared error
(MSE) loss.

This approach suffers from the following drawbacks. First, the softmax
operator assumes that every target point matches a source point, which does
not hold in partial-to-partial registration. Second, it allows multiple target
points to correspond to the same source point. Third, in practice, the use of
an SVD layer, if not carefully designed, can make the training process
unstable [29, 30, 31]. To address this, below, we introduce an OT layer that
explicitly reflects the underlying correspondence problem, and enable it to
handle partially-observed point sets via outlier bins. Furthermore, we improve
training stability by imposing direct supervision on the output of the OT
layer.

###  3.3 Optimal Transport Layer

Given the score matrix $\mathcal{S}\in\mathbb{R}^{M\times N}$, we seek to
obtain a partial assignment of the target points to the source points, which
would effectively provide us with correspondences. In this section, we first
introduce the outlier bins that allow us to handle partial point sets, and
then discuss how we compute the partial assignment.

Outlier bins. Because we tackle partial-to-partial registration, some points
in one set will have no corresponding point in the other. To model this,
inspired by [32], which tackles the 2D-2D matching scenario, we adopt a
dustbin design. Specifically, we extend the score matrix $\mathcal{S}$ by one
row and one column so as to form an augmented score matrix
$\mathcal{\bar{S}}$. The values at the newly-created positions in
$\mathcal{\bar{S}}$ are then set to

$\mathcal{\bar{S}}_{i,N+1}=\mathcal{\bar{S}}_{M+1,j}=\mathcal{\bar{S}}_{M+1,N+1}=\alpha,\;\;\forall
i\in[1,M],\;\forall j\in[1,N]\;,$ (3)

where $\alpha\in\mathbb{R}$ is a learnable parameter. The values at the other
indices directly come from $\mathcal{S}$.

Optimal transport. Given the augmented score map $\mathcal{\bar{S}}$, we aim
to find a partial assignment
$\mathcal{\bar{P}}\in\mathbb{R}^{(M+1)\times(N+1)}$, defining correspondences
between the two point sets, extended with the outlier bins. For such an
assignment to be valid, it should ensure that each point in one set has at
most one corresponding point in the other. However, it should allow multiple
points to be assigned to the outlier bins. We express this by constraining
$\mathcal{\bar{P}}$ to belong to the set

$\mathbf{U}(\mathbf{a},\mathbf{b})=\left\\{\mathcal{\bar{P}}\in\mathbb{R}_{+}^{(M+1)\times(N+1)}:\mathcal{\bar{P}}\mathbbm{1}_{M+1}=\mathbf{a}\>\text{and}\>\mathcal{\bar{P}}^{\top}\mathbbm{1}_{N+1}=\mathbf{b}\right\\}\;,$
(4)

where $\mathbf{a}=[\mathbbm{1}_{M}^{\top},N]^{\top}$, and
$\mathbf{b}=[\mathbbm{1}_{N}^{\top},M]^{\top}$, with
$\mathbbm{1}_{M}=[1,1,...,1]^{\top}\in\mathbb{R}^{M}$. Note that, here, we
further make use of Kantorovich’s relaxation, which allows the elements of
$\mathcal{\bar{P}}$ to take any non-negative value instead of constraining
them to be in $\\{0,1\\}$. More detail on this relaxation can be found in
Section 2.3 of [28]. The resulting set $\mathbf{U}(\mathbf{a},\mathbf{b})$ is
bounded and defined by $N+M+2$ equality constraints, forming a convex
polytope.

Then, from the optimal transport theory [33, 28], the probability matrix
$\mathcal{\bar{P}}$ can be obtained by solving

$\min_{\mathcal{\bar{P}}\in\mathbf{U}(\mathbf{a},\mathbf{b})}\langle\mathcal{\bar{S}},\mathcal{\bar{P}}\rangle-\lambda
E(\mathcal{\bar{P}})\;,$ (5)

where $\langle\cdot,\cdot\rangle$ is the Frobenius dot product and $E(\cdot)$
is an entropy regularization term defined as
$E(\mathcal{\bar{P}})=-\displaystyle\sum_{i,j}\mathcal{\bar{P}}_{i,j}\left(\log(\mathcal{\bar{P}}_{i,j})-1\right)$.

Log-domain Sinkhorn’s algorithm. The standard way to solve (5) consists of
making use of Sinkhorn’s algorithm [34, 28]. In concurrent work, this
algorithm has been employed within deep networks for 2D-2D matching with a
similar dustbin design to ours [35], and for 2D-3D matching assuming that
every point has a correspondence [36]. In practice, we found that using the
original version of this algorithm within a deep network, as in [36], suffers
from numerical overflow, preventing training to converge. To address this, we
adopt the log-domain computation suggested in Section 4.4 of [28]. The
resulting algorithm is provided in Algorithm 1, where the matrix operator
$logsumexp(\mathcal{A})=log(exp(\mathcal{A}_{1,1})+...+exp(\mathcal{A}_{i,j})+...+exp(\mathcal{A}_{M,N}))$.
Since all operations performed by this algorithm are differentiable, the
training errors can be backpropagated to the rest of the network.

Input : $\mathcal{\bar{S}}$, $\mathbf{a}=[\mathbbm{1}_{M}^{\top},N]^{\top}$,
$\mathbf{b}=[\mathbbm{1}_{N}^{\top},M]^{\top}$,
$\mathbf{f}_{0}=\mathbb{0}_{M+1}$, $\mathbf{g}_{0}=\mathbb{0}_{N+1}$,
$\lambda$ and $k$ (nb. of iters)

Output : Assignment matrix $\mathcal{\bar{P}}$

Init : $\mathcal{\bar{S}}=-\mathcal{\bar{S}}/\lambda$

while  _$l\leq k$_ do

$\mathbf{f}^{(l+1)}=\lambda
log(\mathbf{a})-\mathit{logsumexp}(\mathcal{\bar{S}}+\mathbbm{1}_{N+1}\mathbf{g}^{(l)\top})$

$\mathbf{g}^{(l+1)}=\lambda
log(\mathbf{b})-\mathit{logsumexp}(\mathcal{\bar{S}}+\mathbf{f}^{(l+1)}\mathbbm{1}_{M+1}^{\top})$

end while

$\mathcal{\bar{P}}_{i,j}=\mathit{exp}(\mathcal{\bar{S}}_{i,j}+\mathbf{f}_{i}+\mathbf{g}_{j})$
;

$\triangleright$ transfer back to the original domain.

Algorithm 1 Log-domain Sinkhorn’s algorithm.

###  3.4 Loss Function

To train our network, we propose to rely on direct supervision of the output
of the OT layer. To this end, let
$\mathcal{\bar{M}}\in\\{0,1\\}^{(M+1)\times(N+1)}$ be the matrix of ground-
truth correspondences, with a 1 indicating a correspondence between a pair of
points. Such correspondences can be estimated using the ground-truth
transformation matrix, as discussed in more detail in Section 4.1. We then
express our loss function as the negative log-likelihood of the assignment,
given by

$\mathcal{L}(\mathcal{\bar{P}},\mathcal{\bar{M}})=\frac{-\sum\limits_{i=1}^{M+1}\sum\limits_{j=1}^{N+1}(\log\mathcal{\bar{P}}_{i,j})\mathcal{\bar{M}}_{i,j}}{\sum\limits_{i=1}^{M+1}\sum\limits_{j=1}^{N+1}\mathcal{\bar{M}}_{i,j}}\;,$
(6)

where the denominator normalizes the loss value so that different training
samples containing different number of correspondences have the same influence
in the overall empirical risk.

###  3.5 Obtaining the Rigid Transformation at Test Time

At test time, we compute correspondences from the output $\mathcal{\bar{P}}$
of the OT layer. To this end, we search for the largest value in each row of
$\mathcal{\bar{P}}$. If the resulting index is a valid point, we treat the
pair as a correspondence; if the largest value occurs in the outlier bin, we
then discard the current sample. With the resulting set of matches, we obtain
the rigid transformation by solving the Procrustes problem. We refer the
reader to the supplementary material for the details of this step.

##  4 Experiments

Model | MSE(R) | RMSE(R) | MAE(R) | MSE(t) | RMSE(t) | MAE(t)
---|---|---|---|---|---|---
ICP | 1134.552 | 33.683 | 25.045 | 0.0856 | 0.293 | 0.250
FGR | 126.288 | 11.238 | 2.832 | 0.0009 | 0.030 | 0.008
Go-ICP | 195.985 | 13.999 | 3.165 | 0.0011 | 0.033 | 0.012
PointNetLK | 280.044 | 16.735 | 7.550 | 0.0020 | 0.045 | 0.025
DCP(v2) | 45.005 | 6.709 | 4.448 | 0.0007 | 0.027 | 0.020
PRNet | 10.235 | 3.199 | 1.454 | 0.0003 | 0.016 | 0.010
Ours | 0.107 | 0.328 | 0.0521 | 3.384e-06 | 0.00183 | 0.000281
Table 1:  Partial-to-partial registration with unseen objects. Model | MSE(R) | RMSE(R) | MAE(R) | MSE(t) | RMSE(t) | MAE(t)
---|---|---|---|---|---|---
ICP | 1217.618 | 34.894 | 25.455 | 0.086 | 0.293 | 0.251
FGR | 98.635 | 9.932 | 1.952 | 0.0014 | 0.038 | 0.007
Go-ICP | 157.072 | 12.533 | 2.940 | 0.0009 | 0.031 | 0.010
PointNetLK | 526.401 | 22.943 | 9.655 | 0.0037 | 0.061 | 0.033
DCP(v2) | 95.431 | 9.769 | 6.954 | 0.0010 | 0.034 | 0.025
PRNet | 24.857 | 4.986 | 2.329 | 0.0004 | 0.021 | 0.015
PRNet* | 15.624 | 3.953 | 1.712 | 0.0003 | 0.017 | 0.011
Ours | 0.127 | 0.357 | 0.0696 | 3.953e-06 | 0.00198 | 0.000406
Table 2:  Partial-to-partial registration with unseen categories.

We first compare our approach to the state of the art in the same settings
used in [5, 6], and then evidence its effectiveness in the more challenging
scenario where self-occluded points in one set are discarded, as with, e.g., a
LiDAR or depth sensor, and on real scans. Finally we provide an ablation study
of the different components of our approach.

###  4.1 Partial-to-Partial Registration on ModelNet40

ModelNet40. This dataset [37] contains 12,311 CAD models spanning 40 object
categories. As in [6], we split it into 9,843 training and 2,468 testing
samples. The point clouds in this dataset are normalized in the range [-1, 1]
on each axis. During training, we sample 1024 points from the original point
clouds to define the source point set $\mathcal{X}$. We then sample a random
rigid transformation along each axis, with rotation in
$[0^{\circ},45^{\circ}]$ and translation in $[-0.5,0.5]$, and apply it to the
point set $\mathcal{X}$ to obtain the target point set $\mathcal{Y}$. To
simulate partial scans of $\mathcal{X}$ and $\mathcal{Y}$, as in [6], we
randomly select one point in $\mathcal{X}$ and one point in $\mathcal{Y}$
independently, and only keep their 768 nearest neighbors.

Evaluation metrics. We report the mean squared error (MSE), root mean squared
error (RMSE), and mean absolute error (MAE) between the predicted
transformation and the ground truth, treating rotation and translation
separately. For the rotation, we evaluate these errors in degrees.

Implementation details. We implemented our network using Pytorch [38] and
trained it from scratch. We used the Adam optimizer [39] and set the learning
rate to $10^{-3}$. We relied on mini-batches of size 20 and trained the
network for 40,000 iterations. For the OT layer, we used $k=50$ iterations and
set $\lambda=1$. Training was performed on two NVIDIA 2080Ti GPUs.

To build the ground-truth assignment matrix $\mathcal{\bar{M}}$, we
transformed $\mathcal{X}$ using the ground-truth transformation $\mathcal{T}$
to get a transformed point set $\tilde{\mathcal{X}}$. We then computed the
pair-wise Euclidean distance matrix between $\tilde{\mathcal{X}}$ and
$\mathcal{Y}$, which we thresholded to 0.05 to obtain a correspondence matrix
$\mathcal{M}\in\\{0,1\\}$. Then we augmented $\mathcal{M}$ with an extra row
and column acting as outlier bins to get $\mathcal{\bar{M}}$. The points
without any correspondence were treated as outliers, and the corresponding
positions in $\mathcal{\bar{M}}$ were set to one. While this strategy does not
guarantee a bipartite matching, in practice, we found that one-to-many
correspondences occurred only very rarely. Strict bipartite matchings could be
obtained via forward-backward check, but we believe that the benefit of such
an operation would be outweighed by its computational cost.

Partial-to-Partial registration with unseen objects. As a first experiment, we
tackle the case where all categories are observed during training, but new
objects from these categories are provided at test time. To this end, we train
on all 9,843 samples and test on the remaining 2,468 ones, as in [6]. In Table
1, we compare our OPRNet with ICP, FGR [3], PointNetLK [4], DCP(v2) [5] and
PRNet [6]. OPRNet outperforms all the competitors by a large margin in all the
metrics.

Partial-to-Partial registration with unseen categories. To test the
generalization ability of our method, we perform an unseen-categories
experiment similar to that in [6]. Specifically, we split the ModelNet40
objects into separate training and testing categories. However, instead of
training our model on ShapeNetCore [40] and testing it on the resulting
ModelNet40 testing data, as done in [6] for PRNet, we restrict ourselves to
the much smaller amount of ModelNet40 training data. For the comparison to be
fair, we thus report the results of [6] trained on the same data as us as
_PRNet_ , and that of the model trained on ShapeNetCore as _PRNet*_. In any
event, as shown in Table 2, our OPRNet outperforms all the baselines, whether
trained on the same data as us or on a larger dataset.

Model | MSE(R) | RMSE(R) | MAE(R) | MSE(t) | RMSE(t) | MAE(t)
---|---|---|---|---|---|---
ICP | 1229.670 | 35.067 | 25.564 | 0.0860 | 0.294 | 0.250
FGR | 764.671 | 27.653 | 13.794 | 0.0048 | 0.070 | 0.039
Go-ICP | 150.320 | 12.261 | 2.845 | 0.0008 | 0.028 | 0.029
PointNetLK | 397.575 | 19.939 | 9.076 | 0.0032 | 0.0572 | 0.032
DCP(v2) | 47.378 | 6.883 | 4.534 | 0.0008 | 0.028 | 0.021
PRNet | 18.691 | 4.323 | 2.051 | 0.0003 | 0.017 | 0.012
Ours | 4.232 | 2.057 | 0.677 | 7.145e-05 | 0.00845 | 0.00270
Table 3:  Partial-to-partial registration with unseen point clouds containing Gaussian noise. Model | MSE(R) | RMSE(R) | MAE(R) | MSE(t) | RMSE(t) | MAE(t)
---|---|---|---|---|---|---
ICP* | 1152.602 | 33.950 | 19.651 | 0.0208 | 0.144 | 0.116
FGR* | 25.890 | 5.088 | 1.275 | 0.00183 | 0.0428 | 0.0109
DCP(softmax) | 0.796 | 0.892 | 0.151 | 5.307e-05 | 0.00728 | 0.00134
Ours | 0.116 | 0.341 | 0.0902 | 7.369e-06 | 0.00271 | 0.000790
Table 4:  Registration with self-occluded point sets. Note that ICP and FGR
failed in several cases, leading to unreasonably high errors. The * therefore
indicates that we report the average over the cases where MSE < 3000\.

Partial-to-Partial registration with unseen objects with gaussian noise. To
study the sensitivity of our method to noise, we perform an experiment in the
same setting as the first one, but adding Gaussian noise to the data. As in
[5, 6], the noise was sampled from $\mathcal{N}(0,0.01)$ and clipped to
$[-0.05,0.05]$. As shown in Table 3, our method still outperform all the
competitors.

###  4.2 Registration with Self-occluded Point Sets

We now turn to the case where we account for self-occlusion in one of the
point sets. This mimics the scenario where the point set is captured by a
sensor such as LiDAR or a depth camera. This is particularly challenging,
since, in essence, only roughly half of the points are observed.

Self-occluded dataset. For this experiment, we use the auto-aligned ModelNet40
dataset [37, 41]. This dataset contains the 40 object categories’ mesh models,
which we use to render depth maps given camera viewpoints and intrinsic
parameters. Specifically, we treat the full point clouds corresponding to the
meshes as source sets and the resulting depth maps as target sets. We use the
look-at method111https://www.scratchapixel.com/lessons/mathematics-physics-
for-computer-graphics/lookat-function to place the camera, and set the
distance between the camera and the center of the object to be 3. We then
randomly sample the elevation and azimuth in $[22.5^{\circ},67.5^{\circ}]$.

Implementation details. In the training process, we render the depth maps
online. We set the number of source points to be 1024, and that of target
points to be 512. We use the same hyper-parameter values as in the previous
experiments, except for the number of training iterations, which we restrict
to 20,000 because online depth rendering increases training time, which we
keep below 11 hours.

Figure 2:  Alignments and correspondences obtained with our approach in the
self-occlusion scenario. We plot the point sets (source in orange, target in
blue) after alignment, but shifted horizontally for clarity. The green lines
in the bottom row depict the correspondences found with our approach.

Registration results. We compare our approach with ICP and FGR. However, we do
not compare it against PRNet, because we do not have access to the four 32GB
GPUs required to train it. Furthermore, because the rendering process was
implemented using Pytorch3D, which relies on Pytorch1.5, the runtime of the
SVD layer increased significantly, to the point of requiring 416 hours to
train DCP. To circumvent this, we therefore replaced the SVD layer + MSE loss
of DCP with a softmax layer and the same NLL loss as us, and refer to this
baseline as DCP(softmax). Note that, as will be shown in Section 4.4,
replacing SVD + MSE with softmax + NLL improves accuracy, and we therefore
believe that our comparison is fair. The results of this comparison are
provided in Table 4, where we again consistently outperform the competitors.
In Fig. 2, we visualize qualitative results of our approach. Note that the
correspondences and alignments we obtain are highly accurate.

###  4.3 Partial-to-Partial Registration of Real Scans

Figure 3:  Registering real scans. We successfully align the source (orange)
and target (blue) sets.

We now demonstrate the use of our approach on the Stanford Bunny dataset [42],
which contains real scans of objects and was also used in [6]. We process the
data in the same way as in Section 4.1. Unlike in [6], where PRNet had to be
fine-tuned to handle such real scans, we directly apply our OPRNet trained for
the ModelNet40 unseen object experiments. As shown by Fig. 3, our model yields
accurate alignments, thus evidencing the good generalization ability of our
approach.

###  4.4 Method Analysis

Points |  ICP |  FGR |  DCP |  PRNet |  Ours
---|---|---|---|---|---
512 |  0.018 |  0.032 |  0.013 |  0.051 |  0.033
1024 |  0.028 |  0.059 |  0.018 |  0.075 |  0.052
2048 |  0.054 |  0.132 |  0.053 |  0.141 |  0.112
Table 5:  Inference time (in seconds) Model | MSE(R) | RMSE(R) | MAE(R) | MSE(t) | RMSE(t) | MAE(t)
---|---|---|---|---|---|---
DCP(SVD) | 550.782 | 23.468 | 17.827 | 0.0161 | 0.126 | 0.0981
DCP(softmax) | 22.923 | 4.787 | 2.811 | 0.00291 | 0.0539 | 0.0345
OT | 2.526 | 1.589 | 0.602 | 0.000158 | 0.00336 | 0.00550
Table 6:  Analysis of the effectiveness of our OT layer. DCP(SVD) indicates
that the rigid transformation is computed within the network, as in DCP. By
contrast, DCP(softmax) uses the NLL loss directly on the score map, but
without relying on optimal transport to compute the correspondences. Our OT
layer yields a significant accuracy boost in all the metrics.

Inference time. In Table 5, we compare the inference time of our method to
that of the baselines for different point set sizes. For this comparison to be
fair, we ran all methods on a desktop computer with an Intel(R) Core(TM)
i7-7700K CPU @ 4.20GHz, an Nvidia GTX 2080 Ti GPU, and 32GB memory. For the
non learning-based methods (ICP and FGR), we used their Open3D [43]
implementation. For the learning-based ones (DCP and PRNet), we used Pytorch
implementations. Our method yields comparable runtimes to FGR, while being
slightly faster than PRNet.

Effectiveness of the OT layer. To evaluate the benefits of our OT layer, we
compare its results with two alternatives. The first one consists of relying
on the SVD layer used in DCP. Specifically, this layer takes the DGCNN
features to build a score map, feeds the score map to a softmax, and then
solves the Procrustes problem using SVD to obtain the rigid transformation, on
which the loss function is computed. The second alternative consists of
directly imposing supervision on the output of the above-mentioned softmax,
i.e., avoiding the SVD computation but not relying on optimal transport to
obtain the correspondences. For these experiments, we used the same setup as
for partial-to-partial registration with unseen categories, except that we
removed the Transformers from feature embedding network of all the methods.
The reason for this modification is that, with such Transformers, training
DCP(SVD) requires at least one 32GB GPU, which we don’t have access to. As
shown in Table 6, while direct supervision of the score maps (softmax) yields
much more robust results than the use of an SVD layer, our OT layer yields a
further significant boost in all metrics.

##  5 Conclusion

We have introduced the first one-shot learning-based partial-to-partial
registration method. At the heart of our approach lies an optimal transport
layer enabling our network to handle partial correspondences between the two
input point sets. Our method does not require a complicated training strategy,
yet yields more accurate results than other approaches that do. Furthermore,
it has a reasonably low memory footprint and generalizes well to unseen data,
making it broadly applicable. In the future, we will aim to deploy it for
LiDAR-based registration, with a particular focus on the task of removing
debris in low Earth orbit.

##  6 Broader Impact

Our work tackles the general problem of partial-to-partial 3D registration. As
such, and because our approach yields high accuracy and generalizes to unseen
data, many applications can potentially benefit from it. In particular,
partial-to-partial registration is commonly used to reconstruct indoor scenes
[17], with potentially high impact in the gaming and mixed reality industry.
Furthermore, it can be used for self-localization in outdoor scenes, thus
greatly facilitating autonomous navigation, as illustrated by the KITTI
registration dataset [44]. This, we believe, can have a high societal impact,
to which we will aim to actively contribute by deploying our approach in the
context of space debris removal. From an ethical perspective, we do not
foresee major risks, because point clouds carry very little personal
information, such as gender or ethnicity. While the security risks might be
higher, e.g., in the presence of adversarial attacks, we expect adversarial
training to significantly mitigate them.

##  7 Registration as a Procrustes Problem

Given $n$ correspondences $(\mathbf{x}_{i},\mathbf{y}_{i})$ between the source
point set $\mathcal{X}$ and the target point set $\mathcal{Y}$, our goal is to
compute a rigid transformation, encoded via a rotation matrix $\mathbf{R}\in
SO(3)$ and a translation vector $\mathbf{t}\in\mathbb{R}^{3}$, that aligns
each $\mathbf{x}_{i}$ to its corresponding $\mathbf{y}_{i}$. This is achieved
by minimizing the mean-squared error

$\underset{\mathbf{R}\in
SO(3),\mathbf{t}}{\mathrm{argmin}}\frac{1}{n}\sum_{i=1}^{n}\lVert\mathbf{R}\mathbf{x}_{i}+\mathbf{t}-\mathbf{y}_{i}\rVert^{2}\;.$
(7)

To do so, we first discard the relative translation of the point sets by
computing their center of mass

$\bar{\mathbf{x}}=\frac{1}{n}\sum\limits^{n}_{i}\mathbf{x}_{i}\;,\bar{\mathbf{y}}=\frac{1}{n}\sum\limits^{n}_{i=1}\mathbf{y}_{i}\;,$
(8)

and removing them from each point as, e.g., $\mathbf{x}_{i}-\bar{\mathbf{x}}$.
Then, the rotation matrix can be obtained by minimizing the mean-squared error

$\displaystyle\underset{\mathbf{R}}{\mathrm{argmin}}$
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\lVert\mathbf{R}(\mathbf{x}_{i}-\bar{\mathbf{x}})-(\mathbf{y}_{i}-\bar{\mathbf{y}})\rVert^{2}$
(9) $\displaystyle s.t.\;\;\mathbf{R}^{\top}\mathbf{R}=\mathbf{I}$ (10)

which is a Procrustes problem. To solve it, we first compute the cross-
covariance matrix

$\mathbf{H}=\sum^{n}_{i=1}(\mathbf{x}_{i}-\bar{\mathbf{x}})(\mathbf{y}_{i}-\bar{\mathbf{y}})^{\top}\;.$
(11)

Then, we use singular value decomposition (SVD) to obtain
$\mathbf{H}=\mathbf{U}\mathbf{S}\mathbf{V}^{\top}$. The rotation matrix and
translation vector can then be computed in closed-form as

$\mathbf{R}=\mathbf{V}\Sigma\mathbf{U}^{\top},\;\;t=-\mathbf{R}\bar{\mathbf{x}}+\bar{\mathbf{y}}\;,$
(12)

where $\Sigma$ is a diagonal matrix, where the last value is defined as
$sign(\det(\mathbf{U}\mathbf{V}^{\top}))$ and the others are set to 1. This
guarantees that the determinant of $\mathbf{R}$ is positive, thus making it a
valid rotation matrix. For further details please check [45, 46].

## References

  * [1] P. Besl and N. Mckay.  A method for registration of 3d shapes.  PAMI, 14(2):239–256, February 1992.
  * [2] Jiaolong Yang, Hongdong Li, Dylan Campbell, and Yunde Jia.  Go-icp: A globally optimal solution to 3d icp point-set registration.  TPAMI, 38(11):2241–2254, 2015.
  * [3] Qian-Yi Zhou, Jaesik Park, and Vladlen Koltun.  Fast global registration.  In ECCV, pages 766–782. Springer, 2016.
  * [4] Yasuhiro Aoki, Hunter Goforth, Rangaprasad Arun Srivatsan, and Simon Lucey.  Pointnetlk: Robust & efficient point cloud registration using pointnet.  In CVPR, pages 7163–7172, 2019.
  * [5] Yue Wang and Justin M Solomon.  Deep closest point: Learning representations for point cloud registration.  In ICCV, pages 3523–3532, 2019.
  * [6] Yue Wang and Justin M Solomon.  Prnet: Self-supervised learning for partial-to-partial registration.  In NIPS, pages 8812–8824, 2019.
  * [7] Aleksandr Segal, Dirk Haehnel, and Sebastian Thrun.  Generalized-icp.  In In Robotics: Science and Systems, 2009.
  * [8] Sofien Bouaziz, Andrea Tagliasacchi, and Mark Pauly.  Sparse iterative closest point.  In Computer graphics forum, volume 32, pages 113–123. Wiley Online Library, 2013.
  * [9] François Pomerleau, Francis Colas, Roland Siegwart, et al.  A review of point cloud registration algorithms for mobile robotics.  Foundations and Trends® in Robotics, 4(1):1–104, 2015.
  * [10] Szymon Rusinkiewicz and Marc Levoy.  Efficient variants of the icp algorithm.  In Proceedings Third International Conference on 3-D Digital Imaging and Modeling, pages 145–152. IEEE, 2001.
  * [11] David M Rosen, Luca Carlone, Afonso S Bandeira, and John J Leonard.  Se-sync: A certifiably correct algorithm for synchronization over the special euclidean group.  The International Journal of Robotics Research, 38(2-3):95–125, 2019.
  * [12] Haggai Maron, Nadav Dym, Itay Kezurer, Shahar Kovalsky, and Yaron Lipman.  Point registration via efficient convex relaxation.  ACM Transactions on Graphics (TOG), 35(4):1–12, 2016.
  * [13] Gregory Izatt, Hongkai Dai, and Russ Tedrake.  Globally optimal object pose estimation in point clouds with mixed-integer programming.  In Robotics Research, pages 695–710. Springer, 2020.
  * [14] Andrew E. Johnson and Martial Hebert.  Using spin images for efficient object recognition in cluttered 3d scenes.  TPAMI, 21(5):433–449, 1999.
  * [15] Radu Bogdan Rusu, Nico Blodow, Zoltan Csaba Marton, and Michael Beetz.  Aligning point cloud views using persistent feature histograms.  In IROS, pages 3384–3391. IEEE, 2008.
  * [16] Radu Bogdan Rusu, Nico Blodow, and Michael Beetz.  Fast point feature histograms (fpfh) for 3d registration.  In ICRA, pages 3212–3217. IEEE, 2009.
  * [17] Andy Zeng, Shuran Song, Matthias Nießner, Matthew Fisher, Jianxiong Xiao, and Thomas Funkhouser.  3dmatch: Learning local geometric descriptors from rgb-d reconstructions.  In CVPR, pages 1802–1811, 2017.
  * [18] Marc Khoury, Qian-Yi Zhou, and Vladlen Koltun.  Learning compact geometric features.  In ICCV, pages 153–161, 2017.
  * [19] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola.  Deep sets.  In NIPS, pages 3391–3401, 2017.
  * [20] C.R. Qi, H. Su, K. Mo, and L.J. Guibas.  Pointnet: Deep learning on point sets for 3d classification and segmentation.  In CVPR, 2017.
  * [21] C.R. Qi, L. Yi, H. Su, and L.J. Guibas.  Pointnet++: Deep hierarchical feature learning on point sets in a metric space.  In NIPS, 2017.
  * [22] Y. Wang, Y. Sun, Z. Liu, S. Sarma, M. Bronstein, and J.M. Solomon.  Dynamic graph cnn for learning on point clouds.  In ACM Transactions on Graphics (TOG), 2019.
  * [23] Yangyan Li, Rui Bu, Mingchao Sun, Wei Wu, Xinhan Di, and Baoquan Chen.  Pointcnn: Convolution on x-transformed points.  In NIPS, pages 820–830, 2018.
  * [24] Matan Atzmon, Haggai Maron, and Yaron Lipman.  Point convolutional neural networks by extension operators.  In ACM Transactions on Graphics (TOG), 2018.
  * [25] Bruce D Lucas, Takeo Kanade, et al.  An iterative image registration technique with an application to stereo vision.  In IJCAI. Vancouver, British Columbia, 1981.
  * [26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In NIPS, pages 5998–6008, 2017.
  * [27] Cédric Villani.  Optimal transport: old and new, volume 338.  Springer Science & Business Media, 2008.
  * [28] Marco Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In NIPS, pages 2292–2300, 2013.
  * [29] Zheng Dang, Kwang Moo Yi, Yinlin Hu, Fei Wang, Pascal Fua, and Mathieu Salzmann.  Eigendecomposition-free training of deep networks with zero eigenvalue-based losses.  In ECCV, pages 768–783, 2018.
  * [30] Z. Dang, K. M. Yi, Y. Hu, F. Wang, P. Fua, and M. Salzmann.  Eigendecomposition-free training of deep networks for linear least-square problems.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 16, 2020.
  * [31] Wei Wang, Zheng Dang, Yinlin Hu, Pascal Fua, and Mathieu Salzmann.  Backpropagation-friendly eigendecomposition.  In NIPS, pages 3156–3164, 2019.
  * [32] Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich.  Superpoint: Self-supervised interest point detection and description.  In CVPR Workshops, pages 224–236, 2018.
  * [33] Gabriel Peyré, Marco Cuturi, et al.  Computational optimal transport.  Foundations and Trends® in Machine Learning, 11(5-6):355–607, 2019.
  * [34] Richard Sinkhorn and Paul Knopp.  Concerning nonnegative matrices and doubly stochastic matrices.  Pacific Journal of Mathematics, 21(2):343–348, 1967.
  * [35] Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich.  Superglue: Learning feature matching with graph neural networks.  In CVPR, pages 4938–4947, 2020.
  * [36] Liu Liu, Dylan Campbell, Hongdong Li, Dingfu Zhou, Xibin Song, and Ruigang Yang.  Learning 2d–3d correspondences to solve the blind perspective-n-point problem.  In arXiv Preprint, 2020.
  * [37] Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao.  3d shapenets: A deep representation for volumetric shapes.  In CVPR, pages 1912–1920, 2015.
  * [38] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer.  Automatic differentiation in pytorch.  In ICLR, 2017.
  * [39] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In ICLR, 2015.
  * [40] Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al.  Shapenet: An information-rich 3d model repository.  arXiv preprint arXiv:1512.03012, 2015.
  * [41] Nima Sedaghat, Mohammadreza Zolfaghari, Ehsan Amiri, and Thomas Brox.  Orientation-boosted voxel nets for 3d object recognition.  In BMVC, 2016.
  * [42] Greg Turk and Marc Levoy.  Zippered polygon meshes from range images.  In Proceedings of the 21st annual conference on Computer graphics and interactive techniques, pages 311–318, 1994.
  * [43] Qian-Yi Zhou, Jaesik Park, and Vladlen Koltun.  Open3D: A modern library for 3D data processing.  arXiv:1801.09847, 2018.
  * [44] Andreas Geiger, Philip Lenz, and Raquel Urtasun.  Are we ready for autonomous driving? the kitti vision benchmark suite.  In CVPR, pages 3354–3361. IEEE, 2012.
  * [45] Peter H Schönemann.  A generalized solution of the orthogonal procrustes problem.  Psychometrika, 31(1):1–10, 1966.
  * [46] David W Eggert, Adele Lorusso, and Robert B Fisher.  Estimating 3-d rigid body transformations: a comparison of four major algorithms.  Machine vision and applications, 9(5-6):272–290, 1997.