Aggregation and Purification: Dual Enhancement Network for Point Cloud Few-shot Segmentation

1 Introduction

Point cloud semantic segmentation, a fundamental task in computer vision with wide applications in autonomous driving, robotics, etc, has achieved conspicuous advancements with the development of deep learning [Long et al., 2015; Sun et al., 2023c; Wang et al., 2024; Sun et al., 2023d; Luo et al., 2023; Mai et al., 2024b]. This progress can be attributed to well-established datasets and elaborately designed algorithms. However, the inherent category sensitivity of the fully supervised segmentation models restricts their capacity to predefined training categories. In pursuit of the human-like intelligence of recognizing new classes with only a glance, point cloud few-shot segmentation [Zhao et al., 2021b] (PCFSS) is proposed to extend the segmentation model to novel categories without extensive labeled data collection and timeconsuming model retraining. Specifically, PC-FSS achieves the segmentation of query point clouds in new classes, which were previously unseen during the training process, using only a handful of annotated support samples.

∗Corresponding author

In previous literature, prototypical learning has emerged as the mainstream paradigm to tackle the PC-FSS. In specific, these methods [Ning et al., 2023; Zhao et al., 2021b] concentrate information of the support samples into a set of representative prototypes to guide the classification of query points. Despite yielding promising results, the frequently encountered significant scene discrepancies result in the prototypes derived solely from support samples struggling to furnish query points classification with essential cues.

As one of the principal challenges of PC-FSS, the discrepancies between the support and query samples can be systematically classified into the following two types: intrasemantic diversity and semantic inconsistency. The former primarily originates from categories that concurrently appear in both support and query samples. (As the foreground category "sofa" and background category "desk" in Figure 1). Although the support prototypes containing co-occurrent information are capable of aiding in the segmentation of the query counterparts, they simultaneously risk introducing cognitive biases due to intra-class diversity such as variety of scale, pose, and so on. While the semantic inconsistency primarily exists in backgrounds including distinct semantics (As the "door" and "bed" in Figure 1). The resulting inconsistent support prototypes may be redundant or even disruptive to the classification of query points, and the query background points with the semantics that are absent in support samples may also be erroneously assigned to the foreground prototypes. To make matters worse, the negative impact of scene discrepancies is inevitably amplified by inbuilt low-data regimes of few-shot segmentation, leading to suboptimal results.

A series of works attempt to resolve the issue by calibrating the support prototypes with query context [Ning et al., 2023; Zhu et al., 2023; Zhang et al., 2023]. While these methods can improve the query-awareness of prototypes in certain scenarios, the strategy of indiscriminately aggregating query information without accounting for different types of discrepancies tends to be highly sensitive to variations in scene composition. On the other hand, existing methodologies predominantly focus on unidirectional updating support prototypes, overlooking the crucial aspect that query features, as the primary entities of segmentation, also necessitate enhancement to be more task-aware.

After an in-depth analysis of the above issues, we design a novel Dual Enhancement Network (DENet), including a mutual enhancement aggregation and a consistent purification strategy, to coherently tackle different kinds of scene discrepancies by reconciling the intrinsic tension between the support prototypes and query point features in a bidirectional manner. To deal with the intra-semantic diversity, in the mutual aggregation module (MAM), we enable bidirectional feature communication between support prototypes and query features via a sequence of customized attention layers. Compared to the vanilla cross-attention adopted in previous works [Ning et al., 2023], the feature interaction within MAM not only reduces intra-semantic diversity from both support and query aspects, but also endows the primary entities of segmentation, i.e., the query features with taskawareness and clean separation of background/foreground semantics. To deal with the semantic inconsistency, we elegantly design the consistent purification strategy (CPS) to explicitly eliminate ambiguous prototypes by comparing them with robust holistic features. Then the refined prediction from the filtered prototypes is adopted as the pseudo label to impose an additional regularization loss, which encourages the aggregation process of MAM to concentrate more on semantic consistent features. Moreover, the CPS can be employed iteratively, and its inherently parameter-free nature enables it to function as a plug-and-play module in the inference phase.

