[论文解读] Transformer-based out-of-distribution detection for clinically safe segmentation

Mark S. Graham, Petru-Daniel Tudosiu+10 more|arXiv (Cornell University)|May 21, 2022

Machine Learning in Healthcare被引用 9

一句话总结

该论文提出了一种基于3D变换器的分布外（OOD）检测方法，用于临床图像分割，采用VQ-GAN压缩头颈部CT扫描，并利用变换器估计数据似然。该方法成功识别了远距离和近距离OOD扫描，图像似然与分割性能高度相关，从而可在分割前可靠地过滤掉有问题的输入。

ABSTRACT

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. Code is available at github.com/marksgraham/transformer-ood.

研究动机与目标

• 解决深度学习模型在分布外（OOD）医学图像上做出自信错误预测的临床风险。
• 评估现有不确定性估计方法（如dropout、集成）在3D CT分割中远距离和近距离OOD场景下的表现。
• 开发一种稳健的OOD检测系统，在处理前过滤掉可能导致分割失败的数据。
• 探究基于VQ-GAN潜在表示上变换器估计的图像似然是否与分割网络性能相关。
• 在包括非头颈部CT、MRI和受损扫描在内的多样化OOD数据上验证该方法。

提出的方法

• 使用VQ-GAN将3D头颈部CT图像压缩到离散潜在空间，降低维度以实现高效建模。
• 训练变换器以建模潜在标记序列的自回归特性，实现输入图像的精确似然估计。
• 将每个输入体积的对数似然作为分布一致性度量。
• 将似然得分用作OOD风险的代理指标，低似然值表示可能存在数据分布偏移。
• 在将图像输入下游分割网络前，将似然阈值设定为预过滤器。
• 使用蒙特卡洛dropout训练并评估分割网络，以分析其性能与似然得分的相关性。

实验结果

研究问题

• RQ1 标准不确定性估计方法（如dropout、集成）在3D分割中应用于远距离和近距离OOD医学图像时是否失效？
• RQ2 基于VQ-GAN潜在表示的变换器似然估计器能否有效检测远距离和近距离OOD的3D医学图像？
• RQ3 通过变换器在VQ-GAN潜在空间上估计的图像似然是否与分割网络性能（如假阳性率）存在强相关性？
• RQ4 所提出的OOD检测方法是否能防止分割网络对异常输入做出自信错误的预测？
• RQ5 在VQ-GAN中引入感知损失和对抗损失如何影响OOD检测性能？

主要发现

• 所有标准不确定性估计方法（如dropout、集成）在OOD数据上均失效，在远距离和近距离OOD情况下均产生自信的错误预测。
• 所提出的VQ-GAN + 变换器方法在区分分布内头颈部CT与10种不同OOD类别（包括非头颈部CT、MRI和受损扫描）方面实现了100% AUC。
• 图像对数似然与分割网络性能强相关：似然值低于-7000的输入具有显著更高的假阳性率。
• 似然值接近分布内数据的受损扫描（如噪声σ=0.01，AUC=0.49）被分割网络处理得更好，表明其存在细微损伤。
• 该方法避免了传统生成式OOD检测中对高度OOD样本分配高似然值的常见失败模式，表明其在鲁棒性方面优于先前方法。
• 空间似然图可定位损伤源（如颅骨剥离、噪声），实现对问题图像区域的可解释性检测。