PVT: Point-Voxel Transformer for Point Cloud Learning

To leverage the success of convolutional neural networks on 2D images, many researchers have endeavored to project 3D point clouds directly onto voxels and perform 3D convolution for information aggregation [6, 21]. However, the memory and computational consumption of such full-voxel-based models increases cubically with respect to the voxel’s resolution. To tackle this shortcoming, O-CNN [43], OctNet [33], and kd-Net [17] constructed tree structures for 3D voxels to bypass the invalid computation of the empty voxels. Although such methods are efficient in data structuring, geometric details are lost during the projection because multiple bucketing points into the same voxel leads to indistinguishable features for learning.

Compared with most voxel-based 3D models, our model is more efficient and effective for the following reasons: 1) Instead of using 3D convolutions, we propose SWA to handle the sparsity characteristic of voxels that have a linear computational complexity with respect to voxel resolution and bypasses the invalid computation of empty voxels. 2) As we employ RA in the point domain, an inherent advantage is that we can avoid feature loss during voxelization, maintaining the geometric details of a point cloud.

Instead of voxelization, it is possible to make a neutral network that consumes directly on point clouds. [29] proposed PointNet, the pioneering work that learns directly on sparse and unstructured point clouds. Inspired by PointNet, previous works introduced sophisticated neural modules to learn per-point features. These models can be generally classified as: 1) neighbouring feature pooling [63, 16], 2) graph message passing [45, 59, 42], and 3) attention-based or Transformer-based models [34, 60, 55, 14, 65, 9].

Owing to the sparsity and irregularity of point clouds, the methods that directly consume points have achieved state-of-the-art performance. However, the cost of data structuring of such methods has become the computation burden, especially on the large-scale point clouds dataset [27, 51]. In this study, we handle this shortcoming by using the SWA module.

[2] proposed a neural machine translation method with an attention mechanism, in which attention weight is computed through the hidden state of an RNN. Then [25] further proposed self-attention to visualize and interpret sentence embeddings. Subsequent works employed self-attention layers to replace some or all the spatial convolution layers, such as Transformer for machine translation [40]. [7] proposed the bidirectional transformers (BERT), which is one of the most powerful models in the NLP field.

Given the success of self-attention and Transformer architectures in NLP, researchers have applied them to vision tasks [15, 32, 62]. For instance, [8] proposed an image recognition network, ViT, which directly applied a Transformer architecture on image patches and achieved better performance than the traditional convolutional neural networks. [26] recently introduced Swin Transformer to incorporate inductive bias for spatial locality, hierarchy and translation invariance.

Recently, plenty of researchers have attempted to explore Transformer-based architectures for point cloud learning. [9] proposed PT¹ to extract global features by introducing the dot-product SA mechanism. [14] proposed offset-attention to calculate the offset difference between the SA features and the input features by element-wise subtraction. Moreover, [19] proposed PT² to build local vector attention in neighborhood point sets and achieved significant progress. However, they wasted a high percentage of the total time on structuring the irregular data, which becomes the efficiency bottleneck. In this work, we study how to solve this shortcoming, so as to design an efficient Transformer-based architecture for point cloud analysis.

There also exists some Transformer-based architecture for various point cloud processing tasks. For instance, [11] proposed the SE(3)-Transformer, a variant of the self-attention module for 3D point clouds and graphs, which is equivariant under continuous 3D roto-translations. Late, [57] and [49] have attempted to explore Transformer-based architectures for point cloud completion and achieved significant performance on all completion tasks. Recently, [58] proposed Point-BERT to unleash the scalability and generalization of Transformers for 3D point cloud representation learning. Extensive experiments also demonstrate that Point-BERT significantly improves the performance of standard point cloud Transformers. Different from these Transformers, in this paper, we study how to design an efficient and high-performance Transformer architecture for point cloud learning, making it suitable for edge devices with limited computational resources or real-time autonomous driving scenarios.

This branch aims to effectively capture local information, which can bypass expensive sampling and neighbor points querying. Specifically, it contains three steps: voxelization, feature aggregation, and devoxelization.

The voxelized and devoxelized methods map the input point cloud $\mathcal{P}$ to a new set of voxel features $V$ and transform the voxel-wise features back to the point-wise features $F_{local}$ (Readers can refer to [27] for more details). In this work, we apply the standard Transformer architecture [40] to perform feature aggregation on regular 3D voxels, which can improve accuracy significantly. However, the standard Transformer architecture conducts global-SA which leads to quadratic complexity with respect to the number of voxels, making it unsuitable for many 3D vision problems requiring dense prediction, such as semantic segmentation.

