SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haobo Jiang $^{1}$ , Mathieu Salzmann $^{2,3*}$ , Zheng Dang $^{2}$ , Jin Xie $^{1*}$ , and Jian Yang $^{1*}$
$^{1}$ PCA Lab, Nanjing University of Science and Technology, China
$^{2}$ CVLab, EPFL, Switzerland
$^{3}$ ClearSpace, Switzerland
{jiang.hao.bo, csjxie, csjyang}@njust.edu.cn
{zheng.dang, mathieu.salzmann}@epfl.ch
In this paper, we introduce an SE(3) diffusion model-based point cloud registration framework for 6D object pose estimation in real-world scenarios. Our approach formulates the 3D registration task as a denoising diffusion process, which progressively refines the pose of the source point cloud to obtain a precise alignment with the model point cloud. Training our framework involves two operations: An SE(3) diffusion process and an SE(3) reverse process. The SE(3) diffusion process gradually perturbs the optimal rigid transformation of a pair of point clouds by continuously injecting noise (perturbation transformation). By contrast, the SE(3) reverse process focuses on learning a denoising network that refines the noisy transformation step-by-step, bringing it closer to the optimal transformation for accurate pose estimation. Unlike standard diffusion models used in linear Euclidean spaces, our diffusion model operates on the SE(3) manifold. This requires exploiting the linear Lie algebra  $\mathfrak{se}(3)$  associated with SE(3) to constrain the transformation transitions during the diffusion and reverse processes. Additionally, to effectively train our denoising network, we derive a registration-specific variational lower bound as the optimization objective for model learning. Furthermore, we show that our denoising network can be constructed with a surrogate registration model, making our approach applicable to different deep registration networks. Extensive experiments demonstrate that our diffusion registration framework presents outstanding pose estimation performance on the real-world TUD-L, LINEMOD, and Occluded-LINEMOD datasets. The code is available at https://github.com/JiangHB/DiffusionReg.
Accurately estimating the 6D pose of an object, comprising both its spatial position and orientation, is a critical task in the field of computer vision, and has been extensively applied in various domains such as robotics grasping [10, 52, 68], augmented reality [32, 33], and autonomous navigation [7, 16, 58]. While substantial efforts have been dedicated to developing methods for 6D pose estimation based on RGB or RGB-D data [26, 42, 41, 52, 38, 63, 49], the advancements of 3D sensors (such as Kinect and LiDAR) and 3D registration techniques has promoted the emergence of point cloud registration-based pose estimation as a promising direction.
Haobo Jiang, Jin Xie, and Jian Yang are with PCA Lab, Key Lab of Intelligent Perception and Systems for High-Dimensional Information of Ministry of Education, and Jiangsu Key Lab of Image and Video Understanding for Social Security, School of Computer Science and Engineering, Nanjing University of Science and Technology, China.
Nevertheless, the state-of-the-art methods for object-level 3D registration [44, 62, 15] primarily focus on synthetic data and struggle to yield precise registration on real-world 6D pose estimation datasets such as TUD-L [20], LINEMOD [18], and Occluded-LINEMOD [6]. Unlike the controlled geometric structures, manually defined transformations, and artificial partial overlap present in synthetic data, real-world pose estimation datasets present significant challenges, including large rotations and translations, natural noise interference, and severe occlusions, considerably exacerbating the registration difficulty. Prior attempts to tackling these challenges include [11], which introduces a general match normalization layer to regularize the feature distribution of the source and model point clouds thus facilitating feature matching, and [60, 48], which design RANSAC-based outlier rejection strategy to enhance the registration robustness. While these initial attempts constitute significant steps towards making object-level point cloud registration practical, much work remains to be done to fully address all the aforementioned real-world challenges.
In this paper, inspired by the remarkable success of diffusion models in generative AI [19, 45, 46, 47], we propose to adapt them to the 3D registration field, and introduce an SE(3) diffusion model-based registration method for robust 6D object pose estimation in the real world. Our approach formulates 3D registration as a denoising diffusion process in SE(3) (special Euclidean group), aiming to progressively refine the pose of the source point cloud to achieve a precise alignment with the model point cloud. Training our model then involves two key operations: An SE(3) diffusion process and an SE(3) reverse process. The SE(3) diffusion process progressively converts the optimal transformation between the source and model point clouds into a noise one by continuously injecting perturbation transformations. Conversely, the SE(3) reverse process learns a denoising network to gradually refine the noise transformation to the optimal one. Our diffusion model operates on the SE(3) manifold, which poses the challenge of extending the diffusion/reverse formulas from linear Euclidean space to a nonlinear manifold. To address this challenge, we exploit the linear Lie algebra  $\mathfrak{se}(3)$  associated with SE(3) to conduct linear diffusion/reverse computations and map the  $\mathfrak{se}(3)$  results to SE(3) to obtain the desired diffused/reverse transformations.
To effectively train our denoising network, we exploit a Bayesian formalism and derive a registration-specific variational lower bound as optimization objective for model optimization. Furthermore, we reformulate our denoising network with a surrogate registration model, which makes our approach applicable to different deep registration methods such as [54, 62]. During the inference stage, the learned denoising network progressively refines an identity transformation to approach the optimal one, given the source and model point clouds as conditioning signals.
Compared to previous registration methods, our diffusion registration framework offers the following two advantages. First, the diffusion process generates a diverse set of poses for the source point cloud, allowing for more comprehensive model training. This increased pose diversity facilitates the model's ability to handle a wider range of rotation/translation and improves its generalization capabilities. Second, by incorporating the Bayesian posterior in the reverse process, our approach effectively guides the update of the source point cloud pose at each reverse step. This guidance mitigates the risk of getting trapped in local optima, leading to more robust and accurate pose estimation results. Our empirical results provide substantial evidence supporting these two advantages.
To summarize, our main contributions are as follows: (i) We introduce a novel SE(3) diffusion model-based 3D registration framework for robust 6D object pose estimation, where the optimal transformation is estimated via a progressive denoising process. (ii) To train our denoising network, we follow a Bayesian approach and establish a 3D registration-specific variational lower bound as optimization objective for our SE(3) diffusion model. Furthermore, we reformulate the denoising network with a surrogate registration model, enabling the integration of different deep registration models into our framework. (iii) To the best of our knowledge, we are the first to successfully adapt the diffusion model from linear Euclidean space to the SE(3) point cloud registration task for 6D object pose estimation. Our extensive experiments on real-world datasets, including TUD-L [20], LINEMOD [18] and Occluded-LINEMOD [6], confirm the effectiveness of our framework.
Point Cloud Registration. Point cloud registration aims to estimate the rigid transformation between a pair of point clouds. Existing methods can be roughly divided into optimization-based techniques and deep learning-based ones. Iterative Closest Point (ICP) [4] alternately searches for the closest correspondences and estimates transformation until convergence. Go-ICP [61] uses the branch-and-bound algorithm to improve the robustness of ICP to initialization. Robust ICP [64] designs a
Figure 1: Probabilistic graphical model of our SE(3) diffusion model-based registration framework.
Welsch's function-based robust metric to obtain an outlier-robust alignment evaluation for optimization. Sparse ICP [5] reformulates ICP using sparsity-inducing norms. Many other ICP variants have been proposed, such as[8, 43, 14, 2, 17], exhibiting promising registration performance. Nevertheless, the current main research effort focuses on deep learning models. In this context, DCP [54] constructs pseudo correspondences for SVD-based transformation estimation using deep feature similarity. PRNet [55] leverages keypoint identification and Gumbel softmax for establishing more reliable correspondences. PointNetLK [1] and FMR [21] introduce the Lucas & Kanada (LK) algorithm [3] and inverse compositional (IC) algorithm [3] into deep models for transformation refinement via feature alignment. CEMNet [23] uses a planning-based cross-entropy method for transformation optimization. RPMNet [62], RGM [15] and RIENet [44] leverage Sinkhorn optimization, neighborhood structure consistency, and graph matching, respectively, for robust outlier removal. Many other deep registration models such as [28, 9, 36, 30, 29, 67] have been developed and achieve impressive registration performance. However, these registration models mainly focus on synthetic datasets, such as ModelNet40 [56]. As shown in [11], they still yield limited precision on real-world 6D object pose estimation datasets such as TUD-L [20], LINEMOD [18], and Occluded-LINEMOD [6].
6D Object Pose Estimation. Estimating the 6D pose (orientation and position) of objects has gained increasing attention in recent years. Early approaches directly estimated object pose from extracted image features through regression or classification. PoseCNN [57] decouples 6D pose estimation into center-based translation regression and quaternion-based rotation regression. SSD-6D [24] extends the single-shot object detector to cover the full 6D object space for pose inference. Trabelsi et al. [51] propose a multi-attentional network for iterative pose refinement using the appearance and flow information. DeepIM [31] and LatentFusion [37] learn pose estimation by minimizing the error between the rendered model and the observation. More recent research has focused on a two-stage estimation pipeline, where 2D keypoints are extracted to solve for 6D pose using the PnP algorithm [13]. Rad et al. [42] take the predicted 2D projections of the corners of 3D bounding boxes as the keypoints, while Zhao et al. [65] manually designate keypoints over the surface of the object model. PVNet[41] establishes a pixel-wise voting network for keypoint estimation using predicted pixel-wise voting vectors. Other methods [6, 40, 50, 52, 49, 52] also present promising pose estimation performance. Recently, thanks to the advances in 3D registration techniques, the community has witnessed a growing interest in 3D registration-based 6D object pose estimation, which recovers the object's 6D pose by estimating the rigid transformation between the object (source) and the model point clouds. This is what we achieve here, but, in contrast to most existing methods which focus on synthetic data, we tackle the challenging scenario of working with real point clouds.
3.1 Revisiting the Euclidean Diffusion Model
Diffusion models, as generative models, aim to generate new data by progressively denoising noisy inputs [19, 45, 46, 47]. Their training phase involves a diffusion process and a reverse process. The diffusion process successively injects Gaussian noise into the data sample  $\mathbf{x}_0 \sim p_{data}$  ( $p_{data}$  indicates the distribution of the training data) so as to gradually transform  $\mathbf{x}_0$  into the noise data  $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  (standard Gaussian distribution), thus forming a Markov chain  $\mathbf{x}_0 \to \mathbf{x}_1 \to \dots \to \mathbf{x}_T$ . As demonstrated in [19], the random variable  $\mathbf{x}_t \sim q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) \coloneqq \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t \mathbf{I})$  can also be expressed in a closed form  $\mathbf{x}_t \sim q(\mathbf{x}_t \mid \mathbf{x}_0)$ , which can be formulated as follows:
$$
\mathbf {x} _ {t} = \sqrt {\bar {\alpha} _ {t}} \mathbf {x} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \varepsilon , \varepsilon \sim \mathcal {N} (\mathbf {0}, \mathbf {I}), \tag {1}
$$
Here, the diffusion coefficients  $\bar{\alpha}_{t} = \prod_{s=0}^{t} \alpha_{s} = \prod_{s=0}^{t} (1 - \beta_{s})$ , and  $\beta_{s}$  indicates the noise coefficient determined by a linear schedule [19] or a cosine schedule [35]. Then, the reverse process learns a denoising network, a parameterized normal distribution  $p_{\theta}(\mathbf{x}_{t-1} \mid \mathbf{x}_t) := \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_{\theta}(\mathbf{x}_t, t), \beta_t \mathbf{I})$ , to progressively denoise the noisy data  $\mathbf{x}_T$  into the clean one  $\mathbf{x}_0$ , forming a reverse Markov chain  $\mathbf{x}_T \to \mathbf{x}_{T-1} \to \dots \to \mathbf{x}_0$ . Here,  $\boldsymbol{\mu}_{\theta}(\mathbf{x}_t, t)$  indicates the parameterized mean of the normal distribution. The following variational lower bound of the log likelihood over the training data is then derived as the optimization objective for training the denoising network:
$$
\mathbb {E} _ {\mathbf {x} _ {0} \sim p _ {d a t a}} [ \log p _ {\theta} (\mathbf {x} _ {0}) ] \geq \mathbb {E} _ {q} [ \sum_ {t > 1} \mathrm {D} _ {\mathrm {K L}} (q (\mathbf {x} _ {t - 1} | \mathbf {x} _ {t}, \mathbf {x} _ {0}) \| p _ {\theta} (\mathbf {x} _ {t - 1} | \mathbf {x} _ {t})) - \log p _ {\theta} (\mathbf {x} _ {0} | \mathbf {x} _ {1}) ]. \tag {2}
$$
Based on Bayes' formula, the random variable  $\mathbf{x}_{t-1}$  following the posterior distribution  $q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$  in Eq. 2 can be represented as
$$
\mathbf {x} _ {t - 1} = \frac {\sqrt {\bar {\alpha} _ {t - 1}} \beta_ {t}}{1 - \bar {\alpha} _ {t}} \mathbf {x} _ {0} + \frac {\sqrt {\alpha_ {t}} (1 - \bar {\alpha} _ {t - 1})}{1 - \bar {\alpha} _ {t}} \mathbf {x} _ {t} + \sqrt {\tilde {\beta} _ {t}} \varepsilon , \tag {3}
$$
where the random variable  $\varepsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})$  and the variance scale  $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t - 1}}{1 - \bar{\alpha}_t}\beta_t$ .
3.2 SE(3) Diffusion Registration Model for 6D Object Pose Estimation
In the context of 6D object pose estimation based on point clouds, our objective is to determine the rigid transformation between a partially-scanned source point cloud  $\mathcal{X} = \{\mathbf{x}_i\in \mathbb{R}^3\}_{i = 1}^N$  and a complete model point cloud  $\mathcal{M} = \{\mathbf{m}_j\in \mathbb{R}^3\}_{j = 1}^M$  so as to align their overlapping regions precisely. This transformation encompasses a rotation matrix  $\mathbf{R}\in SO(3)$  and a translation vector  $\mathbf{t}\in \mathbb{R}^3$ , representing the orientation of the object and its spatial position, respectively. To obtain a partial source point cloud, following [11], we mask an input depth map to restrict it to the object, and convert the resulting depth values into 3D points using the known camera intrinsic parameters. By contrast, the model point cloud is generated by uniformly sampling a complete mesh model. We denote the optimal transformation  $\mathbf{H}_0\in SE(3)$  and the identity transformation  $\mathbb{H}\in SE(3)$  as
$$
\mathbf {H} _ {0} = \left[ \begin{array}{l l} \mathbf {R} & \mathbf {t} \\ \mathbf {0} ^ {\top} & 1 \end{array} \right], \quad \mathbb {H} = \left[ \begin{array}{l l} \mathbf {I} & \mathbf {0} \\ \mathbf {0} ^ {\top} & 1 \end{array} \right], \tag {4}
$$
both of which are represented by  $4 \times 4$  homogeneous transformation matrices.
Inspired by the progressive generation of diffusion models, we adopt the concept of denoising diffusion to tackle the point cloud registration task, specifically in the SE(3) space, and thus propose an SE(3) diffusion model-based 3D registration framework for robust 6D object pose estimation. While [27, 53] have introduced variations of the diffusion model on SE(3) or SO(3) manifolds, they are not suitable for our 3D registration task. The training phase of our framework thus involves an SE(3) diffusion process and an SE(3) reverse process. Given a pair of source and model point clouds, the former continuously disturbs their optimal transformation by injecting noise into it, while the latter aims to learn a denoising network to progressively convert the noisy transformation to the optimal one. As such, during inference, the learned denoising network can be leveraged to recover the transformation between the source and model point clouds through a progressive denoising process. The training phase and inference phase of our algorithm are detailed in Algorithms 1 and 2. Below, we provide a detailed explanation of our SE(3) diffusion process and SE(3) reverse process.
3.2.1 SE(3) Diffusion Process
In accordance with the standard diffusion model described in Sec. 3.1, our SE(3) diffusion process progressively disturbs the optimal transformation  $\mathbf{H}_0$  of point clouds  $\mathcal{X}$  and  $\mathcal{M}$  into a noisy transformation  $\mathbf{H}_T$  by injecting perturbation transformations (noise) into it, thus forming a diffusion Markov chain:  $\mathbf{H}_0\rightarrow \mathbf{H}_1\rightarrow \dots \rightarrow \mathbf{H}_T$ . However, our SE(3) diffusion process features two critical differences from the conventional one: (i) Our diffusion process operates on the nonlinear SE(3) manifold, unlike the standard diffusion process which acts in linear Euclidean space. Consequently, the linear diffusion operations provided in Eq 1 cannot be directly applied to our model. (ii) The standard diffusion process assumes the noise variable  $\mathbf{x}_T\sim \mathcal{N}(\mathbf{0},\mathbf{I})$  to be centered around the zero vector. However, in our model for 3D registration, we require the noise transformation  $\mathbf{H}_T$  to be centered around the identity transformation  $\mathbb{H}$ . To account for these differences while generating the
Algorithm 1 Training phase	Algorithm 2 Inference phase
repeat
Sample X, M, H0 ~ pdata
Sample t ~ Uniform{1,...,T}), Ht = Exp(γ√1-αtε)F(√αt; H0, H)
RegisterˆHt→0 = fθ(T(X, Ht), M)
Optimize loss Lt = loss(ˆHt→0, H0Ht-1)
until converged	for t = T,..., 1 do
RegisterˆHt→0 = fθ(Xt, M), where Xt = T(X, Ht)
Estimate Ht-1 = Exp(λ0 Log (ˆHt→0Ht) + λ1 Log (Ht))
end for
Return: H0
desired diffusion Markov chain, we propose to interpolate the transformations between  $\mathbf{H}_0$  and  $\mathbb{H}$  on the SE(3) manifold, incorporating noise perturbations. Formally, our interpolation-based SE(3) diffusion formula formulates  $\mathbf{H}_t\sim q(\mathbf{H}_t\mid \mathbf{H}_0)$  at time step  $t$  ( $1\leq t\leq T$ ) as
$$
\mathbf {H} _ {t} = \underbrace {\operatorname {E x p} \left(\gamma \sqrt {1 - \bar {\alpha} _ {t}} \varepsilon\right)} _ {\text {P e r t u r b a t i o n}} \underbrace {\mathcal {F} \left(\sqrt {\bar {\alpha} _ {t}} ; \mathbf {H} _ {0} , \mathbb {H}\right)} _ {\text {I n t e r p o l a t i o n}}, \varepsilon \sim \mathcal {N} (\mathbf {0}, \mathbf {I}), \tag {5}
$$
In essence, this formalism utilizes an interpolation function  $\mathcal{F}(\sqrt{\bar{\alpha}_t};\mathbf{H}_0,\mathbb{H})$  to derive an intermediate transformation that lies between the optimal transformation  $\mathbf{H}_0$  and the identity transformation  $\mathbb{H}$ . Subsequently, we augment this interpolated transformation with a randomly-sampled perturbation transformation (noise), denoted as  $\mathrm{Exp}(\gamma \sqrt{1 - \bar{\alpha}_t}\varepsilon)$ , to yield the diffused transformation  $\mathbf{H}_t$ . Below, we will introduce these two components in detail.
Transformation Interpolation Function  $\mathcal{F}(\sqrt{\bar{\alpha}_t};\mathbf{H}_0,\mathbb{H})$ . The interpolation function  $\mathcal{F}:SE(3)\times SE(3)\times [0,1]\to SE(3)$  aims to compute an intermediate transformation between the optimal transformation  $\mathbf{H}_0$  and the identity one  $\mathbb{H}$ . Following the setting of diffusion coefficients in Eq. 1, we set the interpolation weight on transformation  $\mathbf{H}_0$  at time step  $t$  to  $\sqrt{\bar{\alpha}_t}$  and, consequently, the interpolation weight on  $\mathbb{H}$  is  $1 - \sqrt{\bar{\alpha}_t}$ . As the time step  $t$  increases, the interpolation weight  $\sqrt{\bar{\alpha}_t}$  on  $\mathbf{H}_0$  decreases gradually, resulting in a transition of the interpolated transformation from  $\mathbf{H}_0$  to  $\mathbb{H}$ . However, due to the non-linearity of the SE(3) manifold, direct linear interpolation methods, such as weighted averages, cannot be applied to interpolate intermediate transformations. To address this, we exploit the Lie algebra  $\mathfrak{se}(3)$  associated to SE(3). The Lie algebra  $\mathfrak{se}(3)$  is a linear 6D vector space corresponding to the tangent space of SE(3) at the identity transformation. Transferring any transformation in SE(3) to a 6D vector element in  $\mathfrak{se}(3)$  can be achieved using the logarithm map  $\mathrm{Log}:SE(3)\to \mathbb{R}^6$ ; conversely, the exponential map  $\mathrm{Exp}:\mathbb{R}^6\to SE(3)$  transforms a 6D vector in  $\mathfrak{se}(3)$  back to the SE(3) manifold. Therefore, we can first project the SE(3) transformation to  $\mathfrak{se}(3)$ , perform linear interpolation in this tangent space, and then convert the interpolated vector back to SE(3) to obtain the interpolated transformation. Formally, our transformation interpolation function  $\mathcal{F}$  is thus expressed as
$$
\mathcal {F} \left(\sqrt {\bar {\alpha} _ {t}}; \mathbf {H} _ {0}, \mathbb {H}\right) = \operatorname {E x p} \left(\left(1 - \sqrt {\bar {\alpha} _ {t}}\right) \cdot \operatorname {L o g} \left(\mathbb {H} \mathbf {H} _ {0} ^ {- 1}\right)\right) \mathbf {H} _ {0}, \tag {6}
$$
which starts by calculating the relative transformation  $\mathbb{H}\mathbf{H}_0^{-1}$  from the optimal transformation  $\mathbf{H}_0$  to the identity transformation  $\mathbb{H}$ , and mapping it to the linear Lie algebra  $\mathfrak{se}(3)$  through the logarithmic map. In this linear space, we then scale the vector  $\mathrm{Log}(\mathbb{H}\mathbf{H}_0^{-1})$  by the interpolation weight  $1 - \sqrt{\overline{\alpha}_t}$ , and map this weighted vector back into the corresponding weighted relative transformation using the exponential map. The weighted relative transformation quantifies the offset of the interpolated transformation relative to  $\mathbf{H}_0$ . Finally, multiplying the optimal transformation  $\mathbf{H}_0$  by the computed transformation offset yields the desired interpolated transformation at time step  $t$ .
Perturbation Transformation  $\mathrm{Exp}(\gamma \sqrt{1 - \bar{\alpha}_t}\varepsilon)$ . Following the standard diffusion formula in Eq. 1, we introduce random noise (i.e., a perturbation transformation) into the interpolated transformation  $\mathcal{F}(\sqrt{\bar{\alpha}_t};\mathbf{H}_0,\mathbb{H})$  at each time step to randomize our SE(3) diffusion process. As indicated in Eq. 1, conventional Euclidean diffusion models typically draw such noise from a Gaussian distribution in Euclidean space. However, formulating a Gaussian distribution on the SE(3) manifold is non-trivial. To tackle this issue, we utilize the Lie algebra once more. Specifically, we randomly sample a 6D noise vector  $\varepsilon \in \mathbb{R}^6$  from  $\mathcal{N}(\mathbf{0},\mathbf{I})$  over  $\mathbb{R}^6$ , which can be interpreted as an element in  $\mathfrak{se}(3)$ . Then, we scale the noise vector by the factor  $\gamma \sqrt{1 - \bar{\alpha}_t}$  as in Eq. 1 to control the magnitude of the perturbation at different time steps. Finally, this scaled noise vector is converted back to SE(3)