A lightweight capsule network via channel-space decoupling and self-attention routing
Yifan Guo,Sulan Zhang,Chunmei Zhang,Hongli Gao,Huajie Li
DOI: https://doi.org/10.1007/s11042-024-18861-1
IF: 2.577
2024-03-14
Multimedia Tools and Applications
Abstract:Compared to traditional convolutional neural networks (CNNs), the Capsule network (CapsNet), due to its capsule-based design that aligns better with the principle of human neurons, possesses stronger representation ability by capturing potential spatial structural relationships among different parts of an entity. However, the transformation of neurons into capsules and the iterative routing mechanism result in a considerable computational burden, which is the main drawback of CapsNet. Additionally, the fully connected decoder network exhibits significant reconstruction errors on more complex datasets (e.g., CIFAR-10), which adversely affects the model's classification performance. To this end, this paper proposes a Lightweight Capsule Network (DSA-CapsNet) based on channel-space decoupling and self-attention routing. First, a set of residual blocks is employed to construct the residual extraction layer, where the deep features are decoupled to respectively model the correlations within channel and space, thereby reducing the number of parameters and generating initial capsules simultaneously. Secondly, a self-attention routing algorithm is introduced between capsule layers to effectively handle fewer capsules and then allow for stacking more layers. Lastly, a deconvolution decoder module is used as a better reconstruction method, replacing the fully connected decoder in CapsNet. Through the evaluation of four benchmark datasets, DSA-CapsNet drastically reduces the number of parameters and runtime while exhibiting better classification results. Particularly, on the CIFAR-10 dataset, DSA-CapsNet achieves a 75.38 reduction in parameters compared to the original CapsNet, with a 25.71 increase in classification accuracy.
computer science, information systems, theory & methods,engineering, electrical & electronic, software engineering