[논문 리뷰] Fast Unsupervised Brain Anomaly Detection and Segmentation with Diffusion Models

Walter Hugo Lopez Pinaya, Mark S. Graham외 13명|arXiv (Cornell University)|2022. 06. 07.

Generative Adversarial Networks and Image Synthesis인용 수 6

한 줄 요약

이 논문은 건강한 뇌 MRI 및 CT 스캔의 잠재 표현에 적용된 노이즈 제거 확산 확률 모델(Denosing Diffusion Probabilistic Models, DDPMs)을 사용하여 빠르고 비지도 학습 기반의 뇌 이상 징후 탐지 및 분할 방법을 제안한다. DDPM의 역방향 노이즈 제거 과정을 활용하여 잠재 공간 내의 이질성을 식별함으로써, 상태의 최신 자동회귀형 트랜스포머 모델과 경쟁 가능한 성능을 달성하면서도 추론 시간을 1시간 이상에서 10분 이내로 단축시켜 실시간 응용에 대한 임상적 타당성을 확보한다.

ABSTRACT

Deep generative models have emerged as promising tools for detecting arbitrary anomalies in data, dispensing with the necessity for manual labelling. Recently, autoregressive transformers have achieved state-of-the-art performance for anomaly detection in medical imaging. Nonetheless, these models still have some intrinsic weaknesses, such as requiring images to be modelled as 1D sequences, the accumulation of errors during the sampling process, and the significant inference times associated with transformers. Denoising diffusion probabilistic models are a class of non-autoregressive generative models recently shown to produce excellent samples in computer vision (surpassing Generative Adversarial Networks), and to achieve log-likelihoods that are competitive with transformers while having fast inference times. Diffusion models can be applied to the latent representations learnt by autoencoders, making them easily scalable and great candidates for application to high dimensional data, such as medical images. Here, we propose a method based on diffusion models to detect and segment anomalies in brain imaging. By training the models on healthy data and then exploring its diffusion and reverse steps across its Markov chain, we can identify anomalous areas in the latent space and hence identify anomalies in the pixel space. Our diffusion models achieve competitive performance compared with autoregressive approaches across a series of experiments with 2D CT and MRI data involving synthetic and real pathological lesions with much reduced inference times, making their usage clinically viable.

연구 동기 및 목표

• 정상 이상 징후 데이터가 필요한 비지도 학습 기반의 신속하고 정확한 뇌 병변 탐지에 대한 임상적 필요를 충족시키기 위해.
• 자기회귀형 트랜스포머의 한계, 즉 높은 추론 시간, 누적 오류, 의료 영상 분석에서의 단방향 순서 편향을 극복하기 위해.
• 잠재 공간에서의 확산 모델의 빠름과 안정성을 활용하여 뇌 영상에서 실시간 이상 징후 탐지 및 분할을 가능하게 하기 위해.
• 확산 모델이 트랜스포머 기반 방법과 비슷하거나 이를 초월하는 이상 징후 탐지 성능을 달성하면서도 훨씬 더 빠른 속도를 제공할 수 있음을 입증하기 위해.

제안 방법

• 방법은 2D MRI 및 CT 영상을 저차원 잠재 공간으로 압축하여 구조적 정보를 유지하는 변동형 오토인코더(Variational Autoencoder, VAE)를 사용한다.
• 정상 뇌 해부학적 분포를 학습하기 위해, 건강한 잠재 표현 전용으로 노이즈 제거 확산 확률 모델(Denosing Diffusion Probabilistic Model, DDPM)을 훈련한다.
• 이상 징후는 역방향 노이즈 제거 과정 중 복원 오차를 측정하여 탐지한다: 높은 오차는 정상성에서의 이질성을 나타낸다.
• 샘플링 단계를 500에서 50으로 줄이면서도 영상 품질과 성능을 유지하기 위해, DDIM(Denoising Diffusion Implicit Models)를 사용하여 역과정을 가속화한다.
• 잠재 공간에서의 픽셀 단위 복원 오차를 계산하고 이를 픽셀 공간으로 복원함으로써 이상 징후 분할을 수행한다.
• 해상도와 추론 속도를 제어하기 위해 고정된 다운샘플링 인자(f=4 또는 f=8)를 사용하며, 다양한 데이터셋에서 성능을 평가한다.

실험 결과

연구 질문

• RQ1 잠재 공간에서의 확산 모델이 뇌 영상에서 최신 자동회귀형 트랜스포머와 비교해 경쟁 가능한 이상 징후 탐지 성능을 달성할 수 있는가?
• RQ2 DDIM 샘플링을 사용할 경우, 탐지 정확도를 떨어뜨리지 않고도 추론 시간을 크게 단축시킬 수 있는가?
• RQ3 이상 징후 레이블 없이도, 합성 및 실제 병변 병변을 2D CT 및 MRI 스캔에서 탐지하고 분할할 수 있는가?
• RQ4 이 방법은 해상도가 증가함에 따라 어떻게 스케일링되며, 자기회귀 모델에서 흔히 발생하는 오류 누적 문제를 피할 수 있는가?
• RQ5 이 방법은 특히 뇌내 출혈 탐지와 같이 시간이 중요한 상황에서 실시간 임상 적용에 충분히 빠른가?

주요 결과

• 제안된 DDPM 기반 방법은 KCH 데이터셋에서 DICE 스코어 0.471을 기록하여 최고 성능을 낸 트랜스포머 앙상블(0.435)을 능가했으며, 추론 시간도 크게 단축시켰다.
• DDIM 샘플링을 사용한 50단계(f=4)로 추론 시간을 100개 슬라이스 기준 324초(약 6분 이내)로 단축시켰고, 트랜스포머 앙상블의 4,907초(80시간 이상) 대비 약 15배 빠르게 성능을 달성했다.
• CROMIS 데이터셋에서 DDIM(f=8)를 사용한 결과, 추론 시간 12초에 DICE 스코어 0.286를 기록하여 트랜스포머 기반 베이스라인(0.205)과 f-AnoGAN(0.146)을 모두 능가했다.
• 이 방법은 다양한 데이터에 대해 뛰어난 강건성을 보였으며, KCH에서 최고의 DICE 스코어 0.471, WMH에서 0.298을 기록하여 실제 병변 분할에서 모든 베이스라인을 능가했다.
• DDIM를 사용함으로써 표준 DDPM 대비 10배 빠른 속도 향상을 달성하면서도 성능은 유사하게 유지되어 임상 적용에 있어 핵심 요소로 입증되었다.
• 특히 고해상도 환경에서 자기회귀 모델이 증가하는 계산 비용 문제로 어려움을 겪는 상황에서도, 이 방법은 성능과 추론 속도 면에서 항상 뛰어난 우월성을 보였다.