[논문 리뷰] InceptionCapsule: Inception-Resnet and CapsuleNet with self-attention for medical image Classification

Elham Sadeghnezhad, Sajjad Salem|arXiv (Cornell University)|2024. 02. 03.

COVID-19 diagnosis using AI인용 수 4

한 줄 요약

이 논문은 의료 영상 분류를 위한 하이브리드 딥 러닝 모델인 InceptionCapsule을 제안한다. 이 모델은 Inception-ResNet, 캡슐 네트워크, 그리고 자기주의(self-attention)를 결합한 것으로, 사전 훈련된 ImageNet 가중치를 활용하고 중간층 특징을 활용하여 풍부한 벡터 표현을 생성함으로써 분류 정확도를 향상시킨다. 이로 인해 BUSI with GT에서 98.88%의 정확도를 기록했으며, 5개 클래스 분류 과제에서 Kvasir에서는 97.62%의 정확도를 달성하였다.

ABSTRACT

Initial weighting is significant in deep neural networks because the random selection of weights produces different outputs and increases the probability of overfitting and underfitting. On the other hand, vector-based approaches to extract vector features need rich vectors for more accurate classification. The InceptionCapsule approach is presented to alleviate these two problems. This approach uses transfer learning and the Inception-ResNet model to avoid random selection of weights, which takes initial weights from ImageNet. It also uses the output of Inception middle layers to generate rich vectors. Extracted vectors are given to a capsule network for learning, which is equipped with an attention technique. Kvasir data and BUSI with the GT dataset were used to evaluate this approach. This model was able to achieve 97.62 accuracies in 5-class classification and also achieved 94.30 accuracies in 8-class classification on Kvasir. In the BUSI with GT dataset, the proposed approach achieved accuracy=98.88, Precision=95.34, and F1-score=93.74, which are acceptable results compared to other approaches in the literature.

연구 동기 및 목표

• 의료 영상 분류를 위한 깊은 신경망에서 무작위 가중치 초기화로 인한 불안정성과 과적합 문제를 해결하기 위해.
• 중간층의 Inception-ResNet 레이어에서 유래한 풍부한 계층적 벡터 표현을 생성함으로써 특징 표현을 향상시키기 위해.
• 동적 특징 가중치 부여를 위한 자기주의를 갖춘 캡슐 네트워크를 통해 분류 성능을 향상시키기 위해.
• Kvasir 및 BUSI with GT와 같은 다양한 의료 영상 기준 데이터셋에서 모델을 평가하여 강건성과 일반화 능력을 입증하기 위해.
• 주목할 만한 주의, 잔차 학습, 캡슐 기반 특징 학습의 융합을 통해 다중 클래스 의료 영상 분류에서 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• 무작위 초기화를 방지하기 위해 사전 훈련된 ImageNet 가중치로 Inception-ResNet 기반 모델을 초기화함으로써 전이 학습을 활용한다.
• Inception-ResNet 아키텍처의 중간층에서 깊은 특징을 추출하여 풍부하고 구분력 있는 벡터 표현을 생성한다.
• 추출된 특징 벡터를 캡슐 네트워크에 입력하여 계층적 공간 관계를 모델링하고 일반화 능력을 향상시킨다.
• 중요한 특징을 동적으로 강조하고 노이즈를 억제하기 위해 캡슐 네트워크 내부에 자기주의 메커니즘을 통합한다.
• 공간 일관성을 바탕으로 의미 있는 캡슐로 특징을 그룹화하기 위해 캡슐 레이어에서 '의견에 의한 라우팅' 메커니즘을 적용한다.
• 과적합을 줄이고 강건성을 향상시키기 위해 레이블 스무딩을 적용한 교차 엔트로피 손실을 사용하여 엔드 투 엔드로 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 Inception-ResNet를 캡슐 네트워크와 자기주의와 융합함으로써 의료 영상 분류에서 특징 표현을 향상시킬 수 있는가?
• RQ2 Inception-ResNet의 중간층 특징을 사용하는 것이 의료 영상 작업에서 표준 특징 추출 방식보다 더 높은 성능을 내는가?
• RQ3 캡슐 네트워크에 자기주의를 통합함으로써 의료 영상 데이터셋에서 분류 정확도와 강건성에 어떤 영향을 미치는가?
• RQ4 제안된 InceptionCapsule 모델이 Kvasir 및 BUSI with GT 데이터셋에서 기존 최신 기술 수준의 방법들을 초월할 수 있는가?
• RQ5 ImageNet에서의 전이 학습이 의료 영상 분류에서 과적합을 줄이고 수렴을 가속화하는 데 얼마나 기여하는가?

주요 결과

• InceptionCapsule 모델은 5개 클래스 분류 과제에서 Kvasir 데이터셋에서 97.62%의 정확도를 기록하여 다중 클래스 의료 영상 분류에서 뛰어난 성능을 입증하였다.
• 8개 클래스 분류 과제에서 Kvasir 데이터셋에서 모델은 94.30%의 정확도를 달성하여 복잡한 다중 레이블 시나리오로의 확장성을 보였다.
• BUSI with GT 데이터셋에서 모델은 98.88%의 정확도, 95.34%의 정밀도, 93.74%의 F1 점수를 기록하여 많은 기존 방법들을 능가하는 성능을 보였다.
• 캡슐 네트워크에 자기주의를 통합함으로써 특징의 구분 능력이 크게 향상되었으며, 특히 복잡하거나 노이즈가 많은 의료 영상에서 두드러졌다.
• 사전 훈련된 Inception-ResNet 가중치의 사용은 훈련의 불안정성을 줄이고 수렴 속도를 가속화하여 더 나은 일반화에 기여하였다.
• 모델은 다양한 의료 영상 작업에서 강건한 성능을 보였으며, 주목, 잔차 학습, 캡슐 기반 특징 학습의 융합이 효과적임을 확인하였다.