Capsule Network was proposed by Geoffrey Hinton and Sara Sabour, aiming to characterize both the existence probability of features and their spatial attributes through vectorized representations, thereby compensating for the limitations of traditional Convolutional Neural Networks in modeling spatial hierarchical relationships. Compared with scalar neurons, capsules are capable of simultaneously representing whether a feature exists as well as its pose, orientation, and other spatial information, which provides advantages in tasks such as image classification and object recognition.

In recent years, research on capsule networks has mainly focused on network architecture design, routing mechanism improvements, and the expansion of application scenarios. From the perspective of network structure, researchers have attempted to enhance the feature representation capability of models by stacking multiple capsule layers or introducing different forms of capsule modules [6]. Regarding routing mechanisms, dynamic routing serves as a core component of capsule networks. Through multiple iterations, it aggregates information from lower-level capsules to higher-level capsules and is considered a key factor contributing to the modeling capability of capsule networks. For example, [4] introduces sparse attention routing to improve both efficiency and accuracy.

However, as the number and dimensionality of capsules increase, the dynamic routing mechanism requires iterative computation over a large number of capsule pairs, resulting in a significant growth in both model parameters and computational complexity. This issue limits the practical deployment of capsule networks, especially in scenarios with constrained computational resources.

To address the large parameter scale and high computational cost of capsule networks, existing studies have explored lightweight improvements from perspectives such as network architecture design and routing strategies. Among these approaches, parameter compression and weight sharing are common techniques for reducing model complexity. By sharing transformation parameters among multiple capsules, the number of transformation matrices can be effectively reduced, thereby alleviating the parameter growth problem as the capsule scale increases. For instance, [3] proposed the concept of shared-weight capsules, where all primary capsules share a single transformation matrix when mapping to a higher-level capsule. This approach significantly reduces the number of parameters and lowers computational complexity.

In addition, some studies attempt to simplify the dynamic routing process. One approach is to reduce the number of routing iterations. For example, [7] replaces traditional dynamic routing with a non-iterative self-attention routing mechanism, thereby reducing both computation and parameter requirements. Another direction focuses on optimizing routing computation through parallelization. For instance, [8] proposes a parallelized dynamic routing strategy that reduces computational complexity and demonstrates significant advantages over the traditional CapsNet in terms of hardware resource utilization, inference speed, and energy consumption. Although these methods improve operational efficiency to some extent, they may also affect the precision of agreement modeling between capsules.

Overall, existing lightweight research has made certain progress in reducing the number of parameters, but limitations remain in controlling computational complexity. Some methods compress parameters without significantly reducing the computational overhead in the routing stage. Therefore, how to simultaneously balance parameter scale and computational efficiency while maintaining model performance remains a key issue in lightweight capsule network research.

The Attention Mechanism has been widely applied in convolutional neural networks and sequence modeling tasks as an effective method for feature selection and information weighting [9]. Its core idea is to assign different weights to different features, guiding the model to focus on information that is more important for the current task, thereby improving overall representation capability.

In recent years, some studies have explored integrating attention mechanisms into capsule networks to enhance the ability of capsules to perceive critical information. Such methods typically introduce attention modules either during the feature extraction stage or the capsule aggregation stage. By weighting lower-level features or capsule outputs, these approaches highlight important regions or key capsules. For example, [10] employs a multi-branch attention mechanism to facilitate information transmission between capsule layers.

Although the attention mechanism can improve the feature representation capability of capsule networks to a certain extent, its introduction is often accompanied by additional parameters and computational overhead, which conflicts with the goal of lightweight capsule networks. Therefore, how to effectively integrate attention mechanisms with capsule network architectures while controlling model complexity—so that useful feature representations are enhanced without significantly increasing computational burden—remains an open research problem.

To address the issues of large parameter scale and high computational complexity in traditional Capsule Network architectures, this paper proposes a lightweight improved model based on the classical capsule network framework. The complete architecture of the proposed model is illustrated in Figure 1. The overall structure still follows the fundamental design of convolutional feature extraction, the PrimaryCaps layer, and the higher-level capsule classification layer, in order to preserve the capability of capsule networks to model spatial hierarchical relationships and pose information. On this basis, targeted improvements are made to the original architecture from three aspects: model parameter organization, feature representation enhancement, and dynamic routing computation strategy.

