Exploring advanced techniques for enhancing CapsNet: custom squashing, pretraining, and routing in large-scale unbalanced data

Baseline CapsNet theory

The concept of CapsNet was initially proposed by Sabour, Frosst & Hinton (2017). These networks represent images in a whole vector format, which allows them to encode internal properties, including the pose of entities within an image. In contrast to CNNs, which rely on pooling for output routing, CapsNet aims to preserve information to achieve equivariance, particularly in handling viewpoint changes. This preservation is facilitated through dynamic routing, which replaces the pooling mechanism. Lower-level capsules representing specific features are hierarchically routed to parent capsules to capture part-whole relationships via linear transformations. This approach is based on the concept of inverse graphics, which suggests that the neural system deconstructs images into their inherent hierarchical properties. A capsule is a group of neurons that performs a diverse range of internal computations and subsequently encodes the results of these computations into an n-dimensional vector. The vector is the output of the capsule. The length of this output vector indicates the vector’s probability and direction, indicating specific properties about the entity. Using initial convolutional layers in capsule networks permits the reuse and replication of learned knowledge across different parts of the receptive field. An iterative Dynamic Routing algorithm is employed to determine the inputs to the capsules. The output of each capsule is then compared with the actual production of the higher-level capsules. In the case of a match between the outputs, the coupling coefficient between the two capsules is increased. Let $\mathrm{i}$ represent a lower-level capsule and $\mathrm{j}$ represent a higher-level capsule. The prediction vector is calculated as follows in Eq. (1):

(1) $$\hat{\mathrm{u}}(\mathrm{j}|\mathrm{i})=\mathrm{W}\mathrm{i}\mathrm{j}\mathrm{u}\mathrm{i}.$$The $\mathrm{W}\mathrm{i}\mathrm{j}$ trainable weighting matrix and the $\mathrm{u}\mathrm{i}$ Output pose vector from the ( $\mathrm{i}-\mathrm{t}\mathrm{h}$) capsule to the ( $\mathrm{j}-\mathrm{t}\mathrm{h}$) Capsules are employed in this context. The coupling coefficients are calculated using a SoftMax function as follows in Eq. (2):

(2) $$Cij=\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{M}\mathrm{a}\mathrm{x}\left(bij\right)={\displaystyle \frac{\mathrm{e}\mathrm{x}\mathrm{p}\left(bij\right)}{{\sum }_k\mathrm{e}\mathrm{x}\mathrm{p}\left(bik\right)}.}$$The log probability of capsule i coupled with capsule j, denoted by bij, is initialized with zero values. The total input to capsule j is a weighted sum of the prediction vectors, calculated as follows in Eq. (3):

(3) $$sj=\sum _i\mathrm{C}\mathrm{i}\mathrm{j}\hat{\mathrm{u}}j|i.$$

In capsule networks, the length of the output vector is employed to represent the probability for the capsule. Consequently, a non-linear activation function, the squashing function, is used. The squashing function is defined as follows in Eq. (4):

(4) $$Vj=Squash\left(sj\right)={\displaystyle \frac{{\Vert \mathrm{s}\mathrm{j}\Vert }^2}{1+{\Vert \mathrm{s}\mathrm{j}\Vert }^2}{\displaystyle \frac{sj}{\Vert \mathrm{s}\mathrm{j}\Vert }.}}$$

The dynamic routing algorithm can update the $\mathrm{c}\mathrm{i}\mathrm{j}$ Values in each iteration. In this case, the objective is to optimize the Vj vector. In the dynamic routing algorithm, the $\mathrm{b}\mathrm{i}\mathrm{j}$ The Vector is updated in every iteration according to Eq. (5):

(5) $$bij=bij+Vj\hat{\mathrm{u}}j|i.$$

Dataset

