Meta-Sampler: Almost-Universal yet Task-Oriented Sampling for Point Clouds

Earlier pursuit of 3D computer vision tasks is mainly focused on grid-like representations of voxel volumes, as the mature convolutional neural networks (CNNs) can be directly extended to such data representation and easily introduce inductive biases such as translational equivariance [7, 38]. However, voxel volume representation has the ingrained drawback of being uniform and low-resolution with densely compacted empty cells consuming vast amount of computational resources. Recently, point-based networks have attracted wide attention with the emergence of PointNet/PointNet++ [34, 35]. These architectures pioneered the learning of per-point local features, circumvented the constraint of low resolution and uniform voxel representations, and hence introduced significant flexibilities and inspired a plethora of point-based architectures [13, 27, 29, 40, 45]. A number of point cloud based tasks [14] including classification [12, 16, 28, 30, 43, 48], segmentation [3, 19, 25, 18, 17], registration [2, 11, 22], reconstruction [9, 31], and completion [6, 46] are extensively investigated. Nevertheless, few arts targeted the fundamental component of point cloud sampling in this deep learning era.

Point cloud sampling, a basic component in most point-based neural architectures, is usually used to refine the raw inputs and improve computational efficiency for several downstream tasks. Widely-adopted point cloud sampling methods include RS, FPS [35, 33, 19], and IDIS [13]. A handful of recent works began to explore advanced and sophisticated sampling schemes [32, 5, 44]. Nonetheless, despite the remarkable progress in point cloud sampling, these methods are task-agnostic and rather universal, lacking awareness of the important features which a particular task may require.

Recently, Dovrat et al. [8] proposed a learnable, data-driven sampling strategy by imposing a specific task loss to enforce the sampler in learning specific-related features for a particular task. Later, Lang et al. [26] extended the learning approach by introducing a differentiable relaxation to minimise the training and inference accuracy gap for the sampling operation. Nevertheless, by introducing an additional task loss, sampling ultimately becomes constrained and prone to overfitting on a specific task model instead of the task itself. Additionally, this also requires significant extra training to fit a particular goal.

Our meta-sampler hopes to bring the best of both worlds: being task-oriented yet as universal as possible. Instead of directly overfitting onto a task model, we focus on how to learn a task through incorporating a meta-learning algorithm. By introducing a better approach of learning a particular task through joint-training, our pretrained meta-sampler be rapidly fine-tuned to any task, making it almost-universal while easing the computational efficiency to which sampling is targeting in the first place.

Meta-learning, the process of learning the learning algorithm itself, has shown to be applicable to several challenging computer vision scenarios such as few-shot and semi-supervised classification [10, 36], shape reconstruction [39], and reinforcement learning [15, 23] due to its capacity for fast adaptation.

Finn et al. [10] proposed one of the most representative meta-learning methods, termed model-agnostic meta-learning (MAML), that allows the model to quickly adapt to new tasks in different domains such as classification, regression, and reinforcement learning. Later, Antoniou et al. [1] further improved the MAML learning scheme, making the learning more generalisable and stable. Recently, Huang et al. [21] proposed MetaSets, which aims to meta-learn the different geometry of point clouds so that the model can generalize to classification tasks performed on different datasets. In contrast and being analogous to the standard meta-learning problem, our proposed meta-sampler focuses on universal point cloud sampling for different tasks, aiming to achieve fast adaptation to reduce computation efficiency through a training strategy extended from [10, 1]. Our fast adaption is not just across tasks within the meta-training, but also across models, datasets and unforeseen tasks.

The goal of this paper is to develop a learning-based sampling module $f_{\theta}$ with trainable parameters $\theta$, which takes in a point cloud with $m$ points and outputs a smaller subset of $n$ points ($m>n$). Apart from the objective of SampleNet to learn task-specific sampling (i.e., particularly suitable for a single task such as shape classification or reconstruction), we take a step further and aim to propose an universal pretrained model, which can be rapidly adapted to a set of different tasks $S_{T}=\{T_{i}\}_{i=1}^{K_{T}}$. Formally, we define the ideal adaptation of sampling to a specific task $T_{i}$ as capable of achieving satisfactory performance by integrating the sampling module into a set of known networks $S_{A_{i}}=\{A_{i,j}\}_{j=1}^{K_{A_{i}}}$ (trained with unsampled point clouds of $m$ points to solve task $T_{i}$. ). We split $S_{A_{i}}$ into $S_{A_{i}}^{train}$ and $S_{A_{i}}^{test}$ (i.e., task networks used during training are disjoint to the ones for testing) to make sure that our evaluation on $f_{\theta}$ is fair and not overfitting to task models instead of the task itself. Note that while $S_{A_{i}}^{train}$ is available during training, the weights are frozen when learning our sampler as suggested by [26].

