This Probably Looks Exactly Like That: An Invertible Prototypical Network

TL;DR

ProtoFlow addresses the semantic gap in prototypical networks by learning prototypical distributions over a latent space with an invertible normalizing flow backbone. By modeling class prototypes as a Gaussian Mixture in latent space and leveraging an exact inverse $f^{-1}$, it provides faithful data-space visualizations and calibrated uncertainty estimates while maintaining competitive predictive accuracy. The approach achieves state-of-the-art joint generative-predictive performance across diverse datasets and offers richer interpretability through prototype distributions, heatmaps, and prototypical parts, reinforced by a diversity loss and prototype pruning. This work has practical impact for interpretable AI, enabling more trustworthy, data-efficient reasoning in vision tasks and beyond.

Abstract

We combine concept-based neural networks with generative, flow-based classifiers into a novel, intrinsically explainable, exactly invertible approach to supervised learning. Prototypical neural networks, a type of concept-based neural network, represent an exciting way forward in realizing human-comprehensible machine learning without concept annotations, but a human-machine semantic gap continues to haunt current approaches. We find that reliance on indirect interpretation functions for prototypical explanations imposes a severe limit on prototypes' informative power. From this, we posit that invertibly learning prototypes as distributions over the latent space provides more robust, expressive, and interpretable modeling. We propose one such model, called ProtoFlow, by composing a normalizing flow with Gaussian mixture models. ProtoFlow (1) sets a new state-of-the-art in joint generative and predictive modeling and (2) achieves predictive performance comparable to existing prototypical neural networks while enabling richer interpretation.

Figures (20)

• Figure 1: Existing prototypical networks rely on prototypical points and limiting means of visualizing prototypes. Our proposed approach, ProtoFlow, enables the learning of prototypical distributions with inverses, enabling their exact and efficient visualization---ProtoFlow inherently enables richer prototype interpretation.
• Figure 2: Examples of whole-image (left) and prototypical parts (right) explanations.
• Figure 3: An overview of the proposed ProtoFlow model that composes a normalizing flow and Gaussian mixture models. In the normalizing${\mathcal{X}} \xrightarrow{\text{\tiny $f$}} {\mathcal{Z}}$ direction, the invertible composition $f = f_k \circ \dots \circ f_1$pulls back the structured latent density ${p_{\mathcal{Z}}}$ to the complex data density ${p_{\mathcal{X}}}$. The generating${\mathcal{X}} \xleftarrow{\text{\tiny $f^{\hbox{$-1$}}$}} {\mathcal{Z}}$ direction generates points $\tilde{{\bm{x}}} \sim {p_{\mathcal{X}}}$ implicitly by pushing forward samples $\tilde{{\bm{z}}} \sim {p_{\mathcal{Z}}}$ from the latent distribution along the inverse $f^{\hbox{$-1$}}$.
• Figure 4: \ref{['subfig:query']}"This looks like that"-style explanations of bird (top row) and automobile (bottom row) image classification decisions. Rather than using training samples as prototypes, ProtoFlow learns prototype distributions directly over the latent space, leading to the "bird/car-adjacent" images in the third column. The fourth column shows the most-likely dataset image part for each prototypical distribution. \ref{['subfig:mean-cropped']} The mean point image of the bird prototype with a bird-like figure segmented from the background. \ref{['subfig:cifar-examples']} Human-picked images from CIFAR-10 that qualitatively match this prototype image.
• Figure 5: Mean points (center) and generated samples (periphery) from prototype distributions learned on CIFAR-10 with consistency loss and a truncation value of $1$.