BM cell dataset (Matek et al., 2021), the dataset is available to the public, and there is no personal information in this, includes various hematological disease cell microscopic images consisting of more than 170,000 de-identified, expert-annotated cells from bone marrow smears of 961 patients stained using the May-Grünwald-Giemsa/Pappenheim stain. The institutional review board at the Munich Leukemia Laboratory (MLL) has given its permission to use this dataset. Sample microscopic images from the dataset are shown in Fig. 2.

Pre-processing

As seen in Fig. 3, imbalances in the distribution of some BM cell dataset classes cannot reach the expected target. SMOTE (Chawla et al., 2011) Adaptive Synthetic (ADASYN) (He et al., 2008), SMOTEBoost (Chawla et al., 2003), and DataBoostIM (Guo & Viktor, 2004). Sampling approach techniques are used to avoid these imbalances.

In this study, SMOTE, which is preferred over other sampling approximation techniques (Sci-Hub, 2019; Maldonado, López & Vairetti, 2019). It is used to create synthetic instances of the minority class by interpolating between class instances of an imbalanced dataset. Elreedy & Atiya (2019) using Eq. (6) (Reza & Ma, 2019; Juanjuan et al., 2007). In Eq. (6), $t$ it is the actual minority value. $k$ is the number of nearest neighbors of the actual value? $n$ is the nearest neighbor value $\left(n-t\right)$ does Euclidean distance obtain the difference? $random\left(0,1\right)$ is the number value that allows the addition of different feature values. In this study, $k$ = 5 for each minority class.

(6) $$\mathrm{s}\mathrm{m}\mathrm{o}\mathrm{t}\mathrm{e}\left(\mathrm{t},\mathrm{k}\right)=\mathrm{t}+\mathrm{r}\mathrm{a}\mathrm{n}\mathrm{d}\mathrm{o}\mathrm{m}\left(0,1\right)\ast \left(\mathrm{n}-\mathrm{t}\right).$$

Minority classes were rounded to 5,000 images, and classes with over 5,000 images were rounded up. This over-sampling technique improved overall classification performance and representation of the minority classes. Figure 4 shows the change in data points after the SMOTE process. In addition, all images are rescaled using bicubic interpolation (Wikipedia).

The primary reason that the classification performance is better is that SMOTE was used to balance the dataset before training. The models do not automatically deal with class imbalance. This technique is a non-adaptive interpolation algorithm and uses polynomial techniques to sharpen and enlarge digital images.

Baseline CapsNet model

The input images of the baseline CapsNet model were resized to 32 × 32 pixels to reduce the training time and then designed as shown in Fig. 5.

The proposed model takes three arguments: the input image, the number of classes, and the number of routing iterations. After input images are preprocessed as in previous steps, they are fed to the convolutional layer to extract low-level features with filters of 256, kernel size 9, and strides of 1. Then, these features are grouped into primary capsules with filters of 8 × 32, kernel of 9 × 9, and stride of 1 to reduce the model size, each representing a part or aspect of an object. Each capsule in the convolutional layer corresponds to a capsule in the primary capsule layer. A routing-by-agreement approach boosts learning capability and captures the relationships between different parts of an object. Each capsule in a higher-level layer sends its output to capsules in the layer above based on the agreement (compatibility) between their outputs. The routing process consists of iteratively updating the connection weights between capsules based on the match between their outputs. Thus, dynamic routing makes the capsules in higher layers focus on the most relevant capsules in the layer below; the ReLU activation function, a non-linear function in deep neural networks, is used to reduce dimension. The output of the capsule network is obtained by measuring the length of the output vectors of the capsules in the top layer. This length indicates the probability that a given class or object is present in the input image. Standard squash is applied to squash the vectors of the primary capsules as a non-linear activation function. It aims to ensure that short vectors are reduced to a length close to zero and long vectors are reduced to a length close to 1. Then, adjusting the hyperparameter, the learning rate is initially determined to be 0.001, and the Adam (Kingma & Ba, 2014) optimization algorithm is used to change it dynamically. The model is run with a decay ratio of 0.9 and a routing of 3 for 100 epochs. The CapsNet model train accuracy is calculated at 96.99% on the BM dataset. Next, two powerful methods are proposed to enhance further the most important cores in the structure of the baseline CapsNet, such as squash, dynamic routing, and the feature extractor part.