To recap, our primary contributions are concluded as follows: (1) We propose the Dual Enhancement Network (DENet) to tackle different types of scene discrepancies in a coherent and synergistic framework. (2) We introduce the mutual aggregation module (MAM) to alleviate intrasemantic diversity from both the support and query aspects, and the plug-and-play consistent purification strategy (CPS) to mitigate the impact of semantic-inconsistency. (3) Extensive experiments show that DENet significantly outperforms previous SOTAs. The systematic analysis of scene discrepancies in PC-FSS also sheds light on future research from a more comprehensive perspective.

2 RelatedWork

3D Point Cloud Semantic Segmentation. With the recent advances in deep neural networks [Sun et al., 2021; Landrieu and Simonovsky, 2018; Sun et al., 2023a; Mai et al., 2023; Sun et al., 2023b; Wang et al., 2023a; Luo et al., 2024; Pan et al., 2023; Mai et al., 2024a; Xiong et al., 2024], point cloud semantic segmentation has gained significant attention as a solution to address the challenges posed by dense pixellevel annotations. Recently, many deep learning based approaches [Hu et al., 2020; Huang et al., 2018; Lai et al., 2022; Wu et al., 2019; Li et al., 2018; Qi et al., 2017b; Wang et al., 2019; Ye et al., 2018; Zhao et al., 2021a] are proposed to address 3D point cloud semantic segmentation under fully supervised conditions. PointNet [Qi et al., 2017a] introduces a pioneering point-based method to segment raw point clouds directly and extract point features with MLPs. Despite its simplicity and efficiency, PointNet falls short in capturing essential local information from neighboring points. DGCNN [Wang et al., 2019] introduces EdgeConv to address this issue, capturing local point cloud structures through a dynamic k-NN graph for long-distance point intersection analysis. In this work, DGCNN serves as the backbone of our feature extractor to extract local structural characteristics and semantic features. Building upon this, we introduce a novel approach aimed at addressing the challenges of 3D point cloud semantic segmentation in a few-shot learning context.

Point Cloud Few-shot Semantic Segmentation. Few-shot 3D Point Cloud Semantic Segmentation transitions traditional 3D point cloud semantic segmentation into a few-shot context [Wang et al., 2022; Wang et al., 2023b], equipping the model with the capability to segment novel classes using minimal support data. AttMPTI [Zhao et al., 2021b] innovatively applied an attention-aware multi-prototype transductive inference approach to few-shot 3D point cloud semantic segmentation, establishing a foundational method in this domain for classifying new categories using a limited number of annotated examples. However, AttMPTI did not fully address the gap issue between support and query. Subsequent works [Ning et al., 2023; Zhu et al., 2023; Zhang et al., 2023] have utilized information from the query to enhance the adaptability of support prototypes to the context of the query. Building upon these prior efforts, our method takes a step further by addressing the discrepancies between support and query from two perspectives: intra-semantic diversity and semantic inconsistency.

3 Method

3.1 ProblemDefinition

Point cloud few-shot semantic segmentation aims to endow the segmentation model with the ability to rapidly generalize to novel categories with only a few set of labeled samples. Following previous works [Zhao et al., 2021b; Ning et al., 2023], we adopt the widely used episodic metatraining paradigm [Vinyals et al., 2016]. Specifically, considering two non-intersection sets Dtrain and Dtest, which possess non-overlapping target categories (Ctrain ∩ Ctest = ∅). A series of episodes are sampled from Dtrain to train our model, and then we test the trained model on the test set Dtest. Each episode, manifests as an N-way K-shot segmentation task within point cloud data. And each episode is composed of a support set denoted as S = {(I n,k S ,M n,k S )}, where I is point cloud, M is mask, k ∈ {1, · · · ,K}, n ∈ {1, · · · ,N}, and a query set Q = {(IQ,MQ)}. During each episode in Nway K-shot scenario, the model is provided with input data comprising {I n,k S ,M n,k S , IQ}, where I n,k S and M n,k S represent the point cloud and mask of the support set, respectively, and IQ represents the query point cloud. The model predict the segmentation result for IQ, which is compared to MQ as the ground truth during validation. During the training phase, the model is trained to infer the segmentation of the query set IQ by leveraging the information provided in the support set S, guided by the ground truth mask MQ. In the testing phase, the generalization capabilities of the model are assessed using tasks derived from Dtest. To simplify the description, we describe our approach in a 1-way 1-shot setting.

3.2 Overview of Dual Enhancement Network

