SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Prototypical Variational Autoencoder for Few-shot 3D Point Cloud Object Detection
Weiliang Tang*
wltang21@cse.cuhk.edu.hk
Biqi Yang*
bqyang@cse.cuhk.edu.hk
Xianzhi Li
xzli@hust.edu.cn
Pheng-Ann Heng
pheng@cse.cuhk.edu.hk
Yunhui Liu
yhliu@mae.cuhk.edu.hk
Chi-Wing Fu
Department of CSE and SHIAE
cwfu@cse.cuhk.edu.hk
Few-Shot 3D Point Cloud Object Detection (FS3D) is a challenging task, aiming to detect 3D objects of novel classes using only limited annotated samples for training. Considering that the detection performance highly relies on the quality of the latent features, we design a VAE-based prototype learning scheme, named prototypical VAE (P-VAE), to learn a probabilistic latent space for enhancing the diversity and distinctiveness of the sampled features. The network encodes a multi-center GMM-like posterior, in which each distribution centers at a prototype. For regularization, P-VAE incorporates a reconstruction task to preserve geometric information. To adopt P-VAE for the 3D object detection framework, we formulate Geometric-informative Prototypical VAE (GP-VAE) to handle varying geometric components and Class-specific Prototypical VAE (CP-VAE) to handle varying object categories. In the first stage, we harness GP-VAE to aid feature extraction from the input scene. In the second stage, we cluster the geometric-informative features into per-instance features and use CP-VAE to refine each instance feature with category-level guidance. Experimental results show the top performance of our approach over the state of the arts on two FS3D benchmarks. Quantitative ablations and qualitative prototype analysis further demonstrate that our probabilistic modeling can significantly boost prototype learning for FS3D.
3D point cloud object detection is a fundamental scene understanding task, aiming at recognizing and localizing objects in point-based scenes. Recent deep-learning methods are mostly fully supervised, requiring extensive labeled data to support the model training. Such a data-hungry setting is, however,
* Equal contributions to the works. This work was supported by Shenzhen Portion of Shenzhen-Hong Kong Science and Technology Innovation Cooperation Zone under HZQB-KCZYB-20200089. This work was supported by the Hong Kong Centre for Logistics Robotics, the Research Grants Council-General Research Fund (No. 14201620). This work was supported by InnoHK of the Government of Hong Kong via the Hong Kong Centre for Logistics Robotics. This work was supported by project MMT-p2-21 of the Shun Hing Institute of Advanced Engineering (SHIAE).
Figure 1: (a) The basic structure of our proposed Prototypical VAE, where we learn a GMM-like distribution for feature sampling that can be regularized by 3D reconstruction tasks. (b) At the scene level, each point cloud scene can be reconstructed by composing local geometric components, we call them geometric-informative prototypes. (c) At the object level, objects of the same category can be represented in a unified fashion by a class prototype, named the class-specific prototype.
impractical in real applications, considering the heavy and costly loads to densely annotate the data. Inspired by humans' capability of recognizing novel objects by learning from limited data, Few-shot Learning (FSL) [1, 2, 3, 4, 5] has been proposed to enable neural networks to handle novel data by training on large labeled data of base classes and small labeled data of novel classes.
Prototype learning is an effective approach to achieve FSL [2, 6, 7, 8, 9, 10, 11]. It has been applied successfully to various 2D vision tasks, e.g., object detection [12, 13, 14], semantic segmentation [15, 16, 17, 18, 19, 20], and instance segmentation [21, 22]. In prototype learning, some principles are first defined to separate the prototypes, e.g., by category [23, 17, 13, 14, 16, 21], by partial similarity [19, 11], or by intrinsic characteristics [12, 15, 17, 18, 20]. These prototypes can be viewed as clustering centers, such that we can group the latent features based on our pre-defined principles. Importantly, the prototypes provide transferable knowledge that allows us to generalize a model for handling novel classes, thereby facilitating representation learning for downstream tasks.
In this work, we exploit prototype learning for Few-Shot 3D Point Cloud Object Detection, named FS3D. We aim to predict accurate 3D bounding boxes on objects of novel classes, given only a small amount of labeled novel-class exemplars. To address FS3D with prototype learning, Zhao et al. [10] propose a pioneering work. Yet, their learning scheme lacks fine-level supervision, as the intermediate features are simply averaged to update the prototypes, which are then used to augment features for sequential detection. There are two fundamental issues. First, without strong regulation on prototype learning, the prototypes can lose substantial 3D information and become less geometric-informative. Second, due to the data imbalance, prototypes of novel classes are clearly underrepresented than those of the base classes, so the latent space can easily overfit the base ones.
To address the above issues, we propose a variational autoencoder approach particularly designed for prototype learning, named Prototypical VAE (abbr. P-VAE). As Fig. 1(a) illustrates, instead of directly learning the features, our P-VAE learns the distribution parameters, by which we can construct a Gaussian Mixture Model(GMM)-based posterior for sampling features from the probabilistic latent space. For issue 1, we propose to perform reconstruction tasks with an encoder-decoder architecture in our P-VAE framework to better preserve the geometric information in the intermediate features. For issue 2, our P-VAE allows us to sample diverse and varied features controlled by a multi-center distribution. Hence, even with limited novel samples, our approach can still gather richer information to generate good-quality representative prototypes for the novel classes.
To enhance the object detection task with the proposed P-VAE, we formulate two crucial extensions: Geometric-informative Prototypical VAE (GP-VAE) and Class-specific Prototypical VAE (CP-VAE). As Fig. 1 shows, we can view real-world 3D objects from two aspects. In a broader sense, all the objects, regardless of their class labels, can be reasoned by a set of local geometric components; see Fig. 1(b). For GP-VAE, each component is a geometric-informative prototype, such that a given point cloud scene can be constructed by combining these prototypes. Narrowing the scope to each object, objects of the same category have similar shapes, while objects of different categories look more different; see Fig. 1(c). Hence, we propose CP-VAE to obtain class-specific prototypes, such that each prototype can serve as a strong shape prior to guiding the object-level feature learning. We incorporate the two modules in our object detection pipeline, leveraging GP-VAE to account for scene-level feature extraction and CP-VAE for object-level feature refinement; see Fig. 2.
We conduct extensive experiments on two FS3D benchmarks, FS-ScanNet and FS-SUNRGB [10]. The results in Sec. 4.2 show that our method significantly outperforms the SOTA approaches in various few-shot settings. More analysis on prototypes in Sec. 4.3 demonstrate the effectiveness of leveraging VAE for prototypical learning. To summarize, our contributions are listed as follows:
- We propose the Prototypical VAE (P-VAE) to learn multiple distributions centering at prototypes, such that it can construct a probabilistic latent space that allows us to effectively sample diverse features with higher controllability and flexibility.
- To harness P-VAE for object detection, we extend it for geometric-informative and class-specific prototypical VAEs (i.e., GP-VAE & CP-VAE) based on our observations.
- Results demonstrate the superiority of our new approach for FS3D and further studies analyze the prototypes and explain the effectiveness in detail; see Sec. 4 and Appendix.
Variational Autoencoder (VAE) is first introduced by Kingma et al. [24] to learn a generative model for constructing a distribution-based latent space. It provides a probabilistic representation learning scheme with theoretical support. Bornschein et al. [25] assume that features follow a multicenter prior distribution, which encourages unsupervised deep clustering in the embedding space. Norouzi et al. [26] select exemplars from the training data and use them to model the embedding space with GMM. To strengthen the robustness, [27, 28, 29] balance between reconstruction authenticity and generation diversity; [30, 31, 32] aims at preventing VAE models from collapse; [33, 34, 35] explore usages of different regularizations. VAE is widely adopted in 3D generation tasks, e.g., point cloud generation [36, 37, 38] and scene reconstruction [39, 40]. However, it has not been fully explored for 3D perception tasks, especially for the challenging FS3D task.
Few-shot 2D Image Object Detection (FS2D) methods can be classified into two categories. The first family aims to address the data imbalance problem, for example, by fine-tuning with balanced data [41], augmenting objects by different scales [42], augmenting novel samples with generative networks [43], learning reasonable and unbiased latent space [44], etc. The second family focuses on learning support (base)-novel (query) correlations. Kang et al. [45] reweight the query features given the support features; Han et al. [46] fuse the support and query features by spatial alignment; Fan et al. [47] and Yan et al. [48] exploit deep support information to guide the region proposal network and the detection heads; some other works incorporate support information learned from the backbone to guide the feature extraction of the query images, using transformers [49, 50, 51, 52, 53] or graph convolution networks [49]. An intuitive solution is to adopt these FS2D strategies to tackle FS3D, as some of our comparison methods in Sec. 4.2. However, their detection results (see Tab. 1 and 2) are far worse than our P-VAE. The reason is that 3D point clouds are intrinsically more disordered and complex than pixel-grid 2D images, thus naturally lacking abundant data with full annotations.
FSL for 3D Point Cloud is under-explored, as we discussed before. As a pioneer work, Fan et al. [54] tackle the one-shot point cloud classification, exploring base-novel sample relations with graph convolution. Cheraghian et al. [55] introduce a transductive zero-shot learning method with a triplet loss to obtain a better classifier. Zhao et al. [11] incorporate FSL for 3D semantic segmentation with a transformer for multi-prototype transductive inference. To our best knowledge, Prototypical VoteNet [10] is the first and only work on FS3D that harnesses prototype learning. However, their experimental results, refer to the indoor scenes in Fig. 3, are not satisfactory. First, it lacks a strong regulation for prototype learning, like our reconstruction task, to preserve latent geometric information. Second, the quality of their novel-class prototypes is impaired due to data imbalance, whereas we propose a new VAE-based approach to sample features with significant diversity.
In this section, we first define the FS3D task and give a modular overview of our pipeline in Sec. 3.1. Then, we present the formulation of our novel Prototypical VAE (P-VAE) to learn a probabilistic-distribution-based feature space in Sec. 3.2. Building upon the P-VAE, we develop (i) Geometric-informative Prototypical VAE (GP-VAE) (Sec 3.3) and (ii) Class-specific Prototypical VAE (CP-VAE) (Sec 3.4), which are customized for learning prototypes of varying geometric components and varying object categories, respectively.
Figure 2: Illustration of our pipeline. We incorporate two crucial tasks to facilitate the framework to learn representative prototypes: a scene-level reconstruction task with GP-VAE and an object-level reconstruction task with CP-VAE.
3.1 Problem Definition and Framework Overview
In the FS3D task, the full class set  $\mathbb{C}$  is split into a base set  $\mathbb{C}_{\text{base}}$  and a novel set  $\mathbb{C}_{\text{novel}}$ , where  $\mathbb{C}_{\text{base}} \cap \mathbb{C}_{\text{novel}} = \emptyset$ . Each base class in  $\mathbb{C}_{\text{base}}$  has plenty of annotated objects, while the amount of labeled samples is very limited for each novel class in  $\mathbb{C}_{\text{novel}}$ . The FS3D annotation for an instance of class- $i$  is  $(pos, dim, c_i)$ , denoting the box center point, 3D box dimensions, and the object class label. Following the FSL definition, the  $K$ -shot setting means only  $K$  annotated objects can be used for training in each novel class. Our goal is to enable a model pre-trained on sufficiently-annotated  $\mathbb{C}_{\text{base}}$  samples to generalize over new classes using only  $K$  labeled objects per class in  $\mathbb{C}_{\text{novel}}$ .
Fig. 2 shows the overall pipeline. Built upon the detection framework, we perform two reconstruction tasks during the training to regulate the prototype learning and model the probabilistic latent space:
(i) Scene-level reconstruction, in which we incorporate our proposed GP-VAE to learn geometric prototypes  $\{geo\}$  with a variational encoder-decoder architecture. From the learnt distribution, we sample features  $\{z\}$  for decoding and further refine  $\{z\}$  to  $\{z'\}$  with the geometric prototypes for down-stream object detection (see Sec. 3.3); and
(ii) Object-level reconstruction, in which we adopt a voting-based object encoder to cluster the scene-level features  $\{z^{\prime}\}$  into object-level features and incorporate the proposed CP-VAE to learn class-specific prototypes  $\{cls\}$  as distribution centers. Similar to (i), we sample features  $\{w\}$  for decoding and output calibrated features  $\{w^{\prime}\}$  (see Sec. 3.4).
In the end, we use multiple feed-forward network heads, which are fed with  $\{w'\}$ , to generate instance-level 3D box positions and dimensions as the detection results. During the inference, we can skip the reconstruction decoders and adopt the pre-trained prototypes for forward data flow.
3.2 Prototypical Variational Autoencoder
Our proposed Prototypical VAE (P-VAE) is built upon a standard VAE as its backbone to learn high-level prototypes. Referring to Fig. 1(a), the overall architecture consists of an encoder and a decoder.  $X$  is the variable that represents our data and the encoded latent space follows a certain prior distribution  $p(Z)$ . The encoder  $f$  parameterized by  $\theta_{f}$  takes the point cloud  $x\in \mathbb{R}^{N_x\times 3}$  as the input and generates parameters  $\mu \in \mathbb{R}^{N_z\times d}$  and  $\sigma \in \mathbb{R}^{N_z\times d}$ , where  $N_{z}$  is the number of parameters  $\{\mu_i,\sigma_i\}$  and  $d$  is the latent space dimension. We adopt them to obtain a set of predicted posteriors  $q_{\theta_f}(Z_i|X)\sim \mathcal{N}(\mu_i,\sigma_i)$ , where  $i = 1\dots N_z$ .
Based on these  $N_z$  normal distributions, we can sample  $N_z$  new features  $\{z_i|z_i\in \mathbb{R}^d\}_{i = 1}^{N_z}$ . Then, we adopt a decoder  $g$ , which is parameterized as  $\theta_{g}$ , taking each  $z_{i}$  to reconstruct a point cloud with  $N_{p}$  points  $p_i\in \mathbb{R}^{N_p\times 3}$ . Finally, we combine these reconstruction results to form our point cloud prediction  $p$ . Given the ground-truth point cloud  $\hat{p}$ , the optimization goal for P-VAE is
$$
\underset {\theta_ {f}, \theta_ {g}} {\arg \min } [ \mathrm {C D} (p, \hat {p}) + \frac {1}{N _ {z}} \sum_ {i = 1} ^ {N _ {z}} \mathrm {K L} \left(q _ {\theta_ {f}} \left(Z _ {i} \mid X\right) \| p (Z)\right) ], \tag {1}
$$
where CD denotes the Chamfer Distance that supervises the reconstruction task and KL denotes the Kullback-Leibler Divergence to regularize the posteriors  $q_{\theta_f}(Z_i|X), i = 1\dots N_z$  to stay close to the latent space prior  $p(Z)$ . By far, we can learn this distribution space in a standard VAE term.
Importantly, we focus on obtaining prototypes based on latent distributions. Denoting  $\{pro_{i}|pro_{i}\in$ $\mathbb{R}^d\}_{i = 1}^{N_{pro}}$  as  $N_{pro}$  prototypes, our goal is to use them as clustering centers to group the latent embeddings. Hence, we define  $p(Z)$  following an approximate GMM, which is composed of  $N_{pro}$  independent normal distributions that each centers at  $pro_{i}$  and has variance  $I$ :
$$
p (Z = z) = \sum_ {i = 1} ^ {N _ {p r o}} \frac {\exp \left(- \frac {1}{2} \| z - p r o _ {i} \| _ {2} ^ {2}\right)}{A (2 \pi) ^ {d / 2}} \mathbb {1} \left\{i = \underset {i} {\arg \min } \| z - p r o _ {i} \| _ {2} \right\}. \tag {2}
$$
Compared with the standard GMM, our approximate GMM is designed with two strategies. First, we adopt a fixed weight  $\frac{1}{A}\left(\int_{\mathbb{R}^d}p(Z = z)dz = 1\right)$  instead of distinctive learned weights for these normal distributions, such that each prototype can be equally accessed. This design is for imitating a real-world situation, where all the geometric components share the same probability to be observed among common geometric structures. Second, we adopt a hard regularization  $\mathbb{1}$  in Eq. (2), such that each sampled feature  $z_{i}$  can be directly assigned to its nearest prototype. In this way, our intermediate features are more organized, since they are discriminatively grouped by the prototypes. By introducing pro, we modify the posterior distribution  $q_{\theta_f}(Z_i = z\| X) = A\exp (-\frac{1}{2}\| z - \mu_i\| _2^2)\mathbb{1}\{\| z - pro_j\| _2^2\leq$ $\| z - pro_k\| _2^2\}$ , where  $A$  is the normalization term and  $pro_{j}$  is the closest prototype to  $\mu_{i}$ .
With Eq. (2), we can optimize the KL term of Eq. (1) and rewrite Eq. (1) as
$$
\underset {\theta_ {f}, \theta_ {g}} {\arg \min } [ \mathrm {C D} (p, \hat {p}) + \frac {1}{2} (\sum_ {i = 1} ^ {N _ {p r o}} [ \| \sigma_ {i} \| _ {2} ^ {2} - \| \log \sigma_ {i} ^ {2} \| _ {1} ] + \underbrace {\arg \min } _ {\text {p r o}} [ \sum_ {i = 1} ^ {N _ {p r o}} \sum_ {\mu_ {j} \in B _ {i}} \| p r o _ {i} - \mu_ {j} \| _ {2} ^ {2} ]) ], \tag {3}
$$
where  $B_{i} = \{\mu_{k}||\mu_{k} - pro_{i}\|_{2}\leq \| \mu_{k} - pro_{j}\|_{2}, j = 1\dots N_{pro}\}$ , indicating the set of features that are assigned to prototype  $pro_{i}$  by least Euclidean distance. Please see Appendix for detailed derivations.
To find the optimal prototypes  $\{pro_{i}\}_{i = 1}^{N_{pro}}$  is to solve the 'Term Pro' in Eq. (3), which can be viewed as a clustering task to seek  $N_{pro}$  clustering centers, such that the mean distance between all pairs of the sampled features to their closest prototypes is minimized. We adopt the K-means algorithm updated with Exponential Moving Average (EMA) to enhance the training efficiency. Overall, at the  $t^{th}$  training iteration,  $pro_{i}$  is estimated as
$$
p r o _ {i} ^ {t + 1} \leftarrow \alpha p r o _ {i} ^ {t} + (1 - \alpha) \frac {\sum_ {\mu_ {k} \in B _ {i}} \mu_ {k}}{| B _ {i} |}, B _ {i} = \left\{\mu_ {k} \| \mu_ {k} - p r o _ {i} \| _ {2} \leq \| \mu_ {k} - p r o _ {j} \| _ {2}, \forall j = 1, \dots , N _ {p r o} \right\}, \tag {4}
$$
where  $\alpha$  is the smoothness hyperparameter and all the prototypes are initialized with the Gaussian random initialization before the training. After we obtain the new prototypes in an iteration, we take them back to Eq. (3) to get our loss for training the network:
$$
\mathcal {L} = \mathrm {C D} (p, \hat {p}) + \frac {1}{2} \left(\sum_ {i = 1} ^ {N _ {p r o}} \left[ \| \sigma_ {i} \| _ {2} ^ {2} - \| \log \sigma_ {i} ^ {2} \| _ {1} \right] + \sum_ {i = 1} ^ {N _ {p r o}} \sum_ {\mu_ {j} \in B _ {i}} \| p r o _ {i} - \mu_ {j} \| _ {2} ^ {2}\right), \tag {5}
$$
where the first term is the reconstruction loss and the last term is the clustering loss for encouraging each sampled feature to stay close to its associated nearest prototype.
3.3 Scene Reconstruction with Geometric Prototypical VAE (GP-VAE)
The detection framework starts by extracting features from the input scene. Here, we incorporate GP-VAE to enhance the feature quality. We reconstruct the whole point cloud scene with GP-VAE, and obtain geometric-informative prototypes denoted as  $\{geo_i\}_{i=1}^{N_{geo}}$ . These prototypes represent basic geometric components of common real-world 3D objects.
Please refer to Fig. 2 for GP-VAE. We use the PointNet++ [56] as our encoder  $f_{scene}$ . It takes a raw point cloud  $x \in \mathbb{R}^{N_x \times 3}$  as input and outputs parameters  $\{\mu_i, \sigma_i\}_{i=1}^{N_z}$  to construct the posterior normal distributions  $\mathcal{N}(\mu_i, \sigma_i), i = 1 \dots N_z$ . We sample  $\{z_i\}_{i=1}^{N_z}$  from the learned distributions and build a decoder  $g_{scene}$  to reconstruct the original point cloud scene. For  $g_{scene}$ , we adopt the architecture
of FoldingNet [57], which transforms a plane into a 3D object surface. Specifically, we use furthest point sampling (FPS) to obtain  $N_{p}$  points on a  $0.4m \times 0.4m$  plane  $B \in \mathbb{R}^{N_p \times 2}$ . For each  $z_{i}$ , we repeat the feature  $N_{p}$  times then concatenate them with plane  $B$ , followed by two 2-layer MLPs to generate  $p_{i} \in \mathbb{R}^{N_{p} \times 3}$ . We use Eq. (4) to obtain the geometric prototypes  $\{geo_{i}\}_{i=1}^{N_{geo}}$ . We combine the  $N_{z}$  predicted point clouds to construct the scene-level result  $p \in \mathbb{R}^{N_{z}N_{p} \times 3}$ , then use Eq. (5) to supervise the scene-level reconstruction, where the ground truth is the input point cloud  $x$ .
Instead of directly using the sampled features  $\{z_i\}_{i=1}^{N_z}$  to cluster the per-object features, we calibrate  $\{z_i\}_{i=1}^{N_z}$  with the geometric prototypes via a cross attention layer. This calibration enriches the information diversity of each feature by not only its nearest prototype but also other geometric prototypes. For brevity, we use geo  $\in \mathbb{R}^{N_{geo} \times d}$  as the united representation of the prototypes:
$$
z _ {i} ^ {\prime} = \operatorname {S o f t m a x} \left(\frac {\mathbf {Q} \left(z _ {i}\right) \mathbf {K} ^ {\top} (g e o)}{\sqrt {d}}\right) \mathbf {V} (g e o), \tag {6}
$$
where  $\mathbf{Q}$ ,  $\mathbf{K}$ , and  $\mathbf{V}$  are three MLPs to generate the query, key, and value in the attention framework. In this way, we can obtain robust embedding features  $\{z_i'\}_{i=1}^{N_z}$ .
3.4 Object Reconstruction with Class-specific Prototypical VAE
Given the features  $\{z_i^{\prime}\}_{i = 1}^{N_z}$  generated in Sec. 3.3, we use CP-VAE to reconstruct each object and obtain class-specific prototypes as shown in Fig. 1(c). We build a voting-based [58] encoder  $f_{object}$ . For detection purpose, we need per-object features  $\{\hat{z}_i\}_{i = 1}^{N_{ins}}$ , where  $N_{ins}$  is the number of the predicted instances. Therefore, we use a 1-layer MLP to predict the feature offsets and point offsets from  $\{z_i^{\prime}\}_{i = 1}^{N_z}$ , then adopt FPS on the shifted points to obtain  $N_{ins}$  votes. For each vote, we collect a maximum number of 16 shifted features within  $0.3\mathrm{m}$  from the vote, then use max pooling on these features to generate the associated object-level feature  $\hat{z}_i$ . Another built MLP is then fed with  $\{\hat{z}_i\}_{i = 1}^{N_{ins}}$  to output parameters  $\{\mu_i,\sigma_i\}_{i = 1}^{N_{ins}}$ , which can be used to construct the posterior truncated normal distributions  $\mathcal{N}(\mu_i,\sigma_i)$  where  $i = 1\dots N_{ins}$ . We sample  $\{w_i\}_{i = 1}^{N_{ins}}$  and build a FoldingNet decoder  $g_{object}$  (Sec. 3.3) for object-level reconstruction. For each  $w_i$ ,  $g_{object}$  outputs an unique 3D object  $p_i$ .
According to the sampling and grouping mechanisms in feature extraction, we reversely map each  $w_{i}$  to a set of corresponding points in  $x$  (we will detail this operation in the Appendix), which constructs our reconstruction target  $\hat{p}_i$ . Therefore, we modify the reconstruction term of Eq. (5) to