Weakly supervised Unet: an image classifier which learns to explain itself

Abstract

BackgroundExplainability is a major limitation of current convolutional neural network (CNN) image classifiers. A CNN is required which supports its image-level prediction with a voxel-level segmentation. MethodsA weakly-supervised Unet architecture (WSUnet) is proposed to model voxel classes, by training with image-level supervision. WSUnet computes the image-level class prediction from the maximal voxel class prediction. Thus, voxel-level predictions provide a causally verifiable saliency map for the image-level decision. WSUnet is applied to explainable lung cancer detection in CT images. For comparison, current model explanation approaches are also applied to a standard CNN. Methods are compared using voxel-level discrimination metrics and a clinician preference survey. ResultsIn test data from two external institutions, WSUnet localised the tumour precisely at voxel-level (Precision: 0.93 [0.93-0.94]), achieving superior voxel-level discrimination to the best comparator (AUPR: 0.55 [0.54-0.55] vs. 0.36 [0.35-0.36]). Clinicians preferred WSUnet predictions in most test instances (Clinician Preference Rate: 0.72 [0.68-0.77]). ConclusionsWSUnet is a simple extension of the Unet, which facilitates voxel-level modelling from image-level labels. As WSUnet supports its image-level prediction with a causative voxel-level segmentation, it functions as a self-explaining image classifier. O_FIG O_LINKSMALLFIG WIDTH=200 HEIGHT=194 SRC="FIGDIR/small/507144v1_ufig1.gif" ALT="Figure 1"> View larger version (35K): org.highwire.dtl.DTLVardef@52e3ccorg.highwire.dtl.DTLVardef@1e981f0org.highwire.dtl.DTLVardef@151f31eorg.highwire.dtl.DTLVardef@13046a2_HPS_FORMAT_FIGEXP M_FIG Graphical Abstract The weakly-supervised Unet converts voxel-level predictions to image-level predictions using a global max-pooling layer. Thus, loss is computed at image-level. Following training with image-level labels, voxel-level predictions are extracted from the voxel-level output layer. C_FIG FundingAuthors acknowledge funding support from the UK Research & Innovation London Medical Imaging and Artificial Intelligence Centre; Wellcome/Engineering and Physical Sciences Research Council Centre for Medical Engineering at Kings College London [WT 203148/Z/16/Z]; National Institute for Health Research Biomedical Research Centre at Guys & St Thomas Hospitals and Kings College London; National Institute for Health Research Biomedical Research Centre at Guys & St Thomas Hospitals and Kings College London; Cancer Research UK National Cancer Imaging Translational Accelerator [C1519/A28682]. For the purpose of open access, authors have applied a CC BY public copyright licence to any Author Accepted Manuscript version arising from this submission. HIGHLIGHTSO_LIWSUnet is a weakly supervised Unet architecture which can learn semantic segmentation from data labelled only at image-level. C_LIO_LIWSUnet is a convolutional neural network image classifier which provides a causally verifiable voxel-level explanation to support its image-level prediction. C_LIO_LIIn application to explainable lung cancer detection, WSUnets voxel-level output localises tumours precisely, outperforming current model explanation methods. C_LIO_LIWSUnet is a simple extension of the standard Unet architecture, requiring only the addition of a global max-pooling layer to the output. C_LI