Method 1: pretrained-CapsNet, custom squash, and modified dynamic

The proposed architecture, shown in Figs. 6, and 7, is pre-trained on ImageNet with the CapsNet architecture to extract features from the BM cell images after resizing the image to a 75 × 75 scale. These features are then fed into the primary convolution layer of the capsule layers with a filter size of 8 × 32, a stride of 2, and a kernel size of 1. The output of the Conv2D layer is reshaped to match the desired shape of the primary capsule layer, followed by the application of the enhanced equation of the squash function. The standard Squash function used in capsule networks compresses vector magnitudes to values between 0 and 1, while preserving direction as shown in Eq. (4) above. However, this approach is limited by its reliance on the L2 norm alone, which can result in vanishing gradients especially for vectors near their maximum length. It also struggles with minimal vectors, sensitivity to input scale and slow convergence. To address these limitations, we present an improved squash function that incorporates layer normalization and a scaling factor α.

This enhancement stabilizes learning by independently computing mean and variance for each capsule, centering, and scaling activations. The scaling factor is computed based on the squared L2 norm of the input vectors divided by 5. $\alpha {\Vert \frac{{\displaystyle \mathrm{s}\mathrm{j}}}{5}\Vert }^2$, using scaling $\mathrm{\alpha }$, controls the intensity of squashing, and promotes scale invariance. In addition, layer normalization subtracts the meaning from the input vectors to center them around the zero mean, reducing variance and improving stability. Improved robustness is a further advantage of the enhanced function; it reduces sensitivity to parameter addition and internal covariate shift and enhances gradient flow and capsule network expressiveness. Using mean and variance instead of the standard function, layer normalization captures more information than simply the input vector’s squared norm. This better function speeds convergence, improves slopes, and lowers bias. It also allows more complex relationships in the data to be captured. The final equation combines scaling and normalization, providing a comprehensive approach to addressing the limitations of the standard squash function and improving the performance of CapsNet in tasks such as image classification and object detection. The enhanced squash is represented as:

(7) $$\mathrm{E}\mathrm{n}\mathrm{h}\mathrm{a}\mathrm{n}\mathrm{c}\mathrm{e}\mathrm{d}\mathrm{s}\mathrm{q}\mathrm{u}\mathrm{a}\mathrm{s}\mathrm{h}=\mathrm{s}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{e}\ast \mathrm{v}\mathrm{e}\mathrm{c}\mathrm{t}\mathrm{o}\mathrm{r},={\displaystyle \frac{\mathrm{\alpha }{\Vert {\displaystyle \frac{\mathrm{s}\mathrm{j}}{5}}\Vert }^2}{1+\mathrm{\alpha }{|\Vert {\displaystyle \frac{\mathrm{s}\mathrm{j}}{5}}\Vert }^2}{\displaystyle \frac{sj-\mu }{\Vert \mathrm{s}\mathrm{j}\Vert \sqrt{o^2+\epsilon }}.}}$$

(8) $$\mu ={\displaystyle \frac{1}{d}\sum _{i=1}^{d}si,j}$$

(9) $${\text{Ơ}}^2={\displaystyle \frac{1}{d}\sum _{i=1}^{d}si,j-{\mu }^2.}$$Here, $\mathrm{s}\mathrm{j}$ is the capsule input vector, $\mu$ and $o^2$ are the mean and variance of the vector components, respectively, and $\epsilon$ is a small constant added for numerical stability. α = 0.5 was chosen through extensive experimentation as it produced the best results across multiple datasets. This value balances vector compression with sufficient dynamic range to improve convergence and model robustness. Including layer normalization ensures that each capsule’s output is centered and scaled individually, reducing internal covariate shift and promoting scale invariance. Compared to the original function, this enhancement improves training stability and speeds up convergence.

