[論文レビュー] Self-Attention Capsule Networks for Object Classification

Assaf Hoogi, Brian Wilcox+2 more|arXiv (Cornell University)|Apr 29, 2019

Advanced Neural Network ApplicationsComputer Science参考文献 22被引用数 20

ひとこと要約

本稿では、畳み込み層とCapsule層の間に自己注意機構を統合することで、特徴選択と空間相関モデリングを向上させる、自己注意Capsuleネットワーク（SACN）という新規アーキテクチャを提案する。SACNは、Capsule処理の前に顕著な画像領域に注目することで、複雑な医療画像および自然画像データセットにおいて、ベースラインのCapsNets や最新のモデル（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.

研究の動機と目的

• 複雑でごちゃついた画像を処理する際のCapsNetsの限界と高い計算コストを解消すること。
• Capsule処理の前に顕著な画像領域を特定することで、特徴選択を改善すること。
• アノテーションデータが限られた小規模で多様な医療データセットにおける性能向上を図ること。
• 自己注意で強化された浅いCapsuleネットワークアーキテクチャを用いることで、計算負荷を低減しながら高い精度を維持すること。
• 医療画像および自然画像の両ドメインにわたる一般化性能を示すこと。

提案手法

• 畳み込み層とプライマリーキャプセル層の間に自己注意機構を中間層として統合する。
• クエリ、キー、値の行列を用いて、不要な画像領域を抑制する注意スコアを計算する。
• スケーリングされたドット積注意を適用し、顕著な特徴を強調する注意マップを生成する。
• 注意マップをプライマリーキャプセル層で処理することで、空間的関係性とオブジェクトの姿勢情報の保持を図る。
• 計算コストを低減しながら性能を維持するため、浅いCapsuleネットワークアーキテクチャを採用する。
• Adam最適化法と重み減衰を用いた交差エントロピー損失を用いて、エンドツーエンドで学習する。

実験結果

リサーチクエスチョン

• RQ1 自己注意機構は、複雑な画像分類のためのCapsuleネットワークにおける特徴選択を改善できるか？
• RQ2 CapsNetsに自己注意を統合することで、背景のばらつきが著しい医療画像データセットにおける性能が向上するか？
• RQ3 SACNは、ベースラインのCapsNet、ResNet-18、DenseNet-40と比較して、計算コストを低減しつつも、より高い精度を達成できるか？
• RQ4 SACNは、自然画像および複雑な医療スキャンを含む多様なデータセットに一般化できるか？
• RQ5 注意機構は、オクルージョン、低コントラスト、画像の損傷が生じる状況でも耐障害性を向上させるか？

主な発見

• MNISTでは99.5%の分類精度を達成し、最新の最先端性能と同等であった。
• SVHNでは、ベースラインのCapsNetと比較して2.4%の精度向上を達成し、現実世界の挑戦的データに対して有効であることが示された。
• CIFAR-10では、ベースラインのCapsNetと比較して3.5%の精度向上を達成し、自然画像分類への強い一般化能力を示した。
• LiTS（CT肝病変）を含む医療データセットにおいて、SACNはベースラインのCapsNet、ResNet-18、DenseNet-40と比較して、精度および標準偏差の両面で顕著に優れた性能を示した。
• SACNは、ベースラインのCapsNetsと比較して計算負荷を低減し、複雑なデータセットにおいてGPUのメモリ不足エラーを回避した。
• 低コントラストや異種の病変を有する複雑で多様な医療画像において、より高い耐障害性と安定性を示した。