White Aggregation and Restoration for Few-shot 3D Point Cloud Semantic Segmentation

Few-shot 3D Point Cloud Semantic Segmentation (FS-PCS) aims to predict per-point labels for an unlabeled point cloud, given only a few labeled examples. To extract representations from the limited labeled set, existing methods have constructed prototypes with Farthest Point Sampling (FPS). However, we found that this convention results in performance instability due to its sensitivity to FPS-induced variations, while the prototype generation process remains underexplored in the field. This motivates us to investigate deterministic prototype generation method based on attention mechanism. Despite its potential, we found that vanilla attention module suffers from the distributional gap between prototypical tokens and support features. To overcome this, we provide a simple approach, White Aggregation and Restoration Module (WARM), which resolves the misalignment by wrapping cross-attention with whitening and coloring transformations. Specifically, whitening aligns the features to tokens before the attention process, and coloring subsequently restores the original distribution to the attended tokens. This design enables robust attention, thereby generating prototypes that capture the semantic relationships in support features. WARM achieves state-of-the-art performance with a significant margin on the S3DIS dataset, and competitive performance on the ScanNet dataset. Further experiments demonstrate its effectiveness in deterministic prototype generation. Code is publicly available at: https://github.com/JiyunIm00/WARM.git

Jiyun Im, SuBeen Lee, Miso Lee, Jae-Pil Heo · 2025