Upsampling Autoencoder for Self-Supervised Point Cloud Learning

In recent years, deep learning on 3D point cloud has attracted increasing attention of researchers. As 3D points are irregularly sampled and have quite a few special properties such as irregularity and permutation invariance, the traditional neural networks, i.e., convolutionl neural networks (CNNs) in the 2D field, cannot be directly applied to point clouds. Therefore, previous works attempt to convert point clouds into a regular grid structure and apply 3DCNN to feature learning [84, 73, 67]. However, the computational consumption and memory of 3DCNN increase cubically with the increase of voxel’s resolution [41].

Another mainstream method is directly operating the unordered point clouds [4, 39, 81, 22, 83, 62, 78, 9, 55]. For instance, Qi et al. [48] pioneered this series of works and proposed a novel neutral architecture, PointNet, which can directly handle irregular and unordered 3D points by stacking multi-layer perceptrons (MLPs) and capturing global shape with max-pooling. Later, Qi et al. [49] proposed PointNet++ to overcome the disadvantage that PointNet fails to capture local structures by developing a hierarchical grouping architecture at different set abstraction level. Subsequently, PointCNN [35], PointConv [68], RSCNN [40] and DGCNN [65] also focus on local structures of point cloud and further improve the quality of captured features. However, these models mentioned above are supervised and rely on large-scale labeled point cloud dataset. In this paper, we propose a novel unsupervised framework, namely UAE for point cloud analysis.

There are several prior attempts [23, 1, 3, 18, 24, 76] on learning shape-specific invariant representations of a point cloud based on unsupervised reconstruction models. These methods discover valuable information of 3D point cloud by reconstructing the input data, which has been shown to effectively learn shape-specific invariant representations. For example, FoldingNet [76] proposed an end-to-end auto-encoder to get the codeword that can represent the high-dimensional embedding of a point cloud, and the fully-connected decoder was replaced with the folding-based decoder. GraphTER [19] proposed a novel unsupervised learning method to capture intrinsic patterns of point cloud structure under both global and local transformations. However, because of lacking sufficient semantic supervision, this kind of methods are insufficient to understand the high-level semantic information of point clouds.

Witnessing the great success of self-supervised on computer vision tasks, many researchers have endeavored to apply it to the field of self-supervised point cloud analysis. Xie et al. [72] pioneered this line of works by proposing PointContrast to build representations of scene-level point cloud which relies on the complete 3D reconstruction of a scene with point-wise correspondences between the different views of a point cloud. However, these point-wise correspondences need to post-process the input data by registering the different depth maps into a single 3D scene. Later, Zhang et al.[80] proposed DepthContrast to side-step the necessity of registered point clouds or correspondences, by considering each depth map as an instance and discriminating between them, even if they come from the same scene. Yet, most of them require careful treatment of negative key samples by either relying on large batch sizes [72, 80], memory banks or customized mining strategies [17] to find the negative key samples. Additionally, the performance of these methods critically depends on the choice of data augmentations [72, 21, 19, 29], limiting the applicability of the approach. As a comparison, in this work, we investigate a useful pretext task for learning shape-specific invariant representations.

Some researchers have also exploited the self-supervised pre-training model, for example, OcCo [61] is investigated by using occlusion completion as a pre-training task to learn an initialization for point cloud models. Compared with OcCo, our model has two advantages: firstly, our UAE does not require extra data-preprocessing to generate occluding point clouds based on different viewpoints; and secondly, our method uses point cloud upsampling as a pre-training task and design a novel joint loss function to capture better shape information by enforcing the upsampled points uniformly distributed on the underlying shape surface.

Self-supervised learning has attracted increasing attention since the cost of accurate annotated dataset is extraordinarily expensive [42, 45, 53, 59]. There are several classes of methods for learning representations such as clustering [6, 7], GANs [15, 43], pretext tasks [14, 44] etc. Recently, contrastive learning [46, 16, 71, 58, 26, 11] have been proposed to learn the unsupervised representation of 2-dimensional natural images. It can be considered as a pretext task where the goal is to maximize the representational similarity between positive key samples and dissimilarity between negative key samples for the input query. The positive key sample is generated with a random data augmentation module, which, given the input, generates a pair of random views of the input. Other inputs in the same batch are often used as negative key samples. Generally, contrastive and related methods strongly depend on data augmentation and the sampling of negative pairs [25, 20], which limit their applicability. In contrast to these methods, we pre-train the UAE on a simple but meaningful upsampling objective, aiming to capture high-level semantic information within a simple and effective framework.

