PCRDiffusion: Diffusion Probabilistic Models for Point Cloud Registration

Most traditional methods need a good initial transformation and converge to the local minima near the initialization point. One of the most profound methods is the Iterative Closest Point (ICP) algorithm [37], which begins with an initial transformation and iteratively alternates between solving two trivial subproblems: finding the closest points as correspondence under current transformation, and computing optimal transformation by SVD [38] based on identified correspondences. Though ICP can complete a high-precision registration, it is susceptible to the initial perturbation and easily prone to local optima. Thus, variants of ICP have been proposed [39, 40, 41, 42, 43], and they can improve the defects of ICP and enhance the registration performance [44]. Generalized-ICP [39] integrates a probabilistic module into an ICP-style framework to enhance the algorithm’s robustness. Go-ICP [40] combines the ICP that integrated the maximum correntropy criterion [45] with the Branch-and-Bound scheme and searches for an optimal solution in the 3D motion space. Besides, Trimmed ICP [46] introduces the least trimmed squares into each part of the ICP pipeline to tackle the partial overlap problem. However, all these methods retain a few essential drawbacks. Firstly, they depend strongly on the initialization. Secondly, it is difficult to integrate them into the deep learning pipeline as they lack differentiability. Thirdly, explicit estimation of corresponding points leads to quadratic complexity scaling with the number of points [47], which can introduce significant computational challenges.

In order to address the aforementioned problems, learning-based methods have made significant advancements in recent years, which are usually divided into correspondence-based methods [14, 15, 16, 17] and correspondence-free methods [9, 10, 11, 12, 13]. Specifically, in the correspondence-free methods, they are often necessary to seek differences between global features and require them to be sensitive to posture. PointNetLK [9] incorporates the Lucas & Kanade (LK) algorithm the PointNet [31] to iteratively align the input point clouds. In PCRNet [48], the LK algorithm is substituted with a fully-connected network and the model improves robustness by regressing. Correspondence-based methods need to find the correspondences between point clouds through precise matching and inlier estimation. IDAM [49] introduces a dual stage point elimination method to facilitate the generation of partial correspondences. However, the efficacy of IDAM remains limited under conditions of low overlapping. EGST [50] construct reliable geometric structure descriptors to extract correspondences and is less sensitive to outliers. OPRNet [51] resorts to the Sinkhorn algorithm [52] tailored to partial-to-partial registration. RPMNet [53] utilizes 3D local patch descriptor [54] to find correspondences in unorganized point clouds. However, all methods rely on updating source point cloud to further refine the predicted transformation parameters, as shown in Figure 2. In this paper, we aim to push forward the development of the point cloud registration pipeline further with diffusion probabilistic models and no iteration strategy is required to update source point cloud during training and inference.

Diffusion models [18, 30, 55, 56] have emerged as the cutting-edge family of deep generative models, representing a class of highly advanced deep generative models. They have broken the long-time dominance of generative adversarial networks (GANs) [57] in the challenging task of image synthesis [58, 55, 56] and have also shown potential in computer vision [23, 24, 59]. Specifically, in 3D computer vision, there has been a recent surge of research using generative models for point cloud generation or completion [60, 21]. These models are employed to infer missing parts and reconstruct complete shapes and hold significant implications for various downstream tasks, such as 3D reconstruction, augmented reality, and scene understanding [61, 22, 62]. [60] involves conceptualizing point clouds as particles within a thermodynamic system, utilizing a heat bath to facilitate diffusion from the original distribution to a noise distribution. The point diffusion-refinement model [22] is introduced by conditional denoising diffusion probabilistic models to generate a coarse completion from partial observations. Furthermore, this model establishes a point-wise mapping between the generated point cloud and the ground truth. While diffusion models have achieved great success in generation tasks, their potential for discriminative tasks has yet to be fully explored. Currently, there are some pioneering works that apply diffusion models to image segmentation [29] and object detection [63]. However, despite the considerable interest in this idea, there hasn’t been a successful application of diffusion models to point cloud registration solutions before, and its progress lags significantly behind image processing. To the best of our knowledge, this is the first work that adopts a diffusion model for point cloud registration.

