[논문 리뷰] K-Diag: Knowledge-enhanced Disease Diagnosis in Radiographic Imaging

Chaoyi Wu, Xiaoman Zhang외 3명|arXiv (Cornell University)|2023. 02. 22.

AI in cancer detection인용 수 7

한 줄 요약

K-Diag는 시각적 표현 학습 과정에 영역 지식을 통합하고 대조 학습을 통해 의료 지식 그래프를 통합함으로써 지식 기반 프레임워크를 제안한다. 또한 가속 가능한 프롬프트 벡터를 사용하여 시각적 표현을 적응시킨다. 이는 장수익성 및 제로샷 질병 인식에서 최신 기술 성능을 달성하며, 제로샷 평가에서 106개의 알려지지 않은 방사선 영상 소견 중 79개에서 AUC >0.600을 기록한다.

ABSTRACT

In this paper, we consider the problem of disease diagnosis. Unlike the conventional learning paradigm that treats labels independently, we propose a knowledge-enhanced framework, that enables training visual representation with the guidance of medical domain knowledge. In particular, we make the following contributions: First, to explicitly incorporate experts' knowledge, we propose to learn a neural representation for the medical knowledge graph via contrastive learning, implicitly establishing relations between different medical concepts. Second, while training the visual encoder, we keep the parameters of the knowledge encoder frozen and propose to learn a set of prompt vectors for efficient adaptation. Third, we adopt a Transformer-based disease-query module for cross-model fusion, which naturally enables explainable diagnosis results via cross attention. To validate the effectiveness of our proposed framework, we conduct thorough experiments on three x-ray imaging datasets across different anatomy structures, showing our model is able to exploit the implicit relations between diseases/findings, thus is beneficial to the commonly encountered problem in the medical domain, namely, long-tailed and zero-shot recognition, which conventional methods either struggle or completely fail to realize.

연구 동기 및 목표

• 기존의 감독 학습 기반 접근 방식이 질병 레이블을 상호 독립적이고 수직적인 것으로 간주하여 내재된 의미적 관계를 忽略하는 한계를 해결하기 위해.
• 관련된 의료 개념 간의 공유 지식을 활용하여 장수익성 질병 인식의 일반화 성능을 향상시키기 위해.
• 시각적 표현 학습 과정에 도메인 지식을 통합하여 기존에 볼 수 없었던 질병에 대한 제로샷 진단을 가능하게 하기 위해.
• 질병 쿼리와 시각적 특징 간의 교차 어텐션 메커니즘을 통해 설명 가능한 진단 프레임워크를 개발하기 위해.

제안 방법

• 의료 지식 그래프에서 대조 학습을 통해 질병 및 소견 개념을 연속적이고 관계 인식 가능한 공간에 임bedding하는 지식 인코더를 훈련한다.
• 시각적 인코더는 고정된 지식 인코더 파라미터를 사용하여 훈련되며, 가속 가능한 프롬프트 벡터 모듈을 통해 하류 작업에 효율적으로 적응시킨다.
• Transformer 기반의 질병 쿼리 모듈이 질병 이름(쿼리)과 시각적 특징 간의 교차 어텐션을 수행하여 설명 가능한 어조 가중 예측을 생성한다.
• 모델은 일항 벡터가 아닌 지식 그래프에서 유도된 연속적 레이블 임베딩을 사용하여 병변 간 의미적 관계를 활용할 수 있다.
• 모델은 시각적 특징과 지식 기반 질병 기술 간의 정렬을 위해 대조 목표를 사용하여 엔드 투 엔드로 훈련된다.
• 추론 시에 직접적으로 알려지지 않은 질병 이름을 지식 공간에 인코딩함으로써 소수 샘플 및 제로샷 추론을 모두 지원한다.

Figure 1: In the conventional training scheme (left), manual annotations are often converted into discrete one-hot vectors, that are orthogonal in the embedding space, thus ignoring the implicit relations between labels. While in our proposed knowledge-enhanced classification framework (right), the

실험 결과

연구 질문

• RQ1 시각적 표현 학습 과정에 의료 도메인 지식을 통합하면 방사선 영상에서 장수익성 질병 인식 성능 향상에 기여하는가?
• RQ2 지식 기반 프레임워크는 훈련 중에 볼 수 없었던 질병에 대한 제로샷 진단을 가능하게 하는가?
• RQ3 일항 인코딩과 비교하여 연속적이고 지식 기반 레이블 표현 방식이 질병 간 의미적 관계를 얼마나 잘 포착하는가?
• RQ4 지식 기반 모델이 다양한 해부학적 구조와 다양한 애너테이션 세부 수준을 가진 데이터셋 간에 얼마나 잘 일반화되는가?
• RQ5 질병 쿼리 모듈의 어텐션 메커니즘이 모델 예측에 대해 설명 가능한 통찰을 제공할 수 있는가?

주요 결과

• 세 개의 X-ray 데이터셋에서 K-Diag는 일항 레이블을 사용한 기준 모델 대비 ResNet 기준 2.53%, ViT 기준 2.15%의 AUC 향상을 기록했다.
• 11개의 공개 흉부 X-ray 데이터셋을 통합한 데이터셋에서 훈련한 결과, K-Diag는 기준 모델 대비 평균 AUC가 2.14% 높아져 더 나은 데이터 활용도를 보였다.
• PadChest에서의 제로샷 평가에서 K-Diag는 106개의 알려지지 않은 방사선 영상 소견 중 79개에서 최소 AUC 0.600을 달성했으며, 이는 기존 감독 학습 접근 방식으로는 달성할 수 없는 결과였다.
• 제로샷 설정에서 14개의 소견에 대해 AUC가 최소 0.800 이상, 46개의 소견에 대해 최소 0.700 이상을 기록하여 희귀 또는 새로운 병변에 대한 강력한 일반화 능력을 보였다.
• 대상 클래스 데이터를 늘리지 않고도 클래스 다양성(알려지지 않은 질병)을 추가함으로써 모든 11개 데이터셋에서 성능 향상을 기록했으며, 특히 GoogleNIH와 SIIM-ACR와 같은 소규모 데이터셋에서 두드러진 성능 향상을 보였다.
• 절단 실험을 통해 지식 인코더와 프롬프트 학습 구성 요소가 필수적임을 확인했으며, 이 중 하나를 제거할 경우 성능이 크게 떨어졌다.

Figure 2: Overview of the knowledge-enhanced disease diagnosis workflow. The knowledge encoder (left) is first trained to learn a neural representation of the medical knowledge graph via contrastive learning, and then used to guide the visual representation learning in our knowledge-enhanced classif