Given a point cloud $\mathcal{X}\in\mathbb{R}^{N\times 3}$, we first subsample $\mathcal{X}$ into a lower resolution subset $\mathcal{P}\in\mathbb{R}^{rN\times 3}$, where $r$ denotes that subsampling ratio. There exists a variety of sampling methods: 1) Random Sampling, where the probability of sampling each point in the point cloud follows a uniform distribution; 2) Farthest Point Sampling, where each point to be sampled is as far away as possible from points in the sampled set; 3) Local Sampling, where the points are sampled from a local part of the point cloud. In this paper, we propose to randomly sample a subset of points at a low proportion uniformly (e.g., $r=0.125$), which largely eliminates redundancy, thus creating a rather challenging task that cannot be easily handled by extrapolation from subsampled neighboring points [25]. This highly sparse point cloud is more conducive for us to design an effective encoder to capture high-level semantic information.

The encoder $\varphi(\cdot)$ takes the coordinates of the subsampled point clouds $\mathcal{P}$ as input and outputs high-dimensional features. In contrast to previous self-supervised methods, where the global shape representation is learned by the encoder, we perform point-wise feature extraction on subsampled point cloud ${\mathcal{P}}$, where each point can learn its representation by upsampling decoding the original point cloud to reveal the local structure around it. These representations will gather global shape information about the original point cloud as we randomly subsample points into different subsets (or groups) through training iterations while capturing the local structures under point-wise upsampling decoding.

Generally, any deep learning-based network that takes point clouds as the input and outputs high-dimensional features can be adopted as the encoder $\varphi(\cdot)$ of UAE. However, as 3D points have some special properties such as irregularity and permutation invariance, we can not directly leverage the convolutional neural network (CNN). Therefore, we deploy the dynamic graph CNN (DGCNN [65]) to overcome this limitation. In particular, we adopt the EdgeConv layer in DGCNN as our basic feature extraction block. By performing EdgeConv, our encoder can aggregate the features of neighbor points to the center point and update the feature of the center point.

Specifically, given an input subsampled points $\mathcal{P}=\{p_{i}\}_{i=1}^{rN}\subseteq\mathbb{R}^{C}$, we first apply K-Nearest Neighbor (KNN) algorithm to construct a local graph in the feature space. The KNN algorithm based on the feature space can efficiently and effectively find two points with the most similar semantics, i.e., the points on the two wings of the airplane, which have similar semantics. Then, the EdgeConv layer performs feature transformation on the $rN$ points and output the updated point features of $p_{i}$. The output feature of a point $p_{i}$ is

$$f_{i}=\max_{p_{j}\in\mathcal{N}(p_{i})}{\rm ReLU}(h_{\theta}(p_{i},p_{i}-p_{j})),$$

(1)

where $p_{j}\in\mathcal{N}(p_{i})$ denotes that point $p_{j}$ belongs to the neighborhood of point $p_{i}$, $h_{\theta}$ denotes that the learnable parameters of MLP.

By stacking four EdgeConv layers, each point is aggregated with local neighboring features hops away. Furthermore, as the subsampling operation enlarges the distance between neighboring points, performing EdgeConv for the subsampled points will capture longer-range dependencies and thus facilitates to extract high-level semantic information that is invariant under random sampling.

The decoder $\phi(\cdot)$ takes the point-wise features as input and outputs the upsampled point cloud $\tilde{\mathcal{P}}$. To upsample a point cloud, Li et al. proposed PU-Net [77] to duplicate the point features and then employ separate MLPs to process each copy independently. However, the expanded features would be too similar to the inputs, so affecting the upsampling quality. In recent years, wang et al. proposed MPU [66] to break a $16\times$ upsampling network into four successive $2\times$ subnets to progressively upsample points in multiple steps. MPU preserves superior upsampling details but usually requires a more complex training process.