Window Attention: To obtain efficient modeling power, we propose to compute SA within local 3D windows, which are arranged to evenly partition the voxel space in a non-overlapping manner. Assume that each local 3D window contains $W\times W\times W$ voxels and $R$ denotes the voxel resolution. The computational complexity of a global-SA and a window-based one on $R^{3}$ voxels are:

	$\displaystyle\Omega(Global\textbf{-}SA)=4R^{3}D^{2}+2(R^{3})^{2}D$		(1)
	$\displaystyle\Omega(Window\textbf{-}SA)=4R^{3}D^{2}+2B^{3}R^{3}D$		(2)

where the former is quadratic to the number of voxels $R^{3}$, the latter is linear when $W$ is fixed, $D$ is the dimension of features. In summary, global-SA computation is generally unaffordable for a large voxel resolution, whereas the local window-SA is scalable.

Sparse Window Attention: Different from 2D pixels densely placed on an image plane, non-empty voxels only account for a small proportion of total voxels. Inspired by [13, 54], we design Sparse Window Attention to handle the sparsity characteristic of voxels. As shown in Figure 2, for each querying index $v_{i}$, we first use Window Attention to determine all neighboring voxel indices in the attention field, and then adopt a Hash table to get the hashed integer neighboring voxel indices. Last, a GPU-based Rule Book stores the hashed voxel indices as keys, and the corresponding indices for the non-empty voxel array as values. Finally, we can perform Window Attention to gather the coarse-grained local features. Note that all the steps can be conducted in parallel on GPUs by assigning each querying voxel $v_{i}$ a separate CUDA thread.

The SWA lacks information interaction across windows, which may limit the representation power of our model. Thus, we extend the shifted 2D window mechanism of Swin Transformer [26] to 3D windows for the purpose of introducing cross-window information interaction while maintaining the efficient computation of non-overlapping SWA. Readers can refer to Swin Transformer for more details

As shown in Figure 1, with the cyclic shifted window partitioning approach, the step of feature aggregation of this branch can be described as follows:

The voxel branch gathers the neighborhood information with low resolution. However, in order to capture long-range dependencies, low-resolution voxel branch alone is limited. To this end, we attempt to employ the following self-attention on the entire point cloud for global context aggregation, which is computed as:

	$\displaystyle F_{sa}$	$\displaystyle=softmax(QK^{T})V,\quad F_{sa}\in\mathbb{R}^{N\times D}$		(3)
	$\displaystyle F_{global}$	$\displaystyle=MLP(F_{sa})+F,\quad F_{global}\in\mathbb{R}^{N\times D}$		(4)

where (Q, K, V) $\in\mathbb{R}^{N\times D}$ is generated by shared linear transformations and the input features $F$ and are all ordered independent. Moreover, softmax and weighted sum are both permutation-independent operators. Thus, the self-attention computation is permutation-invariant, making it well-suited to handle the irregular, disordered 3D points.

Relative Attention: Self-attention mentioned above fails to incorporate relative position representations in its structure, whereas such ability is very important to 3D visual tasks. For example, the absolute coordinates of the same object can be completely different with rigid transformations. Therefore, injecting relative position representations are generally more robust. In this study, we design our relative-attention to embed relative position representations, which are not well studied in prior point cloud learning works.

First of all, by embedding RPR into the scaled dot-product self-attention module, Eq 3 can be re-formulated as:

$\displaystyle F_{ra}=softmax(QK^{T}+B)V$

(5)

where $B\in\mathbb{R}^{N\times N}$ is the relative representations bias.

Suppose that the original point cloud with $N$ points is denoted by $\mathcal{P}=\{p_{i}\}_{i=1}^{N}\subseteq\mathbb{R}^{3}$. We compute the relative position $B$ as follows:

$\displaystyle B_{p_{i},p_{j},m}=p_{i,m}-p_{j,m},\quad m\in\{x,y,z\}.$

(6)

To map relative coordinates to the corresponding position encoding, we maintain three learnable look-up tables $t_{x},t_{y},t_{z}\in\mathbb{R}^{L\times N\times N}$ corresponding to the x, y and z axis, respectively. As the relative coordinates are continuous floating-point numbers, we uniformly quantize the range of $B_{p_{i},p_{j},m}$ into $L$ discrete parts and map the relative coordinates $B_{p_{i},p_{j},m}$ to the indices of the tables as:

