[논문 리뷰] Dynamic Routing on Deep Neural Network for Thoracic Disease Classification and Sensitive Area Localization

Yan Shen, Mingchen Gao|arXiv (Cornell University)|2018. 08. 17.

COVID-19 diagnosis using AI참고 문헌 15인용 수 4

한 줄 요약

이 논문은 흉부 X선에서 흉부 질환 분류 및 국소화를 위한 밀도 신경망 내 1×1 컨볼루션 레이어에 통합된 새로운 동적 라우팅 메커니즘을 제안한다. 라우팅-협의성 기반 기법을 컨볼루션 특징 맵으로 확장하고 Grad-CAM을 활용해 해석 가능성을 높임으로써, 더 적은 층 수로도 최신 기준 AUC 성능을 달성하며, 표준 컨볼루션 기반 모델 대비 평균 정확도에서 15% 향상된 국소화 정확도를 확보한다.

ABSTRACT

We present and evaluate a new deep neural network architecture for automatic thoracic disease detection on chest X-rays. Deep neural networks have shown great success in a plethora of visual recognition tasks such as image classification and object detection by stacking multiple layers of convolutional neural networks (CNN) in a feed-forward manner. However, the performance gain by going deeper has reached bottlenecks as a result of the trade-off between model complexity and discrimination power. We address this problem by utilizing the recently developed routing-by agreement mechanism in our architecture. A novel characteristic of our network structure is that it extends routing to two types of layer connections (1) connection between feature maps in dense layers, (2) connection between primary capsules and prediction capsules in final classification layer. We show that our networks achieve comparable results with much fewer layers in the measurement of AUC score. We further show the combined benefits of model interpretability by generating Gradient-weighted Class Activation Mapping (Grad-CAM) for localization. We demonstrate our results on the NIH chestX-ray14 dataset that consists of 112,120 images on 30,805 unique patients including 14 kinds of lung diseases.

연구 동기 및 목표

• 모델 복잡도와 분류 능력 간의 상충 관계로 인한 딥 컨볼루션 신경망의 성능 저하 문제를 해결하기 위해.
• 1×1 컨볼루션 레이어에 새로운 동적 라우팅 메커니즘을 도입하여 NIH ChestX-ray14 데이터셋에서 흉부 질환 분류 정확도를 향상시키기 위해.
• 감정 영역 국소화를 위한 고해상도 Grad-CAM 히트맵을 생성함으로써 모델의 해석 가능성을 향상시키기 위해.
• 동일 환자 이미지가 각 세트 간에 유출되지 않도록 공식적인 훈련/검증/테스트 분할 프로토콜을 따르며 공정한 평가를 보장하기 위해.
• 1×1 컨볼루션 레이어에 적용된 동적 라우팅이 더 적은 층 수로도 우수한 특징 선택 및 분류 성능을 달성할 수 있음을 입증하기 위해.

제안 방법

• 밀도 블록 내 1×1 컨볼루션 레이어에 동적 라우팅을 확장하여, 반복적 협의를 통해 라우팅 계수를 학습함으로써 관련 특징 맵을 선택한다.
• 라우팅을 효율적으로 계산하기 위해 커널 기법을 사용하여 훈련 및 추론 시간을 단일 레이어 수준에 가깝게 줄인다.
• 다운샘플링 블록은 7×7 및 3×3 최대 풀링을 거친 후 1×1 컨볼루션과 2×2 평균 풀링을 적용하여 밀도 라우팅 레이어로 이어진다.
• 32×32 해상도의 주요 캡슐에서 Grad-CAM을 생성하여 국소화를 위해 공간적 세부 정보를 유지한다. 이는 4×4 평균 풀링 레이어 이전 단계에서 수행된다.
• 라우팅 메커니즘은 밀도 레이어 내 특징 맵 간 및 최종 분류 레이어의 주요 캡슐과 예측 캡슐 간에 작동한다.
• 표준 교차 엔트로피 손실을 사용하여 엔드 투 엔드로 학습하며, 추론 시에는 학습 가능한 계수를 가진 최종 라우팅 반복을 사용한다.

실험 결과

연구 질문

• RQ1 표준 컨볼루션 대비 1×1 컨볼루션 레이어에 적용된 동적 라우팅이 흉부 X선 영상 분류 성능 향상에 기여하는가?
• RQ2 밀도 레이어 연결에 라우팅을 확장함으로써 특징 표현력과 모델 정확도가 향상되는가?
• RQ3 저해상도 주요 캡슐에서 생성된 Grad-CAM이 여전히 질환 관련 영역을 정확히 국소화할 수 있는가?
• RQ4 공식 데이터 분할을 사용할 때, 제안된 모델이 최신 기준 모델 대비 AUC 및 정확도 측면에서 어떻게 성능을 내는가?
• RQ5 Grad-CAM을 통한 모델의 해석 가능성은 방사선의학 전문의가 앵커링한 경계 상자와 얼마나 일치하는가?

주요 결과

• 제안된 모델은 이전 모델들보다 훨씬 적은 층 수로 NIH ChestX-ray14 데이터셋에서 최신 기준 AUC 성능을 달성한다.
• 1×1 컨볼루션 캡슐을 표준 1×1 컨볼루션으로 대체한 경우 평균 분류 정확도가 15% 감소하여, 동적 라우팅의 우수성을 입증한다.
• 32×32 주요 캡슐에서 생성된 Grad-CAM 히트맵은 질환 영역을 정확히 국소화하며, Atelectasis의 경우 69.77%의 겹침 비율(IoBB ≥ 0.1)과 Cardiomegaly의 경우 83.33%를 기록한다.
• Atelectasis의 경우, 한 환자에서는 상좌측 폐 병변을 정확히 국소화했고, 다른 환자에서는 상우측 병변을 정확히 식별하여 임상 앵커링 결과와 일치한다.
• Atelectasis의 국소화 정확도는 IoBB ≥ 0.25 조건에서 45.34%로 측정되었고, Cardiomegaly의 경우 82.77%로 지표상의 경계 상자와 강한 일치를 보였다.
• 모델는 전역 평균 풀링 이전에 Grad-CAM을 생성함으로써 정확도를 희생시키지 않은 채 높은 해석 가능성을 유지한다.