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 onlya few labeled examples. To extract representations from the limited la-beled set, existing methods have constructed prototypes with FarthestPoint Sampling (FPS). However, we found that this convention resultsin performance instability due to its sensitivity to FPS-induced varia-tions, while the prototype generation process remains underexplored inthe field. This motivates us to investigate deterministic prototype gen-eration method based on attention mechanism. Despite its potential, wefound that vanilla attention module suffers from the distributional gapbetween prototypical tokens and support features. To overcome this, weprovide a simple approach, White Aggregation and Restoration Module(WARM), which resolves the misalignment by wrapping cross-attentionwith whitening and coloring transformations. Specifically, whitening alig-ns the features to tokens before the attention process, and coloring subse-quently restores the original distribution to the attended tokens. This de-sign enables robust attention, thereby generating prototypes that capturethe semantic relationships in support features. WARM achieves state-of-the-art performance with a significant margin on the S3DIS dataset, andcompetitive performance on the ScanNet dataset. Further experimentsdemonstrate its effectiveness in deterministic prototype generation. Codeis publicly available at: https://github.com/JiyunIm00/WARM.git