Spatial Attention-Based Capsule Networks With Guaranteed Group Equivariance

Abstract

Some capsule networks (CapsNets) reported lately aim to enforce capsule poses and descriptors to be equivariant and invariant respectively by adding extra loss functions as regularization but without providing rigorous proof. To address this problem, a group equivariant spatial attention mechanism (GSA) is proposed to rigidly guarantee the equivariance with mathematical proof while enhancing the spatial information in capsule poses. In addition, to alleviate the computation burden associated with the conventional routing algorithm, group poolings are developed to generate the descriptors and poses of capsules, which contribute greatly to preserving the invariance and equivariance of CapsNets. With the proposed components of GSA and group poolings, a new attentive CapsNet, namely spatial attentive group equivariant CapsNets (SAGE-CapsNets), is constructed in this paper. To validate the invariance and equivariance of SAGE-CapsNets, we conduct experiments involving classification, semantic segmentation, and visualization. The results obtained from these experiments provide empirical evidence of the effectiveness of our proposed approach. Note to Practitioners—This paper is motivated by the problem that existing affine transformations in the real world generally degrade the performance of neural networks in vision tasks like image classification, segmentation, and detection. While conventional capsule networks help to alleviate this problem by learning invariant spatial relationships between features, their robustness to affine transformations is shown through empirical results without rigorous proof. To tackle this issue, we propose a novel capsule network with equivariant components, including group spatial attention and group pooling layers. These components are rigorously proven to be equivariant and greatly contribute to the model’s robustness against affine transformations. Moreover, for practical applications, our proposed attention mechanism improves model performance without significantly increasing computation. Additionally, group pooling preserves model equivariance while reducing computation overhead. As a result, our computation-saving model can be applied to real-world vision applications that require robustness to affine transformations, such as bearing fault diagnosis and facial recognition.