[논문 리뷰] HydraViT: Adaptive Multi-Branch Transformer for Multi-Label Disease Classification from Chest X-ray Images

Şaban Öztürk, Mehmet Y. Turali외 1명|arXiv (Cornell University)|2023. 10. 09.

COVID-19 diagnosis using AI인용 수 5

한 줄 요약

HydraViT는 다중 병변 흉부 X선(CXR) 영상 분류를 위한 새로운 적응형 다중 브랜치 트랜스포머 모델을 제안한다. 이 모델은 복합-공간 인코더와 트랜스포머 기반의 맥락 인코더를 조합하여 장거리 상관관계를 포착하고, 학습된 가중치를 가진 다중 브랜치 출력 모듈을 통해 개별 병변과 동시에 발생하는 병변에 대한 민감도를 향상시킨다. 이는 기존 방법보다 평균 1.2% 높은 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Chest X-ray is an essential diagnostic tool in the identification of chest diseases given its high sensitivity to pathological abnormalities in the lungs. However, image-driven diagnosis is still challenging due to heterogeneity in size and location of pathology, as well as visual similarities and co-occurrence of separate pathology. Since disease-related regions often occupy a relatively small portion of diagnostic images, classification models based on traditional convolutional neural networks (CNNs) are adversely affected given their locality bias. While CNNs were previously augmented with attention maps or spatial masks to guide focus on potentially critical regions, learning localization guidance under heterogeneity in the spatial distribution of pathology is challenging. To improve multi-label classification performance, here we propose a novel method, HydraViT, that synergistically combines a transformer backbone with a multi-branch output module with learned weighting. The transformer backbone enhances sensitivity to long-range context in X-ray images, while using the self-attention mechanism to adaptively focus on task-critical regions. The multi-branch output module dedicates an independent branch to each disease label to attain robust learning across separate disease classes, along with an aggregated branch across labels to maintain sensitivity to co-occurrence relationships among pathology. Experiments demonstrate that, on average, HydraViT outperforms competing attention-guided methods by 1.2%, region-guided methods by 1.4%, and semantic-guided methods by 1.0% in multi-label classification performance.

연구 동기 및 목표

• 다중 병변 흉부 X선(CXR) 영상에서 병변의 크기, 위치, 동시 발생의 변동성을 다루는 데 어려움이 있는 다중 레이블 병변 분류 문제를 해결한다.
• 기존의 컨volutional 네트워크(CNN)와 주의력 강화 모델이 병변 간 장거리 맥락과 동시 발생 패턴을 포착하는 데 한계를 보이는 문제를 해결한다.
• 개별 병변 레이블을 별도로 학습하고 병변 동시 발생 관계를 공동으로 모델링하는 다중 브랜치 출력 모듈을 통합하여 분류 성능을 향상시킨다.
• 자동화된 정확한 흉부 병변 진단을 위한 계산 효율적이면서도 높은 민감도를 지닌 딥 러닝 프레임워크를 개발한다.

제안 방법

• CXR 영상에서 초기 특징 맵을 추출하기 위해 복합-공간 인코더를 사용하고, 그 다음으로 영상 패치 간 장거리 상관관계를 모델링하기 위해 트랜스포머 기반 맥락 인코더를 적용한다.
• 트랜스포머 인코더의 자기주의 메커니즘을 활용하여 임무에 중요한 병변 관련 영역에 적응적으로 집중함으로써, 작거나 산산이 흩어진 병변에 대한 민감도를 향상시킨다.
• 14개 병변 레이블 각각에 대응하는 전용 헤드와 병변 간 동시 발생 관계를 모델링하기 위한 집합 헤드를 포함한 다중 브랜치 출력 모듈을 구현한다.
• 학습 중 개별 레이블 성능과 동시 발생 민감도를 균형 있게 유지하기 위해 다중 브랜치 출력 간의 학습된 적응형 가중치를 도입한다.
• 클래스 불균형 문제를 완화하고 일반화 성능을 향상시키기 위해 레이블별 손실 항목과 공동 손실 항목을 결합하고 가중치를 학습시킨다.
• 해석 가능성 확보를 위해 GradCAM을 적용하여 정확한 병변 영역에 주의가 집중되었는지 검증한다.

실험 결과

연구 질문

• RQ1 표준 CNN에 비해 병변이 작거나 산재한 CXR 영상에서 하이브리드 복합-트랜스포머 백본이 장거리 맥락 모델링을 얼마나 향상시킬 수 있는가?
• RQ2 적응형 가중치를 가진 다중 브랜치 출력 모듈이 동시에 개별 병변과 동시 발생 병변을 최적화함으로써 다중 레이블 CXR 분류 성능을 향상시키는가?
• RQ3 HydraViT는 최신 기술 수준의 주의력 유도, 영역 유도, 의미 유도 방법들과 비교하여 평균 및 레이블별 분류 성능에서 어떤가?
• RQ4 제안된 아키텍처가 주의 시각화를 통해 정렬 정확도와 모델의 해석 가능성에 얼마나 기여하는가?

주요 결과

• HydraViT는 CXR 데이터셋에서 경쟁 기반의 주의력 유도 방법보다 다중 레이블 분류 성능에서 평균 1.2% 높은 성능을 달성한다.
• 병변 레이블 14개 전반의 평균 F1 점수에서 지역 유도 방법보다 1.4% 향상되었고, 의미 유도 방법보다 1.0% 향상되었다.
• 예측 레이블과 진짜 레이블 간의 일치도가 뛰어나며, 진짜 병변에 대한 평균 예측 점수는 비현존 레이블보다 현저히 높게 나타났다(예: 상위 5개 예측에서 0.4 대비 0.1).
• 절단 실험 결과, 트랜스포머 기반 맥락 인코더와 다중 브랜치 출력 모듈이 성능 향상에 각각 기여하며, 특히 후자는 순차적 단일 브랜치 학습 대비 약 14배 빠른 학습 속도를 제공한다.
• GradCAM 열맵의 시각적 점검 결과, HydraViT가 병변 관련 영역을 효과적으로 강조하고 있음을 확인하여 주의 집중의 강력한 정확도를 입증한다.
• 희귀 병변이나 복잡한 동시 발생 패턴을 가진 병변 클래스에서도 높은 성능를 유지하여 일반화 능력 향상과 편향 감소를 시사한다.