To achieve the dual objectives of high accuracy and rapid convergence, we must first carefully evaluate the best training strategy to better learn each individual task, and then design a training strategy which is adaptive to multiple tasks. We build our sampler $f_{\theta}$ based on the previous state-of-the-art learnable sampling network — PointNet-based SampleNet architecture [34, 26] — and then introduce our training technique in a bottom-up manner.

For an individual task $T_{i}$, we hope that the $f_{\theta}$ learns to sample the best set of points $\forall A\in S_{A_{i}}$.

The conventional way of training the SampleNet uses one single frozen network $A^{\prime}$ as $S_{A}^{train}$ for training by defining a sampling task loss $\mathcal{L}_{ST_{i}}$ targeting $T_{i}$ as:

where $\mathcal{L}_{T_{i}}$ is the loss when pretraining $A^{\prime}$. We refer to this configuration as single-model, single-task training. As mentioned previously, this training method, having accomplished promising results in several tasks, still exhibits a large accuracy discrepancy between the results on $A^{\prime}$ and $S_{A}^{test}$. In other words, even though $A^{\prime}$ is frozen during the training of SampleNet, the sampling stage is overfitted onto the task network instead of the task itself.

To alleviate the issue of model-wise overfitting, we extend (1) and create a joint-training approach for a single task. Specifically, we take a set of weight-frozen models $\{A_{i,j}\}_{j=1}^{k},1<k<<K_{A_{i}}$ as $S_{A_{i}}^{train}$ and compute $\mathcal{L}_{ST_{i}}$ as:

It is critical to understand that all the frozen task models are under inference mode (i.e., not significantly sabotaging computation power) and that only a very small number of task models (easily obtainable online or by self-training with different random initial weights) would bring significant improvements to the sampler’s performance. We further show in Section 4.2 that a very small $k>1$ allows the sampling network generalise better across $S_{A_{i}}$, as $S_{A_{i}}^{train}$ becomes a vicinity distribution rather than a single specific instance to $S_{A_{i}}$.

In addition to the joint $\mathcal{L}_{ST_{i}}$, we also update the weights with a simplification loss comprising the average and maximum nearest neighbour distance and a projection loss to enforce the probability of projection over the points to be the Kronecker delta function located at the nearest neighbour point (identical to the SampleNet loss [26]).

Instead of restricting ourselves to a single task (e.g., classification), we consider whether training the sampler over multiple tasks could lead to our vision of an almost-universal sampler. Broadly, we aim to extend the sampler beyond multi-model to multi-task, such that given any task $T_{i}\in S_{T}$, where $S_{T}$ is a set of tasks, a good initial starting point could be achieved for the sampler. In this way, adapting or fine-tuning to a particular task (which may even be beyond the known set) will be rapid and cheap.

To tackle this, we draw inspiration from the MAML framework and propose a meta-learning approach for rapidly adaptive sampling [10]. In essence, we aim to utilise the set of $S_{A_{i}}^{train}$ to mimic the best gradients in learning a particular task $T_{i}$ for meta-optimisation, such that given any task $T_{i}\in S_{T}$ or even beyond the known set of tasks, the MAML network can quickly converge within a few iterations and without additional training of the task networks.

The joint-training procedure discussed in the previous section motivates that a particular task is better solved with a set of task networks instead of just one — we transfer this idea to the meta-optimisation such that the sampler is adaptive to a number of tasks instead of just one. Formally, we first optimise the adaptation of $f_{\theta}$ to $T_{i}\in S_{T}$ by updating the parameters $\theta$ to $\theta^{\prime}_{i,j}$ for every $A_{i,j}$ through the gradient update:

where $\alpha$ is the step size hyperparameter. Similar to MAML, we can directly extend the single gradient update into multi-gradient updates to optimise the effectiveness of $\theta^{\prime}_{i,j}$ on $T_{i}$.

With the inner update (3), we then follow the meta-optimisation procedure through a stochastic gradient descent:

where $\beta$ is the meta step size hyperparameter that could either be fixed or accompanied with annealings. Note that we apply the single task loss in the inner update (3) but sum all losses from all weights to resemble a task in the meta-update (4). Section 4.3 shows that our meta-optimisation design is sufficient in learning tasks for rapid adaptation. Simplification and projection losses are also directly optimised at this stage. They are however directly updated rather than included in the meta-update fashion as they are task-agnostic.

