[Paper Review] Attentional Prototype Inference for Few-Shot Segmentation

Haoliang Sun, Xiankai Lu+5 more|arXiv (Cornell University)|May 14, 2021

Advanced Neural Network Applications57 references4 citations

TL;DR

This paper proposes Attentional Prototype Inference (API), a probabilistic latent variable framework for few-shot image segmentation that models object prototypes as distributions to handle uncertainty and intra-class variation. By jointly inferring probabilistic prototypes and attention maps via variational Bayesian inference, API achieves state-of-the-art performance on four benchmarks, outperforming existing prototype-based methods in both 1-shot and 5-shot settings.

ABSTRACT

This paper aims to address few-shot segmentation. While existing prototype-based methods have achieved considerable success, they suffer from uncertainty and ambiguity caused by limited labeled examples. In this work, we propose attentional prototype inference (API), a probabilistic latent variable framework for few-shot segmentation. We define a global latent variable to represent the prototype of each object category, which we model as a probabilistic distribution. The probabilistic modeling of the prototype enhances the model's generalization ability by handling the inherent uncertainty caused by limited data and intra-class variations of objects. To further enhance the model, we introduce a local latent variable to represent the attention map of each query image, which enables the model to attend to foreground objects while suppressing the background. The optimization of the proposed model is formulated as a variational Bayesian inference problem, which is established by amortized inference networks. We conduct extensive experiments on four benchmarks, where our proposal obtains at least competitive and often better performance than state-of-the-art prototype-based methods. We also provide comprehensive analyses and ablation studies to gain insight into the effectiveness of our method for few-shot segmentation.

Motivation & Objective

• Address the limitations of deterministic prototype-based methods in few-shot segmentation, which struggle with ambiguity, noise, and intra-class variation due to limited labeled data.
• Improve generalization by modeling object prototypes as probabilistic distributions rather than fixed vectors, capturing uncertainty and intra-class diversity.
• Introduce a local latent variable for attention maps to enhance focus on foreground objects and suppress background interference.
• Formulate the optimization as a variational inference problem using an evidence lower bound (ELBO), enabling joint learning of prototypes and attention.
• Demonstrate the effectiveness of the framework through extensive experiments and ablation studies across diverse few-shot segmentation benchmarks.

Proposed method

• Define a global latent variable representing the probabilistic prototype of each object category, modeled as a distribution to capture uncertainty and intra-class variation.
• Introduce a local latent variable for each query image to represent the attention map, enabling dynamic focus on foreground regions.
• Formulate the optimization as a variational Bayesian inference problem using amortized inference networks to approximate posterior distributions.
• Derive a new evidence lower bound (ELBO) tailored for few-shot segmentation, enabling joint estimation of prototype and attention posteriors.
• Use Monte Carlo sampling to approximate the segmentation output, aggregating multiple stochastic predictions for robust and complete segmentation maps.
• Apply stochastic gradient descent with a single sample during training to balance computational cost and generalization, while using 5–15 samples at inference for improved accuracy.

Experimental results

Research questions

• RQ1 Can modeling object prototypes as probabilistic distributions improve robustness to label ambiguity and intra-class variation in few-shot segmentation?
• RQ2 How does incorporating a learnable attention mechanism as a local latent variable enhance foreground localization and segmentation accuracy?
• RQ3 To what extent does variational inference with a jointly optimized ELBO objective improve generalization in low-data regimes?
• RQ4 How does Monte Carlo sampling over stochastic prototypes and attention maps affect segmentation quality and inference efficiency?
• RQ5 What is the contribution of each component (probabilistic prototype, attention map, variational inference) to the overall performance, as shown through ablation studies?

Key findings

• API achieves at least competitive and often superior performance compared to state-of-the-art prototype-based methods on four benchmarks: Pascal-5i, MS-COCO, FSS-1000, and LIDC-IDRI.
• On the medical imaging dataset LIDC-IDRI, the method demonstrates strong statistical properties, with cross-energetic distance results indicating faithful handling of intrinsic image ambiguities.
• Ablation studies confirm that both the probabilistic prototype and the variational attention mechanism significantly contribute to performance gains.
• Monte Carlo sampling improves segmentation completeness and noise reduction, with performance saturating at 15 samples, though inference is efficient with 5 samples in practice.
• The model produces diverse, plausible segmentation hypotheses, which could support clinical decision-making by providing uncertainty-aware predictions.
• Inference time increases slightly with attention (less than 0.1s under 5-shot with M=15), and the model maintains acceptable speed with minimal sampling at test time.