ProtoMask: Segmentation-Guided Prototype Learning

TL;DR

ProtoMask addresses explainability gaps in prototype-based image classifiers by tying prototypes to object parts through segmentation-guided multi-view representations. It extends ProtoPNet with a multi-view pipeline where segmentation masks generate views that feed embeddings, which are evaluated against prototypes via a distance-based similarity and a max-pooling strategy for classification. A multi-term loss encourages accurate classification, well-separated latent spaces, diverse prototypes, and sparsity, implemented through a two-stage training procedure that includes prototype projection. Across CUB200-2011, Stanford Dogs, and Stanford Cars, ProtoMask achieves competitive accuracy and improved explainability metrics, though its performance is sensitive to segmentation mask quality and cropping strategies.

Abstract

XAI gained considerable importance in recent years. Methods based on prototypical case-based reasoning have shown a promising improvement in explainability. However, these methods typically rely on additional post-hoc saliency techniques to explain the semantics of learned prototypes. Multiple critiques have been raised about the reliability and quality of such techniques. For this reason, we study the use of prominent image segmentation foundation models to improve the truthfulness of the mapping between embedding and input space. We aim to restrict the computation area of the saliency map to a predefined semantic image patch to reduce the uncertainty of such visualizations. To perceive the information of an entire image, we use the bounding box from each generated segmentation mask to crop the image. Each mask results in an individual input in our novel model architecture named ProtoMask. We conduct experiments on three popular fine-grained classification datasets with a wide set of metrics, providing a detailed overview on explainability characteristics. The comparison with other popular models demonstrates competitive performance and unique explainability features of our model. https://github.com/uos-sis/quanproto

Figures (4)

• Figure 1: ProtoMask operates on multiple views of the same image, as illustrated on the left side. The design consists of a CNN backbone model, which extracts a feature vector $z$ for each view separately. A similarity score for each prototype feature vector pair is computed. The score of the nearest feature vector per prototype is then used in the classification layer $h$.
• Figure 2: Distribution of the mean segmentation mask size (top) and mean mask object overlap (bottom) per image across all datasets for the SAM2 and SLIT segmentation method.
• Figure 3: Prototype Relevance Propagation (PRP) visualization of one of the top-5 most activated prototypes on a representative test image across all models.
• Figure 4: a) SAM2 and SLIT example segmentations across datasets. b) Examples of uncropped input views. The illustrated examples were used in the projection step of our model, establishing a 1:1 relation between a prototype and its nearest input view.