# REM: Routing Entropy Minimization for Capsule Networks

Riccardo Renzulli
University of Turin
Computer Science Department
<EMAIL_ADDRESS>Enzo Tartaglione
LTCI, Télécom Paris
Institut Polytechnique de Paris
<EMAIL_ADDRESS>Marco Grangetto
University of Turin
Computer Science Department
<EMAIL_ADDRESS>
###### Abstract

Capsule Networks ambition is to build an explainable and biologically-inspired
neural network model. One of their main innovations relies on the routing
mechanism which extracts a parse tree: its main purpose is to explicitly build
relationships between capsules. However, their true potential in terms of
explainability has not surfaced yet: these relationships are extremely
heterogeneous and difficult to understand.
This paper proposes REM, a technique which minimizes the entropy of the parse
tree-like structure, improving its explainability. We accomplish this by
driving the model parameters distribution towards low entropy configurations,
using a pruning mechanism as a proxy. We also generate static parse trees with
no performance loss, showing that, with REM, Capsule Networks build stronger
relationships between capsules.

##  1 Introduction

Capsule Networks (CapsNets) [21, 6, 9] were recently introduced to overcome
the shortcomings of Convolutional Neural Networks (CNNs). CNNs loose the
spatial relationships between its parts because of max pooling layers, which
progressively drop spatial information [21]. Furthermore, CNNs are also
commonly known as “black-box” models: most of the techniques providing
interpretation over the model are post-hoc: they produce localized maps that
highlight important regions in the image for predicting objects [22]. CapsNets
attempt to preserve and leverage an image representation as a hierarchy of
parts, carving-out a parse tree from the networks. This is possible thanks to
the iterative routing mechanism [21] which models the connections between
capsules. This can be seen as a parallel attention mechanism, where each
active capsule can choose a capsule in the layer above to be its parent in the
tree [21]. Therefore, CapsNets can produce explainable representations encoded
in the architecture itself [21] yet can be still successfully applied to a
number of applicative tasks[26, 17, 1].

Figure 1: With REM, a Capsule Network builds stronger relationships with lower
entropy between capsules.

However, understanding what really happens inside a CapsNet is still an open
challenge. Indeed, for a given input image, there are too many active (and
each-other coupled) capsules, making the routing algorithm connections still
difficult to understand, as the coupling coefficients typically have each-
other similar values (Fig. 1), non exploiting the routing algorithm potential
and making architectures without routing performing comparably [3].
Furthermore, backward and forward passes of a CapsNet come at an enormous
computational cost, since the number of trainable parameters is very high. For
example, the CapsNet model deployed on the MNIST dataset by Sabour et al.[21]
is composed by an encoder and a decoder part. The full architecture has $8.2$M
of parameters. Do we really need that amount of trainable parameters to
achieve competitive results on such a task? Recently, many pruning methods
were applied to CNNs in order to reduce the complexity of the networks,
enforcing sparse topologies [24, 14, 13]: is it possible to tailor one of
these approaches with not only the purpose of lowering the parameters, but
aiding the model’s interpretability?

This work introduces REM (Routing Entropy Minimization) for CapsNets, which
aims to improve the explainabilty of the routing algorithm of CapsNets.
Pruning can effectively reduce the overall entropy of the connections of the
parse tree-like structure encoded in a CapsNet, because in low pruning regimes
it removes noisy couplings which cause the entropy to increase considerably.
We collect the coupling coefficients studying their frequency and cardinality,
observing lower per-class conditional entropy on these: the pruned version
adds a missing explicit prior in the routing mechanism, grounding the coupling
of the unused primary capsules disallowing fluctuations under the same
baseline performance on the validation/test set. This implies that the parse
trees are more stable for the pruned models, providing higher explainability
on the relevant features (per-class) selected from the routing mechanism.

This paper is organized as follows: in Section 2 we introduce some of the
basic concepts of CapsNets and their related works, in Section 3 we describe
our techinque called REM, in Section 4 we investigate the effectivess of our
method testing it on many datasets. The last two sections discuss the
limitations of our approach and the conclusion of our work.

##  2 Background and Related Work

Figure 2: Capsule Network architecture (encoder network).

This section first describes the fundamental aspects of CapsNets and their
routing algorithm introduced by [21]. Then, we review the literature
especially related to how introduce sparsity into CapsNets. The notation used
in our work is summarized in Table 1.

Symbols | Meaning
---|---
$\Xi$ | Cardinality of the dataset
$M,N$ | Spatial dimensions of the PrimaryCaps layer
$T$ | Primary capsules types
$I$ | Primary capsules
$J$ | Output (digit) capsules
$\hat{y}^{\xi}$ |  Target class for the $\xi$-th sample
$y^{\xi}$ |  Predicted class for the $\xi$-th sample
$\theta$ | Generic parameter for the model
$\boldsymbol{u}_{i}$ |  Pose of the $i$-th primary capsule (normalized)
${\boldsymbol{s}}_{j}$ |  Pose of the $j$-th output capsule (non-normalized)
$\boldsymbol{v}_{j}$ |  Pose of the $j$-th output capsule (normalized)
$c_{ij}^{\xi}$ |  Coupling coefficient for the $\xi$-th input linking the $i$-th primary capsule to the $j$-th output capsule
$\boldsymbol{c}_{-j}$ |  Couplings to the same $j$-th output caps
$\boldsymbol{W}$ | Transformation matrix
$\mathbb{1}_{\hat{y}^{\xi}}$ |  One-hot encoding for the target class $y^{\xi}$
$\Theta(\cdot)$ | One-step function
$q_{K}(\cdot)$ |  Uniform quantizer on $K$ levels
$\tilde{x}$ |  Quantized representation of $x$
$\mathbb{P}(x)$ |  Probability of $x$
$\mathbb{H}_{j}$ |  Entropy of the couplings toward $j$-th output capsule
$\mathbb{H}$ | Entropy of the couplings for the CapsNet
Table 1: Summary of the notation.

