DiffRef3D: A Diffusion-based Proposal Refinement Framework for 3D Object Detection

3D object detectors can be broadly categorized into two methods: point-based methods and voxel-based methods. Point-based methods (Qi et al. 2017a, b) take the raw point clouds as inputs and employ permutation invariant operations to aggregate the point features. Voxel-based methods (Zhou and Tuzel 2018; Yan, Mao, and Li 2018) divide point clouds into a regularly structured grid and apply convolutional operations to extract features. Two-stage models employ the above methods to serve as a region proposal network (RPN), introducing the proposal refinement stage that further improves the prediction quality.

PointRCNN (Shi, Wang, and Li 2019) performs foreground segmentation and generates initial proposals using a PointNet++ (Qi et al. 2017b) backbone, and then extracts RoI features from the point clouds inside the proposals for refinement. PV-RCNN (Shi et al. 2020a) proposes a Voxel Set Abstraction module to encode the voxel-wise feature volumes into a small set of keypoints, then aggregates to the RoI-grid points to take advantage of both point-based and voxel-based methods to refine proposals. Voxel R-CNN (Deng et al. 2021) follows a typical pipeline of anchor-based two-stage models, using SECOND (Yan, Mao, and Li 2018) as RPN and refine proposals with the voxel features pooled for RoI. CenterPoint (Yin, Zhou, and Krahenbuhl 2021) predicts proposals from object centers computed from class-specific heatmap then point features are extracted from 3D centers of each face of the proposals for refinement. CT3D (Sheng et al. 2021) embeds and decodes proposals into effective proposal features based on neighboring point features for box prediction by channel-wise transformer.

There are special cases of two-stage models which utilize additional proposal refinement modules in order to adopt the cascade detection paradigm. 3D Cascade RCNN (Cai et al. 2022) iteratively refines proposals through cascade detection heads while considering point completeness score. CasA (Wu et al. 2022) progressively refines proposals by cascade refinement network while aggregating features from each stage by cascade attention module. The previous cascade detection paradigm boosts performance by leveraging the advantages of multiple predictions such as ensemble, however it lacks flexibility in sampling steps and increases memory usage for additional detection heads. In contrast, iterative sampling of diffusion-based models exploits advantages without increasing memory usage for each sampling step and sampling steps are adjustable without additional training.

The diffusion framework (Dhariwal and Nichol 2021; Ho, Jain, and Abbeel 2020) has shown remarkable performance in the field of image generation. Notably, DDIM (Song, Meng, and Ermon 2020) expedites the sampling process through the utilization of a more efficient class of iterative implicit probabilistic models. Drawing attention to its denoising capabilities and performance, there have been exploding recent efforts to apply it to perception tasks (Shan et al. 2023; Zou et al. 2023) beyond generation tasks. DiffusionDet (Chen et al. 2022) first adopt the diffusion model to the 2D object detection problem by denoising random boxes to object bounding boxes via the diffusion process. DiffusionInst (Gu et al. 2022) utilizes additional mask branches and instance-aware filters upon DiffusionDet to address the instance segmentation task, and the following works (Lai et al. 2023; Wang et al. 2023) show competitive performance on semantic segmentation. However, diffusion-based perception models have been mostly explored in the image domain. Through DiffRef3D, we demonstrate the first application of the diffusion framework to a 3D perception task that involves point cloud input.

A proposal refinement module in two-stage 3D object detectors takes the initial proposals of the RPN as input and exploits the information of the RoIs to refine them. Formally, the refinement module takes a proposal $X^{P}$ and outputs the box residual $\hat{x}$ and the confidence score $\hat{c}$:

The RoI feature pooling operation $\textrm{RoIPool}(\cdot)$ extracts proposal feature $f^{P}$ from $X^{P}$. $\textrm{Det}(\cdot)$ represents the detection head, which consists of a regression branch producing $\hat{x}$ and a classification branch producing $\hat{c}$. $\hat{x}$ is supervised with the target residual $x^{gt}$ between the proposal and its target object box $X^{T}$, i.e., $x^{gt}=X^{T}\ominus X^{P}$ where $\ominus$ represents box encoding functions (Yan, Mao, and Li 2018), and the target for $\hat{c}$ is calculated with the intersection over union (IoU) between $X^{P}$ and $X^{T}$.

During training, a small amount of Gaussian noise is gradually added to the data sample $x_{0}$ through a Markovian chain of the forward diffusion process. The forward diffusion process is formulated as:

where $\beta_{t}$ represents a noise variance schedule (Ho, Jain, and Abbeel 2020; Nichol and Dhariwal 2021). For brevity, let $\alpha_{t}=1-\beta_{t}$, $\bar{\alpha_{t}}=\prod_{i=1}^{t}\alpha_{i}$, and a noisy sample $x_{t}$ at timestep $t,0\leq t\leq T,$ is derived as follows:

