Uncertainty-aware Prototype Learning with Variational Inference for Few-shot Point Cloud Segmentation

2026-03-20

AI Summary

This work addresses the challenges of rigid prototype representations and poor generalization in few-shot 3D point cloud semantic segmentation, which stem from scarce annotations. To overcome these limitations, the authors propose UPL, a novel approach that introduces variational inference into prototype learning for the first time, yielding uncertainty-aware probabilistic prototype representations. A dual-stream prototype refinement module is designed to jointly leverage information from both support and query samples, enabling dynamic optimization of class prototypes. Evaluated on ScanNet and S3DIS benchmarks, the method achieves state-of-the-art performance, significantly enhancing model robustness and generalization while providing reliable uncertainty estimates.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Computer Vision: Segmentation

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Security and Privacy: Data transparency and provenance
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

Few-shot 3D semantic segmentation aims to generate accurate semantic masks for query point clouds with only a few annotated support examples. Existing prototype-based methods typically construct compact and deterministic prototypes from the support set to guide query segmentation. However, such rigid representations are unable to capture the intrinsic uncertainty introduced by scarce supervision, which often results in degraded robustness and limited generalization. In this work, we propose UPL (Uncertainty-aware Prototype Learning), a probabilistic approach designed to incorporate uncertainty modeling into prototype learning for few-shot 3D segmentation. Our framework introduces two key components. First, UPL introduces a dual-stream prototype refinement module that enriches prototype representations by jointly leveraging limited information from both support and query samples. Second, we formulate prototype learning as a variational inference problem, regarding class prototypes as latent variables. This probabilistic formulation enables explicit uncertainty modeling, providing robust and interpretable mask predictions. Extensive experiments on the widely used ScanNet and S3DIS benchmarks show that our UPL achieves consistent state-of-the-art performance under different settings while providing reliable uncertainty estimation. The code is available at https://fdueblab-upl.github.io/.

Problem

• few-shot segmentation
• point cloud
• uncertainty modeling
• prototype learning
• 3D semantic segmentation

Innovation

• Uncertainty-aware
• Prototype Learning
• Variational Inference
• Few-shot Segmentation
• Point Cloud