Chimeric U-Net – Modifying the standard U-Net towards explainability

Abstract

Healthcare guided by semantic segmentation has the potential to improve our quality of life through early and accurate disease detection. Convolutional Neural Networks, especially the U-Net-based architectures, are currently the state-of-the-art learning-based segmentation methods and have given unprecedented performances. However, their decision-making processes are still an active field of research. In order to reliably utilize such methods in healthcare, explainability of how the segmentation was performed is mandated. To date, explainability is studied and applied heavily in classification tasks. In this work, we propose the Chimeric U-Net, a U-Net architecture with an invertible decoder unit, that inherently brings explainability into semantic segmentation tasks. We find that having the restriction of an invertible decoder does not hinder the performance of the segmentation task. However, the invertible decoder helps to disentangle the class information in the latent space embedding and to construct meaningful saliency maps. Furthermore, we found that with a simple k-Nearest-Neighbours classifier, we could predict the Intersection over Union scores of unseen data, demonstrating that the latent space, constructed by the Chimeric U-Net, encodes an interpretable representation of the segmentation quality. Explainability is an emerging field, and in this work, we propose an alternative approach, that is, rather than building tools for explaining a generic architecture, we propose constraints on the architecture which induce explainability. With this approach, we could peer into the architecture to reveal its class correlations and local contextual dependencies, taking an insightful step towards trustworthy and reliable AI. Code to build and utilize the Chimeric U-Net is made available under:

https://github.com/kenrickschulze/Chimeric-UNet---Half-invertible-UNet-in-Pytorch