Point Cloud Recognition with Position-to-Structure Attention Transformers

Our proposed PS-Former model is a Transformer-based architecture, as seen in Figure 2. The overall goal of the PS-Former is to learn an effective representation for downstream tasks such as classification and segmentation. The input is $N\times 3$ coordinate data which will first go through an input embedding module to project the data from the low-dimension space to the high-dimension space. Two self-attention layers are adopted to make points interact with each other to obtain more developed features before condensation. Following are our two key designs: Condensation Layer and Position-to-Structure Attention Layer. The features are concatenated together for further recognition tasks.

The vanilla attention mechanism [39] is as follows:

where $Q$ and $K$ are query and key respectively, $C$ is the feature dimensions and $A$ is the final attention matrix. However, in the 3D space of points, a more meaningful way to consider the correlation between points is the euclidean distance instead of the dot product. Thus, we introduce a new attention mechanism based on euclidean distance in our model. We also keep this characteristic in the high dimensional space as well to facilitate an easier learning process. Moreover, we use a similar normalization method as PCT[13] and we’ll see later in the Condensation Layer (section 3.3) it will be a key design for our algorithm. Our final attention equation is as following:

A common challenge in point cloud recognition is exploiting rich local and global structures. Methods such as PointNet++[34] and PCT[13] have adopted a combination of sampling—Farthest Point Sampling (FPS)—and grouping methods to tackle this difficulty. Our approach extracts the local structures through a two-step process based entirely on attention: sampling and feature construction. The generation of these features through attention create a much richer and more focused local neighborhood context as compared to the more handcrafted approach of PCT and PointNet++. Furthermore, the proposed method proves to be a natural integration into the Transformer network. The Condensation Layer also compresses information for less computational complexity. The schema of the Condensation Layer can be viewed in Figure 3.

We select $T$ points from the point cloud to serve as center points of local neighborhoods. The traditional Farthest Point Sampling strategy does not identify the most important neighborhoods because its main objective, to maximize inter-point distance, does not motivate meaningful center points. Since attention is a learned measure of affinity and an integral part to the Transformer, we sample the center points based on the row sum of the attention matrix. We apply the softmax operation in columns at first, which can be understood as each point’s contribution to each other. The row sum can be seen as how much information each point will take. We make the choice of the center points based on this computed row sum. Note that we don’t simply choose the top-k points with the largest sum since it will likely cause points to crowd together. The detailed algorithm is shown in Algorithm 1.

After we get the center points, we use the strategy which will be described in section 3.4 to obtain the structure features for each point from the relative position features. Moreover, the position features are updated using the vanilla attention mechanism.

Our key design is to decouple the position and structure information and allow the position feature to gradually enrich the structure feature. To achieve this, we design the Position-to-Structure Attention Layer as presented in Figure 4. The structure information is acquired using the relative position features which can be seen as the neighborhood structure features around one specific point. In our Position-to-Structure Attention Layer, there are three main processes to update the features which can be seen in Figure 5. Each point’s position feature is updated through the weighted sum along with its structure feature, while the relative position feature is also used to enrich the structure feature.

In particular, the attention matrix is computed by both the position features and the structure features as following:

Both the position features and structure features are updated through attention while structure features are also influenced by top-k related points according to the attention matrix.

where $\mathcal{K}_{1},\mathcal{K}_{2}$ are mapping functions: $R^{d}\rightarrow R^{d}$, $\circ$ is the concatenate operation, $F_{\text{pos}}$ is the position features for each point, $F_{\text{str}}$ is the structure features for each point and $F_{\text{gather}}$ is the features gathered using the relative position features.

Dataset and implementation. A primary assay for point cloud classification, ModelNet40[44] is a 3D point cloud dataset composed of 12,311 CAD models from 40 different object classes. To maintain consistency with other studies, the same train-test split is used as the given dataset: 9,843 training samples and 2,468 testing samples. In processing the data, we randomly select 1,024 points, as many other methods do. It is important to note that we discard information regarding surface normals and solely use the 3D coordinates. Following PointNet[32] convention, we augment the training data with random translation, scaling and dropout. While testing, augmentation is omitted and the test set is evaluated without any voting strategy. The model is trained with the Adam optimizer [18] with its learning rate and weight decay both set at 1e-4.

Results and analysis. We achieved a competitive result on the ModelNet40 dataset, as seen in Table 1. The best overall accuracy achieved by the PS-Former without the assistance of normals or voting ensembles was 93.9%. For the experiment, we investigated the results of our novel Condensation Layer (shown in Figure 6). We observe that our Condensation Layer identifies points of interests much more commonly in central areas, while the Farthest Point Sampling (FPS) strategy identifies points in an even distribution in the point cloud. We hypothesize that our Condensation Layer favors more information rich centers that hold high attention values with respect to its neighboring points; the Farthest Point Sampling technique likely produces an even collection of points due to it’s objective to maximize point-to-point distance. With intuitively chosen attention-based centers, our model is able to outperform FPS-based methods such as PointNet++ and PCT. Also, with the novel Position-to-Structure Attention and Condensation Layer, our PS-Former outperforms other Transformer-based methods.

