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

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

Advanced Neural Network Applications참고 문헌 22인용 수 20

한 줄 요약

이 논문은 복합적인 의료 및 자연 이미지 데이터셋에서 특히 뛰어난 분류 정확도와 강건성을 달성하면서 기존 CapsNets 및 최신 모델인 ResNet-18과 DenseNet-40에 비해 계산 부담을 줄이는, 컨볼루션층과 캡슐층 사이에 자기주의 메커니즘을 통합한 새로운 아키텍처인 Self-Attention Capsule Networks(SACN)을 제안한다. 자기주의 메커니즘을 통해 캡슐 처리 이전에 주목할 만한 이미지 영역에 집중함으로써, 특징 선택과 공간 상관관계 모델링을 향상시킨다.

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.

연구 동기 및 목표

• 복잡하고 혼잡한 이미지를 다루는 데에서 CapsNets의 한계와 높은 계산 비용을 해결하기 위해.
• 캡슐 처리 이전에 주목할 만한 이미지 영역을 식별하여 특징 선택을 향상시키기 위해.
• annotation 데이터가 제한된 소규모이고 다양한 의료 데이터셋에서 성능을 향상시키기 위해.
• 자기주의로 강화된 얕은 캡슐 아키텍처를 통해 계산 부담을 줄이면서도 높은 정확도를 유지하기 위해.
• 의료 및 자연 이미지 도메인 전반에서의 일반화 능력을 입증하기 위해.

제안 방법

• 컨볼루션층과 주요 캡슐층 사이에 자기주의 메커니즘을 중간층으로 통합한다.
• 쿼리, 키, 밸류 행렬을 사용하여 관련이 없는 이미지 영역을 억제하는 자기주의 점수를 계산한다.
• 스케일된 도트곱 자기주의를 적용하여 주요 특징을 강조하는 자기주의 맵을 생성한다.
• 자기주의 맵을 주요 캡슐층을 통해 처리하여 공간적 관계와 객체 자세 정보를 유지한다.
• 계산 비용을 줄이기 위해 얕은 캡슐 네트워크 아키텍처를 사용한다.
• Adam 옵timizer와 가중치 감소를 사용하여 전체 네트워크를 엔드 투 엔드로 훈련한다.

실험 결과

연구 질문

• RQ1 자기주의가 캡슐 네트워크의 특징 선택에 복잡한 이미지 분류에서 향상시킬 수 있는가?
• RQ2 CapsNets에 자기주의를 통합하면 배경 변동성이 높은 의료 영상 데이터셋에서 성능 향상이 이루어지는가?
• RQ3 SACN은 계산 비용을 줄이며 기존 CapsNet, ResNet-18, DenseNet-40보다 더 높은 정확도를 달성할 수 있는가?
• RQ4 SACN은 자연 이미지와 복잡한 의료 영상 스캔을 포함한 다양한 데이터셋에서 일반화 능력을 보여주는가?
• RQ5 자기주의 메커니즘이 가림, 저대비 및 이미지 손상 상황에서 강건성을 향상시키는가?

주요 결과

• SACN은 MNIST에서 99.5%의 분류 정확도를 달성하여 최신 기술 수준의 성능을 보였다.
• SVHN에서 SACN은 기존 CapsNet 정확도를 2.4% 향상시켜 실제 도메인에서의 도전적인 데이터에서의 효과를 입증했다.
• CIFAR-10에서 SACN은 기존 CapsNet 정확도를 3.5% 향상시켜 자연 이미지 분류에 대한 강력한 일반화 능력을 보였다.
• 의료 데이터셋, 특히 LiTS(CT 간병변)에서 SACN은 기존 CapsNet, ResNet-18, DenseNet-40보다 정확도와 표준편차 측면에서 뚜렷한 승리를 거두었다.
• SACN은 기존 CapsNets에 비해 계산 부담을 줄였고, 복잡한 데이터셋에서 GPU 메모리 부족 오류를 방지했다.
• 저대비 및 이질적인 병변이 있는 복잡하고 이질적인 의료 영상에서 높은 강건성과 안정성을 보였다.