Capsule Networks Fundamentals. CapsNets group neurons into capsules, namely
activity vectors, where each capsule accounts for an object of one of its
parts. Each element of these vectors accounts for different properties of the
object such as its pose and other properties like color, deformation, etc.
[21]. The magnitude of a capsule stands for the probability of existence of
that object in the image [21]. Fig. 2 shows a standard architecture of a
CapsNet with two capsule layers, PrimaryCaps and DigitCaps (also called
OutputCaps). The poses of $L$-th capsules $\boldsymbol{u_{i}}$, called primary
capsules, are built upon convolutional layers. In order to compute the poses
of the capsules of the next layer $L+1$, an iterative routing mechanism is
performed. Each capsule $\boldsymbol{u_{i}}$ makes a prediction
$\hat{\boldsymbol{u}}_{{j|i}}$, thanks to a transformation matrix
$\boldsymbol{W}_{ij}$, for the pose of an upper layer capsule $j$

$\hat{\boldsymbol{u}}_{{j|i}}=\boldsymbol{W}_{{ij}}\boldsymbol{u}_{i}.$ (1)

Then, the total input ${\boldsymbol{s}}_{j}$ of capsule $j$ of the DigitCaps
layer is computed as the weighted average of votes
$\hat{{\boldsymbol{u}}}_{j|i}$

${\boldsymbol{s}}_{j}=\sum\limits_{i}c_{ij}\hat{{\boldsymbol{u}}}_{j|i},$ (2)

where $c_{ij}$ are the coupling coefficients between a primary capsule $i$ and
an output capsule $j$. The pose ${\boldsymbol{v}}_{j}$ of an output capsule
$j$ is then defined as the normalized “squashed” $s_{j}$

${\boldsymbol{v}}_{j}=\textrm{squash}({\boldsymbol{s}}_{j})=\frac{\|{\boldsymbol{s}}_{j}\|^{2}}{1+\|{\boldsymbol{s}}_{j}\|^{2}}\frac{{\boldsymbol{s}}_{j}}{\|{\boldsymbol{s}}_{j}\|}.$
(3)

So the routing algorithm computes the poses of output capsules and the
connections between capsules of consecutive layers. The coupling coefficients
are computed dynamically by the routing algorithm and they are dependent on
the input [21]. The coupling coefficients are determined by a “routing
softmax” activation function, whose initial logits $b_{ij}$ are the log prior
probabilities the $i$-th capsule should be coupled to the $j$-th one

$c_{ij}=\textrm{softmax}(b_{ij})=\frac{e^{b_{ij}}}{\sum\limits_{k}.e^{b_{ik}}}\\\
$ (4)

At the first step of the routing algorithm they are equals and then they are
refined by measuring the agreement between the output ${\boldsymbol{v}}_{j}$
of the $j$-th capsule and the prediction $\hat{{\boldsymbol{u}}}_{j|i}$ for a
given input. The agreement is defined as the scalar product
$a_{ij}={\boldsymbol{v}}_{j}\cdot\hat{{\boldsymbol{u}}}_{j|i}$. At each
iteration, the update rule for the logits is

$b_{ij}\leftarrow b_{ij}+a_{ij}.$ (5)

The steps defined in (2), (3), (4), (5) are repeated for the $t$ iterations of
the routing algorithm. In order to train the network, Sabour  _et al._ [21]
replaced the cross entropy loss with the margin loss.

###  2.1 Related work

Capsule Networks. They were first introduced by Sabour  _et al._ [21] and
since then a lot of work has been done, both to improve the routing mechanism
and to build deeper models. Regarding the routing algorithm, Hinton  _et al._
[6] replace the dynamic routing with Expectation-Maximization, adopting matrix
capsules instead of vector capsules. Wang  _et al._ [25] model the routing
strategy as an optimization problem. Li  _et al._ [11] use master and aide
branches to reduce the complexity of the routing process. Hahn  _et al._ [5]
incorporates a self-routing method such that capsule models do not require
agreements anymore. De Sousa Ribeiro  _et al._ [2] replace the routing
algorithm with variational inference of part-object connections in a
probabilistic capsule network, leading to a significant speedup without
sacrificing performance. Since the CapsNet model introduced by Sabour  _et
al._ [21] is a shallow network, several works attempted to build deep
CapsNets. Rajasegaran  _et al._ [18] propose a deep capsule network
architecture which uses a novel 3D convolution based dynamic routing algorithm
aiming at improving the performance of CapsNets for more complex image
datasets. Gugglberger  _et al._ [4] introduce residual connections to train
deeper capsule networks.