As illustrated in Figure 2, the proposed DENet follows the prototypical learning paradigm and mainly focuses on dealing with the misalignment between the support prototypes P = {(PS bg ,PS fg )} originate from support point features FS ∈ RN×C and query point features FQ ∈ RN×C . Note that we adopt the same prototype generation and transductive inference process as done in [Zhao et al., 2021b] for a fair comparison without claiming any contribution. DENet synergistically integrates two components, i.e., Mutual Aggregation Module (MAM) and Consistent Purification Strategy (CPS), to handle different types of discrepancies. MAM enables the bidirectional feature aggregation to bridge the intrasemantic gap while preserving the foreground-background discriminability (Section 3.3). CPS mitigates the negative impact of semantic inconsistency by eliminating ambiguous prototypes. (Section 3.4). The details are as follows.

3.3 The Mutual AggregationModule

Due to the significant intra-class diversity between query and support samples, the local prototypes derived solely from support samples fail to comprehensively capture class information, thereby proving inefficient as directors for the classification of query points. We design the MAM to promote a more holistic object pattern understanding of both the support prototypes and query features, thus reducing the impact of the intra-semantic gap. MAM consists of three different attention modules, namely the support-to-prototype masked attention module, the query-to-prototype alignment attention module, and the prototype-to-query attention module.

Support-to-Prototype Masked Attention. Since prototypes are obtained by averaging small portions of point features [Zhao et al., 2021b], only local and limited information is contained. we adopt masked attention layers to aggregate support features to the corresponding foreground or background prototypes to endow them with a holistic perception of the category, such as geometric structures. The standard cross-attention process is formulated as

where √C is a scaling factor, Wq,Wk,Wv are linear projections of query, key, and value in the same dimension C. However, the soft nature of attention makes the process noisy and less dependable - foreground prototypes might have small weights distributed in the background features and vice versa. Inspired by [Cheng et al., 2022], we deploy masked attention to aid the clean separation of semantics between foreground and background. In practice, we force foreground prototypes to focus on the foreground features of the support sample while directing background prototypes to interact with the background support features, formally:

where MS ∈ {0, -∞}np×N determines whether the query p is allowed(= 0) or not allowed (= -∞) to attend to the i-th feature, defined as:

In this way, Pfg S and Pbg S are restricted to absorbing contextual information from the foreground and background respectively. The resulting prototypes P, equipped with task awareness, also effectively avoid cross-category interference.

Query-to-Prototype Alignment Attention. To enhance the semantics perception of prototypes from the query perspective, thereby reducing the gap between them, we aggregate query category knowledge into P via cross-attention layers. Considering that the query mask is not accessible, we introduce an alignment matrix A to filter out implausible matching. Formally:

where √C is a scaling factor, Wq,Wk,Wv are linear projections of query, key, and value in the same dimension C, same as before but with different weights. The alignment matrix A is obtained by

The alignment matrix ensures the correspondence in feature aggregation, and the query-to-prototype alignment attention significantly reduces the intra-semantic diversity between the prototype and query features.

Prototype-to-Query Attention. In addition to modulating the support prototypes to approach query samples, we deem that the query features, as the main entity for the segmentation task, can also be calibrated toward the support prototypes, forming a bidirectional alignment. To achieve that, prototype-to-query attention is included to aggregate discriminative prototype features. Specifically, we employ crossattention with residual connections to update query features with prototypes:

The reversed cross-attention layer injects task-related information into query features, enabling a more consistent matching process for label prediction.

3.4 The Consistent Purification Strategy

Similar to AttMPTI [Zhao et al., 2021b], we use the graph relationship between bi-directional enhanced support prototypes ˜P and query features ˜FQ to derive the initial prediction results for each point in the query sample using a label propagation algorithm, denoted as M 0 pred ∈ {0, 1}1×N .

However, due to the presence of semantic inconsistency, some of the support prototypes may be redundant or even disruptive to the label propagation. That is, some background points of the query may inappropriately correspond to the support foreground prototypes, and vice versa. Therefore, we introduce a consistent purification strategy, which evaluates and filters the support prototypes to reduce the negative impact of semantic inconsistency.

To delve into the correlation between support prototypes and query samples, we first apply mask average pooling (MAP) to the initial prediction of query features M 0 pred to obtain the holistic foreground and background information of the query, denoted as Qfg and Qbg , respectively.

