AutoCaps-Zero: Searching for Hardware-Efficient Squash Function in Capsule Networks

Capsule networks (CapsNets) offer distinct advantages over conventional convolutional neural networks (CNNs) by introducing the concept of a capsule. Specifically, this innovation achieves both rotational invariance and spatial awareness, making CapsNets a powerful tool in the field of machine learning. However, this breakthrough comes with an increased level of computational complexity. In our comprehensive experimental analysis of CapsNets, we meticulously inspected its various components and identified the squash function as the main computational bottleneck. To address this challenge, In this paper, we adapts the principles of neural architecture search (NAS) and introduces AutoCaps-Zero, a framework that automatically searches the hardware-efficient squash function to reduce model inference time. Meanwhile, CapsNet models incorporating the searched squash function have exhibited excellent performance across datasets of various sizes, while retaining robust features that make them resistant to adversarial attacks. Besides, these models maintain high performance even on challenging datasets like multiMNIST. Particularly, our experimental results demonstrate that the squash function searched by AutoCaps-Zero reduces the execution time of the squash function itself by approximately 68 %. Consequently, deploying the searched squash function on our benchmark models can reduce the end-to-end graphic processing unit (GPU) inference time by up to 34%. Overall, with the searched function, the CapsNet code will be released.