Inspired by the PU-GAN [33], we design a novel upsampling decoder to expand the point features, which is mainly composed of feature up and feature down blocks. As shown in Figure 2, we first upsample the point features $F_{in}$ (after a MLP layer) by $\frac{1}{r}\times$ to generate $F_{up}$ and downsample it to generate $F_{down}$; then, instead of directly constructing the original point cloud $\mathcal{X}$, we adopt residual learning to regress the per-point feature offset by calculating the difference between $F_{up}$ and $F_{down}$; Ultimately, we feed them into a feature up block and a MLP layer to restore the original point cloud $\mathcal{X}$. Such a strategy that utilizes feature offset to self-correct the expanded features has two advantages: firstly, it facilitates the production of fine-grained features while avoiding tedious multi-step training; secondly, the features of the same object can be completely different with respect to rigid transformations. In the following, we detail the design choices of feature up and feature down blocks.

Feature up block. To upsample the point features $\frac{1}{r}$ times, we adopt the commonly-used variation expansion operator [33, 34] by duplicating $F_{in}$ with $\frac{1}{r}$ copies and concatenating with a regular 2D grid. However, such operator may introduce redundant information or extra noise [34]. To rectify these problems, we propose to make use of the offset-attention [22] as the global refinement unit by considering the overall shape structure. The reason behind is that, compared with the widely-used self-attention unit [33, 34, 79], the offset-attention is generally more robust because it works by replacing the attention feature with the offset between the input of self-attention module and attention feature. The pipeline of this block is illustrated in Figure 3 (top).

Feature down block. To downsample the expanded features, we propose a novel GCN-based feature down block, illustrated in Figure 3 (bottom). Given the expanded features $F_{up}\in\mathbb{R}^{N\times D}$, our method works in two steps. First, we reshape $F_{up}$ of shape $N\times D$ to $rN\times\frac{1}{r}\times D$ and use one layer of EdgeConv to downsample $F_{up}$ to shape $rN\times D$ using learnable parameters. Second, we feed them into a set of MLPs to regress the point features $F_{down}$.

In contrast to previous works [33, 77, 66], our feature down block leverages GCNs, which are common modules for feature extraction. To the best of our knowledge, we are the first to introduce a GCN-based feature downsampling block. Our GCN design choice stems from the fact that GCNs enable our feature down block to encode spatial information from point neighborhoods and learn new features from the latent space rather than simply using Convs.

To encourage the upsampled points to be similar to the original point cloud and uniformly distributed on the underlying shape surface, a novel joint loss function was proposed to train our UAE in an end-to-end fashion. Next, we will detail the design of this loss function.

Reconstruction loss. Our UAE reconstructs the original point cloud $\mathcal{X}$ by predicting coordinate for each unsampled point. Each element in the decoder’s output is a vector of coordinate values representing a point’s spatial location. We formulate our objective function to encourage the geometric consistency between $\tilde{\mathcal{P}}$ and $\mathcal{X}$ in the point space:

$$\mathcal{L}_{CD}(\tilde{\mathcal{P}},\mathcal{X})=\frac{1}{|\tilde{\mathcal{P}}|}\sum_{\tilde{p}\in\tilde{\mathcal{P}}}\min_{x\in\mathcal{X}}||\tilde{p}-x||_{2}+\frac{1}{|\mathcal{X}|}\sum_{x\in\mathcal{X}}\min_{\tilde{p}\in\tilde{\mathcal{P}}}||x-\tilde{p}||_{2},$$

(2)

where $\mathcal{L}_{CD}(\cdot)$ means the Chamfer Distance (CD) to measure the average closest point distance between two point sets, $||\cdot||_{2}$ denotes the L2 distance between two points.

Repulsion loss. Since the reconstruction loss alone can not ensure that the upsampled points over the underlying object surfaces will be uniformly distributed, which is important for capturing high-level shape information. To solve this issue, we propose to use the repulsion loss [77] as the constraint term of uniformity distribution, which is represented as:

