[论文解读] Self-Attention Capsule Networks for Object Classification

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

Advanced Neural Network Applications参考文献 22被引用 20

一句话总结

本文提出自注意力胶囊网络（SACN），一种新颖的架构，通过在卷积层与胶囊层之间引入自注意力机制，以增强特征选择和空间相关性建模。通过在胶囊处理前聚焦于显著的图像区域，SACN在复杂医学图像和自然图像数据集上实现了更高的分类准确率和鲁棒性，同时相比基线胶囊网络及SOTA模型（如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在处理复杂、杂乱图像时的局限性以及高计算成本的问题。
• 通过在胶囊处理前识别显著图像区域，提升特征选择能力。
• 在标注数据有限的小型多样化医学数据集上提升性能。
• 通过结合自注意力机制的浅层胶囊网络架构，在保持高准确率的同时降低计算负载。
• 在医学图像与自然图像两个领域均展示出良好的泛化能力。

提出的方法

• 在卷积层与主胶囊层之间引入自注意力机制作为中间层。
• 使用查询（query）、键（key）和值（value）矩阵计算注意力分数，以抑制无关图像区域。
• 应用缩放点积注意力机制，生成突出主要特征的注意力图。
• 将注意力图输入主胶囊层，以保留空间关系与目标姿态信息。
• 采用浅层胶囊网络架构，降低计算成本同时保持性能。
• 使用交叉熵损失函数，配合Adam优化器与权重衰减，进行端到端训练。

实验结果

研究问题

• RQ1 自注意力机制能否提升胶囊网络在复杂图像分类中的特征选择能力？
• RQ2 在背景变化大、复杂度高的医学影像数据集中，将自注意力机制集成到CapsNets中是否能提升性能？
• RQ3 SACN能否在计算成本更低的前提下，实现优于基线CapsNet、ResNet-18与DenseNet-40的准确率？
• RQ4 SACN在自然图像与复杂医学影像等多样化数据集上是否具备良好的泛化能力？
• RQ5 注意力机制是否能在遮挡、低对比度及图像损坏等情况下提升模型的鲁棒性？

主要发现

• 在MNIST数据集上，SACN实现了99.5%的分类准确率，达到SOTA水平。
• 在SVHN数据集上，SACN将基线CapsNet的准确率提升了2.4%，证明其在复杂真实世界数据上的有效性。
• 在CIFAR-10数据集上，SACN将基线CapsNet的准确率提升了3.5%，展现出在自然图像分类任务中的强大泛化能力。
• 在医学数据集（包括LiTS，即CT肝脏病变数据集）中，SACN在准确率与标准差方面均显著优于基线CapsNet、ResNet-18与DenseNet-40。
• 与基线CapsNets相比，SACN显著降低了计算负载，避免了在复杂数据集上出现GPU显存不足的问题。
• 该模型在低对比度、异质性病变的复杂、异构医学图像中表现出更高的鲁棒性与稳定性。