[论文解读] 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）和β-VAE在医学影像中进行无监督异常定位的问题，比较了重建误差、ELBO梯度、KL梯度以及基于能量的投影方法。研究发现，像Rec-grad这样的基于梯度的预测器在脑部MRI中优于重建误差和能量投影方法，且性能对VAE超参数（尤其是β-VAE中的β）极为敏感。

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与β-VAE框架下，多种异常定位预测器（重建误差、ELBO梯度、KL梯度、基于能量的投影）在医学影像中的有效性。
• 探究基于能量的投影方法（在自然图像中表现良好）是否能有效泛化到具有挑战性的医学影像任务中，如脑肿瘤定位。
• 分析β-VAE超参数（特别是KL权重β）对异常定位性能的影响。
• 通过集成学习评估多种预测器之间的互补性，旨在提升单个方法无法达到的定位精度。

提出的方法

• 在正常T2脑部MRI图像（BraTS2018）上训练VAE与β-VAE，以学习健康组织的数据分布。
• 计算多种异常得分：像素级重建误差、ELBO和KL散度相对于输入的梯度，以及基于能量的投影重建误差。
• 通过迭代能量最小化优化重建结果，使其更贴近正常流形，从而提升正常区域的保真度。
• 使用逻辑回归对多个预测器（Rec-Error、KL-grad、Rec-grad）进行集成，以提升定位性能。
• 在不同潜在空间维度和β值下，使用测试数据的AUROC评估所有方法的性能。
• 通过比较不同VAE配置的性能，评估模型超参数的敏感性。

实验结果

研究问题

• RQ1 与标准重建误差相比，基于能量的投影是否在医学影像中显著提升异常定位性能？
• RQ2 KL梯度预测器的鲁棒性如何随不同VAE架构和超参数（特别是β-VAE中的β）变化？
• RQ3 在脑部MRI异常定位中，哪种异常检测预测器——重建误差、ELBO梯度、KL梯度或Rec-grad——能获得最高的AUROC？
• RQ4 当通过集成学习组合时，多种预测器在多大程度上能相互补充？
• RQ5 β-VAE中的潜在空间维度和β值如何影响不同异常定位方法的性能？

主要发现

• 基于能量的投影提升了正常区域的重建准确性，但未超越基于梯度的方法，其AUROC为0.861，低于Rec-grad的0.900。
• KL梯度预测器并非始终稳健；其性能严重依赖于VAE超参数，尤其是β值和潜在空间大小。
• 对于β-VAE，提高β值普遍提升了多数预测器的性能，β=10时表现最佳，尽管组合方法（Combi）在高β值下性能下降。
• Rec-grad预测器取得了最高的AUROC（0.905），优于重建误差（0.856）和基于能量的投影（0.861）。
• 通过Rec-Error、KL-grad和Rec-grad进行集成学习得到的AUROC为0.903，略低于仅使用Rec-grad的性能，表明各预测器之间互补性有限。
• 标准VAE在高β设置下，正常区域的重建质量较差，表明模型未能充分学习正常流形，从而限制了投影方法的收益。