Diffusion Models with Implicit Guidance for Medical Anomaly Detection (MICCAI 2024)

세줄요약

• AnoDDPM 을 medical image (Brain MRI & wrist X-ray) 에 응용
• 영상 전체를 정상 복원 하지 말고, anomalous 한 부분만 집중적으로 복원
• Noise 의 type 은 데이터셋마다 성능이 다름

논문 링크: https://papers.miccai.org/miccai-2024/paper/1315_paper.pdf

(사진 출처: MVTec 데이터셋)

컴퓨터 비전에서 anomaly detection은 이미지에서 이상한 부분을 찾아내는 것이 목적입니다. 위 사진에서 보다시피 초록색은 멀쩡한 사진이고, 빨간색 안에 있는 사진들이 불량 사진인데, 이걸 사람이 일일이 검수하기 힘드니 인공지능이 알아서 이상한 부분 있으면 찾아내라는겁니다.

Anomaly detection 은 아직 사람이 인공지능보다 훨씬 잘하는데, 그 이유는

• 보통 normal dataset이 많음 (class imbalance)
• Anomaly가 어떤 방식으로 출현할지 모름 (unpredictable)
• 정상과 비정상의 경계가 모호함 (borderline cases)

즉, 인공지능을 어떻게 훈련해야할지도 정의하기가 어렵기 때문입니다.

Normative modeling

한가지 보편적인 접근 방식은 정상 분포 자체를 학습시키는 것입니다. 어떤 모델이 정상 데이터에 대한 분포를 학습하고 나면, 어떤 임의의 비정상 데이터에 대해 평소와 다른 출력을 낼 것이라는 것이 기본 개념이죠.

기본적인 학습 구조는 다음과 같습니다.

1. Autoencoder 학습

(자료 출처: Building Autoencoders in Keras)

먼저, normal data 만을 활용하여 autoencoder 을 학습합니다. 이 autoencoder 은 임의의 정상 이미지를 입력 받고, 그 이미지를 latent space 로 압축한 뒤 다시 복원하는 것을 학습하는데, 이를 통해 어떤 데이터가 들어오더라도 정상 데이터처럼 복원하는 것이 목적입니다.

2. Abnormal data reconstruction

이제 위에서 활용한 autoencoder 에 데이터를 넣어서 복원해 본 뒤, 픽셀별로 원래 데이터와의 차이 (reconstruction error) 을 계산합니다.

(자료 출처: Takahiro Nakao, Journal of Digital Imaging 2021)

Chest X-Ray 데이터에 적용해 본 결과입니다. Normal 데이터는 이상한 점이 별로 없어서 reconstruction error 이 전반적으로 낮은 반면, abnormal 데이터는 오른쪽에 병변이 빨갛게 칠해져 있는 것을 확인할 수 있습니다.

Diffusion Models with Implicit Guidance for Medical Anomaly Detection

이제 본격적으로 논문 이야기를 할 준비가 되었습니다. 이 논문의 포인트부터 정리하겠습니다:

• Brain MRI 영상, wrist X-ray 영상에서 anomaly detection 을 수행합니다.
• Diffusion model 을 활용해서 reconstruction 을 수행합니다.
• Intermediate mask 를 생성 (implicit guidance)해서 anatomical coherence 를 개선합니다.

Diffusion model 에 대한 설명은 생략하고, 기존의 diffusion 모델이 왜 문제가 돼서 "implicit guidance" 를 제안하게 되었을까요?

(출처: 원 논문)

먼저, anomaly map의 생성 과정부터 설명드리겠습니다. 앞서 말한 autoencoder과 비슷하게, MRI 영상을 latent space로 압축하는 과정이 필요합니다. 하지만, diffusion model의 encoder은 영상을 압축하지 않고, noise 를 점진적으로 추가하여 굉장히 noisy 하게 만듭니다. 즉, 어떤 MRI 영상이든 일단 거의 안보이는 수준으로 만들고, 그걸 다시 denoising 하는거죠.

이렇게 훈련시킨 DDPM (노랑색) 의 healthy reconstruction 을 보시면, 처음 영상 (Reverse, 빨간색) 과 아예 다르게 생긴 사람의 뇌가 생성된 것을 확인할 수 있습니다. 왜 이럴까요?

Diffusion model 이 영상에 noise를 추가하는 과정에서, 원 이미지의 특징마저 훼손시켜버린 것입니다.

그 결과, normal 영상을 reconstruction 하는 과정에서 아예 다른 사람처럼 생긴 영상을 만들어버렸고, 이 때문에 anomaly map 이 굉장히 noisy해졌습니다. 그럼 어떡할까요? 저자들은 이 문제를 해결하기 위해서 영상 전체에 대해 reconstruction 을 수행하는 것이 아니라, anomaly 가 있는 부분만 reconstruction 하는 아이디어를 꺼냅니다.

Anomaly Score via Intermediate masks