We first introduce notations utilized throughout this paper. Given two point clouds: source point cloud $\mathcal{P}=\{{\textbf{p}_{i}}\in\mathbb{R}^{3}\mid{i=1,\dots,N}\}$ and template point cloud $\mathcal{Q}=\{{\textbf{q}_{j}}\in\mathbb{R}^{3}\mid{j=1,\dots,M}\}$, where each point is represented as a vector of $(x,y,z)$ coordinates. Point cloud registration task aims to estimate a rigid transformation $\{\textbf{R},\textbf{t}\}$ which accurately aligns $\mathcal{P}$ and $\mathcal{Q}$, with a 3D rotation $\textbf{R}\in SO(3)$ and a 3D translation $\textbf{t}\in\mathbb{R}^{3}$. The transformation can be solved by:

where $\mathcal{H}^{*}$ is the set of ground-truth correspondences between $\mathcal{P}$ and $\mathcal{Q}$. In this paper, for the convenience of formalization, we merge rotation and translation into a single transformation denoted as G. Note that the meaning of G varies depending on the specific context. In the correspondence-free methods, G represents a 7-dimensional vector composed of a rotation quaternion and a translation. In the correspondence-based methods, G represents a transformation matrix as follow:

As a result, our objective becomes: finding the rigid-body transformation $\textbf{G}\in{SE}(3)$ which best aligns $\mathbf{P}_{\mathrm{T}}$ and $\mathbf{P}_{\mathrm{S}}$ such that:

where we use shorthand $(\cdot)$ to denote transformation of $\mathbf{P}_{\mathrm{S}}$ by rigid transform G.

The denoising diffusion probabilistic model is a generative model where generation is modeled as a denoising process, in a specific form for point cloud registration:

The denoising process begins with standard Gaussian noise and continues until an optimal transformation is achieved. Specifically, a sequence of transformation variables denoted as $\textbf{G}_{T}$, $\textbf{G}_{T-1}$, $\ldots$, $\textbf{G}_{0}$ is generated through the denoising process, where each variable exhibits a progressively reduced level of noise, where $\textbf{G}_{T}$ is sampled from a standard Gaussian prior, and $\textbf{G}_{0}$ represents the final transformation.

We provide a detailed review of the formulation of diffusion models. Starting from a data distribution $\textbf{G}_{0}\sim q\left(\textbf{G}_{0}\right)$, we define forward process $q$ which produces data samples $\textbf{G}_{1},\textbf{G}_{2},\ldots,\textbf{G}_{T}$ by gradually adding Gaussian noise at each timestep $t$. In particular, the added noise is scheduled according to a variance schedule $\beta_{1},\ldots,\beta_{T}$:

A notable property of the forward process is that it admits sampling $\textbf{G}_{t}$ at an arbitrary timestep $t$ knowing $\textbf{G}_{0}$ in convenient closed-form evaluation: using the notation $\alpha_{t}:=1-\beta_{t}$ and $\tilde{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}$, we have

The forward process variances $\beta_{t}$ can be learned by reparameterization [64], or held constant as hyperparameters, and expressiveness of the reverse process is ensured in part by the choice of Gaussian conditionals in $p_{\theta}\left(\textbf{G}_{t-1}\mid\textbf{G}_{t}\right)$, because both processes have the same functional form when $\beta_{t}$ are small [30]. In this work, we fix the forward process variances $\beta_{t}$ to constants.

Furthermore, we define reverse process $p_{\theta}(\textbf{G}_{0:T})$ as a Markov chain with learned Gaussian transitions starting at a standard Gaussian prior $p\left(\textbf{G}_{T}\right)=\mathcal{N}\left(\textbf{G}_{T};\mathbf{0},\mathbf{I}\right)$:

which aims to invert the noise corruption process. Since calculating $q\left(\textbf{G}_{t-1}\mid\textbf{G}_{t}\right)$ (ground truth reverse process) exactly should depend on the entire data distribution, we can approximate $q\left(\textbf{G}_{t-1}\mid\textbf{G}_{t}\right)$ using a neural network $f_{\theta}$ with parameter $\theta$, which is optimized to predict a mean $\bm{\mu}_{\theta}\left(\textbf{G}_{t},t\right)$ and a diagonal covariance matrix $\bm{\Sigma}_{\theta}\left(\textbf{G}_{t},t\right)$. Intuitively, the forward process $q\left(\textbf{G}_{t}\mid\textbf{G}_{t-1}\right)$ can be seen as gradually injecting more random noise to the data, with the reverse process $p_{\theta}\left(\textbf{G}_{t-1}\mid\textbf{G}_{t}\right)$ learning to progressively remove noise to obtain realistic samples by mimicking the ground truth reverse process $q\left(\textbf{G}_{t-1}\mid\textbf{G}_{t}\right)$.