After squashing vectors, batch normalization (Santurkar et al., 2018) is applied with an impulse of 0.8 to maintain standardized output values. A dense layer called GN-Caps, consisting of 160 neurons, is initialized with He-normal and zero initializers for the kernel and bias, respectively, and processes the normalized output. Next, the routing algorithm is applied. Subsequent routing iterations include SoftMax activation and a dense layer of 160 neurons. The output vectors are then squashed. The element-wise multiplication of the squashed SoftMax output and the lower-level capsules (initially predicted vectors) is used to compute the routing weights between the lower and upper-level capsules and to further improve routing performance. The result is then passed through a CPReLU. (This extends to PReLU (Krstić et al., 2021)). We introduce the CPReLU activation function to the dynamic routing process. Unlike standard ReLU or PReLU, CPReLU uses routing predictions (e.g., SoftMax outputs) as a conditional parameter for adapting negative activations.

This allows the model to adjust how it handles negative activations based on the incoming data.

(10) $$\mathrm{S}\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{d}\mathrm{a}\mathrm{r}\mathrm{d}\mathrm{P}\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}\left(\mathrm{x}\right)=\mathrm{M}\mathrm{a}\mathrm{x}\left(0,\mathrm{x}\right)+\mathrm{\alpha }\cdot \mathrm{M}\mathrm{i}\mathrm{n}\left(0,\mathrm{x}\right).$$

CPReLU replaces the static slope parameter α with a dynamic ( $\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{d}\mathrm{i}\mathrm{c}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}\mathrm{s}$), enabling instance-specific adaptation. The CPReLU function uses the following enhanced equations:

(11) $$f\left(\mathrm{x},\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{d}\mathrm{i}\mathrm{c}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}\right)=\mathrm{M}\mathrm{a}\mathrm{x}\left(0,\mathrm{x}\right)+\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{d}\mathrm{i}\mathrm{c}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}\mathrm{s}\times \mathrm{M}\mathrm{a}\mathrm{x}\left(0,-\mathrm{x}\right).$$

