COTReg: Coupled Optimal Transport based Point Cloud Registration

The best-known traditional registration algorithm is Iterative Closest Point (ICP) [24], which alternates between rigid motion estimation and the correspondences searching [25, 26] by solving a $L_{2}$-optimization. However, ICP converges to spurious local minima due to the non-convexity. Based on this algorithm, many variants have been proposed. For example, the Levenberg-Marquardt ICP [27] uses a Levenberg-Marquardt algorithm to produce a transformation by replacing the singular value decomposition with gradient descent and Gaussian-Newton approaches, accelerating convergence while ensuring high accuracy. Go-ICP [28] solves the point cloud alignment problem globally by using a Branch-and-Bound (BnB) optimization framework without prior information on correspondence or transformation. RANSAC-like algorithms are widely used for robust finding of the correct correspondences for registration [29]. FGR [30] optimizes a Geman-McClure cost-induced correspondence-based objective function in a graduated non-convex strategy and achieves high performance. TEASER [31] reformulates the registration problem as an intractable optimization and provides readily checkable conditions to verify the optimal solution. However, most existing methods still face challenges when point clouds contain a mixture of noise, outlier, and partial overlap [32]. In contrast, the learning-based methods tend to have strong robustness [33].

Deep point cloud registration methods can be commonly classified into correspondences-free methods and correspondences-based methods [32]. The core idea of correspondence-free registration approaches [34, 35] is to estimate the transformation by minimizing the difference between global features extracted from two input point clouds. However, the global features rely on the whole points of a point cloud, leading to a major limitation that correspondences-free approaches are inapplicable to real scenes with partial overlaps [11, 32]. On the contrary, correspondence-based methods extract the per-point or per-patch embeddings of the two input point clouds to generate a mapping and estimate a rigid transformation [32]. For instance, DCP [25] employs DGCNN [36] for feature extraction and an attention module to generate soft matching pairs. FCGF [20] uses $1\times 1\times 1$ kernel to extract compact metric features for geometric correspondence. RGM [16] considers information exchange between two point clouds to extract discriminative features for point-wise matching. Predator [15] and PRNet [1] pay attention to the detection of points in the overlap region and use these features of the detected points to generate matches. DGR [11] proposes a 6-dimensional convolutional network architecture for inlier likelihood prediction and estimates the transformation via a weighted Procrustes module. 3DRegNet [17] uses an inlier prediction model to estimate the inliers. PointDSC [12] integrates pairwise similarity as a constraint to improve correspondence accuracy. RPMNet [37] and [38] propose methods to solve the point cloud partial visibility by integrating the Sinkhorn algorithm into a network to get soft correspondences from local features. Soft correspondences can increase the robustness of registration accuracy. IDAM [39] incorporates both geometric and distance features into the iterative matching process. Some works [13, 14, 40, 41] also focus on alleviating the dependence on the manual annotation in the process of extracting features. Most of these methods generate correspondences by applying the point feature similarities then directly match the highest response [39]. This strategy has two obvious limitations. First of all, some inliers are assigned to outliers, since inliers and outliers are treated equally in the correspondence prediction stage [15]. Second, the one-shot matching may fail because there are multiple possible correspondences due to randomness [39], while the per-point feature itself is not powerful enough to distinguish the geometric properties in point clouds [16]. Based on such observations, some additional constraints should be imposed to obtain high-quality correspondences. Inspired by edge similarity in graph matching [22, 23], the matched pairs between two point clouds hold an equal length constraint in both Euclidean and feature spaces, we propose a coupled optimal transport-based framework that incorporates overlap scores into pointwise matching as well as structural matching to generate more accurate correspondences. We also design a positional encoding scheme that assigns intrinsic geometric properties to per-point features to enhance distinctions among point features in indistinctive regions.

Optimal transport (OT) [21] is a method to exploit the best assignment between two general objects, which is widely applied in many machine learning applications [42]. OT provides a way to predict the correspondences and measure the similarity between two distributions. The distance based on OT is called the Wasserstein distances (WD), which has been used in various computer vision tasks [43]. For instance, Su et al. [44] employed the Wasserstein distance to deal with the 3D shape matching and surface registration problem. Dang et al. [38] applied Wasserstein distances to handle 3D point cloud registration. Gromov-Wasserstein (GWD) [21] extends WD by computing couplings between metric measure spaces. Unlike calculating the distance between two entity sets as in WD, GWD can be utilized to calculate the structure distance within each domain. The GWD has been used for shape analysis [45], graph matching [46], etc. The properties of WD and GWD that focus, respectively, on features and structure information motivate us to utilize WD to build pointwise matching and GWD to build structural matching. Contrasting with previous correspondence prediction algorithms for point-cloud registration, we treat point clouds as probability distributions, i.e., overlap scores are embedded in a specific metric space based on the similarities of features and structures. The correspondence prediction is then used to compute a coupled optimal transport distance between two probability distributions that can endow pointwise and structural matchings with overlap scores to generate more accurate correspondences.