Then, by measuring the average correlation between the support foreground prototypes ˜Pfg and the holistic query foreground feature Qfg , denoted as τ1, as well as the average correlation between the support background prototypes ˜Pbg and the holistic query background feature Qbg , denoted as τ2. We take the minimum of τ1 and τ2 as the holistic relevance of the support prototypes to the query samples, which serves as the prototype purification threshold τ , i.e.τ = min(τ1, τ2).

where the "AVG" denotes the average operation across all prototypes. We base this on the simple but straightforward assumption that the foreground-to-foreground similarity and background-to-background similarity should be higher than the background-to-foreground or foreground-to-background similarity.

where C(i) is a judgment function on the ith prototype that decides whether to keep Pi (C(i)=1) or to filter out the inconsistent Pi (C(i)=0). The purified prototypes ˆP are subsequently used in label propagation along with the query features ˜FQ to get the refined prediction M 1 pred . In addition to the baseline design cross-entropy loss Lce for initial query prediction M 0 pred , the CPS introduces an additional regularization loss to encourage the aggregation process of MAM to focus more on semantically consistent features,

and the overall training objective L is defined as:

It should be noted that the CPS not only introduces supplementary regularization loss during the training phase, but also acts as a plug-and-play module to refine the prediction.

4 Experiments

4.1 Experiment Setup

Datasets. In accordance with previous methods, we evaluate the proposed DANet on two commonly 3D few-shot segmentation benchmarks, S3DIS [Armeni et al., 2016] and ScanNet [Dai et al., 2017]. S3DIS contains 272 point cloud datasets covering indoor areas from three different buildings. It includes a wide range of indoor environments such as lobbies, corridors, offices, and pantries. Each point cloud dataset is annotated with 12 different semantic categories, along with an additional background category to represent various surrounding clutter. ScanNet consists of 1,513 point clouds extracted from 707 different indoor scenes, and provides a rich annotation framework with 20 semantic categories for segmentation and an additional class for unannotated regions. Compared to S3DIS, ScanNet presents a wider variety of room types, including bathroom and living room, with more irregular point cloud scenes. In each dataset, the classes are systematically partitioned into two mutually exclusive subsets, denoted as S0 and S1, ensuring no overlap. These subsets are alternately utilized for cross-validation purposes, with one serving as the test classes Ctest and the other as the training classes Ctrain. Consistent with the preprocessing strategy of previous work, we divide rooms into 1m x 1m blocks, with 2048 points randomly sampled for each block.

EvaluationMetric. Following previous works [Zhao et al., 2021b; Ning et al., 2023], we adopt mean-IoU as our evaluation metrics. In our few-shot setting, the mean-IoU is obtained by averaging the IoU over all testing classes Ctest.

Implementation Details. All models are implemented using the PyTorch framework and trained on one NVIDIA GeForce RTX 3090 GPU. The feature extraction module is pre-trained on the training set Dtrain for 100 epochs with a batch size of 32. We use the Adam optimizer with a learning rate of 0.001 to optimize the pre-trained model. During episodic training, we initialize the feature extraction module by loading the pre-trained weights. The initial learning rate is set to 0.0001 for the feature extraction module and 0.001 for the other modules, and the Adam optimizer is used to update all parameters. The learning rates are halved every 5000 iterations and the batch size is set to 1. Following [Zhao et al., 2021b; Ning et al., 2023], models undergo 40,000 iterations. During each iteration, we randomly sample an episode, utilizing Gaussian noise and arbitrary z-axis rotations to augment the point clouds in both the support and query sets.

4.2 Comparison with State-of-the-Art Methods

We experimented on two popular benchmarks, S3DIS and ScanNet, comparing our method with previous PC-FSS methods in Table 1 and Table 2, respectively. Across various evaluation scenarios on both benchmarks, the proposed DENet consistently outperforms prior advanced approaches, strongly proving the effectiveness of our method. Specifically, on ScanNet, our DENet achieves 63.95% and 71.13% mIoU for 1-way 1-shot and 1-way 5-shot settings. This performance represents an improvement of 3.59% and 1.6% over the most competitive approach, QGE [Ning et al., 2023]. We additionally observed that the enhancements offered by our proposed DENet are more pronounced on the more complex ScanNet dataset as compared to the S3DIS dataset. This can be attributed to the greater complexity of scenes and categories present in the ScanNet dataset, leading to more conspicuous discrepancies between support and query. Consequently, a series of results demonstrate that our DENet, developed through a systematic analysis of scene discrepancies in the PC-FSS task, exhibits commendable performance and robustness in handling complex scene scenarios. Figure 6 displays prediction visualizations on ScanNet dataset, illustrating that our proposed method yields more visually appealing results, intuitively demonstrating the superiority of our approach.