First, a parameter-sharing mechanism is introduced during the transformation process of higher-level capsules. This mechanism reduces redundant linear transformation parameters between lower-level capsules and higher-level capsules, thereby effectively compressing the model size. Second, a lightweight spatial attention mechanism is incorporated at the PrimaryCaps stage to adaptively reweight feature responses at different spatial locations, enhancing the discriminative power of the capsule input features. Finally, during the dynamic routing stage, a score-based routing pruning strategy is adopted. Dynamic routing computation is performed only among a small number of higher-level capsules with high matching scores, thereby reducing the computational overhead of the routing process.

Through these improvements, the proposed model maintains the structural advantages of the original capsule network while effectively reducing the parameter scale and providing a feasible approach for further lowering computational complexity.

Figure 1. WTCaps network model diagram.

In the classical Capsule Network architecture, feature mapping from lower-level capsules to higher-level capsules is achieved through a set of independent linear transformation matrices. Specifically, for each pair consisting of a lower-level capsule and a higher-level capsule, the network learns a dedicated transformation matrix to map the pose vector of the lower-level capsule into the prediction space of the corresponding higher-level capsule. Let the output vector of theiii-th lower-level capsule be:

The prediction vector corresponding to the j higher-level capsule is computed using an independent transformation matrix:

where W i j denotes the transformation matrix between the iii-th lower-level capsule and the jjj-th higher-level capsule. This design implies that each capsule pair requires its own set of parameters. As the number of capsules increases, the number of parameters grows rapidly, leading to significant parameter redundancy.

To address this issue, this paper redesigns the transformation mechanism between capsules by introducing a parameter-sharing strategy based on higher-level capsules. Instead of assigning an independent transformation matrix to each pair of lower-level and higher-level capsules, only one shared linear transformation matrix is learned for each higher-level capsule. All lower-level capsules use this shared parameter when mapping to that higher-level capsule:

where W j ∈ ℝ d ' × d represents the shared transformation matrix corresponding to the jjj-th higher-level capsule. Through this approach, lower-level capsules no longer rely on independent parameters for differentiation, thereby significantly reducing the parameter scale of the model.

After the pose mapping stage, the input vector of a higher-level capsule is obtained by computing the weighted sum of prediction vectors from all lower-level capsules:

where a i j represents the contribution coefficient of the i lower-level capsule to the j higher-level capsule, which measures the degree of agreement between them. Functionally, this coefficient is similar to the routing coefficient in the original capsule network. However, it no longer relies on the coupling modeled by independent transformation matrices; instead, it acts as an independent weighting factor in constructing the input of higher-level capsules.

Finally, the output of the higher-level capsule is obtained through a nonlinear squashing function:

which ensures that the length of the output vector reflects the probability of the existence of the corresponding class.

Through the proposed parameter-sharing mechanism, the model substantially reduces redundant linear transformation parameters between lower-level and higher-level capsules while preserving the core principles of capsule networks.

In the classical Capsule Network, the dynamic routing process computes coupling coefficients between each lower-level capsule and all higher-level capsules, and continuously updates them through multiple iterations. Although this fully connected routing strategy can effectively model the relationship between parts and wholes, its computational complexity grows linearly with the number of lower-level and higher-level capsules, resulting in considerable computational overhead in practical applications. In the PrimaryCaps layer, a large number of lower-level capsules are typically generated. However, many of them contribute little to the final classification result but still fully participate in the routing computation, leading to significant redundancy.

To address this issue, this paper introduces a Top-K routing pruning mechanism in the dynamic routing stage. By pre-screening the matching relationships between lower-level and higher-level capsules, only a subset of routing connections with higher contributions is retained for subsequent dynamic routing computation. In this way, the computational cost of the routing stage can be reduced without significantly affecting model performance. The schematic diagram of the Top-K routing process is shown in Figure 2.