The conditional diffusion model learns a reverse diffusion process, aiming to progressively refine $x_{t}$ back to $x_{0}$ with the guidance of the conditional input $y$. The reverse diffusion process is formulated as follows:

The hypothesis $X^{H}_{t}$ is a secondary inspection region derived by applying the given box residual to the proposal. i.e., $X^{H}_{t}=X^{P}\oplus x_{t}$. As the baseline refinement module extracts the proposal feature $f^{P}$ (Eq. 1), HAM also extracts the hypothesis feature $f^{H}_{t}$ using the same RoI feature pooling operation:

The hypothesis feature is further processed to incorporate the information of $x_{t}$ and $t$ by applying a self-attention block $\textrm{Attn}(Q,K,V)$ and a temporal transformation block $\textrm{TT}(\cdot,t)$.

The self-attention block is a common multi-head attention module (Vaswani et al. 2017) that takes $f^{H}_{t}$ as input for the query, key, and value. Except, we embed $x_{t}$ as a positional encoding to the query. We utilize a shallow MLP $g_{\phi}(\cdot)$ to convert $x_{t}$ as the same dimensional vector with $f^{H}_{t}$. Formally, the attention feature of the hypothesis $a^{H}_{t}$ is computed as:

$a^{H}_{t}$ is a feature vector that represents the information of the hypothesis region with consideration of its relative displacement from the proposal.

The temporal transformation block reweights $a^{H}_{t}$ with respect to $t$ to adjust the effect of the noise signal which gets dominant as timestep increases. Specifically, the output of HAM $h_{t}$ is formulated as follows:

where $W_{t},b_{t}$ are the scaling and shifting factors derived from an MLP $s_{\psi}(\cdot)$:

Finally, $h_{t}$ is aggregated with the proposal feature for prediction:

During the training stage, We set the reconstruction target $x_{0}=x^{gt}$ and sample noisy residual $x_{t}$ by the forward diffusion process. The model is trained to perform the reverse diffusion process of reconstructing the true residual, which is identical to refining a hypothesis to the target object bounding box. We elucidate the methodology for generating hypotheses from proposals and present the overall training process of DiffRef3D in pseudo-code through Algorithm 1.

During the inference stage, DiffRef3D refines the hypothesis generated from random Gaussian noise and follows iterative sampling steps from DDIM (Song, Meng, and Ermon 2020) for progressive refinement. Starting from a random Gaussian noise $x_{T}\sim\mathcal{N}(0,1)$, DDIM progressively denoises Gaussian noise into data samples with step size $\Delta t$, i.e., $x_{T}\rightarrow x_{T-\Delta t}\rightarrow...\rightarrow x_{0}$. The progressive refinement is formulated as follows:

where $\epsilon^{(t)}$ represents noise term added to data sample and $\sigma_{t}$ represents stochastic parameter, respectively. Algorithm 2 describes the inference procedure of DiffRef3D in pseudo-code.

The KITTI benchmark (Geiger, Lenz, and Urtasun 2012) is one of the most popular benchmarks for 3D object detection evaluation. The dataset consists of 7,481 LiDAR samples for training and 7,518 LiDAR samples for testing. Training samples are divided into the train set with 3,712 samples and the val set with the remaining 3,769 samples. We present two performance evaluations for comparative analysis: For evaluation on KITTI val set, we compared DiffRef3D-V, DiffRef3D-PV, and DiffRef3D-T trained with the train set with their respective baselines. For KITTI test set, we trained a model with our framework using randomly selected 5,981 samples, validated with the remaining 1,500 samples, and evaluated the performance on the online test leaderboard. The performance were assessed across three classes: car, pedestrian, and cyclist on three different levels (easy, moderate, and hard). The results on both the val and test sets were evaluated using the mean average precision calculated at 40 recall positions.

We conduct extensive ablation studies on the KITTI val set to verify our design choices of DiffRef3D.

We visualize the hypothesis-driven box prediction process facilitated by the proposal refinement process of DiffRef3D. Figure 3 illustrates iterative sampling steps of DiffRef3D in each timestep. Visualization result implies the necessity of proposals as conditional inputs due to point cloud sparsity and relatively small object size compared to point cloud scenes. Figure 3 (a) shows hypotheses generated around proposal centers tend to prioritize regions with higher object presence, while hypotheses exhibit less realistic box forms since generated from random Gaussian noise. During the iterative sampling, DiffRef3D progressively refines proposals and hypotheses toward object boxes, as demonstrated in Fig. 3 (b) and (c). This result highlights DiffRef3D successfully employing the diffusion process to iteratively denoise hypotheses and refine proposals.

To verify the temporal transformation block considers the effect of the noise signal, we visualize the norm of the scaling factor along every timestep. Results are illustrated in Fig. 4. We affirm that the norm of the scaling factor is larger at a smaller timestep, implying that the refinement module takes more attention to the hypothesis feature at the later refinement process. Therefore, the temporal transformation block controls the importance of hypothesis features before aggregating it with proposal features.