[論文レビュー] Unsupervised anomaly localization using VAE and beta-VAE

Leixin Zhou, Wenxiang Deng+1 more|arXiv (Cornell University)|May 19, 2020

COVID-19 diagnosis using AI参考文献 21被引用数 13

ひとこと要約

本稿では、VAEおよびbeta-VAEを用いて、脳MRIにおける教師なし異常局在化を調査している。再構成誤差、ELBO勾配、KL勾配、エネルギーに基づく投影を比較した。脳MRIでは、再構成誤差やエネルギー投影よりも勾配ベースの予測子（例：Rec-grad）が優れていることが判明し、性能はVAEのハイパーパrameter、特にbeta-VAEにおけるbetaに極めて敏感である。

ABSTRACT

Variational Auto-Encoders (VAEs) have shown great potential in the unsupervised learning of data distributions. An VAE trained on normal images is expected to only be able to reconstruct normal images, allowing the localization of anomalous pixels in an image via manipulating information within the VAE ELBO loss. The ELBO consists of KL divergence loss (image-wise) and reconstruction loss (pixel-wise). It is natural and straightforward to use the later as the predictor. However, usually local anomaly added to a normal image can deteriorate the whole reconstructed image, causing segmentation using only naive pixel errors not accurate. Energy based projection was proposed to increase the reconstruction accuracy of normal regions/pixels, which achieved the state-of-the-art localization accuracy on simple natural images. Another possible predictors are ELBO and its components gradients with respect to each pixels. Previous work claimed that KL gradient is a robust predictor. In this paper, we argue that the energy based projection in medical imaging is not as useful as on natural images. Moreover, we observe that the robustness of KL gradient predictor totally depends on the setting of the VAE and dataset. We also explored the effect of the weight of KL loss within beta-VAE and predictor ensemble in anomaly localization.

研究の動機と目的

• VAEおよびbeta-VAEフレームワーク内での異常局在化予測子（再構成誤差、ELBO勾配、KL勾配、エネルギーに基づく投影）の有効性を評価すること。
• 自然画像で有効なエネルギーに基づく投影が、脳腫瘍局在化のような挑戦的な医療画像処理タスクにも一般化するかを調査すること。
• 特にKL重み（beta）に注目したbeta-VAEのハイパーパrameterが、異常局在化性能に与える影響を分析すること。
• アンサンブル学習を用いて複数の予測子の相乗効果を評価し、個々の手法を上回る局在化精度を達成することを目的とする。

提案手法

• 正常なT2 MRI脳画像（BraTS2018）を用いて、健康な組織のデータ分布を学習するため、VAEおよびbeta-VAEを訓練した。
• 複数の異常スコアを計算した：ピクセル単位の再構成誤差、入力に関するELBOおよびKL発散の勾配、エネルギーに基づく再構成誤差の投影。
• 反復的エネルギー最小化を適用し、正常マニフォールドに再構成を適合させることで、正常領域の忠実度を向上させた。
• ロジスティック回帰を用いて、複数の予測子（Rec-Error、KL-grad、Rec-grad）をアンサンブル化し、局在化性能を向上させた。
• さまざまな潜在空間次元およびbeta値を用いて、テストデータにおけるAUROCを評価した。
• モデルのハイパーパrameterへの感受性を評価するため、異なるVAE設定間での性能を比較した。

実験結果

リサーチクエスチョン

• RQ1 エネルギーに基づく投影は、標準的な再構成誤差と比較して、医療画像における異常局在化に顕著に効果を発揮するか？
• RQ2 KL勾配予測子のロバスト性は、VAEアーキテクチャーやハイパーパrameter、特にbeta-VAEにおけるbetaにどのように依存するか？
• RQ3 再構成誤差、ELBO勾配、KL勾配、Rec-gradの中から、脳MRIにおける異常局在化で最高のAUROCを達成するのはどれか？
• RQ4 アンサンブル学習によって複数の予測子を組み合わせた場合、それらはどの程度相乗効果を発揮するか？
• RQ5 beta-VAEにおける潜在空間次元およびbeta値は、異なる異常局在化手法の性能にどのように影響を与えるか？

主な発見

• エネルギーに基づく投影は正常領域の再構成精度を向上させたが、勾配ベースの手法に劣り、AUROCは0.861であったのに対し、Rec-gradは0.900を記録した。
• KL勾配予測子は一貫したロバスト性を示さず、VAEのハイパーパrameter、特にbetaと潜在空間サイズに強く依存していた。
• beta-VAEでは、betaを増加させることでほとんどの予測子の性能が向上し、beta=10が最良の結果をもたらしたが、Combiの性能は高betaで低下した。
• Rec-grad予測子は最高のAUROC 0.905を達成し、再構成誤差（0.856）およびエネルギー投影（0.861）を上回った。
• Rec-Error、KL-grad、Rec-gradを用いたアンサンブル学習によるAUROCは0.903に留まり、Rec-grad単体よりわずかに低かった。予測子間の相乗効果は限定的であった。
• 標準VAEでは、特に高beta設定において正常領域の再構成品質が著しく低く、モデルが正常マニフォールドを適切に学習できていないことが示唆された。これは、投影手法の利点が制限されることを意味する。