(12) $$f\left(X,Predications\right)=\{\begin{array}{cc}x\hfill & ifx\ge 0\hfill \\ x+Predictions\ast max\left(0,-x\right)\hfill & ifx<0\hfill \end{array}.$$When input (x) is not negative, it returns x, unchanged. If x is negative, it sums predictions scaled by the maximum of 0 and the negative value of input (x).

While PReLU provides a global slope adjustment for negative inputs across the network, conditional PReLU provides a more localized and adaptive adjustment on a per-neuron basis, allowing for greater flexibility and potentially better model fit. We suggest adding a CPReLU activation function to the routing process, which is different from the traditional dynamic routing mechanism in capsule networks, where the capsule outputs are updated without activation-aware modification. In our method, the routing weights are changed not just by capsule agreement but also by softmax-based predictions, which act as conditional coefficients in the CPReLU function. This architecture lets the network change the size of negative activations depending on how sure it is about routing assignments. This makes capsule updates more non-linear and sensitive to input. The CPReLU function lets you change the way it works according to where the data is coming from. This helps reduce dead neurones, speeds up routing convergence, and better captures complex information connections, especially in datasets that are noisy or imbalanced.

The output of the final routing iteration passes through a 32-neuron dense layer before being predicted using a dense layer with the number of neurons equal to the 21 classes of the dataset, with Softmax activation to generate class probabilities. The GN-CapsNet model combines Google Inception V3 (Cloud, 2024). With CapsNet to take advantage of Inception V3’s ability to capture features at multiple scales and the efficiency of CapsNet in handling complex relationships within data. In the VGG-CapsNet model, which combines VGG16 with CapsNet, the VGG-CapsNet model is proposed to exploit the feature extraction capabilities of VGG16, which excels at learning hierarchical features from images and RES-CapsNet is combined ResNet152V2 (Researchgate). With CapsNet. The models were trained using categorical cross-entropy loss and the Adam optimizer, with hyperparameters such as learning decay (0.9), learning rate (0.001), and normalization momentum (0.8) tuned for optimal performance after establishing Pretrained-CapsNet architecture.

Method 2: Fire-CapsNet

The presented architecture is shown in Figs. 8 and 9. The implementation of Fire-CapsNet in the method follows a hierarchical structure with three main components: Base Filter Substructure (BFS): This is the first stage that performs basic feature extraction from the input image.

A convolutional layer with a small kernel size (1 × 1) with a custom Swish activation function using multiple parameters, such as ( $\alpha ,\beta ).$ They are used to extract low-level features such as edges and corners. $\alpha$ (scaling factor): This adjusts the total range of the function’s output. A value of α greater than 1 would amplify positive outputs, while a value between 0 and 1 would compress them. β (Sigmoid Steepness): This controls how sharply the function transitions from negative to positive values. Thus, a higher β produces a steeper transition.

To address the limitations of standard activation functions, such as ReLU, or standard swish which can result in dead neurons and non-smooth gradients, we have developed a customized Swish activation function for Fire-CapsNet architecture.

The custom or enhanced swish is computed as:

(13) $$\mathrm{S}\mathrm{w}\mathrm{i}\mathrm{s}\mathrm{h}\left(x,\beta \right)={\displaystyle \frac{\alpha x}{1+e^{-\beta x}}}$$where σ is the sigmoid function and α and β are trainable or fixed scaling parameters. This modified Swish activation function: preserves gradient flow across layers due to its smooth, non-monotonic shape. It enhances learning stability and generalization, especially in deep networks. It is also more robust than ReLU when applied in Fire modules and capsule layers.

Further, Batch normalization is applied to improve the stability of the training. Second components of structure squeeze net (SNS): This section incorporates Fire modules, a technique for efficient feature extraction. The fire module is a building block of the Squeeze Net (Trending Papers) Architecture. However, the standard fire method has been modified in this method. It starts with a 1 × 1 convolutional layer (squeeze) to reduce the number of input channels (squeeze the input). Batch normalization is applied to normalize the activation of the layer. The Swish activation function (implemented as a custom swish) introduces non-linearity. It then separates into two branches. Expand1 A 1 × 1 convolutional layer to increase the number of channels. Expand3 branch: A 3 × 3 convolutional layer to capture spatial information. The outputs of the Expand1 and Expand3 branches are concatenated along the channel axis. After seven fire modules are used, a residual connection is used to preserve information flow and facilitate the training of deeper networks. For example, after fire4, there is a residual connection from the input (BN1) to the output (fire4), similarly, after the second convolution, a residual connection to fire6. In standard fire models, max pooling is often used for down sampling, which can result in a loss of information. This implementation replaces max pooling with a convolutional layer with a stride greater than 1. While ResNet employs identity or convolution-based skip connections between standard convolution layers, Fire-CapsNet integrates residual connections across modular Fire blocks, enabling lightweight yet effective feature reuse. Additionally, Fire-CapsNet extends residual connections into the capsule routing structure, a design not found in traditional residual networks like ResNet. This not only improves gradient flow but also enhances capsule-based representation learning. This stride reduces the spatial dimensions of the feature maps while preserving more information than max pooling as shown in Table 1.

The third component in this method is CapsNet Substructure (CNS): This core component introduces capsule layers. The primary capsule layer transforms the extracted features into capsules with a filter size of 8 × 32, a stride of 2, and a kernel size of 1. The output of the Conv2D layer is reshaped to match the desired shape of the primary capsule layer. An enhanced squash activation function is then employed to limit the capsule activations between 0 and 1. The digit capsule layer performs routing iterations to learn the relationships between the capsules and predict the image’s class probability. Finally, a deep decoder part is added to the CapsNet model, including a dropout layer with a dropout rate of 0.3 for regularization. The decoder consists of dense layers of 512, 1,024, 2,048, and output neurons, custom swish activation, and batch normalization. It is designed to reconstruct the input images from the capsule outputs, ensuring that the model encodes meaningful information and reduces overfitting. The models were trained using loss margin and the Adam optimizer, with hyperparameters such as learning decay (0.9), learning rate (0.001), and normalization momentum (0.8) tuned for optimal performance after establishing the Fire-CapsNet architecture. The proposed Fire-CapsNet is better than standard CapsNet because it adds lightweight Fire modules from SqueezeNet to make multi-scale feature extraction more efficient. Residual connections are established between Fire blocks and the first layers to improve the flow of gradients and the reuse of features. A custom Swish activation algorithm makes gradients smoother and increases non-linearity. These improvements lead to improved representation, faster convergence, and higher accuracy on substantial, unbalanced datasets.

Evaluation process of models

It analyzed a large collection of BM cell images using the pre-trained CapsNet model. This provides an opportunity to show the effects of the modification methods and assess how well the model classified BM cells. Important measures of classification accuracy covering several facets were used to evaluate the model’s performance: precision, specificity, recall, and F1-score. These measurements were computed with formulas, as will be explained below:

(14) $$\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}\mathrm{n}={\displaystyle \frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{P}}.}$$

