Point Transformer

I Introduction

Processing 3D point sets using deep neural networks has become very popular the past few years. The three-dimensional information has a wide range of applications in autonomous driving and computer vision . However, training neural networks on point sets is not trivial. First, point sets are unordered, thus require the neural network to be permutation invariant. Second, the number of points in the set is usually dynamic and unstructured. Finally, the network needs to be robust against rotation and translation to operate in the metric space, and since the points describe objects, the network needs to capture the spatial relations between the points.

Standard neural architectures, such as convolutional neural networks (CNN), have shown promising results for structured data. For that reason, several point set processing approaches attempt to transform the points into regular representations such as voxel grids or rendered views of the point clouds . However, transforming the point sets leads to loss of shape information as geometric relations between points are removed. Furthermore, these methods suffer from high computational complexity due to the sparsity of the 3D points. To address these limitations, there is another family of approaches that act directly on the point set. The main idea is to process each point individually with a multi-layer perceptron (MLP) and then fuse the representation to a vector of fixed size with a set pooling operation over a latent feature space . Set pooling is a symmetric function that is permutation invariant. Additionally, under certain conditions, set pooling acts as a universal set function approximator . Nevertheless, Wagstaff et al. argue that reducing the latent representation to a vector of fixed length can be impractical since the cardinality of the input set is usually not considered. Thus, the capacity of the vector may not be sufficient enough to capture the spatial relations of the point set which may reduce the overall performance. Therefore, the set pooling mechanism can become a bottleneck for point processing networks.

Our goal and motivation stems from removing the set pooling method and overcoming the aforementioned bottleneck, while still achieving a permutation invariant representation that models the point set relations in terms of object shape and geometric dependencies. Therefore, it is necessary to introduce a symmetric set function that replaces traditional set pooling operations. For that, we adapt the attention mechanism , which was originally introduced for natural language processing, that is used to weight and score sequences (words) based on learned importance. To our understanding, we face a similar problem in 3D point processing, given that we need to relate representations of the input points to capture and describe the object’s shape. Additionally, attention itself does not depend on the input ordering, i.e. it is permutation-invariant, as it is comprised of matrix multiplication and summation only, which makes it well-suited for our problem. However, the output is still unordered, thus, directly processing the output of attention for standard computer vision tasks is not possible. Consequently, our goals can be outlined as follows:

Avoid the bottleneck that can occur while employing set pooling operations .

Present a novel permutation invariant network architecture that adapts the popular and prevalent attention mechanism for 3D point processing.

Demonstrate superior performance compared to traditional set pooling methods to justify the use of attention and reinforce the claims made by Wagstaff et al.

To address these problems, we propose SortNet, a permutation invariant network module, that learns ordered subsets of the input with latent features of local geometric and spatial relations. For that, we learn important key points, which we call top-k selections, that replace the set pooling operation. Since current state-of-the-art methods have shown that aggregating local and global information increases the network’s capabilities of capturing context information , we employ SortNet to generate local features of the point cloud. Moreover, global features of the entire point cloud are related to the sorted local features using local-global attention. Local-global attention attends both feature representations to capture the underlying shape. Since the local features are ordered, the output of local-global attention is ordered and permutation invariant; and thus it can be used for a variety of visual tasks such as shape classification and part segmentation. An overview of our network is outlined in Fig. 1. Since we aim to process 3D point sets using the ideas proposed by the Transformer network architecture , we took inspiration from , and name our network Point Transformer.

Overall, our contributions can be summarized as follows:

We propose Point Transformer, a neural network that uses the multi-head attention mechanism and operates directly on unordered and unstructured point sets.

We present SortNet, a key component of Point Transformer, that induces permutation invariance by selecting points based on a learned score.

We evaluate Point Transformer on two standard benchmarks and show that it delivers competitive results.

II Related Work

Below, we discuss approaches that process 3D points and are related to our work.

Point clouds are irregular and unordered sets of points with a variable amount of elements, thus applying standard neural networks on 3D points is not possible. For that reason, previous approaches rely on transforming the point sets into an ordered representation, such as voxel grids. The metric space is discretized into small regions (voxels), which are labeled as occupied if a point lies inside the voxel. Then, 3D convolutional networks (CNN) can be easily applied to the voxel-based representation . This pre-processing, however, reduces the resolution as multiple points are combined into a single voxel and thus damages important spatial relations of the metric space. Furthermore, voxelization increases the memory requirements and computational complexity due to the sparsity of the 3D points. To address these limitations, multiple extensions have been proposed that try to leverage the sparsity of 3D data , but still fail to process large amounts of input points.