Sparse Capsule Networks. A naive solution to reduce uncertainty within the
routing algorithm is to simply running more iterations. As shown by Paik  _et
al._ [16] and Gu  _et al._ [3], the routing algorithms tends to overly
polarize the link strengths, namely a simple route in which each input capsule
sends its output to only one capsule and all other routes are suppressed. On
the one hand, this behavior is desirable because it makes the routing
algorithm more explainable, by making it possible to extract a parse tree
thanks to this coupling coefficients. On the other hand, running many
iterations is only useful in the case of networks with few parameters, as
demonstrated by Renzulli  _et al._ [20], otherwise the performance will drop.
Rawlinson  _et al._ [19] trained CapsNets in an unsupervised setting, showing
that in unsupervised learning the routing algorithm does not discriminate
among capsule anymore: the coupling coefficients collapse to the same value.
Therefore, they sparsify latent capsule layers activities by masking output
capsules according to a custom ranking function. Kosiorek  _et al._ [9] impose
sparsity and entropy constraints into capsules, but they do not employ an
iterative routing mechanism. Jeong  _et al._ [7] introduced a structured
pruning layer called ladder capsule layers, which removes irrelevant capsules,
namely capsules with low activities. Kakillioglu  _et al._ [8] solve the task
of 3D object classification on point clouds with pruned Capsule Networks.
Their objective was to compress robust capsule models in order to deploy them
on resource-constrained devices.

The main contribution of our work relies on the fact that we regularize and
prune the parameters in a CapsNet as a way to minimize the entropy of the
connections computed by the routing algorithm. In fact, we show that
relationships between objects and their parts in a standard CapsNets described
by Sabour et al. [21] have high entropies. We minimize these entropies so that
we can extract stronger parse trees. This allows us to effectively build
dictionaries upon the input datasets and understand which are the shared
object parts and transformations between different the entities in the images.

##  3 Routing Entropy Minimization

The coupling coefficients computed by the routing mechanism model the part-
whole relationships between capsules of two consecutive capsule layers, as
displayed in Fig. 1. Assigning parts to objects (namely learning how each
object is composed), is a challenging task. One of the main goal of the
routing algorithm is to extract a parse tree of these relationships. Given the
$\xi$-th input of class $j$, an ideal parse tree for a primary capsule $i$
detecting one of the parts of the entity in the input $\xi$ would ideally lead
to

$\boldsymbol{c}_{i-}^{\xi}=\mathbb{1}_{\hat{y}^{\xi}}.$ (6)

This means that the routing process is able to carve a parse tree out of the
CapsNet which explains perfectly the relationships between parts and wholes.
One of the problems of this routing procedure is that there is no constraint
on how strong a parse tree should be. In this section we present our techinque
REM, first showing how to extract a parse tree and then how to build stronger
parse trees. The pipeline of our method is depicted in Fig. 3.

Figure 3: Pipeline of REM.

###  3.1 Parse trees extraction

Once we have a trained CapsNets model, in order to interpret the routing
mechanism, we extract all the possible routing coupling coefficients and build
a parse tree. Towards this end, we want to define a metric which helps us
deciding if the relationships captured by the routing algorithm resemble or
not a parse tree. Therefore, we organize the coupling coefficients into
associative arrays, we can compute the number of occurrences of each coupling
sequence in order to compute the entropy of the whole coupling dictionary to
measure the simplicity of the parse tree. In the next paragraphs, we explain
how to generate these sequences by discretizing the coupling coefficients and
how to create the dictionary.
Quantization. During the quantization stage, we first compute the continuous
coupling coefficients $c_{ij}^{\xi}$ for each $\xi$-th input example. It
should be noticed that $c_{ij}^{\xi}$ are the coupling coefficients obtained
after the forward pass of the last routing iteration. Then, we quantize every
$c_{ij}^{\xi}$ into $K$ discrete levels through the uniform quantizer
$q_{K}(\cdot)$, obtaining

$\tilde{c}_{ij}^{\xi}=q_{K}({c}_{ij}^{\xi}).$ (7)

We choose the lowest $K$ such that the accuracy is not deteriorated. We will
here on refer to CapsNet+Q as trained CapsNet where the coupling coefficients
are quantized.

Extracting the parse tree. Given the quantized coupling coefficients of a
CapsNet+Q, we can extract the parse tree (and create a dictionary of parse
trees) for each class $j$, where each entry is a string composed by the
quantization indices of the coupling coefficients. We will extract the
coupling coefficients $\boldsymbol{\tilde{c}}_{-j}^{\xi}$ between the primary
capsules $I$ and the predicted $j$-th output capsule. Given a dictionary for
the coupling coefficients of a CapsNet+Q, we can compute the entropy for each
class as

$\mathbb{H}_{j}=-\sum\limits_{\xi}\left\\{\mathbb{P}(\boldsymbol{\tilde{c}}_{-j}^{\xi}\mid
y^{\xi}=j)\cdot\log_{2}\left[\mathbb{P}(\boldsymbol{\tilde{c}}_{-j}^{\xi}\mid
y^{\xi}=j)\right]\right\\}$ (8)

where $\mathbb{P}(\boldsymbol{\tilde{c}}_{-j}^{\xi}\mid y^{\xi}=j)$ is the
frequency of occurrences of a generic string $\xi$ for each _predicted_ class
$y^{\xi}$. Finally, the entropy of a dictionary for a CapsNet+Q on a given
dataset is the average of the entropies $H_{j}$ of each class

