[논문 리뷰] Using StyleGAN for Visual Interpretability of Deep Learning Models on Medical Images

Kathryn Schutte, Olivier Moindrot외 3명|arXiv (Cornell University)|2021. 01. 19.

Explainable Artificial Intelligence (XAI)참고 문헌 20인용 수 36

한 줄 요약

논문은 입력 의학 영상이 모델 예측을 바꾸려면 어떻게 변화해야 하는지 시각화하기 위한 StyleGAN 기반 방법을 제안하며, GradCAM보다 더 해석 가능한 설명을 제공합니다. 무릎 X선 사진과 조직 샘플에서의 검증은 유망한 결과를 보여줍니다.

ABSTRACT

As AI-based medical devices are becoming more common in imaging fields like radiology and histology, interpretability of the underlying predictive models is crucial to expand their use in clinical practice. Existing heatmap-based interpretability methods such as GradCAM only highlight the location of predictive features but do not explain how they contribute to the prediction. In this paper, we propose a new interpretability method that can be used to understand the predictions of any black-box model on images, by showing how the input image would be modified in order to produce different predictions. A StyleGAN is trained on medical images to provide a mapping between latent vectors and images. Our method identifies the optimal direction in the latent space to create a change in the model prediction. By shifting the latent representation of an input image along this direction, we can produce a series of new synthetic images with changed predictions. We validate our approach on histology and radiology images, and demonstrate its ability to provide meaningful explanations that are more informative than GradCAM heatmaps. Our method reveals the patterns learned by the model, which allows clinicians to build trust in the model's predictions, discover new biomarkers and eventually reveal potential biases.

연구 동기 및 목표

• 의료 영상에서 해석가능한 AI에 대한 필요성을 제시하여 규제 및 임상 신뢰를 얻는 데 기여한다.
• StyleGAN을 이용해 잠재 공간을 이미지로 매핑하고 잠재 방향을 예측과 연결하는 3단계 방법을 제안한다.
• 잠재 공간 편집이 생물학적으로 타당한 진행 또는 예측 변화의 변화를 생성한다는 것을 보인다.
• GradCAM과 비교하여 더 풍부한 해석 가능성과 잠재 바이오마커 발견의 가능성을 시연한다.

제안 방법

• 의료 영상에서 StyleGAN2를 학습시켜 512 차원의 잠재 공간 W를 학습한다.
• 합성 쌍 (w, G(w))를 이용해 현실 이미지를 W의 해당 잠재 벡터로 매핑하는 인코더 E를 학습한다.
• 잠재 코드에 로지스틱 회귀를 학습시켜 모델 출력을 예측하게 하고, 예측을 변화시키는 W의 방향 α를 얻는다.
• 새로운 이미지 x에 대해 w = E(x)를 계산하고 λ*α만큼 따라가며 G(w + λ*α)를 생성해 업데이트된 예측을 얻는다.
• 생성된 이미지 시퀀스를 통해 시각적 변경이 예측과 어떤 관련이 있는지 해석하고 임상적으로 의미 있는 패턴을 식별한다.

실험 결과

연구 질문

• RQ1 StyleGAN의 잠재 공간 조작이 의학 영상에서 블랙박스 모델의 예측에 대한 설명을 제공할 수 있는가?
• RQ2 학습된 잠재 방향으로의 이동이 예측 변화와 상관관계가 있는 임상적으로 타당한 이미지 진화를 만들어 내는가?
• RQ3 생성된 시각적 진화가 의료 의사 결정 과제에서 GradCAM 히트맵보다 더 유용한가?
• RQ4 이 접근법이 의학 영상 분류기의 바이오마커 또는 잠재적 편향을 드러내는 데 도움이 될 수 있는가?

주요 결과

• X선에서 무릎 골관절염의 경우, 방법은 GradCAM 로컬라이제이션을 넘어 KL 등급이 증가함에 따라 생물학적으로 타당한 진행 특징을 산출한다.
• 무릎 X선에서 predictor f의 검정 AUC는 89%였고 잠재 공간 예측기 tilde f의 검정 AUC는 80%였다.
• 방사선 의사들이 StyleGAN 생성기의 X선 샘플의 현실성을 평가한 평균 Turing 테스트 정확도는 58%였다.
• 조직 샘플에서 원발 전이 탐지의 predictor f는 92%의 검정 AUC를 달성했고 잠재 예측기 tilde f는 95%의 검정 AUC를 달성했다.
• GradCAM 히트맵은 예측 특징이 이미지 전반에 확산될 때 덜 정보적이었고, 잠재 공간 편집은 의미 있는 특징 변화를 강조했다.
• 이 접근법은 임상의 신뢰를 구축하고 새로운 바이오마커를 발견하며 잠재적 편향을 드러내는 데 도움이 될 수 있다.