$$\mathcal{L}_{rep}=\sum^{N}_{i=0}\sum_{\tilde{p_{j}}\in\mathcal{N}(\tilde{p_{i}})}\delta(||\tilde{p_{j}}-\tilde{p_{i}}||)\omega(||\tilde{p_{j}}-\tilde{p_{i}}||),$$

(3)

where $\mathcal{N}(\tilde{p_{i}})$ is the point set of the k-nearest neighbors of point $\tilde{p_{i}}$, $N$ is the number of upsampled points and $||\cdot||$ is the L2-norm. $\delta(m)=-m$ is the repulsion term, which is a decreasing function to penalize $\tilde{p_{i}}$ if $\tilde{p_{i}}$ is located too close to other points in $\mathcal{N}(\tilde{p_{i}})$. We further use the fast-decaying weight function $\omega(m)=e^{-m^{2}/h^{2}}$ ($h$ as the finite support radius [37]) to penalize $\tilde{p_{i}}$ only when it is too close to its neighboring points.

Overall loss function. The overall loss function can be calculated as the weighted sum of all the terms described above:

$$\mathcal{L}_{total}=\alpha\mathcal{L}_{CD}+\beta\mathcal{L}_{rep},$$

(4)

where $\alpha$ and $\beta$ are the weighting factors for the loss functions of reconstruction and repulsion, respectively. Following [77], we set $\alpha=100$ and $\beta=1$.

Pre-training Dataset. We train our UAE on ShapeNet dataset [8], which consists of 57, 448 synthetic 3D CAD models organized into 55 categories with a further 203 subcategories, organized according to WordNet synsets. For pre-training we use the normalized version of ShapeNet, where all shapes are consistently aligned and normalized to fit inside a unit cube.

Architecture. The network architecture used for pre-training is shown in Fig 4. we deploy four EdgeConv layers whose dimensions are [64,64,128,256] and one MLP (648) layer as the encoder. The number of nearest neighbors set to 20 for all EdgeConv layers. After the four EdgeConv layers, we concatenate the output of these layers to get a 64+64+128+256=512 dimensional point cloud features and feed them into one MLP layer with input channels of 512 and output channels of 648. Similar to previous works [33, 50], the upsampling decoder is used to obtain the reconstructed point cloud, where $D$ is 128.

Pre-training Hyperparameters. We uniformly sample 2, 048 points on ShapeNet dataset for our UAE pre-training. We use the Adam optimizer with no weight decay (L2 regularisation). The learning rate is set to 1e-3 initially and is decayed by 0.7 every 10 epochs. We pre-train the models for 120 epochs. The batch size is 32, and the momentum of batch normalization is 0.9. Note that, all the experiments are trained on an NVIDIA RTX3090 GPU.

Evaluation Metrics. For the classification task on ModelNet10 and ModelNet40 datasets, we use the overall accuracy (OA) as the metric. On ShapeNet Part dataset, we evaluate our scheme with part classification accuracy and mean Intersection over-Union (mIoU). For each sample, IoU is computed for each part that belongs to that object category. The mean of all part IoUs is regarded as the IoU for that sample. For point cloud upsampling task, we use the Chamfer distance (CD), Hausdorff distance (HD), and point-to-surface distance (P2F) w.r.t ground truth meshes as evaluation metrics. The smaller the metrics, the better the performance.

Dataset. We utilize ModelNet40 [70] and ModelNet10 [70] for shape classification task. We follow the same data split protocols of PointNet-based methods [65] for these two datasets. For ModelNet40, the train set has 9, 840 models and the test set has 2, 468 models, and the dataset consists of 40 categories. For ModelNet10, 3, 991 models are for fine-tuning and 908 models for testing. It contains 10 categories. We follow the experimental configuration [48]: (1) we uniformly sample 1,024 points from the mesh faces for each model; (2) the point cloud is re-scaled to fit the unit sphere; and (3) the (x,y,z) coordinates and the normal of the sampled points are used in the experiment. During the training process, randomly scaling and perturbing the objects are adopted as the data augmentation strategy in our experiment.

