기술 요약: 의료 이상 탐지를 위한 마스크 가이드 잠재 확산 기반의 재구성-변화 식별 (Reconstruction-Shift Discrimination via Mask-Guided Latent Diffusion for Medical Anomaly Detection)

문제 정의

비지도 학습 기반의 의료 이상 탐지는 정상 이미지만을 사용하여 병리적 소견을 식별하고 영향을 받은 영역을 국소화하는 것을 목표로 한다. 오토인코더(Autoencoder), GAN, 확산 모델(Diffusion Models)과 같은 재구성 기반 방식이 주를 이루고 있으나, 재구성 잔차(reconstruction residuals)를 유일한 이상 신호로 사용하는 데 있어 내재적인 모호함 문제를 겪는다. 표현력이 뛰어난 모델은 병리적 구조까지도 의도치 않게 보존하여 실제 병변에 대한 반응이 약해질 수 있다. 반대로, 양성 해부학적 변이, 영상 노이즈, 또는 획득 차이는 정상 영역에서 큰 재구성 오차를 생성할 수 있다. 이러한 잔차들을 단일 이미지 수준 점수로 집계하는 과정은 병리적 변화와 양성 재구성 아티팩트를 혼동하게 만들어, 실제 이상과 정상적인 변이를 구분하기 어렵게 만든다.

방법론: 판별적 마스크 가이드 확산 (Discriminative Mask-Guided Diffusion, DMD)

저자들은 잔차 기반 국소화(localization)를 보완하기 위해 **재구성-변화 식별(reconstruction-shift discrimination)**로부터 유도된 학습된 이미지 수준 신호를 제안하는 판별적 마스크 가이드 확산 (DMD) 프레임워크를 제안한다. 이 방법은 두 가지 별도의 단계로 작동한다.

1단계: 양자화된 정상 잠재 표현 (Quantized Normal Latent Representation)

정상 이미지에 대해 벡터 양자화 오토인코더(VQ-VAE)를 학습시켜 압축된 잠재 표현을 학습한다.

• 인코더 ($E_\phi$) 및 디코더 ($D_\eta$): 이미지를 잠재 공간으로 매핑하고 재구성한다.
• 벡터 양자화 (Vector Quantization): 공간적 잠재 벡터를 학습 가능한 코드북($C$)의 가장 가까운 항목으로 대체하여 이산 그리드 $z_q$를 생성한다.
• 목적 함수: 재구성 오차와 commitment loss를 최소화한다. 학습 후, 다음 단계를 위한 안정적인 잠재 공간을 보장하기 위해 $E_\phi$, $D_\eta$, $C$를 고정(freeze)한다.

2단계: 마스크 가이드 잠재 확산 + 재구성-변화 식별

이 단계에서는 고정된 인코더/디코더를 사용하여 잠재 확산 모델과 분류기를 공동 학습시킨다.

1. 마스크 가이드 잠재 확산:

   • 국소적인 정사각형 마스크($M$)가 양자화된 잠재 그리드 $z_0$에 적용되어 특정 영역을 섭동(perturbation) 대상으로 선택한다.
   • 순방향 확산 과정(forward diffusion process)은 마스크 처리된 영역에만 가우시안 노이즈를 추가하며, 마스크 처리되지 않은 영역은 깨끗한 컨텍스트(context)로 남겨둔다.
   • 디노이징 네트워크($\epsilon_\omega$)는 노이즈를 예측하고 깨끗한 잠재 상태 $\hat{z}_0$를 재구성한다.
   • 디코더는 이 복원된 잠재 상태를 다시 이미지 공간으로 매핑하여, 원래의 정상 이미지($x$)에 대한 "재구성된 대응물($x_r$)"을 생성한다.
   • 확산 손실($L_{diff}$)은 마스크 처리된 위치에 대해서만 계산된다.

2. 재구성-변화 식별 (Reconstruction-Shift Discrimination):

   • 자가 지도 학습 기반의 분류 작업이 정의된다. 여기서 원래의 정상 이미지 $x$는 클래스 0으로, 마스크 확산에 의해 "변화(shift)"가 유도된 재구성 대응물 $x_r$은 클래스 1로 레이블링된다.
   • 분류기($f_\psi$)는 원래 이미지와 변화된 재구성본을 구별하는 법을 학습한다.
   • 분류 손실($L_{cls}$)은 디노이징 네트워크에 판별적 피드백을 제공하여, 모델이 입력값과는 구별되면서도 해부학적으로 타당한 재구성을 생성하도록 유도한다.
   • 전체 손실은 $L_{stage2} = L_{diff} + \lambda L_{cls}$이다.

이중 수준 추론 (Dual-Level Inference)

테스트 시, DMD는 두 가지 상호 보완적인 출력을 제공한다.