Intermediate mask 의 역할은 예측한 reconstruction과 실제 input 을 비교해서 어떤 부분이 비정상인지 판단하는 것입니다.

이 때, 두 영상 $x, y$에 대한 intermediate mask $m$의 정의는 다음과 같습니다:

• $S_{LPIPS}:$ LPIPS metric, 높을 수록 두 영상의 차이가 큼

계산된 마스크는 영상 전체에 대해 [0, 1]로 normalize 되고, 추가적인 후처리를 거치게 됩니다 (근데 0, 1로 normalize 하면 normal 영상은 전부 anomalous 해지지 않나..?). $m$을 구하고 나면, 이 값에 따라 reconstruction과 input 사이의 interpolation ($m$=1이면 reconstruction only, 0이면 input only) 으로 최종 $x_t$ (timestep $t$의 영상) 을 유추합니다.

(simplified from original equation for clarity)

이렇게 해서, 최종 anomaly score $S$는 모든 harmonization step 에서의 intermediate mask 의 harmonic mean 으로 정의됩니다.

Results

Dataset

이 논문에서는 Brain MRI와 wrist x-ray 를 활용하였으나, Brain MRI 위주로 설명하겠습니다. IXI dataset 과 ATLAS 200 을 활용하였으며, 이 데이터들은 Brain MRI 와 함께 전문가의 병변 annotation 이 함께 제공됩니다.

Performance

논문의 Table 1 과 Fig. 3 을 살펴보겠습니다.

먼저 저자들은 lesion annotation mask를 ground truth로 놓고, anomaly mask 와의 Dice score 을 계산했습니다. Simplex Noise 를 사용했을 때, lesion의 크기에 따라 small 은 11.5, medium 은 39.2, large 는 63.6 으로 lesion의 크기가 커질수록 anomaly detection 의 성능도 올라간 것을 볼 수 있습니다.

Figure 3 에서 gaussian noise 와 simplex noise 의 차이도 주목할만한데, 제일 왼쪽 그림에서 gaussian noise 를 활용한 복원은 모두 normal image 로 복원해버려서 false positive가 많이 발생한 반면, simplex noise 를 활용한 복원은 조금 더 원 image에 truthful 하게 생성한 것을 볼 수 있습니다. 다만, ventricle dilatation 도 어떻게 보면 anomaly 기 때문에, Gaussian 복원이 꼭 문제가 있었나? 에 대한 질문에 더 깊이 생각해 볼 필요가 있습니다.

Insights into noise type

주의: 뇌피셜이 난무할 수 있습니다

원문에선 AnoDDPM에서 사용한 simplex noise 를 응용했습니다. 왜일까요? AnoDDPM원문을 살펴보면 다음과 같은 이야기를 합니다.

Natural images follow a power law distribution of frequencies.

즉, frequency domain 에서 이미지를 바라보았을 때, 영상의 대부분은 low-frequency (smooth background, large shapes) 에 속한다는 것입니다. 하지만, gaussian noise 는 spectral density 가 uniform 해서, anomaly가 대부분 low-frequency에 있을 경우 적합하지 않을 수 있습니다. 이 경우, simplex noise 를 잘 설계하여 power law distribution 을 따르게 할 수 있습니다.

Brain MRI 에서의 결과를 다시 상기해보면, simplex noise 를 활용했을 때 더 정확한 anomaly detection 결과가 나왔는데, 이는 이 brain MRI 데이터셋의 anomaly 는 low-frequency, 즉 큼지막한 구조적 이상이 대부분을 차지했다 라고 해석할 수 있습니다. Simplex noise 가 더 효과적으로 큰 이상을 제거해서, 복원한 영상과 차이가 많이 나게 만든것이죠.

반면에, 저자들이 wrist X-Ray 에 simplex noise를 활용했을 때 오히려 성능이 저하되었음을 보입니다.

Simplex noise 를 사용한 dice score 을 공개하는 대신, supplementary 에 simplex noise 를 활용해서 실패한 case 를 보여줍니다. 초록색으로 칠한 GT annotation 안에 있는 anomaly 가 anomaly map 에는 노랗게 보이지 않는 (normal) 것을 볼 수 있습니다. 이는 wrist X-ray 의 anomaly 는 high-frequency component (fine details) 가 많아서 brain MRI에서 좋은 성능을 보인 simplex noise 가 효과적이지 않았을 것이라고 추측해볼 수 있습니다.

Conclusion

Thor은 의료영상에서 AnoDDPM을 응용한 논문입니다. 이 논문은 정상 데이터만 활용해서 훈련한 diffusion model이 추가적인 훈련 없이 anomaly 를 탐지할 수 있음을 보였고, 실제로 anomaly 가 있는 부분만 집중적으로 복원을 하면서 false positive 를 줄였습니다. 또한, Noise 의 type은 결국 task 에 맞게 설계를 해야 한다는 결과를 보였습니다.