View-based methods: In contrast to building voxel grids, a lot of research has been conducted on rendering point clouds into 2D images, i.e. structured representation of the underlying 3D shape. Then, working with traditional CNNs is possible . Since shape information can be occluded by rendering point clouds from a specific viewpoint, multi-view approaches have been proposed that render multiple images from different angles . Even though images are rendered from different views, the model still fails to capture all geometric and spatial relations. To this day, multi-view approaches achieve impressive results on standard 3D benchmarks. However, the transformation from sparse 3D points into images increases computational complexity as well as required memory.

Shape-based methods: PointNet is a pioneering network architecture that operates directly on 3D point sets, and it is invariant to input point permutations. Therefore, a transformation into a structured representation is no longer necessary. PointNet uses a multi-layer perceptron (MLP) with shared weights that encodes spatial features to each input point separately. Then, a symmetric function, e.g. max pooling, is applied to the latent features to induce permutation invariance and create a global feature representation of the input. PointNet established the de facto standard for point processing that many state-of-the-art approaches still rely on . However, it is not able to encode and capture local information, since the max pooling operation induces permutation invariance, but also destroys local structures and relations of the points in metric space. To address this issue, Qi et al. proposed the improved PointNet++ architecture, a hierarchical model that abstracts the input points with every layer to produce sets with fewer elements. First, centroids of local regions are sampled using hand-crafted algorithms, then local features are encoded to the centroids by exploring the local neighborhood. Thus, allowing the network to capture fine-grained patterns and improving the performance on current datasets. A general approach related to unordered sets was introduced by Zaheer et al. demonstrating the capabilities of pooling operations to induce permutation invariance. Importantly, they prove that the set pooling method is a universal approximator for any set function. In general, problems arise with set pooling when the reduced feature vector lacks the capacity to capture important geometric relations. Our work addresses this limitation with a network topology that encodes the entire point cloud by relating local information with the global shape structure.

Convolutions on Point Clouds: Classic convolutional neural networks require the input data to be ordered, such as images or voxel grids. Since points are unstructured, an active research area is the definition of convolution operations that can operate on irregular 3D point sets such as KPConv , SpiderCNN or PointCNN . These methods achieve state-of-the-art performance on a variety of tasks. However, due to the irregularities of the shape and point density, point convolutions are usually hard to design and the kernel needs to be adapted for different input data .

II-B Attention

Attention itself has its origin in natural language processing . Traditionally, encoder-decoder recurrent neural networks (RNN) were used for machine translation applications, where the last hidden state is used as the context vector for the decoder to sequentially produce the output. The problem is that dependencies between distant inputs are difficult to model using sequential processing. Bahdanau et al. introduced the attention mechanism that takes the whole input sequence into account by taking the weighted sum of all hidden states and additionally, models the relative importance between words. Vaswani et al. improved the attention mechanism by introducing multi-head attention and proposing an encoder-decoder structure that solely relies on attention instead of RNNs or convolutions. Therefore, they reduce the computational complexity. In this work, multi-head attention is the basis for Point Transformer.

Attention with point cloud processing: Neural networks that rely on attention achieved impressive results in machine translation, and were adopted to function on point clouds by utilizing the points as sequences. Vinyals et al. proposed a network that processes unordered sets using attention. They show that the network is able to sort numbers. However, they only focus on generic sets. In contrast, we present an approach that is applied to different point cloud related tasks for capturing shape and geometry information. Recently, Lee et al. proposed Set Transformer, a method that is related to our approach. They adapt the original Transformer network to process unordered sets by using induced points, i.e. trainable parameters of the network, that are attended to the input. Set Transformer focuses on general sets as input. Furthermore, Lee et al. demonstrate that it is applicable to point sets. In our work, Point Transformer is specifically designed to process point clouds and leverage important characteristics of points in metric space such as shape and geometric relations.

Xie et al. propose ShapeContextNet, where they hierarchically apply the shape context approach that acts as a convolutional building block. To overcome the difficulties of manually tuning the shape context parameters, Xie et al. employ self-attention to combine the selection and feature aggregation process into one trainable operation. However, similar to point cloud convolutions, shape context relies on a manual selection of the shape context kernels which is sensitive to the irregularities of point cloud data.

