Model Guidance via Explanations Turns Image Classifiers into Segmentation Models

TL;DR

Heatmaps from explainable AI methods often resemble segmentation maps, enabling weakly supervised segmentation with image-level labels. The authors unify heatmap-based guidance with standard segmentation by unrolling Layer-wise Relevance Propagation (LRP) into an encoder–decoder architecture with weight tying and class-specific decoders, trained with standard segmentation losses. They show formal parallels to conventional segmentation networks and demonstrate competitive performance on PASCAL VOC 2012 across backbones, with notable gains when pixel-level labels are scarce. The approach preserves classifier performance while leveraging image-level labels to improve segmentation, offering a practical path for semi-supervised segmentation with readily pluggable training objectives and code availability.

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}.

Figures (3)

• Figure 1: Unrolled LRP-0 for standard convolutional classifiers is an encoder-decoder CNN. (A) Sketch of architectural building blocks. (B) Decoder weights are tied with corresponding encoder layers. (C) Activation functions are skipped from encoder to decoder. (D) Up-convolutions are tied to the respective pooling operation in the encoder. In (A) we sketch the decoder architecture induced by an individual semantic class i; this decoder is replicated for all classes with shared weights, except for the bottleneck 1x1 convolution, which is class-specific.
• Figure 2: Exemplary unrolled LRP segmentations on the VOC validation set, 500 labelled images scenario. Left: input image; Right: ground truth segmentation; In-between: Evolution of segmentation over training of unrolled LRP, ranging from the first iteration to peak validation mIoU.
• Figure 3: Exemplary unrolled LRP heatmaps on the VOC validation set, 500 labelled images scenario (cf. Fig. \ref{['fig:segmentations']} for corresponding segmentations). Left: input image; Right: ground truth segmentation; In-between: Evolution of heatmaps of the positive classes over training of unrolled LRP, ranging from the first iteration to peak validation mIoU. Note that all heatmaps except for the initial (left-most) one are shown after softmax to render standard blue-white-red visualization informative.