As is mentioned above, only considering the pointwise matching is not sufficient to find the accurate correspondences due to the ambiguous and repeated patterns in the 3D acquisition point clouds. Therefore, we exploit both pointwise and structural matchings to indicate the correspondence between source and target point clouds.

For point-wise matching, the cross-distance matrix $\bm{C}^{pq}\in\mathbb{R}^{N\times M}$ is derived based on the feature space between point cloud $\bm{\mathcal{P}}$ and $\bm{\mathcal{Q}}$ with each element satisfying

where $\mathcal{D}_{f}\left(\bm{f}_{\bm{p}_{i}},\bm{f}_{\bm{q}_{j}}\right)=\|\frac{\bm{f}_{\bm{p}_{i}}}{\|\bm{f}_{\bm{p}_{i}}\|_{2}}-\frac{\bm{f}_{\bm{q}_{j}}}{\|\bm{f}_{\bm{q}_{j}}\|_{2}}\|_{2}$ represents the distance between $\bm{f}_{\bm{p}_{i}}$ and $\bm{f}_{\bm{q}_{j}}$.

To construct the structural matching, we first calculate the discrepancy (structure) matrix $\bm{C}^{p}\in\mathbb{R}^{N\times N}$ of the point pairs within $\bm{\mathcal{P}}$ in both Euclidean and feature spaces with the elements related to $\{\bm{p}_{i},\bm{p}_{k}\}$ satisfying

where $\mathcal{D}_{e}\left(\cdot,\cdot\right)$ is a function related to distances between two points in Euclidean space satisfying $\mathcal{D}_{e}\left(\bm{p}_{i},\bm{p}_{k}\right)=2\tanh(\|\bm{p}_{i}-\bm{p}_{k}\|_{2})$. $\lambda\in[0,1]$ is a hyperparameter controlling the contribution of feature and coordinate information. Similarly, the elements of $\bm{C}^{q}\in\mathbb{R}^{M\times M}$ related to $\{\bm{q}_{j},\bm{q}_{l}\}$ for $\bm{\mathcal{Q}}$ satisfy

We jointly exploit the pointwise and structural matchings equipped with overlap scores to indicate the correspondences between source and target point clouds. This leads to an optimization problem that can be solved by a coupled optimal transport method (called coupled optimal transport since it contains two types of distance, i.e., WD and GWD).

where $\bm{1}_{n}$ denotes an $n$-dimensional all-one vector. $\xi_{1}$ and $\xi_{2}$ are two non-negative hyperparameters controlling the pointwise and structural matchings, respectively. If $\xi_{1}>0$ and $\xi_{2}=0$, then it only depends on the pointwise matching, and if $\xi_{1}=0$ and $\xi_{2}>0$, it only considers the structural matching. When $\xi_{1}>0$ and $\xi_{2}>0$, it allows the framework to effectively consider both pointwise and structural matchings for better correspondence predictions by sharing the assignment matrix $\Gamma$. After obtaining $\Gamma$, we can apply some methods, such as SVD and RANSAC, to estimate the transformation.

Next, we introduce the Wasserstein distance-based pointwise matching and the Gromov-Wasserstein distance-based structural matching, respectively.

Our Wasserstein distance-based pointwise matching method can be regarded as a variant of the nearest neighbor search with an additional bijectivity constraint that enforces global matching consistency. Our model is appropriate for partial registration, since the overlap scores have been introduced. Given two point clouds $\bm{\mathcal{P}}$ and $\bm{\mathcal{Q}}$, with their associated features $\bm{\mathcal{F}}_{p}$ and $\bm{\mathcal{F}}_{q}$, overlap scores $\bm{\mu}_{p}$ and $\bm{\mu}_{q}$, the assignment matrix $\Gamma$ can be estimated based on the following Theorem.

Please refer to the proof of this theorem in Appendix A.