Specifically, under the shared transformation matrix setting, a lower-level capsule obtains a prediction vector for the j higher-level capsule after linear transformation. Based on the consistency relationship between the prediction vector and the current output vector of the higher-level capsule, the matching score between them can be defined as:

This score essentially measures the degree of support that the iii-th lower-level capsule provides to the jjj-th higher-level capsule, and it is consistent in form with the agreement term used to update routing logits in the original dynamic routing algorithm.

In the Top-K routing pruning strategy, routing coefficients are not computed and updated for all capsule pairs. Instead, for each lower-level capsule, the top K candidate higher-level capsules with the highest matching scores are selected from all higher-level capsules, forming a candidate set K i . Only when j ∈ K i does the lower-level capsule participate in the dynamic routing computation of the corresponding higher-level capsule. All other routing connections are directly pruned for the current sample.

Based on this mechanism, the normalization process of routing coefficients is restricted within the Top-K subspace:

where b i j denotes the routing logits, whose update rule still follows the iterative procedure of the original dynamic routing algorithm. It can be seen that the Top-K routing pruning strategy does not alter the core principle of dynamic routing; instead, it explicitly constrains the routing search space, thereby reducing unnecessary interactions between irrelevant capsules.

Through this mechanism, the number of higher-level capsules involved in the routing computation for each lower-level capsule is reduced to K. As a result, the computational cost of operations such as vector inner products and weighted summations in the routing stage is significantly decreased. Meanwhile, since the retained routing connections correspond to higher matching scores, the pruning strategy can also suppress noisy routing to some extent, which helps improve the stability of the dynamic routing process.

Figure 2. TOP-K routing diagram.

Under the shared transformation matrix and the Top-K routing pruning mechanism, the number of parameters and the computational complexity in the higher-level capsule stage are significantly reduced. However, this lightweight design also places higher demands on the quality of feature representations produced by lower-level capsules. If the lower-level capsules contain excessive irrelevant regions or background noise, it may negatively affect the consistency computation of prediction vectors during the subsequent routing process.

To address this issue, this paper introduces a lightweight spatial attention mechanism into the PrimaryCaps layer to adaptively reweight spatial features before generating lower-level capsules. The attention structure is inspired by the spatial attention component of Convolutional Block Attention Module (CBAM) [11]. By modeling the importance of different spatial locations, this mechanism guides the network to focus on more discriminative regional features.

It should be noted that this spatial attention module only performs weighting on feature maps without altering their spatial dimensions or the number of channels. Moreover, it does not introduce additional capsule structures. Therefore, it does not increase the number of capsules nor interfere with the subsequent capsule reshaping and routing mechanisms. The attention-enhanced feature maps are then used to construct lower-level capsule vectors, thereby improving their stability and effectiveness under the shared-weight transformation and Top-K routing pruning settings.

By introducing spatial attention at the lower-level capsule stage, the proposed method maintains the lightweight advantage of the model while effectively alleviating the potential performance degradation caused by parameter sharing and routing pruning.

This study evaluates the classification performance of the proposed model on five publicly available datasets: MNIST, Fashion-MNIST, CIFAR-10, SVHN, and SmallNORB.

MNIST: A grayscale handwritten digit dataset containing classes 0 - 9, with images of size 28 × 28. It consists of 60,000 training images and 10,000 test images.

Fashion-MNIST: A grayscale image dataset of 70,000 fashion products in 10 categories, with images of size 28 × 28.

CIFAR-10: A dataset of 60,000 RGB natural images across 10 categories, with image size 32 × 32. Each category contains 5,000 training images and 1,000 test images.

SVHN: A dataset for digit classification, containing over 600,000 real-world RGB images of printed digits (0 - 9), each of size 32 × 32.

SmallNORB: A dataset of 3D toy objects across 5 categories, with each object captured at 18 different azimuth angles (0˚ - 340˚) and 9 different elevations, with image size 96 × 96.

