STAR-Caps: Capsule Networks with Straight-Through Attentive Routing

Karim Ahmed, Lorenzo Torresani

Abstract: Capsule models have been shown to be powerful models for image classification, thanks to their ability to represent and capture viewpoint variations of an object. However, capsule networks have high computational complexity due to the large cost of their recurrent dynamic routing. This is a major drawback that makes their use for large-scale image classification challenging. In addition, training capsule models can often be sensitive to the number of capsules specified as a hyperparameter. In this work, we propose a capsule-based network called STAR-Caps that exploits a straight-through attentive routing mechanism to address the drawbacks of capsule networks. By utilizing efficient attention modules augmented by differentiable binary routers, the proposed mechanism efficiently estimates the routing coefficients between capsules without recurrence, as opposed to the previous related work. Subsequently, the binary routers that utilize straight-through estimators make dynamic decisions to either connect or disconnect the route between capsules, allowing more stable and faster performance. The experiments conducted on several image classification datasets, including MNIST, SmallNorb, CIFAR-10, CIFAR-100 and ImageNet show that STAR-Caps is superior to the baseline capsule networks, in terms of the classification accuracy as well as the training and inference time.

Code Link: http://bit.ly/2WhO8Ms

CMT Num: 4875