Our method is based on the observations as follows: for $\forall\leavevmode\nobreak\ \bm{p}_{i},\bm{p}_{k}\in\bm{\mathcal{P}}$ and $\forall\leavevmode\nobreak\ \bm{q}_{j},\bm{q}_{l}\in\bm{\mathcal{Q}}$ with their associated features $\bm{f}_{\bm{p}_{i}},\bm{f}_{\bm{p}_{k}}\in\bm{\mathcal{F}}_{p}$ and $\bm{f}_{\bm{q}_{j}},\bm{f}_{\bm{q}_{l}}\in\bm{\mathcal{F}}_{q}$, if $\bm{p}_{i}\rightarrow\bm{q}_{j}$ and $\bm{p}_{k}\rightarrow\bm{q}_{l}$ are correct correspondences, then the distance between $\bm{p}_{i}$ and $\bm{p}_{k}$ should be similar to the distance between $\bm{q}_{j}$ and $\bm{q}_{l}$. Thus, the structural difference in both Euclidean and feature spaces should be small, i.e., $|c^{e}\left(\bm{p}_{i},\bm{p}_{k}\right)-c^{e}\left(\bm{q}_{j},\bm{q}_{l}\right)|$ and $|c^{f}\left(\bm{f}_{\bm{p}_{i}},\bm{f}_{\bm{p}_{k}}\right)-c^{f}\left(\bm{f}_{\bm{q}_{j}},\bm{f}_{\bm{q}_{l}}\right)|$ should be small. Gromov-Wasserstein distance is often applied to find the correspondences between two sets of samples based on their pairwise intra-domain similarity (or distance) matrices. Thus, we can approximately transform the correspondence prediction to a structural matching problem based on the Gromov-Wasserstein distance as the following Theorem.

The proof is available in Appendix B.

Now, we introduce how to solve the problem in (8). For simplicity, we denote a matrix $\bm{H}\left(\bm{C}^{p},\bm{C}^{q},\Gamma\right)\in\mathbb{R}^{N\times M}$ with each element $[\bm{H}\left(\bm{C}^{p},\bm{C}^{q},\Gamma\right)]_{kl}=\sum_{i=1}^{N}\sum_{j=1}^{M}\left(\bm{C}^{p}_{ik}-\bm{C}^{q}_{jl}\right)^{2}\Gamma_{ij}$. We get $\sum_{ijkl}\Gamma_{ij}\Gamma_{kl}\left(\bm{C}^{p}_{ik}-\bm{C}^{q}_{jl}\right)^{2}=\left<\bm{H}\left(\bm{C}^{p},\bm{C}^{q},\Gamma\right),\Gamma\right>$. The problem in Eq. (8) can be rewritten as

Standard optimal transport only allows a meaningful comparison of measures with the same total mass, i.e., $\sum_{i=1}^{N}\bm{\mu}_{p_{i}}=\sum_{j=1}^{M}\bm{\mu}_{q_{j}}$, which does not always satisfy the registration requirement due to multiple correspondences. Following [49], we replace the constraints in Eq. (11b) with soft-marginals ($\mathcal{KL}$ divergence). Optimization in Eq. (11) is then translated into an unconstrained approximate transport problem

where $\tau>0$ is a regularization parameter to adjust the strength of penalization of the soft margins. We use the generalized proximal point method [50] and projected gradient descent to solve the problem in Eq. (12) based on the $\mathcal{KL}$ metric. Following [45, 50], we fix $\Gamma^{(k)}$ at iteration $k+1$ for $k\geq 0$, $\mathcal{KL}\left(\Gamma|\Gamma^{(k)}\right)$ acts as a regularization centered on the previous solution $\Gamma^{(k)}$. The update rule for Eq. (12) at iteration $k+1$ can be written as

with initialization $\Gamma^{(0)}=\bm{\mu}_{p}{\bm{\mu}_{q}}^{\top}$. $\epsilon>0$ is a regularization parameter. $\mathcal{KL}\left(\Gamma|\Gamma^{(k)}\right)$ can be interpreted as a damping term that encourages $\Gamma^{(k+1)}$ not to be very far from $\Gamma^{(k)}$. For $\epsilon$ small enough, $\Gamma^{(k+1)}$ in Eq. (13) converges to the optimal solution of problem in Eq. (11) as $\tau$ increases. Choosing $\epsilon$ trades off convergence speed with closeness to the original transport problem [48]. The solution of the problem in Eq. (13) is based on the following theorem.

Please refer to the proof of the theorem in Appendix C. The problem in Eq. (14) can be solved using the Sinkhorn algorithm [48, 49]. Specifically,