$\mathbb{H}=\frac{1}{J}\sum\limits_{j}\mathbb{H}_{j}.$ (9)

Intuitively, the lower (9), the stronger the connections between capsules are
in the parse tree-like structure carved-out from the routing algorithm. We
also target to obtain the distribution of these coupling coefficients. In
general, we know we have $\Xi\times I\times J$ coupling coefficients for the
full dataset (with potential redundancies). Given the $i$-th primary capsule,
however, we are only interested to $\boldsymbol{c}_{-j}^{\xi}|y^{\xi}=j$. In
this way, we reduce the coupling coefficients space to $I\times J$. We compute
then the average of all the inputs belonging to an object class in order to
output just $I\times J$ coupling coefficients.
Preliminary result. Let us now generate the parse trees on the CapsNet model
trained on the MNIST dataset as in [21], investigating their evolution during
training. To this end, we compute the distribution of the coupling
coefficients on the test set for two models: the model trained after the first
epoch and the one that achieved the lowest loss value on the validation set.

Figure 4: Coupling coefficients distributions of two CapsNets+Q at epochs 1
and 190 on MNIST (test set).

Fig. 4 shows the distributions for the two models and Table 2 reports the
corresponding accuracy and entropy values. We sampled a subset of the
distributions for all the object classes. It can be observed that after the
first epoch CapsNet is clearly far from optimality, both in term of
performance and parse tree explainability: indeed all coupling coefficients
are almost all equals to the value selected at initialization, i.e. $1/J$.
This effect is almost the same when we train a CapsNet with $t=1$ as Gu et al.
[3], where its entropy is exactly zero but capsules are uniformly coupled. On
the other hand, a well trained CapsNet achieves higher generalization but its
entropy is getting higher, so the routing mechanism is difficult to interpret
and to explain. As a matter of fact, we can see in Table 3 the dictionary for
the $\text{CapsNet}_{190}\text{+Q}$: the number of keys for each class is too
high to make the routing algorithm explainable. Notice that the MNIST test set
has 10k images and the total number of keys is 8147. We target a model that
has high generalization but a dictionary where number of keys is low.