$\displaystyle idx_{i,j,m}=\lfloor\frac{B_{p_{i},p_{j},m}+s_{max}}{s_{quad}}\rfloor$

(7)

where $s_{max}$ is the maximum size of point cloud coordinates and $s_{quad}=\frac{2\cdot s_{max}}{L}$ is the quantization size, and $\lfloor\cdot\rfloor$ denotes floor rounding.

We look up the tables to retrieve corresponding embedding with the index and sum them up to obtain the position encoding of

$\displaystyle B_{p_{i},p_{j}}=\sum_{m=1}^{3}t_{m}[idx_{i,j,m}],$

(8)

where $t[idx]$ indicates the $idx$-th entry of the learnable look-up table $t$, and $B_{p_{i},p_{j}}$ means the relative position encoding between $p_{i}$ and $p_{j}$.

External Attention: Although the RA module has demonstrated significant performance on small-scale point clouds, it is unsuitable for large-scale point clouds (e.g., SemanticKITTI [3]) due to its unacceptable $\mathcal{O}(N^{2})$ memory consumption. Therefore, in this work, we perform External Attention computation on large-scale point clouds. External Attention, is a novel attention mechanism based on two external, small, learnable, shared memories, which can be implemented easily by simply using two cascaded linear layers and two normalization layers. It has linear complexity and implicitly considers the correlations between all data samples, making it suitable for large-scale point clouds.

We effectively fuse the outputs of two branches with an addition as they are providing complementary information:

$\displaystyle F^{\prime}=F_{local}+F_{global},\quad F^{\prime}\in\mathbb{R}^{N\times D}$

(9)

The proposed PVT is related to several prior works which includes PVCNN [27], PCT [14], PT¹[9], PT²[65],.

Although we are inspired by the idea of PVCNN [27], our PVT is different in several ways: 1) in the voxel branch, PVCNN uses a 3D convolution to gather local information while we employ SWA computation within each local window, which is highly efficient. Per-layer complexity for different layer types are shown in Table 1, for large-scale point clouds, approximately 10% of the voxels are non-empty ($r=0.1$) [51] which means that our SWA can save significant computation power compared with window-attention. 2) in the point branch, PVCNN uses 1D convolution, which is computation efficient but lacks the global context modeling capability. By performing RA (or EA) computation on the entire points, our method gathers the global modeling power.

Unlike prior Transformer-based 3D models that need to gather the downsampled points and find the corresponding neighbors by using expensive FPS and $k$-NN in point domain, our approach does not require explicitly identify which point is the farthest and what are in the neighboring set. Instead, the voxel branch observes regular data and learns to capture local features using SWA. Additionally, the point branch only needs to perform RA (or EA) on the entire point cloud, which also does not require to find the neighboring points. Thus, our approach obtains highly efficiency than PCT, PT¹ and PT² (see Table 4).

Data. We conduct the classification experiment on the ModelNet40 [47] dataset. ModelNet40 includes 12,311 CAD models from 40 different object classes, in which 9843 models are used for training and the rest for testing. We follow the experimental configuration of PointNet, i.e., for each model, we uniformly sample 1024 points with 3 channels (or 6) of spatial location (and normal) as input; the point cloud is re-scaled to fit the unit sphere.

Setting. We utilize random translation, random anisotropic scaling and random input dropout strategies to augment the input points data during training. During testing, no data augmentation or voting methods were used. For classification on ModelNet40, the SGD optimizer was used for 200 epochs with the batch size 32. We set the initial learning rate to 0.01 and adopt a cosine annealing schedule to adjust the learning rate at every epoch.

Results. The results are shown in Tables 2 and 3, we can see that our PVT outperforms most previous models. Compared with its baselines, such as PCT, PT¹ and PT², PointASNL, our PVT not only achieves state-of-the-art accuracy of 94.1%, but also has the best speed-accuracy trade-off (10$\times$ faster on average). Figure 3 also provides an accuracy plot under equal-epoch setting. As can be seen, our method outperforms all Transformer-based methods, being the fastest and most accurate towards convergence.