1. 이미지 수준 탐지: 분류기 $f_\psi$는 테스트 이미지를 직접 처리하여, 해당 이미지가 정상적인 원본보다는 "재구성 변화(reconstruction shift)"와 유사할 확률을 나타내는 이상 점수 $S_{img}$를 출력한다.
2. 픽셀 수준 국소화: 테스트 이미지를 인코딩하고 양자화한다. 잠재 그리드 전체에 걸쳐 중첩되는 마스크를 적용하고, 디노이저는 각 마스크에 대해 이미지를 재구성한다. 픽셀 수준의 이상 맵은 입력과 다양한 재구성 대응물 사이의 잔차(절대 차이)를 계산하고, 중첩된 영역에 대해 평균을 내어 생성된다.

주요 기여

1. 재구성-변화 식별: 저자들은 마스크 확산 재구성을 통해 이미지 수준의 이상 점수를 학습하는 정상 전용 자가 지도 신호를 식별하고 구현하였다. 이는 병리적 데이터 없이도 전통적인 잔차 기반 국소화를 보완한다.
2. DMD 프레임워크: 국소 잠재 마스킹과 잠재 확산이 이미지 수준 탐지와 픽셀 수준 국소화를 동시에 구동하는 통합된 아키텍처이다. 이 프레임워크는 단일 복원 과정을 사용하여 판별적 신호(분류기를 위한)와 공간적 단서(잔차 맵을 위한)를 모두 생성한다.
3. 최첨단 성능 (State-of-the-Art Performance): 이 방법은 다섯 가지 다양한 의료 영상 데이터셋(뇌 MRI, 유방 초음파, 흉부 방사선 촬영)에서 우수한 성능을 입증하였다.

실험 결과

논문은 DMD를 BraTS2021, IXI/ATLAS 2.0, WMH, BUSI, VinDr-CXR의 다섯 가지 데이터셋에 대해 평가하였다.

• 이미지 수준 탐지: DMD는 모든 데이터셋에서 가장 높은 AUC, AP, F1 점수를 달성했다. 예를 들어, BraTS2021에서 DMD는 AUC **87.2%**를 기록하여 가장 강력한 베이스라인인 MAD-AD를 1.1% 앞질렀으며, F1 점수는 3.1% 향상되었다. VinDr-CXR에서는 AUC **84.3%**에 도달했고, BUSI에서는 AUC 93.8% 및 AP **96.0%**를 달성했다.
• 픽셀 수준 국소화: DMD는 AP$_{pix}$ 및 Dice 점수 측면에서 고전적인 재구성 방식(AE, VAE, f-AnoGAN)과 최신 확산 기반 베이스라인(AnoDDPM, AutoDDPM, MAD-AD)을 일관되게 능가했다. BraTS2021에서 DMD는 **76.3%**의 Dice 점수를 달성하였으며, 이는 차순위 방법인 MAD-AD(69.6%)보다 현저히 높은 수치이다.
• 교차 데이터셋 강건성: BraTS2021으로 학습된 모델을 재학습 없이 WMH 데이터셋에 적용했을 때도 DMD는 강력한 성능을 유지하였으며, 이는 서로 다른 병변 유형 및 획득 특성에 대한 일반화 가능성을 시사한다.
• 정성적 분석: 시각화 결과, DMD는 특히 미세하고 저대조인 이상 부위에서 다른 방법들이 산발적이거나 약한 반응을 보이는 것과 달리, 실제 병변 경계와 더 잘 일치하는 더 선명하고 공간적으로 일관된 이상 맵을 생성함을 보여주었다.
• 절제 연구 (Ablation Studies): 분류기, 판별적 피드백, 사전 학습 단계, 마스킹 전략 또는 잠재 확산 구성 요소 중 하나라도 제거하면 성능 저하가 발생하였으며, 이는 각 구성 요소의 필요성을 확인시켜 준다. 특히 잠재 확산을 제거했을 때 가장 큰 성능 하락이 관찰되었는데, 이는 해부학적 변이를 모델링하는 데 있어 확산의 역할을 강조한다.

의의 및 주장

본 논문은 DMD가 단일 재구성 잔차를 탐지와 국소화 모두에 사용하는 근본적인 모호함을 해결한다고 주장한다. 분류기를 통한 이미지 수준 판별과 잔차를 이용한 픽셀 수준 국소화로 과업을 분리함으로써, 이 프레임워크는 두 접근 방식의 장점을 모두 활용한다.

저자들은 이미지 수준의 신호가 "재구성 변화" 개념을 통해 오직 정상 훈련 데이터로부터 완전히 유도되었으며, 병리적 주석이 필요하지 않다는 점을 강조한다. 본 방법이 최첨단 결과를 달성했음에도 불구하고, 저자들은 이미지 수준의 신호가 고정된 마스킹 정책에 의존한다는 점을 한계로 언급하며, 미지의 병리 및 획득 조건에 대한 전이 가능성에 대해서는 추가적인 광범위한 연구가 필요하다고 조심스럽게 덧붙였다. 향후 연구로는 적응형 변화 구축 및 다기관 평가가 제안되었다.