Let $(\bm{u}^{k},\bm{v}^{k},\bm{a}^{k},\bm{b}^{k})$ be the solution returned at the $k$-th iteration of the algorithm. Here $\bm{a}=B(\bm{u},\bm{v})\textbf{1}_{M}$ and $\bm{b}=B(\bm{u},\bm{v})^{\top}\textbf{1}_{N}$ with $B(\bm{u},\bm{v})=\text{diag}\left(\exp{\left(\frac{\bm{u}}{\epsilon}\right)}\right)\cdot\exp{\left(-\frac{\bm{C}}{\epsilon}\right)}\cdot\text{diag}\left(\exp{\left(\frac{\bm{v}}{\epsilon}\right)}\right)$. Suppose we are at iteration $k+1$ for $k\geq 0$ with a fixed $\bm{v}^{k}$, i.e.,

Multiplying both sides by $\exp\left(\frac{\bm{u}_{i}^{k}}{\epsilon}\right)$, we get

Similarly, with $\bm{u}^{k}$ fixed, we get

We use the pseudocode in Algorithm 1 to illustrate the solution. The inner iterations can be determined by $\epsilon,\bm{C}_{ij}$ and $\max\{M,N\}$ and the proof is similar to Theorem 2 in [49]. In our experiments, we found that when setting $\xi_{1}=1.0$, $\tau=5.0$, $\epsilon=0.001$, $N_{I}=100$ and $N_{O}=20$, it can obtain satisfactory results.

The solution of the optimization problem in Eq. (8) is sought over the space of $N\times M$ permutation matrices. Because of memory constraints and speed limitations, it is not suitable to solve large scale registration problem. To this end, COTReg adopts a hierarchical matching strategy that first establishes superpoint-level correspondences and then predicts point-level correspondences according to superpoint-level matches. The pipeline of our COTReg is illustrated in Figure 3, which is a shared weighted two-stream encoder-decoder network. Given a pair of point cloud $\bm{\mathcal{P}}$ and $\bm{\mathcal{Q}}$, the encoder aggregates the raw points into superpoints $\bar{\bm{\mathcal{P}}}=\{\bar{\bm{p}}_{i}\in\mathbb{R}^{3}|i=1,2,...,\bar{N}\}$ and $\bar{\bm{\mathcal{Q}}}=\{\bar{\bm{q}}_{j}\in\mathbb{R}^{3}|j=1,2,...,\bar{M}\}$, while jointly learning the associated features $\bm{\mathcal{F}}_{\bar{p}}=\{\bm{f}_{\bar{p}_{i}}\in\mathbb{R}^{b}|i=1,2,...,\bar{N}\}$ and $\bm{\mathcal{F}}_{\bar{q}}=\{\bm{f}_{\bar{q}_{j}}\in\mathbb{R}^{b}|j=1,2,...,\bar{M}\}$. The overlap attention block updates the features as $\bar{\bm{\mathcal{F}}}_{\bar{p}}$ and $\bar{\bm{\mathcal{F}}}_{\bar{q}}$, and projects them to coarse level overlap score vectors $\bm{\mu}_{\bar{p}}=\{\bm{\mu}_{\bar{p}_{i}}\in\left[0,1\right]\}_{i=1}^{\bar{N}},\bm{\mu}_{\bar{q}}=\{\bm{\mu}_{\bar{q}_{j}}\in\left[0,1\right]\}_{i=1}^{\bar{M}}$. The updated features and overlap scores are then used to calculate the coarse-level correspondences. Finally, the decoder transforms the superpoint level features and overlap scores to per-point feature descriptors $\bm{\mathcal{F}}_{p}$ and $\bm{\mathcal{F}}_{q}$, and overlap scores $\bm{\mu}_{p}$ and $\bm{\mu}_{q}$, which are used to estimate fine-level correspondences.

Our model is an end-to-end learning framework, using the ground truth correspondences as supervision. The loss function $\mathcal{L}=\mathcal{L}_{C}+\mathcal{L}_{F}+\mathcal{L}_{CO}+\mathcal{L}_{FO}$ is composed of an coarse-level loss $\mathcal{L}_{C}$ for superpoint matching, a point matching loss $\mathcal{L}_{F}$ for point matching, a binary classification loss $\mathcal{L}_{CO}$ for coarse-level overlap scores, and a classification loss $\mathcal{L}_{FO}$ for fine-level overlap scores.