[Paper Review] Self-Attention Capsule Networks for Object Classification

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

Advanced Neural Network Applications22 references20 citations

TL;DR

This paper proposes Self-Attention Capsule Networks (SACN), a novel architecture that integrates a self-attention mechanism between convolutional and capsule layers to enhance feature selection and spatial correlation modeling. By focusing on salient image regions before capsule processing, SACN achieves superior classification accuracy and robustness—especially on complex medical and natural image datasets—while reducing computational load compared to baseline CapsNets and state-of-the-art models like ResNet-18 and 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.

Motivation & Objective

• To address the limitations of CapsNets in handling complex, cluttered images and high computational costs.
• To improve feature selection by identifying salient image regions before capsule processing.
• To enhance performance on small, diverse medical datasets where annotated data is limited.
• To maintain high accuracy while reducing computational load through a shallow capsule architecture enhanced by self-attention.
• To demonstrate generalization across both medical and natural image domains.

Proposed method

• Integrates a self-attention mechanism as an intermediate layer between convolutional and primary capsule layers.
• Uses query, key, and value matrices to compute attention scores that suppress irrelevant image regions.
• Applies scaled dot-product attention to generate an attention map highlighting dominant features.
• Processes the attention map through the primary capsule layer to preserve spatial relationships and object pose information.
• Employs a shallow capsule network architecture to reduce computational cost while maintaining performance.
• Trains end-to-end using cross-entropy loss with Adam optimizer and weight decay.

Experimental results

Research questions

• RQ1 Can self-attention improve feature selection in capsule networks for complex image classification?
• RQ2 Does integrating self-attention into CapsNets enhance performance on medical imaging datasets with high background variability?
• RQ3 Can SACN achieve better accuracy than baseline CapsNet, ResNet-18, and DenseNet-40 with reduced computational cost?
• RQ4 How does SACN generalize across diverse datasets, including natural images and complex medical scans?
• RQ5 Does the attention mechanism improve robustness in the presence of occlusions, low contrast, and image corruptions?

Key findings

• SACN achieved a 99.5% classification accuracy on MNIST, matching state-of-the-art performance.
• On SVHN, SACN improved baseline CapsNet accuracy by 2.4%, demonstrating effectiveness on challenging real-world data.
• On CIFAR-10, SACN improved baseline CapsNet accuracy by 3.5%, showing strong generalization to natural image classification.
• In medical datasets, including LiTS (CT liver lesions), SACN significantly outperformed baseline CapsNet, ResNet-18, and DenseNet-40 in both accuracy and standard deviation.
• SACN reduced computational load compared to baseline CapsNets, avoiding GPU out-of-memory errors on complex datasets.
• The model showed higher robustness and stability on complex, heterogeneous medical images with low-contrast and heterogeneous lesions.