(15) $$\mathrm{S}\mathrm{p}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{f}\mathrm{i}\mathrm{c}\mathrm{i}\mathrm{t}\mathrm{y}={\displaystyle \frac{\mathrm{T}\mathrm{N}}{\mathrm{T}\mathrm{N}+\mathrm{F}\mathrm{P}}.}$$

(16) $$\mathrm{R}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{l}={\displaystyle \frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{N}}.}$$

(17) $$\mathrm{F}1\text{-}\mathrm{s}\mathrm{c}\mathrm{o}\mathrm{r}\mathrm{e}={\displaystyle \frac{2\ast \left(\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}\mathrm{n}\ast \mathrm{R}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{l}\right)}{\mathrm{P}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}\mathrm{n}+\mathrm{R}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{l}}.}$$

Method 1 results

The performance metrics of the pre-trained CapsNet model show exceptional results across different classes of BM cells. The precision, recall, specificity, and F1-score metrics show consistently high values, as shown in Tables 4, 5, 6, and 7. indicating the ability of the models to classify different cell types with high accuracy. Baseline CapsNet performance is less effective than other models in pre-trained models; Classes such as ABE, ART, BAS, BLA, and FGC achieve perfect scores across all metrics, demonstrating the model’s ability to identify these cell types accurately.

The results presented in Table 8 provide a comparative analysis between the baseline CapsNet model and its enhanced versions, VGG-CapsNet and GN-CapsNet, on the BM datasets. The baseline CapsNet achieves an accuracy of 96.99%, requires 100 epochs, and takes approximately 3 h and 22 min to train. In contrast, the improved models show significant improvements in several aspects. All pretrained-CapsNet outperform the baseline model, achieving accuracies of 99.31%, 99.38% and 99.89%, VGG-CapsNet, RES-CapsNet, and GN-CapsNet, respectively. This considerable increase in accuracy highlights the effectiveness of the enhancements incorporated into these models. In addition, all improved models require only 50 epochs to reach convergence, halving the number of epochs needed for the baseline CapsNet. This faster convergence indicates the improved training efficiency introduced by the modifications. In addition, training time is significantly reduced, with VGG-CapsNet completing training in 1 h 20 min, RES-CapsNet in 4 h, and GN-CapsNet completing training in 1 h 33 min, as shown in Table 8. These reductions in training time and improved accuracy highlight the superior performance of the models, as shown in Fig. 10.

