Model Guidance via Explanations Turns Image Classifiers into Segmentation Models

TL;DR

This paper demonstrates that explainability heatmaps can be used as a form of guidance to turn image classifiers into effective segmentation models, especially under weak supervision, by establishing formal parallels and leveraging standard segmentation losses.

Contribution

It introduces a novel approach that unifies explainability heatmaps with segmentation architectures, enabling semi-supervised training with minimal pixel-level labels.

Findings

Differentiable heatmaps are formally similar to encoder-decoder segmentation models.

Heatmap-based models achieve competitive segmentation results.

Weakly supervised training with image-level labels outperforms traditional encoder-decoder models.

Abstract

Heatmaps generated on inputs of image classification networks via explainable AI methods like Grad-CAM and LRP have been observed to resemble segmentations of input images in many cases. Consequently, heatmaps have also been leveraged for achieving weakly supervised segmentation with image-level supervision. On the other hand, losses can be imposed on differentiable heatmaps, which has been shown to serve for (1)~improving heatmaps to be more human-interpretable, (2)~regularization of networks towards better generalization, (3)~training diverse ensembles of networks, and (4)~for explicitly ignoring confounding input features. Due to the latter use case, the paradigm of imposing losses on heatmaps is often referred to as "Right for the right reasons". We unify these two lines of research by investigating semi-supervised segmentation as a novel use case for the Right for the Right Reasons paradigm. First, we show formal parallels between differentiable heatmap architectures and standard encoder-decoder architectures for image segmentation. Second, we show that such differentiable heatmap architectures yield competitive results when trained with standard segmentation losses. Third, we show that such architectures allow for training with weak supervision in the form of image-level labels and small numbers of pixel-level labels, outperforming comparable encoder-decoder models. Code is available at \url{https://github.com/Kainmueller-Lab/TW-autoencoder}.