Implementation Details. The global max pooling and average pooling layer are deployed in our classification head to acquire a 1,296-dimensional global feature vector. Three layers of linear projection with dropout ratio of 50% are used to get the final classification score.

For unsupervised transfer learning, we fix the parameters of encoder and only train the classification head. During training, a random translation in [-0.2, 0.2], a random anisotropic scaling in [0.67, 1.5] and a random input dropout were applied to augment the input data. During testing, no data augmentation or voting methods were used. The batch sizes were 32, 200 training epochs were used and the initial learning rates were 10^-3, with a cosine annealing schedule to adjust the learning rate at every epoch.

For supervised fine-tuning, we fully fine-tuned the pre-trained model on ModelNet40 and ModelNet10 datasets.

Results. The classification results are presented in Table 1. When utilizing DGCNN as the encoder, our method outperforms most of previous unsupervised counterparts and the results on ModelNet10 and ModelNet40 are comparable to certain fully-supervised models. Since the pre-training of the encoder is based on different datasets, the results demonstrate that our framework has a strong generalizability, which is regarded as a significant application of self-supervised representation learning. Notably, most of the classes are unseen by the model during ShapeNet pre-training. Thus, the superior performance further demonstrates that our model has a good ability to generalize to novel classes.

To the best of our knowledge, the most important application of self-supervised learning methods is to make full use of a large number of unlabeled data and boost the performance of supervised learning methods. Thus, we pre-train the encoder with our framework on ShapeNet and fine-tune the weights on downstream classification tasks and compare the results with other supervised fine-tuning methods. As demonstrated in Table 1, the proposed method outperforms all other supervised fine-tuning methods on ModelNet10 and ModelNet40 datasets, which justifies the superiority of our UAE.

Dataset. We use the large-scale 3D dataset ShapeNet Parts [36] as the experiment bed. ShapeNet Parts contains 16,880 models (14,006 models are used for training, and 2874 models are used for testing), each of which is labeled with two to six parts, and the entire dataset has 50 different part labels. We sample 2,048 points from each model as input, with a few pointsets having six labeled parts. We directly adopt the same train–test split strategy similar to DGCNN [65] in our experiment.

Implementation Details. The global max pooling and average pooling layer are also deployed to acquire a 1,296-dimensional global feature vector. And we design a segmentation head to propagate these features to each point hierarchically. Like shape classification task, three layers of linear projection with dropout ratio of 50% are also used to get the final classification score of each point.

We optimize our networks via SGD with batch size 32. The learning rate of unsupervised transfer learning setting starts from 10^-2 and decreases to 10^-4, and the learning rate of supervised fine-tuning setting decays from 10^-1 to 10^-3.

For unsupervised transfer learning, we fix the parameters of encoder and only train the segmentation head. For supervised fine-tuning, we fully fine-tuned the pre-trained model on ShapeNet Part dataset.

Results. As shown in Table 2, the proposed UAE outperforms all other unsupervised point cloud learning models by a large margin, which indicates that our pre-trained model captures more effective semantic information that can transfer well to downstream segmentation tasks. Some results are visualized in Figure 5. We also conduct the object part segmentation experiments under supervised fine-tuning strategy and make comparisons with previous excellent models, such as PointDis, MID-FC and OcCo in Table 2. As shown, our UAE model (Ours-DGCNN) fine-tuned on 100% labeled samples achieves state-of-the-art performance. Compared to the supervised learning framework DGCNN, our pre-trained model (Ours-DGCNN) achieves remarkable performance improvements, which demonstrates the advantage of our pre-training strategy. We can conclude that pre-training with our framework on unlabeled data significantly boosts the performance and can be regarded as a strong initializer for supervised models, which is a critical application of self-supervised learning.

Raw point clouds data acquired from depth cameras or LiDAR scanners are often sparse, noisy and non-uniform, which hinders shape analysis and 3D reconstruction. Point cloud upsampling is thus significant to the subsequent CAD applications such as analysis, rendering or general processing. We evaluate our UAE on the point cloud upsampling task over the PU-GAN’s dataset.