To verify the classification performance and lightweight efficiency of the proposed WTCaps model, comparative experiments were conducted on MNIST, Fashion-MNIST, SVHN, CIFAR-10, and SmallNORB datasets against several representative capsule network models. The experimental results are summarized in Table 1.

The results indicate that WTCaps achieves stable classification performance across all five datasets. Specifically, the error rates on MNIST, SVHN, and CIFAR-10 are comparable to those of current mainstream capsule network methods. In terms of parameter scale, WTCaps contains 4.20 M parameters, significantly fewer than models such as DRCaps, AWRCaps, Aff-CapsNet, and Self-Routing. Meanwhile, WTCaps achieves classification performance that is better than or comparable to other lightweight methods, including CAPCaps and GraCapsnets.

Overall, the experimental results demonstrate that the proposed WTCaps model, based on weight sharing and Top-K routing pruning, effectively reduces model parameter size while maintaining strong classification performance, achieving a reasonable balance between accuracy and model complexity.

Table 1.Classification error rates and parameters of different models on five datasets.

To analyze the impact of each improved module on model performance and complexity, ablation experiments were conducted on the MNIST and CIFAR-10 datasets, as shown in Table 2, with the baseline model being the traditional capsule network model [2]. The results indicate the following: Weight Sharing (WS): Introducing the weight-sharing mechanism reduces the number of model parameters from 6.10 M to 4.30 M. However, a slight decrease in classification accuracy is observed, indicating that while weight sharing effectively compresses the model size, it imposes some constraints on feature representation capability. Top-K Routing Pruning (TR): When only the Top-K routing pruning is applied, the model parameter count remains unchanged, but FLOPs decrease significantly from 552.38 M to 260.15 M, demonstrating a clear reduction in computational complexity. Classification accuracy, however, also decreases slightly, suggesting that excessive pruning of routing connections may result in the loss of discriminative information. Spatial Attention (CA): Introducing the spatial attention mechanism alone has little effect on model parameters or computational complexity, but classification accuracy improves on both datasets, validating the effectiveness of spatial attention in enhancing feature discriminability.

Combining TR and CA leads to a substantial reduction in computational complexity while achieving better classification performance compared to using TR alone, indicating that spatial attention can partially mitigate performance loss caused by routing pruning.

Finally, the full WTCaps model, incorporating WS, TR, and CA, reduces the parameter count and FLOPs to 4.20 M and 200.90 M, respectively, while achieving classification accuracies of 99.67% on MNIST and 92.34% on CIFAR-10. These results demonstrate that the proposed method achieves a favorable balance between model lightweight efficiency and classification performance.

Table 2. Generalization performance from MNIST and CIFAR10 datasets.

To evaluate the robustness of the proposed algorithm to affine transformations, all models were trained on the MNIST dataset and then tested on the affNIST dataset. In this setup, the training images were generated by first padding the original MNIST images (28 × 28) to 40 × 40 pixels and then randomly placing the digits within the 40 × 40 background. This procedure introduces variations in position and scale, allowing assessment of the model’s ability to maintain accurate classification under affine transformations.

Table 3.Generalization performance from MNIST to affNIST dataset.

The test set was generated by first padding the original MNIST images to 40 × 40 pixels, and then applying affine transformations to the digits. These transformations included rotations within ±20˚, shearing within ±45˚, horizontal and vertical scaling in the range of 0.8 - 1.2, and random translations up to 8 pixels in each direction, resulting in the affNIST dataset. For fair comparison, only models that achieved at least 99.2% test accuracy on the original MNIST dataset were evaluated on affNIST. The proposed capsule network model achieved a classification accuracy of 96.69% on affNIST, as shown in Table 3. These results demonstrate that the proposed method achieves state-of-the-art generalization performance on affine-transformed images.

To further evaluate the robustness of the proposed WTCaps model under complex geometric transformations, comparative experiments were conducted on the rotated MNIST and rotated SVHN datasets, with rotation angles set to ranges of ±θ₁ and ±θ₂, respectively, as summarized in Table 4.

