A differentiable Gaussian Prototype Layer for explainable Segmentation

TL;DR

This paper presents a differentiable Gaussian Prototype Layer for explainable segmentation, enabling end-to-end training of neural networks with improved prototype localization and applicability to semi-supervised learning.

Contribution

It introduces a novel gradient-based Gaussian Prototype Layer that allows end-to-end training and enhances explainability in segmentation models.

Findings

Achieves state-of-the-art performance with prototype detection in latent grids.

Provides more precise prototype localization using region proposals.

Facilitates future semi-supervised learning strategies.

Abstract

We introduce a Gaussian Prototype Layer for gradient-based prototype learning and demonstrate two novel network architectures for explainable segmentation one of which relies on region proposals. Both models are evaluated on agricultural datasets. While Gaussian Mixture Models (GMMs) have been used to model latent distributions of neural networks before, they are typically fitted using the EM algorithm. Instead, the proposed prototype layer relies on gradient-based optimization and hence allows for end-to-end training. This facilitates development and allows to use the full potential of a trainable deep feature extractor. We show that it can be used as a novel building block for explainable neural networks. We employ our Gaussian Prototype Layer in (1) a model where prototypes are detected in the latent grid and (2) a model inspired by Fast-RCNN with SLIC superpixels as region proposals. The earlier achieves a similar performance as compared to the state-of-the art while the latter has the benefit of a more precise prototype localization that comes at the cost of slightly lower accuracies. By introducing a gradient-based GMM layer we combine the benefits of end-to-end training with the simplicity and theoretical foundation of GMMs which will allow to adapt existing semi-supervised learning strategies for prototypical part models in future.