4.3 Ablation Study

In this section, we use AttMPTI as our baseline model and conduct extensive ablation studies on S3DIS S0 split under the 1-way 1-shot setting to verify the effectiveness of each component of our DENet, including the mutual enhancement module and the consistent purification strategy.

Effectiveness of Components. As shown in Table 3, both mutual aggregation module (MAM) and consistent purification strategy (CPS) consistency bring a certain performance lift compared with the baseline. (1) With the utilization of MAM, a 5.9% improvement of mIoU can be observed, indicating that reducing intra-semantic diversity between support and query samples can benefit the point cloud few segmentation task. (2) The introduction of CPS achieves further accuracy gains (3.8% in mIoU), primarily attributed to the exclusion of prototypes with ambiguous semantics, which helps to suppress the negative impact of scene discrepancies.

Effectiveness of Mutual Aggregation Module. To perform ablation studies on the components of our Mutual Aggregation module, we used the AttMPTI model as our baseline, while all experiments used the CPS. As shown in Table 4, both the Prototype-Support Masked Attention module and the Prototype-Query Alignment Attention module can improve the accuracy, indicating that optimizing support prototypes can better align them with the query samples. This is consistent with previous methods of updating support prototypes to improve segmentation performance. By further integrating the Query-Prototype Attention module, we observe a performance advance of 1.48% mIoU, attributed to the module's ability to extract consistent information from the query perspective and mitigate the negative effects of the semantic gaps between the query and support point clouds.

Effectiveness of Consistent Purification Strategy. To further explore the applicability of CPS, as described in Table 5, we examine its effectiveness in both the training and testing phases. The introduce of CPS consistently leads to significant performance improvements in both training and testing scenarios. Furthermore, the use of CPS during both training and testing further elevates performance. This is because CPS effectively eliminates ambiguous prototypes and reducing the impact of category inconsistency on model performance.

Hyperparameter Evaluations. Quantitative experiments are conducted to explore the appropriate number of iterations. In Table 6, we compare the performance of different iterations of the CPS. It can be observed that as the number of purification iterations increases, the performance gradually increases and then converges. To achieve a trade-off between accuracy and speed, we finally choose a number of iterations of 3.

Investigation of Dual Enhancement Network. In Figure 5, we employed t-SNE visualizations to concretely illustrate two types of scene discrepancies, as well as the process by which the DENet mitigates them. The intra-semantic diversity can be observed from the orange dots and magenta dots in Figure 5 (a). They represent the foreground features of support and query points belonging to the same category but are not well aligned. The semantic inconsistency is primarily reflected in the intertwining of background interference features with foreground features as shown by the darkgray triangles and lightblue crosses within the red boxes in Figure 5 (a). In the presence of scene discrepancies, support prototypes struggle to effectively guide the classification of query features. The integration of the mutual aggregation module (MAM) markedly enhances the intra-class alignment, we deem this is reasonable as the adaptive feature communication within the customized markedly alleviates intrasemantic diversity in a bidirectional manner while preserving the foreground-background discriminability of corresponding features. The consistent purification strategy (CPS) further eliminates ambiguous prototypes and thus effectively suppresses erroneous label propagation. The MAM and CPS synergistically reinforce each other, thereby constituting the coherent PC-FSS framework DENet, which is capable of addressing different types of scene discrepancies.

5 Conclusion

In this paper, we propose a novel Dual Enhancement Network (DENet) for point cloud few-shot semantic segmentation. By systematically analyze the scene discrepancies in the PC-FSS task, we tailored the MAM module and the CPS module to respectively address intra-semantic discrepancies and semantic inconsistencies issues. Extensive experiments demonstrate the effectiveness of our method.

Contribution Statement

Guoxin Xiong and Yuan Wang contributed equally to this paper.

Acknowledgments

This work was partially supported by the National Defense Basic Scientific Research Program (Grant JCKY20211130B016).