Now, we analyze the computational requirements of PVT and several other baselines by comparing the floating point operations required (FLOPs) and number of parameters (Params) in Table 4. PVT has the lowest memory requirements with only 2.45M parameters, and also puts a low load on the processor of only 1.62G FLOPs, yet delivers highly accurate results of 94.1%. These characteristics make it suitable for deployment on a mobile device. In addition, PT² has attained top accuracy on various point cloud learning benchmarks, but its shortcomings of slow inference time and high computing cost are also obvious. In the summary in Table 4, PVT only spends 6.3% of the total runtime on structuring the irregular data, which is much lower than previous Transformer-based models. Overall, PVT has the best accuracy and the lowest computational and memory requirements compared with its baselines.

Data. The model is also evaluated on ShapeNet Parts [24]. It is composed of a total of 16,880 models (14,006 models are used for training, the rest for testing), each of which has 2 to 6 parts and the whole dataset is labeled in 50 different parts. A total 2048 points are sampled from each model as input and only few point sets have six labeled parts. We directly adopt the same train-test split as PointNet in our experiment.

Setting. The same training setting as in our classification task was adopted. For this task on ShapeNet Part, we train our model using the SGD optimizer for 200 epochs with the batch size 16. The initial learning rate was set to 0.1, with a cosine annealing schedule to adjust the learning rate at every epoch.

Results. Table 5 lists the class-wise segmentation results. Part-average IoU is used to evaluate our model, which is given both overall and for each object category. The results show that our PVT makes an improvement of 2.9% over PointNet. Compared with all baselines, PVT attains the top pIoU with 86.6%.

Visualization. In Figure 4, we illustrate output features from the point and voxel branches respectively, where a warmer color represents larger magnitude. As we can see, the voxel branch captures large, continuous parts while the point-based counterpart captures global shape details (e.g., table legs, airplane wings, and tail).

Data. To further assess our network, we conduct semantic segmentation task on S3DIS dataset [1], which includes 273 million points from six indoor areas of three different buildings. Each point is annotated with a semantic label from 13 classes—e.g., beam, bookcase, chair, column, and window. We follow [38] and [29] and uniformly sampled points from blocks of area size $1m\times 1m$, where each point is represented by a 9D vector (XYZ, RGB, and normalized spatial coordinates).

Setting. During the training process, we generate training data by randomly sampling 4,096 points from each block on-the-fly. Following [45], we utilize Area-5 as the test scene and all the other areas for training. Note that the position and RGB information of points are used to as the input features. The same training setting as in our classification task is adopted, but requires 50 epochs with the batch size 8 to fulfil the experiment.

Results. The results are presented in Tables 6 and Table 8. From Tables 6, we can see that our PVT attains mIoU of 67.3%, which outperforms MLPs-based frameworks such as PointNet [29] and PointNet++ [31], graph-based methods such as DGCNN [45], attention-based models such as PointASNL [53]. Moreover, PVT attains mIoU of 69.2% under 6-fold cross-validation, substantially outperforming most prior models.

Visualization. Fig.5 shows the visualization of PVT’s predictions. The predictions of our network are very close to the ground truth. PVT captures detailed semantic structure in complex 3D scenes, which is important in our network.

Data. To further evaluate our network, we conduct outdoor semantic segmentation task on SemanticKITTI [3] dataset, which includes 43,552 densely labeled LIDAR scans belonging to 21 sequences. Each scan is a large-scale point cloud with $\sim 10^{5}$ points and spanning up to 160×160×20 m in 3D space. Officially, the sequences $00\sim 07$ and $09\sim 10$ (19,130 scans) are used for training, the sequence 08 (4071 scans) for validation, and the sequences $11\sim 21$ (20,351 scans) for online testing. The raw 3D points only have 3D coordinates without color information.

Setting. Based on UNet, we build our backbone network for large-scale point clouds segmentation with a stem, four down-sampling and four up-sampling stages, and the dimensions of these nine stages are 32, 64, 64, 128, 256, 256, 128, 64, and 64, respectively. As for the voxel branch, the voxel resolution is 0.05 m for segmentation experiments. We train PVT for 100 epochs on a single GeForce RTX 3090 GPU with the batch size 8. In addition, the Adam optimizer is employed to minimize the overall loss; the learning rate starts from 0.01 and decays with a rate of 0.5 after every 10 epochs.

Results. As shown in Table 7, PVT outperforms the previous state-of-the-art point-based model BAAF by 5.0% in mIoU with 48× measured speedup. Compared with strong attention-based PointASNL and point-based KPConv, our PVT achieves +18.1% and +6.1% mIoU improvements, with 51× and 25× measured speedup respectively.