The results in Table 9 below demonstrate the performance of different CapsNet architectures. In the CIFAR-10 dataset, the baseline CapsNet achieves an accuracy of 67.81%. VGG-CapsNet achieves an accuracy of 83.58%, showing a significant improvement. Res-CapsNet slightly outperforms with an accuracy of 84.02%, demonstrating the benefits of residual connections in improving feature learning. The highest performance is achieved by GN-CapsNet with an accuracy of 86.54%, highlighting the effectiveness of integrating global and residual learning mechanisms. For the Fashion MNIST dataset, which presents a more challenging, the baseline CapsNet achieves an accuracy of 92.49%, which is quite robust for this dataset. However, the advanced architectures again show significant improvements: VGG-CapsNet reaches 95.22%, Res-CapsNet reaches 95.46%, and GN-CapsNet reaches 95.95%. Finally, the MNIST dataset, known for its simplicity and low complexity, provides a benchmark where even traditional models achieve high accuracy. The baseline CapsNet performs very well, with an accuracy of 99.23%. However, the VGG-CapsNet, Res-CapsNet, and GN-CapsNet architectures perform even better, achieving 99.87%, 99.91%, and 99.93% accuracy, respectively as shown in Fig. 11.

Figure 11 illustrates the accuracy curve for the enhanced CapsNet models on the training data across 50 epochs for GN-CapsNet and VGG-CapsNet and 100 epochs for baseline CapsNet. This curve demonstrates how the model’s classification accuracy evolves as training progresses. Figure 12 provides insights into the model’s dynamics and ability to improve accuracy over successive epochs.

Method 2 results

The high precision, recall, specificity, and F1-scores of the Fire-CapsNet model are illustrated in Table 10 by the performance metrics for different classes. The model’s excellent results on all measures showed that classes such as ABE, ART, BAS, BLA, FGC, and HAC were particularly well-classified.

MNIST dataset in Table 11 achieves the highest accuracy, 99.95%, indicating that the model has learned the patterns in the dataset very well. Fashion-MNIST follows close behind with an accuracy of 95.98%, indicating that the model can generalize similar image recognition tasks well with different categories. Cifar10 has a slightly lower accuracy (90.02%) than the previous two datasets. BM achieves 99.92% accuracy. These results are after 20 epochs for all datasets. It is important to consider that the better classification performance on the BM dataset isn’t just because of proposed model architecture: it is also because SMOTE was used during preprocessing. It is important to note that the optimal number of epochs can vary depending on the complexity of the model, as shown in Figs. 13 and 14.

Comparison of our models with other proposed models and state-of-the-art methods

Table 12 below shows the performance of several proposed enhanced CapsNet models on four different datasets, including MNIST, Fashion-MNIST, CIFAR-10, and a dataset labelled “BM”. Each model is trained on a different number of epochs, ranging from 20 to 100, and the models trained with Method and Method2 consistently outperform the Baseline CapsNet, with Method2 achieving the highest accuracy of 99.91% at 20 epochs on BM, demonstrating its effectiveness in improving model performance. Fashion MNIST dataset and MNIST dataset, all proposed models show higher accuracy compared to the baseline CapsNet. The Fire-CapsNet model trained with Method 2 achieves the highest accuracy of 90.02% with the Cifar10 dataset, demonstrating its effectiveness in handling more complex image datasets.

Table 13 comprehensively compares different CapsNet architectures across multiple datasets, including MNIST, Fashion MNIST, CIFAR-10, and BM. The proposed model, particularly method 3 (Fire-CapsNet), achieves the highest accuracy of 99.95%, outperforming other CapsNet architectures and the different proposed models and CapsNet baselines.

As shown in Table 14, the studies on this BM dataset are limited because it is new and has huge, unbalanced data, but these issues have been solved in our proposed models. Previous methods, such as ResNeXt-50 and SWAV, achieved F1-scores (62–73%) higher than other methods, according to a performance comparison on the BM dataset. Our proposed models increased accuracy significantly, including Enhanced CapsNet architectures; Fire-CapsNet reached 99.92%.