We describe the overall training pipeline of the proposed meta-sampler (Figure 2) as the following:

To comprehensively evaluate the performance of our meta-sampler, we extract a set of representative tasks on 3D point clouds, including shape classification, reconstruction, and shape retrieval. Note that all experiments are conducted on the ModelNet40 [42] (except for ShapeNet [4] used in transferring dataset analysis) to ensure fair evaluation (i.e., without introducing additional information during meta-sampler training). The detailed experimental settings (i.e., task network architecture, task loss) are described as follows:

To justify the effectiveness of the proposed multiple-model training scheme, we present the quantitative comparison of incorporating multiple models training and the traditional SampleNet single-model training strategy in all three individual tasks on the ModelNet40 dataset [42]. We adopt the official train and test splits in this dataset, and follow [8, 26] to pretrain all task networks with the original point clouds (1024 points by default). All the task models are under inference mode during the training of the sampling network. We evaluate our sampling on different sampling ratios calculated as $1024/n$, where $n$ is the number of outputted points from the sampler. Note that for reconstruction and shape retrieval tasks, we were adopting a different dataset and task to prior networks. Much work is required for the adaption and thus we only compare with the previously proposed state-of-the-art SampleNet.

Versatility is a broad term with multiple dimensions requiring evaluation. To fully realize this, we begin with the critical evaluation of our meta-sampler’s adaptiveness on tasks included in the meta-optimisation step. We then extend to the more challenging scenarios of changing model architecture, dataset distribution, and ultimately tasks distinct from the ones used for meta-training. All experiments are conducted with hyperparameters $\alpha$ and $\beta$ set to 1e-3.

As shown in Figure 3, we separately compare the task performance as the training progresses with/without our pretrained meta-sampler for different meta-tasks. It is clear that as the sampling ratio increases (i.e., the task is more challenging), joint-training without meta-sampler starts at lower performance and requires more iterations to converge to a stable result. By contrast, our pretrained meta-sampler allows the network to quickly adapt to the task within one epoch across all sampling ratios and achieves higher accuracies after 10 epochs in most cases.

There are also two intriguing points we would like to address within this empirical study. First, we observe a relatively large fluctuation of the performance for shape classification and shape retrieval — a phenomenon we conjecture to be owing to the distribution shift between training and testing sets. Second, we notice that our meta-sampler not only converges faster, but also pushes the upper bound in some cases. For example, our shape retrieval results at sampling ratio 16 achieved the accuracy of 96.9% (the upper bound of training from scratch is 96.7%) on the 20th epoch (not plotted in the figure). This infers that by learning how to adapt, the sampling model could potentially also be trained to learn better. However, such occurrences do not take place at all times.

The achieved results are plotted in Figure 4. It is apparent that better performance is achieved using our meta-sampler, with a higher starting point and fast convergence speed under both sampling ratios. This further demonstrates the capacity of our meta-sampler in adapting different model architectures. Interestingly, we also notice that the classification accuracy of PointNet++ achieved on sampled points drops significantly compared with that of unsampled point clouds, especially under the sampling ratio of 16. This is likely because FPS is progressively used in each encoding layer. In this case, by adding a SampleNet in front of PointNet++, we are implicitly “double sampling” and leaving very few features for the abstraction layer.

We evaluate our meta-sampler on the point cloud registration — the task of finding the spatial transformation between two point clouds. Here, we follow the standard train-test split of PCRNet [37] to obtain pairs of source and template point clouds with templates rotated by three random Euler angles of $[-45^{\circ},45^{\circ}]$ and translated with a value in the range $[-1,1]$. $\mathcal{L}_{T_{i}}$ is the CD between the source point cloud and the template point cloud with our predicted transformation. We first train three PCRNets to achieve the rotation error of around 7-9 degrees on unsampled point clouds, then freeze the PCRNet weights and perform the proposed joint-training scheme under the conditions with/without the pretrained meta-sampler under the sampling ratio of 32. We ran each setting three times and show the mean and standard deviation of rotational error during training in Figure 6. Even though the task objective (registration), task network (PCRNet), and even the dataset itself (pairs of point clouds from ModelNet40 with extra transformations) are unforeseen during our meta-optimisation, we can still notice two subtle yet solid performance differences adopting our meta-sampler: 1) The pretrained model generally converges faster during the initiation of training and 2) The pretrained model is much more stabilised and improves consistently compared to the model trained from scratch that exhibits a large variance throughout every epoch.