The experimental results show that under standard testing conditions without rotation, all models achieve high classification accuracy. As the rotation range increases, the classification performance of all models declines to varying degrees; however, capsule network-based methods generally exhibit stronger robustness. Compared to ResNet18, WTCaps shows a smaller decrease in accuracy under large-angle rotations, indicating better adaptability to pose variations. In comparison with other capsule network models such as DRCaps, DeepCaps, and CAPCaps, WTCaps maintains comparable classification performance on rotated MNIST and rotated SVHN datasets, and in some rotation ranges, achieves slightly better or equivalent results.

These results indicate that, even with significantly reduced model parameters and computational complexity, WTCaps effectively preserves the inherent robustness of capsule networks to spatial transformations such as rotation.

Table 4.Classification performance of different algorithms on Rotated MNIST and Rotated SVHN.

To evaluate the model’s ability to model and separate overlapping targets in complex scenarios, comparative experiments were conducted on the MultiMNIST dataset using several representative methods. The MultiMNIST dataset contains images with multiple overlapping digits, placing higher demands on the model’s feature disentanglement and object perception capabilities, as shown in Table 5.

The results indicate that traditional CNNs achieve relatively low recognition accuracy on this task, whereas methods incorporating capsule structures and dynamic routing mechanisms significantly improve classification performance. Compared with capsule network models such as Dynamic Routing and Aff-CapsNet, WTCaps achieves comparable classification accuracy to AWRCaps while drastically reducing the number of parameters. Specifically, WTCaps achieves a test accuracy of 95.95% with only 4.20 M parameters.

Although AWRCaps attains slightly higher accuracy, its larger parameter scale makes the model less compact. WTCaps strikes a better balance between model compactness and recognition performance. These results demonstrate that the proposed weight-sharing and Top-K routing pruning mechanisms maintain the capsule network’s strong modeling capability for overlapping objects while effectively reducing model size, highlighting the practical applicability of WTCaps in complex visual scenarios.

Table 5.Classification performance of different algorithms on MultiMNIST.

To further verify the feature representation capability of the proposed WTCaps model during the lightweighting process, experiments were conducted to analyze its reconstruction performance. In capsule networks, a reconstruction branch is introduced after the classification capsules, using the input image as a supervisory signal to reconstruct the original sample. This process provides an intuitive visualization of the model’s ability to capture the pose and structural information of the target object. The reconstruction results are shown in Figure 3.

The experimental results indicate that, compared with baseline capsule networks, WTCaps is still able to reconstruct the overall contours and key structural features of input digits fairly accurately, even under significantly reduced parameters and computational complexity. No noticeable shape distortions or semantic omissions were observed. This demonstrates that the proposed weight-sharing and Top-K routing pruning strategies do not compromise the capsule network’s ability to model spatial hierarchical relationships. These results further confirm that WTCaps maintains strong feature representation and structural reconstruction capabilities while achieving model lightweighting.

Figure 3. Reconstruction results of DRCaps and WTCaps.

To more comprehensively evaluate the characteristics of the proposed algorithm, a comparative experiment on computational efficiency was conducted. Classic capsule network models and recent lightweight capsule network models were tested using input images of the same size and number of channels, and their parameter counts and floating-point operations (FLOPs) were measured to assess computational performance, as shown in Table 6.

The results indicate that traditional capsule networks, such as DRCaps and AWRCaps, maintain good performance but suffer from high parameter counts and FLOPs. In contrast, GraCapsNet and CAPCaps reduce model parameters to some extent through structural simplification, but their computational complexity remains relatively high.

The proposed WTCaps model achieves a more balanced performance in both parameter scale and computational complexity. Specifically, WTCaps contains only 4.20M parameters—significantly fewer than DRCaps and AWRCaps—and reduces the computation to 200.90 MFLOPs, approximately 63.6% lower than DRCaps. These results demonstrate that WTCaps effectively reduces redundant computations in the dynamic routing stage without significantly increasing structural complexity, thereby significantly improving overall computational efficiency while preserving the capsule network’s modeling capabilities.

Table 6.Comparison of the computational efficiency of different algorithms.