The goal of training the reverse diffusion process is to maximize the log-likelihood of the transformation: $\mathbb{E}_{q(\textbf{G}_{0})}\left[-\log{p_{\theta}(\textbf{G}_{0})}\right]$. However, since directly optimizing the exact log-likelihood is intractable, we instead maximize its variational lower bound:

where the inequality is by Jensen’s inequality and the above derivation after the fifth row is based on the Bayes’ rule [30].

In the final objective of Equation 8, as the forward process is fixed and $p(\textbf{G}_{T})$ is defined as a Gaussian prior, $\mathcal{L}_{1}$ does not affect the learning of $\theta$. $\mathcal{L}_{0}$ can be regarded as the reconstruction of the original data, which can be computed by estimating $\mathcal{N}\left(\textbf{G}_{0};\bm{\mu}_{\theta}\left(\textbf{G}_{1},1\right),\mathbf{\Sigma}_{\theta}\left(\textbf{G}_{1},1\right)\right)$ and constructing a discrete decoder. Therefore, the ultimate optimization objective is $\mathcal{L}_{3}$. Based on Bayes’ theorem, it is found that the posterior $q\left(\textbf{G}_{t-1}\mid\textbf{G}_{t},\textbf{G}_{0}\right)$ is a Gaussian distribution as well:

where

The final training objective can be reduced to maximum likelihood given the complete data likelihood with joint posterior $q(\textbf{G}_{1:T}\mid\textbf{G}_{0})$:

How to choice $p_{\theta}\left(\textbf{G}_{t-1}\mid\textbf{G}_{t}\right):=\mathcal{N}\left(\textbf{G}_{t-1};\bm{\mu}_{\theta}\left(\textbf{G}_{t},t\right),\bm{\Sigma}_{\theta}\left(\textbf{G}_{t},t\right)\right)$ is a crucial question. In this paper, we set $\mathbf{\Sigma}_{\theta}\left(\textbf{G}_{t},t\right)=\sigma_{t}^{2}\mathbf{I}$, where $\sigma_{t}^{2}=\tilde{\beta}_{t}$. To represent $\bm{\mu}_{\theta}\left(\textbf{G}_{t},t\right)$, $\frac{1}{\sqrt{\alpha_{t}}}\left(\textbf{G}_{t}-\frac{\beta_{t}}{\sqrt{1-\tilde{\alpha}_{t}}}\bm{\epsilon}\right)$ must be predicted given $\textbf{G}_{t}$. Since $\textbf{G}_{t}$ is available as input to the model $f_{\theta}$, we may choose the parameterization:

where $\bm{\epsilon}_{\theta}$ is the output of $f_{\theta}$. Finally, the training objective is parameterized:

Based on the aforementioned process, the sampling procedure can be expressed as follows:

Note that since point cloud registration is not a purely generative task, we make further adjustments during the implementation process based on the aforementioned theory, such as the training objective.

To demonstrate the generalization of our proposed framework, we conduct separate investigations on both correspondences-based and correspondences-free methods. While the overall algorithmic flow can be shared between the two methods, there are distinct differences in the implementation details of each component. We propose to separate the whole framework into two parts, point cloud encoder and transformation estiamtion decoder, where the former runs only once to extract a global representation or point-wise representation from the raw input point cloud $\mathcal{P}$ and $\mathcal{Q}$, and encodes the noisy transformation calculated by $q\left(\textbf{G}_{t}\mid\textbf{G}_{0}\right)$ in Equation 6. The latter takes these representations as condition to predict object transformation. Notably, for correspondence-free method, the number of input points $N$ is equal to $M$, whereas for correspondence-based method, $N$ is not necessarily equal to $M$.

During training, we first construct the diffusion process from ground-truth transformation to noisy transformation and then train the model to reverse this process. Algorithm 1 provides the pseudo code of PCRDiffsuion training procedure.

For correspondence-based methods, the loss function for joint optimization is identical to that of EGST [50]. It is noteworthy that during the training process, PCRDiffsuion requires only a single pass through the network. Iterative updates by applying the predicted transformation to the source point cloud are unnecessary.

The inference procedure of PCRDiffusion is a denoising sampling process from noise to object transformation $\textbf{G}_{0}$. Starting from transformation sampled in Gaussian distribution, the model progressively refines its predictions, as shown in Algorithm 2.