The Point2Sequence model uses an attention-based sequence-to-sequence network. The approach first extracts local regions and produces local features using an LSTM-based attention module. Using a set pooling method, a global feature vector is generated following the ideas of and . However, it relies on a sequence-to-sequence architecture that tends to be more computational complex than multi-head attention . Furthermore, in contrast to our method, Point2Sequence uses a max-pooling operation to make the network permutation invariant. Yang et al. introduce a network architecture that replaces traditional subsampling methods like furthest point sampling (FPS) with an attention-based selection process using the gumbel-softmax function, which is similar to the proposed SortNet module.

Recently, Tao et. al proposed a multi-head attentional point cloud processing network that uses a rotation invariant representation of point clouds as input. For that, they employ a multi-head attentional convolution layer (MACL) with attention coding. However, their work focuses on designing a rotation invariant network that relies on global max pooling operations, whereas Point Transformer together with SortNet leverages the strengths and advantages of the attention operation to select useful local point structures and relates them to the global shape to induce permutation invariance.

III Fundamentals

III-B Transformer

The Transformer network is an extension of the attention mechanism from Equation (2) that consists of an encoder-decoder structure and introduces multi-head attention. In the following, we explain multi-head attention in detail, as our Point Transformer architecture relies on it.

Instead of employing a single attention function, multi-head attention first linearly projects the queries, keys and values $Q,K,V$ $h$ times to $d_{k},d_{k}$ and $d_{v}$ dimensions, respectively, using separate feed-forward networks to learn relations from different subspaces. Then, attention is applied to each projection in parallel. The output is then concatenated and projected again using a feed-forward network. Thus, multi-head attention can be defined as follows:

that outputs a matrix of dimension $N\times d_{m}$ which order depends on the ordering of $P$. Since the output is not permutation invariant but follows the ordering of the input, Transformer and multi-head attention can not be used directly for point data without further processing. To solve this problem, we introduce our novel Point Transformer architecture that handles unordered point sets.

IV Point Transformer

This section presents Point Transformer, a neural network that operates on point set data and it is based on the multi-head attention mechanism. The network is permutation invariant due to a new module that we name SortNet. Our goal is to explore shape information of the point set by relating local and global features of the input. This is done using cross multi-head attention. To introduce our method, we first give an overview of the complete Point Transformer architecture, which is shown in Fig. 2. Our approach is divided into three parts:

SortNet that extracts ordered local feature sets from different subspaces.

Global feature generation of the whole point set.

Local-Global attention, which relates local and global features.

Consequently, the output of SortNet constitutes one local feature set

Since $\mathcal{Q}$ is an ordered set, it follows that $\mathcal{F}_{m}^{L}$ is ordered as well. To capture dependencies and local features from different subspaces, we employ $M$ separate SortNets. Finally, the $M$ feature sets are concatenated to obtain an ordered local feature set of fixed size

IV-B Global Feature Generation

The second branch of Point Transformer is responsible for extracting global features from the input point cloud. To reduce the total number of points to save computational time and memory, we employ the set abstraction multi-scale grouping (MSG) layer introduced by Qi et al. . We subsample the entire point cloud to $N^{\prime}<N$ points using the furthest point sampling algorithm (FPS) and find neighboring points to aggregate features of dimension $d_{m}$ resulting in a global representation of dimension $N^{\prime}\times d_{m}$. Note that the global feature representation is still unordered since no sorting or set pooling operation was performed.

IV-C Local-Global Attention

The goal of Point Transformer is to relate local and global feature sets, $\mathcal{F}^{L}$ and $\mathcal{F}^{G}$ respectively, to capture shape and context information of the point cloud. After obtaining both feature lists, we employ self multi-head attention $\mathcal{A}^{\text{self}}$ on the local features $\mathcal{F}^{L}$ as well as the global features $\mathcal{F}^{G}$. Then, cross multi-head attention layer $\mathcal{A}^{\text{cross}}$ from Equation (6) is applied such that every global feature is scored against every local feature, thus relating local context with the underlying shape. We call this operation local-global attention $\mathcal{A}^{\text{LG}}$ (see Fig. 2) and define it as follows:

IV-D Complete Model