Model | Accuracy | Entropy
---|---|---
$\text{CapsNet}_{1}\text{+Q}$ | $97.53$ | $1.83$
$\text{CapsNet}_{190}\text{+Q}$ | $99.55$ | $9.41$
Table 2: Classification results (%) and entropy of two CapsNets+Q at epochs 1 and 190 on MNIST (test set). Model | $\mathbf{\textbf{\\#}_{0}}$ | $\mathbf{\textbf{\\#}_{1}}$ | $\mathbf{\textbf{\\#}_{2}}$ | $\mathbf{\textbf{\\#}_{3}}$ | $\mathbf{\textbf{\\#}_{4}}$ | $\mathbf{\textbf{\\#}_{5}}$ | $\mathbf{\textbf{\\#}_{6}}$ | $\mathbf{\textbf{\\#}_{7}}$ | $\mathbf{\textbf{\\#}_{8}}$ | $\mathbf{\textbf{\\#}_{9}}$
---|---|---|---|---|---|---|---|---|---|---
$\text{CapsNet}_{190}\text{+Q}$ | $974$ | $358$ | $968$ | $967$ | $859$ | $856$ | $939$ | $670$ | $965$ | $591$
Table 3: Dictionary sizes for the $\text{CapsNet}_{190}\text{+Q}$ model on
MNIST (test set).

###  3.2 Unconstrained routing entropy

In this subsection we are going to more-formally analyze the distribution of
the coupling coefficients

$c_{ij}=\frac{e^{b_{ij}+\sum_{r=1}^{t}\boldsymbol{v}_{j}^{r}\boldsymbol{u}_{j}\boldsymbol{W}_{ij}}}{\sum_{k}e^{b_{ik}+\sum_{r=1}^{t}\boldsymbol{v}_{k}^{r}\boldsymbol{u}_{k}\boldsymbol{W}_{ik}}}$
(10)

where $t$ indicates the target routing iterations.111for abuse of notation, in
this subsection we suppress the index $\xi$ Let us evaluate the $c_{ij}$ over
a non-yet trained model: as we saw also in Section 3.1, we have

$c_{ij}\approx\frac{1}{J}\ \forall i,j.$ (11)

When updating the parameters, following [3], we have

$\displaystyle\frac{\partial L}{\partial\boldsymbol{W}_{ij}}=$
$\displaystyle\left[\frac{\partial
L}{\partial\boldsymbol{v}_{j}}\frac{\partial\boldsymbol{v}_{j}}{\partial\boldsymbol{s}_{j}}\cdot
c_{ij}+\right.$ $\displaystyle+\left.\sum_{m=1}^{M}\left(\frac{\partial
L}{\partial\boldsymbol{v}_{m}}\frac{\partial\boldsymbol{v}_{m}}{\partial\boldsymbol{s}_{m}}\cdot\boldsymbol{\hat{u}}_{m|i}\frac{\partial
c_{im}}{\boldsymbol{\hat{u}}_{m|i}}\right)\right]\cdot\boldsymbol{u}_{i}$ (12)

where we can have the gradient for $\boldsymbol{W}_{ij}\approx 0$ in a
potentially-high number of scenarios, despite $c_{ij}\neq\\{0,1\\}$. Let us
analyze the simple case in which we have perfect outputs, matching the ground
truth, hence we are close to a local (or potentially the global) minimum of
the loss function:

$\left\|\frac{\partial L}{\partial\boldsymbol{v}_{m}}\right\|_{2}\approx 0\
\forall m.$ (13)

Looking at (4), we see that the right class is chosen, but given the squashing
function, we have as an explicit constraint that, given the $j$-th class as
the target one, we require

$\left\|\boldsymbol{v}_{j}\right\|_{2}\gg\left\|\boldsymbol{v}_{m}\right\|_{2}\
\forall m\neq j$ (14)

on the $\boldsymbol{W}_{ij}$, which can be accomplished in many ways,
including:

  * •

having sparse activation for the primary capsules $\boldsymbol{u}_{i}$: in
this case, we have constant $\boldsymbol{W}_{ij}$ (typically associated to no-
routing based approaches); however, we need heavier deep neural networks as
they have to force sparse signals already at the output of the primary
capsules. In this case, the coupling coefficients $c_{ij}$ are also constant
by definition;

  * •

having sparse votes $\boldsymbol{\hat{u}}_{j|i}$: this is a combination of
having both primary capsules and weights $\boldsymbol{W}_{ij}$ enforcing
sparsity in the votes, and the typical scenario with many routing iterations.

Having sparse votes, however, does not necessarily result in having sparse
coupling coefficients: according to (5), the coupling coefficients are
multiplied with the votes, obtaining the output capsules. The distribution of
the coupling coefficients requires (14) to be satisfied only: if
$\boldsymbol{W}_{ij}$ is not sparsely distributed, we can still have sparse
votes. However, this is the main reason we observe high entropy in the
coupling coefficient distributions (Fig. 4): as the votes
$\boldsymbol{\hat{u}}_{j|i}$ are implicitly sparse (yet also disordered, as we
are not explicitly imposing any structure in the coupling coefficients
distribution), the model is still able to learn but it finds a typical
solution where $c_{ij}$ are not sparse. However, we would like to have
sparsely distributed, recurrent $\boldsymbol{c}_{-j}$, establishing stable
relationships between the features extracted at primary capsules layer.
Minimizing explicitly the entropy term (8) is an intractable problem due to
the non-differentiability of the entropy term and of the quantization step (in
our considered setup) and due to the huge computational complexity to be
introduced at training time. Hence, we can try to implicitly enforce routing
entropy minimization by forcing a sparse and organized structure in the
coupling coefficients. Towards this end, one efficient solution relies in
enforcing sparsity in the $\boldsymbol{W}_{ij}$ representation: enforcing a
vote between the $i$-th primary capsule and the $j$-th output caps to be
exactly zero for any input, according to (10)

$c_{ij}=\frac{1}{\sum_{k}e^{b_{ik}+\sum_{r=1}^{t}\boldsymbol{v}_{k}^{r}\boldsymbol{u}_{k}\boldsymbol{W}_{ik}}}.$
(15)

In this way, having a lower variability in the $c_{ij}$ values (and hence
building more stable relationships between primary and output capsules),
straightforwardly we are also explicitly minimizing the entropy of the
quantized representations for the coupling coefficients. In the next
subsection, we are going to tailor a sparsity technique to accomplish such a
goal.

###  3.3 Enforcing REM with pruning

CapsNets are trained via standard back-propagation learning, minimizing some
loss function like margin loss. Our ultimate goal is to assess to which extent
a variation of the value of some parameter $\theta$ would affect the error on
the network output. In particular, the parameters not affecting the network
output can be pushed to zero in a soft manner, meaning that we can apply an
$\ell_{2}$ penalty term. A number of approaches have been proposed, especially
in the recent years [12, 15, 10]. One recent state-of-the-art approach,
LOBSTER [23] proposes to penalize the parameters by their gradient-weighted
$\ell_{2}$ norm, leading to the update rule

$\displaystyle\theta^{t+1}=$ $\displaystyle\theta^{t}-\eta
G\left[\frac{\partial
L}{\partial\theta^{t}}\right]-\lambda\theta^{t}\textrm{ReLU}\left[1-\left|\frac{\partial
L}{\partial\theta^{t}}\right|\right],$ (16)

where $G\left[\frac{\partial L}{\partial\theta^{t}}\right]$ is any gradient-
based optimization update (for SGD it is the plain gradient, but other
optimization strategies like Adam can be plugged) and $\eta,\lambda$ are two
positive hyper-parameters.
Such a strategy is particularly effective on standard convolutional neural
networks, and easy to plug in any back-propagation based learning system.
Furthermore, LOBSTER is a regularization strategy which can be plugged at any
learning stage, as it self-tunes the penalty introduced according to the
learning phase: for this non-intrusiveness in the complex and delicate routing
mechanism for CapsNets, it resulted a fair choice to enforce REM.

##  4 Experiments and Results

In this section we report the experiments and the results that we performed to
test REM. We first show the results on the MNIST dataset, reporting also how
the entropy and the accuracy values change during training. Then, we test REM
on more complex datasets such as Fashion-MNIST, CIFAR10 and Tiny ImageNet. We
also performed experiments to test the robustness to affine transformations of
CapsNets+REM. Except for Tiny ImageNet, we used the same architectures
configurations and augmentations described in [21].222we have removed the
decoder part of the network, see Sec. 4.3 for more details. We trained models
with five random seeds. We report the classification accuracy (%) and entropy
(averages and standard deviations), the sparsity (percentage of pruned
parameters, median) and the number of keys in the dictionary (median). The
code will be open-source released upon acceptance of the paper, to guarantee
the reproducibility of the results. The experiments were run on a Nvidia
Ampere A40 equipped with 48GB RAM, and the code uses PyTorch 1.9.

###  4.1 Preliminary results on MNIST

In order to assess our REM technique, we analyze in-depth the benefits of
pruning towards REM on the MNIST dataset. Nowadays, despite its outdatedness,
MNIST remains an omni-present benchmark for CapsNets [21, 3, 19].

Figure 5: Accuracy and entropy curves vs pruned parameters on MNIST (test
set).

Fig. 5 shows how the entropy (red line) and classification accuracy (blue
dotted line) changes as the sparsity increases during training. We can see
that at the beginning of the training stage the entropy is low because the
routing algorithm has not learned yet to correctly discriminate the
relationships between the capsules; but at the end of the training process REM
lowers the entropy of the network.

Figure 6: Coupling coefficients distributions for each class of a CapsNets+Q
and a CapsNet+REM on MNIST (test set).

In Fig. 6 we plot the distributions of the coupling coefficients for a
CapsNet+Q and a CapsNet+REM following the method described in Sec. 3.1. We can
see that the distributions of the CapsNet+REM model are sparser that the ones
for the CapsNet+Q model, namely we carve a stronger parse tree while achieving
high generalization. Table 4 shows that there is no performance loss when
pruning a CapsNet, even when REM is applied with its quantization levels.

Model | Accuracy | Sparsity
---|---|---
CapsNet | $99.57\pm_{0.0002}$ | $0$
CapsNet | $99.58\pm_{0.0003}$ | $85.53$
CapsNet+Q | $99.56\pm_{0.0003}$ | $0$
CapsNet+REM | $99.56\pm_{0.0002}$ | $85.53$
Table 4: Results for CapsNets, CapsNets+Q, CapsNets+REM on MNIST (test set). Model | Sparsity | $\mathbf{\textbf{\\#}_{0}}$ | $\mathbf{\textbf{\\#}_{1}}$ | $\mathbf{\textbf{\\#}_{2}}$ | $\mathbf{\textbf{\\#}_{3}}$ | $\mathbf{\textbf{\\#}_{4}}$ | $\mathbf{\textbf{\\#}_{5}}$ | $\mathbf{\textbf{\\#}_{6}}$ | $\mathbf{\textbf{\\#}_{7}}$ | $\mathbf{\textbf{\\#}_{8}}$ | $\mathbf{\textbf{\\#}_{9}}$ | Entropy
---|---|---|---|---|---|---|---|---|---|---|---|---
CapsNet+Q | 0 | $972$ | $489$ | $1007$ | $985$ | $904$ | $856$ | $939$ | $670$ | $968$ | $957$ | $9.53\pm_{0.54}$
CapsNet+REM | 85.53 | $40$ | $46$ | $15$ | $39$ | $12$ | $117$ | $182$ | $14$ | $310$ | $9$ | $4.16\pm_{1.59}$
Table 5: (Quantized) CapsNets number of keys for each class on MNIST (test
set).

Table 5 shows the number of keys in each class for CapsNets+REM and a
CapsNets+Q, namely a CapsNet where the quantization is applied without pruning
the network during training. We also report the number of keys for each class
and the entropy of the dictionary for the same quantized models shown in Table
4. We can see that the dimension of the dictionary for CapsNets+REM is lower
than the one for CapsNets+Q. Also the entropy measure for CapsNets+REM is
lower compared to CapsNets+Q, namely, REM has successufully built a parse tree
between primary and class capsules with stronger relationships. We now show ho
we can improve the explainability of such connections. Each key can be seen as
a cluster of images with shared features or parts. For each key of the
predicted class, we take all the images of the input dataset that belong to
that key, sorted according to the size of that cluster. Then, we overlap all
the images in a cluster. Figure 7 depicts a visual representation of some of
these clusters for a subset of object classes. We can see that a CapsNet+REM
can not only predict which digit is present in an input image, but it can also
show which are the shared object parts (such as different localized parts for
the digit 7 and 9) and their transformations (rotation and translation for the
digit 0) in a cluster.

Figure 7: Each digit is a cluster composed by the input images. Shared object
parts and transformations are extracted with the dictionary of a CapsNet+REM.

###  4.2 Results

We trained and tested CapsNets+REM on more complex datasets such as:

  * •

Fashion-MNIST, 28x28 grayscale images (10 classes);

  * •

CIFAR10, 32x32 RGB images (10 classes);

  * •

Tiny ImageNet, 64x64 RGB images (200 classes).

| F-MNIST | CIFAR10 | T-ImageNet
---|---|---|---
Model | Accuracy |  | Sparsity
---
Entropy | Accuracy |  | Sparsity
---
Entropy | Accuracy |  | Sparsity
---
Entropy
CapsNet+Q | $92.46\pm_{0.002}$ | $0$ | $8.64\pm_{1.15}$ | $78.42\pm_{0.026}$ | $0$ | $6.26\pm_{0.61}$ | $58.50\pm_{0.25}$ | $0$ | $5.18\pm_{0.67}$
CapsNet+REM | $92.62\pm_{0.001}$ | $80.71$ | $4.80\pm_{1.70}$ | $79.25\pm_{0.005}$ | $81.17$ | $4.15\pm_{0.62}$ | $54.02\pm_{0.14}$ | $44.27$ | $3.15\pm_{0.81}$
Table 6: Accuracy, sparsity and entropy results on on Fashion-MNIST, CIFAR10 and Tiny ImageNet (test set). | expanded MNIST | affNIST
---|---|---
Model | Accuracy | Accuracy |  | Sparsity
---
Entropy
CapsNet+Q | $99.22$ | $77.93\pm_{0.005}$ | $0$ | $8.64\pm_{1.15}$
CapsNet+REM | $99.22$ | $81.81\pm_{0.008}$ | $71.26$ | $8.45\pm_{1.10}$
CapsNet+Q | $99.36\pm_{0.0005}$ | $83.14\pm_{0.002}$ | $0$ | $8.45\pm_{0.99}$
CapsNet+REM | $99.48\pm_{0.0002}$ | $85.23\pm_{0.001}$ | $87.32$ | $5.93\pm_{1.39}$
Table 7: Accuracy, sparsity and entropy results on affNIST test set for under
and well-trained models.

Specific setup for Tiny ImageNet. In order to test also in transfer learning
scenarios, we train CapsNets on the Tiny ImageNet dataset, using in this case
only a pretrained ResNet18 model as back-bone and fine-tuning the CapsNet
parameters only. We used 10% of the training set as validation set and the
original validation set as test set.
Discussion. As we can see in Tables 6, a CapsNet+REM has an high percentage of
pruned parameters with a minimal performance loss. So this confirms our
hypothesis that CapsNets are over-parametrized. Furthermore, for the Fashion-
MNIST dataset the entropy of the CapsNet+REM is much lower that the CapsNet+Q,
while for the other datasets the differences between the two approaches is
less evident. This is an expected behavior: for CIFAR10 and Tiny ImageNet,
both CapsNets achieved poor classification results and this introduces noise
on the coupling coefficients, leading to higher uncertainty in the
classification task hence in the features selection.

Robustness to affine transformations. To test the robustness to affine
transformations of CapsNets+REM, we used expanded MNIST: a dataset composed by
padded and translated MNIST, in which each example is an MNIST digit placed
randomly on a black background of 40×40 pixels. We used the affNIST dataset as
test set, in which each example is an MNIST digit with a random small affine
transformation. We tested an under-trained CapsNet with early stopping which
achieved 99.22% accuracy on the expanded MNIST test set as in [21, 3]. We also
trained these models until convergence. We can see in Table 7 that the under-
trained networks entropies are high. Instead, a well-trained CapsNet+REM can
be robust to affine transformations and have a low entropy.

###  4.3 Limitation

A CapsNet is typically composed of an encoder and a decoder part, where the
latter is a reconstruction network with 3 fully connected layers [21]. In the
previously-discussed experiments, we have removed the decoder. The main
limitation of our work arise when computing the entropy of CapsNets trained
with the decoder. We observed that the entropy of a CapsNets+REM is almost the
same as that of a CapsNet+Q shown in Tables 5 and 6. Indeed, when the decoder
is used, the activity vector of an output capsule encodes richer
representations of the input. Sabour et al. [21] introduced the decoder to
boost the routing performance on MNIST by enforcing the pose encoding a
capsule. They also show that, when a perturbed activity vector is feeded to
the decoder, such perturbation affects the reconstruction. So capsules
representations are _equivariant_ , meaning that transformations applied to
the input are described by continuous changes in the output vector. In order
to verify if output capsules of a trained CapsNet+REM without the decoder (so
with low entropy) are still equivariant, we stacked on top of it the
reconstruction network, without training the encoder. We can see in Fig. 8
that the output capsules are still equivariant to many transformations so we
successfully obtained a CapsNet+REM with low entropy but with equivariant
representations encoded into capsules.

Figure 8: MNIST perturbation reconstructions of a freezed CapsNet+REM with a
decoder stacked on top of it.

##  5 Conclusion

This paper investigated and improved the explainability of the routing
algorithm in CapsNets with REM (Routing Entropy Minimization), which drives
the model parameters distribution towards low entropy configurations. We first
showed how to extract the parse tree of a CapsNet by discretizing its
connections and then collecting the possible parse trees in associative
arrays. Standard CapsNets show high entropy in the parse trees structures, as
an explicit prior on the coupling coefficients distribution is missing. We
showed how pruning methods, in low pruning regimes, naturally reduce such
entropy as well as the cardinality over the possible parse trees, testing such
a phenomenon on several datasets. Furthermore, we empirically observe that a
CapsNet+REM model remains robust to affine transformations and its output
capsules are still equivariant. REM opens research pathways towards the
distillation of parse trees and model interpretability, including the design
of a pruning technique specifically-designed for REM.

## References

  * [1] Parnian Afshar, Arash Mohammadi, and Konstantinos N Plataniotis.  Brain tumor type classification via capsule networks.  In 2018 25th IEEE International Conference on Image Processing (ICIP), pages 3129–3133. IEEE, 2018.
  * [2] Fabio De Sousa Ribeiro, Georgios Leontidis, and Stefanos Kollias.  Introducing routing uncertainty in capsule networks.  In Advances in Neural Information Processing Systems, volume 33, pages 6490–6502. Curran Associates, Inc., 2020.
  * [3] Jindong Gu and Volker Tresp.  Improving the robustness of capsule networks to image affine transformations.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
  * [4] Josef Gugglberger, David Peer, and Antonio Rodríguez-Sánchez.  Training deep capsule networks with residual connections.  In Igor Farkaš, Paolo Masulli, Sebastian Otte, and Stefan Wermter, editors, Artificial Neural Networks and Machine Learning – ICANN 2021, pages 541–552, Cham, 2021. Springer International Publishing.
  * [5] Taeyoung Hahn, Myeongjang Pyeon, and Gunhee Kim.  Self-routing capsule networks.  In Advances in Neural Information Processing Systems, volume 32, pages 7658–7667. Curran Associates, Inc., 2019.
  * [6] Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst.  Matrix capsules with EM routing.  In International Conference on Learning Representations, 2018.
  * [7] Taewon Jeong, Youngmin Lee, and Heeyoung Kim.  Ladder capsule network.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3071–3079. PMLR, 09–15 Jun 2019.
  * [8] Burak Kakillioglu, Ao Ren, Yanzhi Wang, and Senem Velipasalar.  3d capsule networks for object classification with weight pruning.  IEEE Access, 8:27393–27405, 2020.
  * [9] Adam Kosiorek, Sara Sabour, Yee Whye Teh, and Geoffrey E Hinton.  Stacked capsule autoencoders.  In Advances in Neural Information Processing Systems, volume 32, pages 15512–15522. Curran Associates, Inc., 2019.
  * [10] Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr.  Snip: Single-shot network pruning based on connection sensitivity.  arXiv preprint arXiv:1810.02340, 2018.
  * [11] Hongyang Li, Xiaoyang Guo, Bo Dai, Wanli Ouyang, and Xiaogang Wang.  Neural network encapsulation.  In Computer Vision – ECCV 2018, pages 266–282, Cham, 2018. Springer International Publishing.
  * [12] Christos Louizos, Max Welling, and Diederik P Kingma.  Learning sparse neural networks through $l\\_0$ regularization.  arXiv preprint arXiv:1712.01312, 2017.
  * [13] Christos Louizos, Max Welling, and Diederik P. Kingma.  Learning sparse neural networks through l_0 regularization.  In Proceedings of the 6th International Conference on Learning Representations, ICLR. OpenReview.net, 2018.
  * [14] Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov.  Variational dropout sparsifies deep neural networks.  In Proceedings of the 34th International Conference on Machine Learning, 2017.
  * [15] Pavlo Molchanov, Arun Mallya, Stephen Tyree, Iuri Frosio, and Jan Kautz.  Importance estimation for neural network pruning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11264–11272, 2019.
  * [16] Inyoung Paik, Taeyeong Kwak, and Injung Kim.  Capsule networks need an improved routing algorithm.  In Proceedings of The Eleventh Asian Conference on Machine Learning, volume 101 of Proceedings of Machine Learning Research, pages 489–502, Nagoya, Japan, 17–19 Nov 2019. PMLR.
  * [17] Mercedes E Paoletti, Juan Mario Haut, Ruben Fernandez-Beltran, Javier Plaza, Antonio Plaza, Jun Li, and Filiberto Pla.  Capsule networks for hyperspectral image classification.  IEEE Transactions on Geoscience and Remote Sensing, 57(4):2145–2160, 2018.
  * [18] J. Rajasegaran, V. Jayasundara, S. Jayasekara, H. Jayasekara, S. Seneviratne, and R. Rodrigo.  Deepcaps: Going deeper with capsule networks.  In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10717–10725, 2019.
  * [19] David Rawlinson, Abdelrahman Ahmed, and Gideon Kowadlo.  Sparse unsupervised capsules generalize better.  CoRR, abs/1804.06094, 2018.
  * [20] Riccardo Renzulli, Enzo Tartaglione, Attilio Fiandrotti, and Marco Grangetto.  Capsule networks with routing annealing.  In Igor Farkaš, Paolo Masulli, Sebastian Otte, and Stefan Wermter, editors, Artificial Neural Networks and Machine Learning – ICANN 2021, pages 529–540, Cham, 2021. Springer International Publishing.
  * [21] Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton.  Dynamic routing between capsules.  In Advances in Neural Information Processing Systems, volume 30, pages 3856–3866. Curran Associates, Inc., 2017.
  * [22] Ramprasaath R. Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra.  Grad-cam: Visual explanations from deep networks via gradient-based localization.  In 2017 IEEE International Conference on Computer Vision (ICCV), pages 618–626, 2017.
  * [23] Enzo Tartaglione, Andrea Bragagnolo, Attilio Fiandrotti, and Marco Grangetto.  Loss-based sensitivity regularization: towards deep sparse neural networks.  CoRR, abs/2011.09905, 2020.
  * [24] Enzo Tartaglione, Skjalg Lepsø y, Attilio Fiandrotti, and Gianluca Francini.  Learning sparse neural networks via sensitivity-driven regularization.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [25] Dilin Wang and Q. Liu.  An optimization view on dynamic routing between capsules.  In ICLR, 2018.
  * [26] Yongheng Zhao, Tolga Birdal, Haowen Deng, and Federico Tombari.  3d point capsule networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1009–1018, 2019.