# Simplified Routing Mechanism for Capsule Networks
## 4 Network Architecture  
In this work, we have used the network architecture proposed by [PERSON] et al. [11] to compare our proposed routing mechanism with other optimization solutions in the field of capsule networks. This capsule network architecture is shown in Figure 6. The original paper used a fixed \(32\times 32\times 1\)-sized input tensor, because they only tested the network efficiency for the MNIST [13] dataset. In contrast, we trained and tested the capsule networks for six fundamentally different datasets in the field of image classification. In our work, the shape of the input layer varies depending on the dataset. We used the following input shapes: \(28\times 28\times 1\), \(48\times 48\times 1\) and \(32\times 32\times 3\). After the input layer, the capsule network architecture consisted of three main components: the first is a convolutional layer, the next is the primary capsule layer, and the last one is the secondary capsule layer.  
The convolution layer contains 256 convolution kernels of size \(9\times 9\) with a stride of 1, and a ReLU (rectified linear unit) [14] activation layer. This convolutional layer generates the main visual features based on intensities for the primary capsule layer.  
The primary capsule block contains a convolutional layer, where both the input and the output are of the size 256. This capsule block also contains a squash layer. In this case, the original squash function (Equation (4)) is used for both implementations. The output of this block contains 32 capsules, where each capsule has 8 dimensions. This capsule block contains advanced features, which are passed onto the secondary capsule block.  
The secondary capsule block has one capsule per class. As mentioned earlier, we worked with several different datasets, so the number of capsules in this capsule block varied, always according to the class number of the dataset: 5, 10 or 43. This capsule block contains the routing mechanism, which is responsible for determining the connection weights between the lower and higher capsules. Therefore, this capsule block represents the main difference between the solution of [PERSON] et al. and our presented method. In this block, we applied our proposed squash function (Equation (11)). The secondary capsule block contains a trainable matrix, called \(W\) (Equation (14)). The shape of the \(W\) matrix, for both solutions, is \(pc\times n\times 16\times 8\), where \(n\) is the number of output classes and \(pc\) depends on the input image shape as follows:  
\[pc=\left\{\begin{array}{ll}32\times 6\times 6,&im_{size}=(28,\ 28)\\ 32\times 8\times 8,&im_{size}=(32,\ 32)\\ 32\times 16\times 16,&im_{size}=(48,\ 48)\end{array}\right. \tag{12}\]  
where \(im_{size}\) is the size of the input image. The routing algorithm was run through \(r=3\) iterations in both cases. The output of this capsule block is a 16-dimensional vector per each class. This means that the block produces \(n\) 16-dimensional capsules, where \(n\) is the number of output classes. The length of the output capsules represents the probability values belonging to the given class.  
Figure 6: The capsule network architecture used in the research, based on work by [PERSON] et al. [11], (green: input, purple: convolutional layer, yellow: primary capsule layer, red: secondary capsule layer, gray: prediction).