To recap, Point Transformer functions as follows: Our architecture is comprised of two independent branches, SortNet for the extraction of local features and a global feature generation module. SortNet constitutes a novel architecture that selects a number of input points based on a learned score from latent features, resulting in $M\cdot K$ ordered feature vectors with dimension $d_{\text{m}}$. In the global feature branch, we employ multi-scale grouping to reduce the total number of points to $N^{\prime}$ while aggregating spatial information. Then, local-global attention is used to relate both spatial signatures, producing a permutation invariant and ordered representation of length $K\cdot M$ with reduced dimension $d^{\prime}_{m}$ (see Fig. 2), which can be used for different tasks such as shape classification or part segmentation. Additionally, we demonstrate the processing chain of our model as a flowchart in Fig. 4.

V Experiments

In this section, we perform two standard evaluations on Point Transformer. We compare our results with approaches that operate directly on 3D point sets , attention-based approaches and methods that use point cloud convolutions . Moreover, we provide a thoughtful analysis and visualizations of the components of our approach. We implement our network in Pytorch where we rely on the RAdam optimizer for all experiments. The weights of each layer are initialized using the popular Kaiming normal initialization method . Our implementation will be made publicly available.

We evaluate Point Transformer on the ModelNet40 dataset and use the modified version by Qi et al. that provides $10.000$ points sampled from the mesh of the CAD model, as well as the normal vectors for each point. The dataset consists of $40$ categories and it is composed of $9843$ training samples and $2468$ test samples. During the training for classification, we augment the input by randomly scaling the shape in the range of $[0.8,1.25]$ and randomly translating in the range of $[-0.1,0.1]$. Additionally, we apply random dropout of the input points as proposed in . For the experiments, we set $N=1024$, $D=6$ ($xyz$ and normals), $d_{m}=512$, $d^{\prime}_{m}=64$, $M=4$ and $K=64$. The results of the shape classification are shown in Table I. Point Transformer outperforms attention-based methods (top part of Table I) and achieves on par accuracy when compared to state-of-the art methods (bottom part of Table I) with a classification accuracy of $92.8\%$.

V-B Point cloud part segmentation

Here, we evaluate Point Transformer on the challenging task of point cloud part segmentation on the ShapeNet dataset , which contains $13.998$ train samples and $2874$ test samples. The dataset is composed of objects from $16$ categories with a total of $50$ part labels. The goal is to predict the class category of every point. To address this task, the network has to learn a deep understanding of the underlying shape. For the part segmentation, we set $M=10$ and $K=16$. Again, we use $xyz$ coordinates with normal vectors ($D=6$) and $N=1024$ input points. For this experiment, we follow the setup of where a one-hot encoding of the category is concatenated to the input points as an additional feature. We report the mean IoU (Intersection-over-Union) in Table I. Finally, we visualize exemplary results of the part segmentation task in Fig. 5.

V-C Network complexity

We examine the network complexity of Point Transformer and perform a comparison to related approaches. The results of this experiment are shown in Table III. We performed all experiments on a Nvidia GeForce 1080Ti. Point Transformer has about $13.5$ million learnable parameters ($51$ MB), which is less when compared to KPConv ($15$ million learnable parameters). However, our model is about $6$ times bigger than PointNet++ and Point2Seq. This is mainly due to the fact that the Transformer model itself has a lot of learnable parameters. For example, one SortNet only has about $10.000$ learnable parameters which shows that SortNet can be incorporated into any existing network architecture without much space requirements and computational overhead, as it only adds about $1.2$ ms of inference time. In many cases, the forward pass of multiple SortNets can additionally be performed in parallel. Even though, Point Transformer has more learnable parameters than, e.g, PointNet++, it still has a faster inference time because multi-head attention blocks are highly optimized and computation is also performed in parallel by employing multiple attention heads. For the computational complexity of the network, an upper bound can be estimated from the most expensive operation, which in our case is the multi-head attention mechanism. The complexity is given by $\mathcal{O}(N^{2}\cdot d_{m})$, thus it scales quadratic with respect to the total number of input points.

V-D Hyperparameter Study

Here, we analyze the effects of different numbers of SortNets in our Point Transformer architecture as well as the amount of Top-K selections on the ModelNet40 dataset . The results are shown in Tab. IV. Furthermore, we present the hyperparameters that were used for the reported results for the classification and the part segmentation task in Tab. V. The parameters follow the notation introduced in Fig. 2 and Fig. 3. The values were found by performing a hyperparameter grid search experiment for the classification and the part segmentation, similar to Tab. IV. We report the set of parameters that achieved the best overall performance. Note, that for the rFF, each value in the parenthesis denotes one layer, where the value represents the feature dimension for that layer.

