[논문 리뷰] Self-Attention Capsule Networks for Image Classification.

Assaf Hoogi, Brian Wilcox외 2명|arXiv (Cornell University)|2019. 04. 29.

Brain Tumor Detection and Classification인용 수 8

한 줄 요약

이 논문은 자기주의 주의(self-attention)를 캡슐 네트워크에 통합하여 특징 선택 및 장거리 의존성 모델링을 향상시키는 새로운 아키텍처인 자기주의 주의 캡슐 네트워크(Self-Attention Capsule Networks, SACN)를 제안한다. 캡슐 라우팅 이전에 주목할 만한 이미지 영역을 강조하기 위해 자기주의 주의를 사용함으로써, SACN은 다양한 데이터셋, 특히 의료 영상 및 자연 이미지 벤치마크에서 기준 모델인 CapsNet, ResNet-18, DenseNet-40보다 분류 정확도와 강인성을 높인다.

ABSTRACT

We propose a novel architecture for object classification, called Self-Attention Capsule Networks (SACN). SACN is the first model that incorporates the Self-Attention mechanism as an integral layer within the Capsule Network (CapsNet). While the Self-Attention mechanism supplies a long-range dependencies, results in selecting the more dominant image regions to focus on, the CapsNet analyzes the relevant features and their spatial correlations inside these regions only. The features are extracted in the convolutional layer. Then, the Self-Attention layer learns to suppress irrelevant regions based on features analysis and highlights salient features useful for a specific task. The attention map is then fed into the CapsNet primary layer that is followed by a classification layer. The proposed SACN model was designed to solve two main limitations of the baseline CapsNet - analysis of complex data and significant computational load. In this work, we use a shallow CapsNet architecture and compensates for the absence of a deeper network by using the Self-Attention module to significantly improve the results. The proposed Self-Attention CapsNet architecture was extensively evaluated on six different datasets, mainly on three different medical sets, in addition to the natural MNIST, SVHN and CIFAR10. The model was able to classify images and their patches with diverse and complex backgrounds better than the baseline CapsNet. As a result, the proposed Self-Attention CapsNet significantly improved classification performance within and across different datasets and outperformed the baseline CapsNet, ResNet-18 and DenseNet-40 not only in classification accuracy but also in robustness.

연구 동기 및 목표

• 복잡한 데이터 처리 및 높은 계산 부담 문제를 보완하기 위해 CapsNet의 한계를 해결한다.
• 주의 메커니즘을 통해 주목할 만한 이미지 영역에 초점을 맞춰 특징 표현을 향상시킨다.
• 복잡한 배경과 다양한 구조를 가진 데이터셋에서의 분류 성능을 향상시킨다.
• 자기주의 주의를 활용해 깊이 있는 특징 추상화를 가능하게 하여 캡슐 네트워크의 얕은 아키텍처를 보완한다.
• 다양한 벤치마크에서 CapsNet, ResNet-18, DenseNet-40과 비교해 뛰어난 강인성과 정확도를 입증한다.

제안 방법

• 컨볼루션 특징 추출 후 전용 레이어로 자기주의 주의 메커니즘을 통합하여 관련이 없는 이미지 영역을 억제한다.
• 주의 맵을 사용해 주요 캡슐 레이어를 안내하여 캡슐 라우팅을 주목할 만한 공간적 특징에 집중시킨다.
• 강조된 영역 내의 공간적 관계와 포즈 벡터를 모델링하기 위해 캡슐 네트워크를 적용한다.
• 주의 기반 특징 선택과 캡슐 라우팅을 결합하여 표현 학습을 향상시킨다.
• 깊은 네트워크를 피하기 위해 특징 추상화에 주의를 기반으로 하는 경량 아키텍처를 설계한다.
• 주의 손실 및 캡슐 손실 구성 요소를 포함한 표준 backpropagation을 사용해 엔드 투 엔드 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 자기주의 주의는 캡슐 네트워크에서 특징 선택 및 장거리 의존성 모델링을 향상시킬 수 있는가?
• RQ2 SACN은 복잡하고 다양한 데이터셋, 특히 혼잡한 배경을 가진 의료 영상에서 어떻게 성능을 발휘하는가?
• RQ3 SACN은 정확도와 강인성 측면에서 기준 모델인 CapsNet, ResNet-18, DenseNet-40보다 어느 정도 뛰어나게 성능을 발휘하는가?
• RQ4 얕은 캡슐 네트워크가 주의 기반 특징 집중을 통해 경쟁적인 성능을 낼 수 있는가?
• RQ5 자기주의 주의의 통합은 모델이 다양한 이미지 클래스와 도메인 간의 일반화 능력을 향상시키는 데 기여하는가?

주요 결과

• SACN은 평가된 모든 데이터셋에서 기준 모델인 CapsNet보다 분류 정확도를 크게 향상시켰다.
• 시험된 벤치마크에서 SACN은 정확도와 강인성 측면에서 ResNet-18과 DenseNet-40을 모두 뛰어넘었다.
• 복잡하고 혼잡한 배경을 가진 의료 영상 데이터셋에서 SACN은 뛰어난 성능을 보였다.
• 자기주의 주의의 통합으로 주목할 만한 특징에 더 잘 집중할 수 있었고, 관련 없는 이미지 영역의 간섭을 줄였다.
• 얕은 캡슐 네트워크임에도 불구하고 주의를 통해 효과적인 보완이 이루어져 강력한 성능을 달성했다.
• 이미지 콘텐츠의 변동성과 배경 복잡성에 대해 모델의 강인성이 향상되었다.