[論文レビュー] 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

ひとこと要約

本論文は、頭部CTスキャンを圧縮するためのVQ-GANと、データの尤度を推定するためのトランスフォーマーを用いた、3次元トランスフォーマーに基づく分布外（OOD）検出手法を提案する。本手法は、遠方OODおよび近方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）の医療画像に対して自信を持って誤った予測を行う臨床的リスクに対処すること。
• 3次元CTセグメンテーションにおける遠方OODおよび近方OODの状況において、従来の不確実性推定手法（例：ドロップアウト、アンサンブル）の有効性を評価すること。
• 処理前にセグメンテーション障害を引き起こす可能性のあるデータをフィルタリングする、堅牢なOOD検出システムの開発。
• VQ-GANの潜在空間上でトランスフォーマーを用いて推定した画像尤度が、セグメンテーションネットワークの性能と相関しているかどうかを調査すること。
• 非頭部CT、MRI、損傷付きスキャンを含む多様なOODデータを用いて、手法の妥当性を検証すること。

提案手法

• 3次元頭部CTボリュームをVQ-GANで離散的潜在空間に圧縮し、効率的なモデリングに適した次元削減を実現する。
• 潜在トークンの自己回帰的シーケンスをモデル化するためのトランスフォーマーを訓練し、入力画像の正確な尤度推定を可能にする。
• 各入力ボリュームの対数尤度を計算し、分布への適合度の指標とする。
• 尤度スコアをOODリスクの代理指標とし、低尤度は分布シフトの可能性を示す。
• 尤度スコアのしきい値処理を、下流のセグメンテーションネットワークへの入力の前処理として統合する。
• モンテカルロドロップアウトを用いてセグメンテーションネットワークを訓練・評価し、尤度スコアとの性能相関を評価する。

実験結果

リサーチクエスチョン

• RQ1 標準的な不確実性推定手法（例：ドロップアウト、アンサンブル）は、3次元セグメンテーションにおける遠方OODおよび近方OODの医療画像に適用した場合に失敗するか？
• RQ2 VQ-GANの潜在空間上でトランスフォーマーに基づく尤度推定器は、遠方OODおよび近方OODの3次元医療画像を効果的に検出できるか？
• RQ3 画像尤度とセグメンテーションネットワークの性能（例：誤検出率）との間に強い相関があるか？
• RQ4 提案手法によるOOD検出は、異常な入力に対して自信を持って誤った予測を行うセグメンテーションネットワークを防げるか？
• RQ5 VQ-GANに知覚的および敵対的損失を組み込むと、OOD検出性能にどのような影響を与えるか？

主な発見

• すべての標準的不確実性推定手法（例：ドロップアウト、アンサンブル）はOODデータに対して失敗し、遠方OODおよび近方OODの両方の状況で、自信を持って誤った予測を出力した。
• 提案手法（VQ-GAN + トランスフォーマー）は、非頭部CT、MRI、損傷付きスキャンを含む10種類の異なるOODクラスに対して、インダストリィの頭部CTとを区別する際、100%のAUCを達成した。
• 画像の対数尤度とセグメンテーションネットワークの性能には強い相関があり、尤度が-7000未満の入力では、誤検出率が顕著に高くなった。
• 分布内データに近い尤度を示す損傷付きスキャン（例：ノイズσ=0.01、AUC=0.49）は、セグメンテーションネットワークによってよりうまく処理された。これは、微細な損傷が認識可能であることを示唆している。
• 本手法は、顕著にOODなサンプルに対して高い尤度を割り当てるという、従来の生成的OOD検出手法の一般的な失敗モードを回避しており、より優れたロバストネスを示している。
• 空間的尤度マップにより、損傷源（例：頭蓋骨除去、ノイズ）が局所化され、問題となる画像領域の説明可能な検出が可能になった。