V-E Point Transformer Design Analysis

We conduct an ablation study to show the influence of each Point Transformer module. Afterward, we qualitatively examine our classification results by visualizing the learned point set regions that contribute to the classification output.

Ablation study of SortNet: We first evaluate Point Transformer using only the SortNet module from Fig. 3 with the classification head from Fig. 2 a). Our aim is to show that the learned scores are based on the importance of points for the classification task. In addition, we want to verify that SortNet selects points that help to understand the underlying shape. Since we cannot explicitly define which are the most important points, we rely on the accuracy score. In detail, we train SortNet based on three different experiments and deliberately set $M=10$ and $K=12$, selecting only a subset of the entire point cloud ($M\cdot K=120$, $N=1024$). In the first experiment, we train SortNet as it is implemented in the Point Transformer pipeline. In the second experiment, we replace the Top-K selection process with the furthest point sampling. Finally, we randomly select $K$ points from the input set instead of the learned Top-K selection. It is important to note, that the last two experiments remove the permutation invariance property. However, we want to show that SortNet performs better than a random selection of points and handcrafted sampling methods. Thus, we rely on random sampling and FPS as baselines. The results are shown in Table II a). With randomly sampled points, SortNet achieves $60.1\%$ classification accuracy. When we apply the FPS to cover most of the underlying shape, the accuracy increases to $74.8\%$, indicating spatial information preservation. Finally, when we use learned Top-K selection, we achieve the highest classification accuracy of $83.4\%$. This empirically shows that SortNet learns to focus on important shape regions.

Ablation study Global Feature Generation: In this ablation study, we compare different sampling methods for the extraction of global features. We rely on the complete Point Transformer pipeline as shown in Fig. 2 and replace the set abstraction (MSG) with different sampling approaches. Again, we evaluate the accuracy of the classification task. The results are presented in Table II b). In the first experiment, we use the complete input point cloud. Then, we sample $N^{\prime}=128$ points using the furthest point sampling, which slightly improves our result by $0.4\%$. When we additionally aggregate features from local regions around the sampled points, i.e. set abstraction with multiscale grouping (MSG) , the accuracy can be further increased to $92.8\%$. This indicates that scoring the local features against every input point makes it harder to find important relations. Additionally, by uniformly selecting fewer points and aggregating local features the network can concentrate on meaningful parts of the underlying shape.

Rotation robustness of SortNet: In this section we evaluate the robustness of SortNet against rotations of the input cloud. For this, we first evaluate Point Transformer on the ModelNet40 test set and randomly rotate the input point cloud. Even though we did not train the network with rotations, we still achieve a classification accuracy of $92.3\%$ compared to $92.8\%$ without rotations. We applied the same input point rotation to PointNet++ and classification accuracy dropped from $91.9\%$ to $88.6\%$. To qualitatively support this claim, we visualize the learned Top-K selections of one SortNet for different rotations in Fig. 6, which shows that SortNet still focuses on the similar local regions even when the input point cloud is rotated.

Visualizations of learned local regions: Here, we show that SortNet focuses on local regions similar to the receptive field of a CNN. For this, we visualize the learned Top-K selections of multiple trained SortNet modules on different models of the same object class in Fig. 7 and Fig. 8. It is apparent, that each SortNet tries to select similar regions even when the shape of the model is slightly different. This, together with the results from the rotational robustness, suggests that SortNet is aware of the underlying shape.

All Top-K selections: As an additional evaluation, we show all selected points of $M=8$ SortNet modules in Fig. 9 for the classification task. We visualize points that were selected from the same SortNet with the same color. It is apparent, that different SortNet modules focus on different parts of the object and in combination, still retain as much as possible of the underlying shape.

VI Conclusion and Future Work

In this work, we proposed Point Transformer, a permutation invariant neural network that relies on the multi-head attention mechanism and operates on irregular point clouds. The core of Point Transformer is a novel module that receives a latent feature representation of the input point cloud and selects points based on a learned score. We relate local features to the global structure of the point cloud, thus exploiting context and inducing shape-awareness. The output of Point Transformer is a sorted and permutation invariant feature list that is used for shape classification and part segmentation. Finally, we show that our point selection mechanism is based on importance for the specified task. As future work, we want to focus on improving the efficiency of the Transformer architecture by implementing recent advances for self-attention, such as .

References