In our framework, the learnable parameters are adjusted by the back propagation algorithm [66]. We implement and evaluate our model with PyTorch on an NVIDIA RTX 3090 GPU. The network is trained with Adam optimizer [67] for 100 epochs on ModelNet40, 50 epochs on 3DMatch and KITTI. The learning rate starts from 0.0001 and decays when loss has stopped improving. The adjustment of the batch size depends on the specific methods used. For correspondence-free method, the batch size is set to 32, while for correspondence-based method, it is set to 8.

We evaluate the proposed method on synthetic datasets ModelNet40 [34], and real-world dataset 3DMatch [35]. ModelNet40 contains 12,308 CAD models of 40 different object categories. Each point cloud contains 2,048 points that randomly sampled from the mesh faces and normalized into a unit sphere. 3DMatch dataset contains 62 indoor subscenes. Some subscenes represent the same scenario with different overlap rates, we ignore this characteristic and treat all subscenes as different scenes. We randomly generate three Euler angle rotations within $[\text{0}^{\circ},\text{45}^{\circ}]$ and translations within $[\text{0},\text{1}]$ on each axis as the rigid transformation during training.

We compare our method to traditional methods: ICP[37] and FGR[68], and recent learning-based methods: PointNetLK [9], PCRNet [48], DeepGMR [69], DCP [70], FMR [10] and RPMNet [53]. We utilize the implementations of ICP and FGR in Intel Open3D [71]. For other baseline, we follow the pipeline provided by original paper. For consistency with previous work, we measure Mean Isotropic Error (MIE) and Mean Absolute Error (MAE). MIE(R) is the geodesic distance in degrees between estimated and ground-truth rotation matrices. It measures the differences between the predicted and the ground-truth rotation matrices. MIE(t) is the Euclidean distance between estimated and ground-truth translation vectors. It measures the differences between the predicted and the ground-truth translation vectors. Here, we give the specific calculation formula of MIE(R) and MIE(t):

Angular measurements are in units of degrees. All metrics should be zero if reference point cloud align to source point cloud perfectly.

We initially investigate the performance of diffusion models in the complete overlapping scenarios on ModelNet40.

Inspired by Gaussian noise, we observe that our method performs well when corresponding relationship is found in the cluttered circumstances. Acquisition of partial point clouds is a common situation in the real-world, therefore we challenge a more extensive and difficult task. We crop the template point cloud $\mathcal{P}$ and the source point cloud $\mathcal{Q}$ respectively, and retain 70% of the points. Noted that, the quantity of points in two point clouds is still consistent after cropping. Moreover, because the operation of cropping is random, some points in two point clouds may lose their correspondence. We verify quantitatively in Table I(d) and experimental results illustrate that our method can still achieve stable results comparing other methods, and the other baseline do not work well.

We conduct experiments on the real-world datasets: 3DMatch (indoor) [35] and KITTI odometry (outdoor) [36]. We utilize 3D object detection dateset in KITTI odometry consists of 14,999 outdoor driving scenarios scanned by Velodyne LiDAR and we use sequences 0-5 for training, 6-7 for validation and 8-10 for testing. 3DMatch dataset contains 62 indoor subscenes. Some subscenes represent the same scenario with different overlap rates, we ignore this characteristic and treat all subscenes as different scenes. For KITTI, the accuracy here refers to the accuracy of inter-frame matching for a single instance, rather than representing the accuracy of the odometry trajectory after matching all point clouds in the sequence.

To be consistent with the input of ModelNet40, each input point cloud is randomly sampled to an average of 2,048 points. The other experimental settings are consistent with previous experiments. Fig. 6 and Fig. 7 show visualization results of KITTI and 3DMatch, respectively. As shown in Table II, we observe that PCRDiffusion can still achieve stable results comparing other methods. PCRDiffusion demonstrates remarkable generalization on real-world datasets, and this framework can still achieve optimal performance without the need for iterative processes.

We conduct extensive ablation studies for a better understanding of the various modules in PCRDiffusion framework. To achieve a fair comparison, we utilize the same experimental settings as Section 4.1 in ablation studies.

In addition, we have uncovered another intriguing characteristic. Upon removing the transformation encoder, our two proposed simplified frameworks degrade to the original iterative-free PCRNet and RPMNet. Remarkably, under the same iterative-free conditions, the utilization of the diffusion model yields substantial performance enhancement. This presents a highly promising practical avenue: integrating the diffusion model as a plug-and-play component into existing point cloud registration networks can yield exceptionally prominent results!