TY - JOUR
T1 - Transformer-based out-of-distribution detection for clinically safe segmentation
AU - Graham, Mark S
AU - Tudosiu, Petru-Daniel
AU - Wright, Paul
AU - Pinaya, Walter Hugo Lopez
AU - Jean-Marie, U
AU - Mah, Yee
AU - Teo, James
AU - Jäger, Rolf H
AU - Werring, David
AU - Nachev, Parashkev
AU - Ourselin, Sebastien
AU - Cardoso, M Jorge
N1 - Funding Information:
MG, PW, WS, SO, PN and MJC are supported by a grant from the Wellcome Trust (WT213038/Z/18/Z). MJC and SO are also supported by the Wellcome/EPSRC Centre for Medical Engineering (WT203148/Z/16/Z), and the InnovateUK-funded London AI centre for Value-based Healthcare. PN is also supported by NIHR UCLH Biomedical Research Centre. YM is supported by a grant from the Medical Research Council (MR/T005351/1). The models in this work were trained on NVIDIA Cambridge-1, the UK’s largest supercomputer, aimed at accelerating digital biology.
Publisher Copyright:
© 2022 M.S. Graham et al.
PY - 2022/5/21
Y1 - 2022/5/21
N2 - In a clinical setting it is essential that deployed image processing systems are robust to the full range of inputs they might encounter and, in particular, do not make confidently wrong predictions. The most popular approach to safe processing is to train networks that can provide a measure of their uncertainty, but these tend to fail for inputs that are far outside the training data distribution. Recently, generative modelling approaches have been proposed as an alternative; these can quantify the likelihood of a data sample explicitly, filtering out any out-of-distribution (OOD) samples before further processing is performed. In this work, we focus on image segmentation and evaluate several approaches to network uncertainty in the far-OOD and near-OOD cases for the task of segmenting haemorrhages in head CTs. We find all of these approaches are unsuitable for safe segmentation as they provide confidently wrong predictions when operating OOD. We propose performing full 3D OOD detection using a VQ-GAN to provide a compressed latent representation of the image and a transformer to estimate the data likelihood. Our approach successfully identifies images in both the far- and near-OOD cases. We find a strong relationship between image likelihood and the quality of a model’s segmentation, making this approach viable for filtering images unsuitable for segmentation. To our knowledge, this is the first time transformers have been applied to perform OOD detection on 3D image data.
AB - In a clinical setting it is essential that deployed image processing systems are robust to the full range of inputs they might encounter and, in particular, do not make confidently wrong predictions. The most popular approach to safe processing is to train networks that can provide a measure of their uncertainty, but these tend to fail for inputs that are far outside the training data distribution. Recently, generative modelling approaches have been proposed as an alternative; these can quantify the likelihood of a data sample explicitly, filtering out any out-of-distribution (OOD) samples before further processing is performed. In this work, we focus on image segmentation and evaluate several approaches to network uncertainty in the far-OOD and near-OOD cases for the task of segmenting haemorrhages in head CTs. We find all of these approaches are unsuitable for safe segmentation as they provide confidently wrong predictions when operating OOD. We propose performing full 3D OOD detection using a VQ-GAN to provide a compressed latent representation of the image and a transformer to estimate the data likelihood. Our approach successfully identifies images in both the far- and near-OOD cases. We find a strong relationship between image likelihood and the quality of a model’s segmentation, making this approach viable for filtering images unsuitable for segmentation. To our knowledge, this is the first time transformers have been applied to perform OOD detection on 3D image data.
KW - out-of-distribution detection
KW - segmentation
KW - Transformers
KW - uncertainty
UR - http://www.scopus.com/inward/record.url?scp=85175475552&partnerID=8YFLogxK
M3 - Conference paper
SN - 2640-3498
SP - 457
EP - 476
JO - Proceedings of Machine Learning Research
JF - Proceedings of Machine Learning Research
T2 - 5th International Conference on Medical Imaging with Deep Learning, MIDL 2022
Y2 - 6 July 2022 through 8 July 2022
ER -