[논문 리뷰] DiffMIC: Dual-Guidance Diffusion Network for Medical Image Classification

Yijun Yang, Huazhu Fu외 3명|arXiv (Cornell University)|2023. 03. 19.

Machine Learning in Healthcare인용 수 4

한 줄 요약

이 논문은 일반적인 의료 영상 분류를 위한 최초의 확산 기반 모델인 DiffMIC을 제안한다. 이 모델은 이중 해상도 조건부 가이던스와 조건별 MMD 정규화를 사용하여 영상의 노이즈를 제거하고 강력한 의미적 표현을 학습한다. DiffMIC은 태반 성숙도 평가, 피부 병변 분류, 당뇨성 망막병변 평가 등 세 가지 다양한 의료 영상 과제에서 최신 기술을 압도적으로 뛰어넘는 성능을 달성한다.

ABSTRACT

Diffusion Probabilistic Models have recently shown remarkable performance in generative image modeling, attracting significant attention in the computer vision community. However, while a substantial amount of diffusion-based research has focused on generative tasks, few studies have applied diffusion models to general medical image classification. In this paper, we propose the first diffusion-based model (named DiffMIC) to address general medical image classification by eliminating unexpected noise and perturbations in medical images and robustly capturing semantic representation. To achieve this goal, we devise a dual conditional guidance strategy that conditions each diffusion step with multiple granularities to improve step-wise regional attention. Furthermore, we propose learning the mutual information in each granularity by enforcing Maximum-Mean Discrepancy regularization during the diffusion forward process. We evaluate the effectiveness of our DiffMIC on three medical classification tasks with different image modalities, including placental maturity grading on ultrasound images, skin lesion classification using dermatoscopic images, and diabetic retinopathy grading using fundus images. Our experimental results demonstrate that DiffMIC outperforms state-of-the-art methods by a significant margin, indicating the universality and effectiveness of the proposed model. Our code will be publicly available at https://github.com/scott-yjyang/DiffMIC.

연구 동기 및 목표

• 다양한 모ality에서 노이즈가 많고 흐릿하며 모호한 의료 영상이 분류 과제에 미치는 영향을 해결하기 위해.
• 확산 확률 모델의 잠재력을 생성 과제를 넘어서 고수준 비전, 특히 의료 영상 분류에 적용하기 위해.
• 스토케스틱한 확산 과정을 통해 영상의 노이즈를 제거하면서도 중요한 해부학적 세부 정보를 유지함으로써 특징 표현을 향상시키기 위해.
• 이중 해상도 조건부 가이던스를 통해 전반적이고 국소적인 사전 지식을 통합하여 분류의 강건성을 향상시키기 위해.
• 조건별로 다른 MMD 정규화를 사용하여 영상 수준과 패치 수준의 특징 간에 분리되고 공유되며 강건한 표현을 학습하기 위해.

제안 방법

• 이중 해상도 조건부 가이던스(DCG) 전략은 각 확산 단계를 전반적 영상 수준과 국소 영역(관심 영역, ROI) 수준의 사전 지식에 기반하여 조절하여 지역적 주의력과 세밀한 특징 학습을 향상시킨다.
• 모델은 전체 영상, 전반적 사전 지식, 국소 사전 지식의 세 입력에 대해 확산 과정을 적용하며, U-Net 아키텍처 내에서 각각의 노이즈 예측 헤드를 갖춘다.
• 전방 확산 과정 중에 조건별 최대 평균 차이(MMD) 정규화를 적용하여 영상 수준과 패치 수준의 잠재 표현 간의 상호정보량을 최대화한다.
• 노이즈 예측은 병합된 노이즈가 포함된 잠재 표현에 대해 평균 제곱 오차(MSE) 손실을 통해 학습되며, MMD 손실은 전반적 및 국소 브랜치의 예측된 노이즈에 적용되어 분포를 일치시킨다.
• 영상 인코더는 특징 임베딩을 추출하며, 이는 확산 U-Net에서 복원된 잠재 특징과 융합되어 최종 분류 로짓을 생성한다.
• 역방향 확산 과정은 특징을 점차적으로 노이즈 제거하며, 추론 시 더 명확하고 클래스 간으로 분리된 표현을 가능하게 하며, t-SNE를 통해 시각화된다.

실험 결과

연구 질문

• RQ1 확산 기반 모델이 노이즈 제거와 의미적 표현 향상으로써 일반적인 의료 영상 분류 성능을 효과적으로 향상시킬 수 있는가?
• RQ2 전반적 및 국소 사전 지식을 모두 활용하는 이중 해상도 조건부 가이던스는 의료 영상에서 주의력과 분류 성능을 어떻게 향상시키는가?
• RQ3 조건별 MMD 정규화는 확산 과정 중에 잠재 공간에서 특징의 분리성과 강건성을 얼마나 향상시키는가?
• RQ4 제안된 DiffMIC 모델은 초음파, 피부 내시경, 망막 영상 등 다양한 의료 영상 모달리티에 일반화되는가?
• RQ5 기존의 CNN과 비전 트랜스포머에 비해 확산 모델 통합은 불균형한 의료 데이터셋에서 정확도와 강건성 측면에서 어떻게 비교되는가?

주요 결과

• 태반 성숙도 평가를 위한 PMG2000 데이터셋에서, DiffMIC은 정확도 0.931과 F1 스코어 0.926을 달성하며, 두 번째로 좋은 성능을 보인 PVT보다 정확도 0.024, F1 스코어 0.024 높게 기록했다.
• 피부 병변 분류를 위한 HAM10000 데이터셋에서, DiffMIC은 두 번째로 좋은 성능을 보인 ProCo보다 정확도 0.019, F1 스코어 0.053 높게 기록했다.
• 당뇨성 망막병변 평가를 위한 APTOS2019 데이터셋에서, DiffMIC은 ProCo보다 정확도 0.021, F1 스코어 0.042 높게 기록하여 불균형 데이터에서 강력한 성능을 보였다.
• 절단 실험 결과, 각 구성 요소인 확산, DCG, MMD 정규화가 성능 향상에 점진적으로 기여하며, 기준 모델(ResNet18) 대비 정확도 0.052 향상된 전체 모델이 가장 우수한 성능을 보였다.
• t-SNE 시각화 결과, 확산 과정이 진행될수록 특징 임베딩이 점점 더 클래스 간으로 분리되는 경향을 보이며, 효과적인 노이즈 제거와 표현 정련이 이루어졌음을 시사한다.
• 모델은 초음파(태반), 피부 내시경(피부 병변), 망막(당뇨성 망막병변) 영상 등 세 가지의 서로 다른 의료 영상 모달리티에 대해 일반화 능력을 보이며, 그 보편성을 확인했다.