Dataset and implementation. We further evaluate our model on ShapeNet[51] for part segmentation. It contains 16,881 shapes with 14,007 as training samples and the remaining 2,874 as testing samples. The dataset contains 50 object parts in total with each object containing no more than 6 object parts. 2048 points are sampled for each point cloud and all the points are assigned with a part label. To classify each point into the correct label, our Condensation Layer’s output token number is set to the input token number which in this case is 2048. After the Position-to-Structure Attention layers, we concatenate the node feature with the whole data feature which is obtained using max pooling as in ModelNet40 classification task. To reduce the computation, in this experiment, we only use the structure feature since this task is more focused on point-wise classification and structure information is more important in this case. We use SGD as our optimizer with the initial learning rate set to 0.1 and the weight decay set to 1e-4. Cosine Annealing Scheduler is used to adjust the learning rate with the minimal learning rate set to 1e-3 [26].

Results and analysis. The IoUs of all the test dataset and each class are presented in Table 2. It can be seen that our model achieves competitive results compared to previous methods. Some of the visualization results are presented in Figure 7.

Dataset and implementation. We also use S3DIS[1] to evaluate our model on scene segmentation. The dataset contains 6 areas and 271 rooms in total. Each data sample contains 4,096 points with each point assigned with a label from 13 classes. Different from ModelNet40 classification and ShapeNet part segmentation, the input of each point contains 9-dimensional features including coordinate, RGB and normal vector information. Due to this change, our model removes subtraction in the gather operation in the Condensation Layer since we want to keep the color and norm information in each point’s structure information while keeping the position information to continue gradually enriching the structure information through the Position-to-Structure Attention Layer. In our training-testing procedure, the training dataset is all the areas except area 5 which leaves area 5 to be the testing dataset. The same optimizer and learning schedule are used as part segmentation which is SGD with the initial learning to 0.1 and Cosine Annealing Scheduler[26].

Results and analysis. The mAcc and mIou of the test dataset are presented in Table 3. Note that since our model is trying to decouple the position and structure features and gradually enrich the structure features through the Position-to-Structure Attention Layer, this task which includes RGB and normal vector input as well as position input is not a direct fit for our task. However, it can still benefit from it given by our results. Furthermore, we present several visualization results on Area 5 in Figure 8.

We analyze the importance of PS-Former’s awareness and distinction between position information and structure information by altering its use and access to such features, as reported in Table 4. We conduct two ablations: 1) removing structure features completely, 2) combining structure features with position features. For the first ablation, the position-only features are fed through vanilla attention. Input dimension is doubled to keep attention dimensions and modeling capability constant. We observe that our model with the proposed Position-to-Structure Attention exhibits a +1.8% improvement over a position-only attention model. The decrease in performance for the first ablation likely comes from the loss of information on local structures. In the second ablation, the structure features are computed once and concatenated with position features to be passed through classical attention. Compared to the ablation, our proposed method performs +1.1% better. The results of the two ablations not only suggest the importance of structure features but also of an awareness of the distinction between structure and position.

In an effort to demonstrate the Condensation Layer’s performance improvements, we conduct ablations on the number of points sampled. As shown in Table 5, our proposed model with the Condensation Layer configured to 512 points performed better than 1024 points, suggesting the efficacy of our attention-based point sampling in creating information rich local structures. This observation coincides with our hypothesis made in Section 4.1. Furthermore, the down-sampling of points provides great relief in terms of computational efficiency because attention is an operation with $\mathcal{O}(n^{2}\cdot d)$ complexity. Through the stated ablations, we verify the utility and desired properties of our proposed Condensation Layer.

Previous research has demonstrated a boost in performance gained from multi-scale features that incorporate both global and local structure[34]. To investigate its effects regarding our PS-Former, we trained a multi-scale PS-Former and applied it to ModelNet40 for classification. The model is comprised of two branches: a PS-Former that samples 256 points and one that samples 1024 points. The features generated by the two are concatenated together after max pooling and fed through the same classification network used for Section 4.1. The model attained a best overall accuracy of 93.6%. We hypothesize that the multi-scale approach does not provide complementary information, as our Condensation Layer and Position-to-Structure Attention layers already serve as effective extractors of structural feature representation.

We also conduct ablation studies on using the cross-attention operation to mediate position and structure feature interaction. After the Condensation Layer each point will have two features: position and structure. We use cross-attention to enrich the structure features from the position features. Besides this, we also conduct two experiments using cross-attention from structure to position and two-way interaction. The results are presented in Table 6. It can be seen clearly that with cross-attention from position features to structures, the result is still decent compared to PCT[13]. However, the application of cross-attention on structure to position, causes the results to go down. This further demonstrates that the weighted sum mechanism is not suited for capturing the structure features from the position features while the decoupling still works. The decrease in performance with structure enriching position validates our design of position to structure.