Dataset. For point cloud upsampling task, we use PU-GAN’s dataset [33] as the experiment bed. The PU-GAN’s dataset has 147 CAD models that were collected from the released datasets of PU-Net [77] and MPU [66], as well as from the Visionair repository [60], covering a rich variety of objects, ranging from simple and smooth models (e.g., Icosahedron) to complex and high-detailed objects (e.g., Statue). Following previous work [33], we randomly select 120 models for training and use the rest for testing.

Implementation Details. We demonstrate a pre-training strategy to evaluate whether the unsupervised pre-training with our model helps improve the performance. We first pre-train our UAE on ShapeNet dataset in an unsupervised fashion, and then employ the learned representation from the pre-trained encoder as an initialization. Eventually, we build our upsampling version of UAE where we only replace the feature extractor of PU-GAN with our pre-trained encoder. We evaluate the effectiveness of our model by comparing the results of PU-GAN with our initialization and other upsampling models in a supervised fashion.

For training, we randomly sample 512 points from each point cloud in the PU-GAN’s dataset and upsample them to 2,048 points. For testing, we randomly sample 2048 points and upsample them to 8196 points. The quality of the upsampled point cloud is measured by the CD, HD and P2F between the original and upsampled point clouds.

Results. We qualitatively and quantitatively compared the unsupervised pre-training with our method with the randomly initialized PU-GAN and other state-of-the-art point set upsampling methods: EAR [28], PU-Net [77], MPU [66] and PU-GCN [50]. Table 3 shows the quantitative comparison results. Our method (Ours (transfer)) achieves the lowest value consistently for most evaluation metrics in an unsupervised manner. Particularly, the supervised fine-tuning with our method (Ours (fine-tune)) outperforms all previous point set upsampling models. The gain is significant in 4$\times$ upsamling ratio which means that our pre-trained model can capture more effective semantic information. Figure 6 shows the point set upsampling results on PU-GAN’s dataset.

The following studies, which are used to investigate the determining factors of our framework, are conducted on both ModelNet40 and ShapeNetPart datasets.

Impact of subsampling strategy. To examine the effectiveness of various subsampling strategies, we conduct a detailed ablation on shape classification and part segmentation tasks. Table 4 presents the results with three types of subsampling methods. We observe that the random sampling achieves the best performance, improving by 0.63%/0.79% on average over Farthest Point Sampling (FPS), and 0.81%/0.58% over local sampling (LS). The reason is that FPS and LS retain too many geometric details (see Figure 7), resulting in the created task being easily solved by extrapolation from neighboring points. In contrast, the task created by random sampling with a low sampling ratio is harder than that of FPS and LS, which provides a more challenging goal for our model. Moreover, random sampling uniformly selects $rN$ points from the original $N$ points. Its computational complexity is $\mathcal{O}(1)$, which is agnostic to the total number of input points. Compared with FPS and LS, random sampling has the highest computational efficiency, regardless of the scale of input point clouds [27]. Based on these characteristics in both computational time and memory, we can conclude that our UAE is suitable for training very large models.

Effect of subsampling ratio. We conduct an ablation study to analyze the setting of subsampling ratio $r$. The results are shown in Table 5. The best performance is achieved when $r$ is set to 12.5%. When the ratio becomes smaller, the lack of key points will result in a performance decline. On the contrary, when the ratio increases, a lot of noises at the local boundary lead to the deformation of feature extraction ability, and then reduces the accuracy of the model. Meanwhile, the task with high subsampling ratio (25% or 50%) can be easily solved by our UAE, which is not conducive to capture high-level semantic information.

Impact of loss function. We also investigate the options of loss functions. The results are shown in Table 6. From this table, we can see that compared with the CD loss, EMD loss achieves better results (+0.13%/+0.18%) because it can better encourage the output points to be located close to the underlying object surfaces [77]. However, in Table 6, the results of the two tasks are close, the time complexity of EMD loss is $\mathcal{O}(N^{2})$ [33], which means that it needs more pretraining time, especially when $N$ is very large. Furthermore, we can see that the performance of the model is significantly improved after adding repulsion loss, especially under the combination of “CD + RL” where the results from the output points are marked semantically and distinguished in spatial position to reduce noise.