Dynamic Routing Between Capsules

TL;DR

This paper introduces a capsule network architecture that uses dynamic routing to improve recognition of overlapping objects, achieving state-of-the-art results on MNIST by modeling hierarchical relationships more effectively.

Contribution

The paper proposes a novel capsule network with dynamic routing that better captures hierarchical pose relationships and improves recognition accuracy over traditional convolutional networks.

Findings

State-of-the-art performance on MNIST

Superior recognition of overlapping digits

Effective hierarchical pose modeling

Abstract

A capsule is a group of neurons whose activity vector represents the instantiation parameters of a specific type of entity such as an object or an object part. We use the length of the activity vector to represent the probability that the entity exists and its orientation to represent the instantiation parameters. Active capsules at one level make predictions, via transformation matrices, for the instantiation parameters of higher-level capsules. When multiple predictions agree, a higher level capsule becomes active. We show that a discrimininatively trained, multi-layer capsule system achieves state-of-the-art performance on MNIST and is considerably better than a convolutional net at recognizing highly overlapping digits. To achieve these results we use an iterative routing-by-agreement mechanism: A lower-level capsule prefers to send its output to higher level capsules whose activity vectors have a big scalar product with the prediction coming from the lower-level capsule.

Tables

Table 1. Figure 3: Sample MNIST test reconstructions of a CapsNet with 3 routing iterations. ( l , p , r ) 𝑙 𝑝 𝑟 (l,p,r) represents the label, the prediction and the reconstruction target respectively. The two rightmost columns show two reconstructions of a failure example and it explains how the model confuses a 5 5 5 and a 3 3 3 in this image. The other columns are from correct classifications and shows that model preserves many of the details while smoothing the noise.

$(l,p,r)$	$(2,2,2)$	$(5,5,5)$	$(8,8,8)$	$(9,9,9)$	$(5,3,5)$	$(5,3,3)$
\pbox 1cmInput
Output

Table 2. Table 1: CapsNet classification test accuracy. The MNIST average and standard deviation results are reported from 3 3 3 trials.

Method	Routing	Reconstruction	MNIST (%)	MultiMNIST (%)
Baseline	-	-	$0.39$	$8.1$
CapsNet	1	no	${0.34}_{\pm 0.032}$	-
CapsNet	1	yes	${0.29}_{\pm 0.011}$	$7.5$
CapsNet	3	no	${0.35}_{\pm 0.036}$	-
CapsNet	3	yes	${0.25}_{\pm 0.005}$	$5.2$

Table 3. Figure 4: Dimension perturbations. Each row shows the reconstruction when one of the 16 16 16 dimensions in the DigitCaps representation is tweaked by intervals of 0.05 0.05 0.05 in the range [ − 0.25 , 0.25 ] 0.25 0.25 [-0.25,0.25] .

Scale and thickness
Localized part
Stroke thickness
Localized skew
Width and translation
Localized part

Table 4. Figure 5: Sample reconstructions of a CapsNet with 3 routing iterations on MultiMNIST test dataset. The two reconstructed digits are overlayed in green and red as the lower image. The upper image shows the input image. L: ( l 1 , l 2 ) subscript 𝑙 1 subscript 𝑙 2 (l_{1},l_{2}) represents the label for the two digits in the image and R: ( r 1 , r 2 ) subscript 𝑟 1 subscript 𝑟 2 (r_{1},r_{2}) represents the two digits used for reconstruction. The two right most columns show two examples with wrong classification reconstructed from the label and from the prediction (P). In the ( 2 , 8 ) 2 8 (2,8) example the model confuses 8 8 8 with a 7 7 7 and in ( 4 , 9 ) 4 9 (4,9) it confuses 9 9 9 with 0 0 . The other columns have correct classifications and show that the model accounts for all the pixels while being able to assign one pixel to two digits in extremely difficult scenarios (column 1 − 4 1 4 1-4 ). Note that in dataset generation the pixel values are clipped at 1. The two columns with the (*) mark show reconstructions from a digit that is neither the label nor the prediction. These columns suggests that the model is not just finding the best fit for all the digits in the image including the ones that do not exist. Therefore in case of ( 5 , 0 ) 5 0 (5,0) it cannot reconstruct a 7 7 7 because it knows that there is a 5 5 5 and 0 0 that fit best and account for all the pixels. Also, in case of ( 8 , 1 ) 8 1 (8,1) the loop of 8 8 8 has not triggered 0 0 because it is already accounted for by 8 8 8 . Therefore it will not assign one pixel to two digits if one of them does not have any other support.

R:$(2,7)$	R:$(6,0)$	R:$(6,8)$	R:$(7,1)$	*R:$(5,7)$	*R:$(2,3)$	R:$(2,8)$	R:P:$(2,7)$
L:$(2,7)$	L:$(6,0)$	L:$(6,8)$	L:$(7,1)$	L:$(5,0)$	L:$(4,3)$	L:$(2,8)$	L:$(2,8)$
R:$(8,7)$	R:$(9,4)$	R:$(9,5)$	R:$(8,4)$	*R:$(0,8)$	*R:$(1,6)$	R:$(4,9)$	R:P:$(4,0)$
L:$(8,7)$	L:$(9,4)$	L:$(9,5)$	L:$(8,4)$	L:$(1,8)$	L:$(7,6)$	L:$(4,9)$	L:$(4,9)$

Equations

$${\mathbf{v}}_j=\frac{||{\mathbf{s}}_j|{|}^2}{1+||{\mathbf{s}}_j|{|}^2}\frac{{\mathbf{s}}_j}{||{\mathbf{s}}_j||}$$

$${\mathbf{s}}_j=\sum _i{c}_{ij}{\hat{\mathbf{u}}}_{j|i},{\hat{\mathbf{u}}}_{j|i}={\mathbf{W}}_{ij}{\mathbf{u}}_i$$

$$c_{ij}=\frac{\exp (b_{ij})}{{\sum }_k\exp (b_{ik})}$$

$$L_k=T_k\max (0,m^{+}-||{\mathbf{v}}_k||{)}^2+\lambda (1-T_k)\max (0,||{\mathbf{v}}_k||-m^{-}{)}^2$$

Full text

1 Introduction

Human vision ignores irrelevant details by using a carefully determined sequence of fixation points to ensure that only a tiny fraction of the optic array is ever processed at the highest resolution. Introspection is a poor guide to understanding how much of our knowledge of a scene comes from the sequence of fixations and how much we glean from a single fixation, but in this paper we will assume that a single fixation gives us much more than just a single identified object and its properties. We assume that our multi-layer visual system creates a parse tree-like structure on each fixation, and we ignore the issue of how these single-fixation parse trees are coordinated over multiple fixations.

Parse trees are generally constructed on the fly by dynamically allocating memory. Following Hinton et al. (2000), however, we shall assume that, for a single fixation, a parse tree is carved out of a fixed multilayer neural network like a sculpture is carved from a rock. Each layer will be divided into many small groups of neurons called “capsules” (Hinton et al. (2011)) and each node in the parse tree will correspond to an active capsule. Using an iterative routing process, each active capsule will choose a capsule in the layer above to be its parent in the tree. For the higher levels of a visual system, this iterative process will be solving the problem of assigning parts to wholes.

The activities of the neurons within an active capsule represent the various properties of a particular entity that is present in the image. These properties can include many different types of instantiation parameter such as pose (position, size, orientation), deformation, velocity, albedo, hue, texture, etc. One very special property is the existence of the instantiated entity in the image. An obvious way to represent existence is by using a separate logistic unit whose output is the probability that the entity exists. In this paper we explore an interesting alternative which is to use the overall length of the vector of instantiation parameters to represent the existence of the entity and to force the orientation of the vector to represent the properties of the entity111This makes biological sense as it does not use large activities to get accurate representations of things that probably don’t exist.. We ensure that the length of the vector output of a capsule cannot exceed $1$ by applying a non-linearity that leaves the orientation of the vector unchanged but scales down its magnitude.

The fact that the output of a capsule is a vector makes it possible to use a powerful dynamic routing mechanism to ensure that the output of the capsule gets sent to an appropriate parent in the layer above. Initially, the output is routed to all possible parents but is scaled down by coupling coefficients that sum to $1$. For each possible parent, the capsule computes a “prediction vector” by multiplying its own output by a weight matrix. If this prediction vector has a large scalar product with the output of a possible parent, there is top-down feedback which increases the coupling coefficient for that parent and decreasing it for other parents. This increases the contribution that the capsule makes to that parent thus further increasing the scalar product of the capsule’s prediction with the parent’s output. This type of “routing-by-agreement” should be far more effective than the very primitive form of routing implemented by max-pooling, which allows neurons in one layer to ignore all but the most active feature detector in a local pool in the layer below. We demonstrate that our dynamic routing mechanism is an effective way to implement the “explaining away” that is needed for segmenting highly overlapping objects.

Convolutional neural networks (CNNs) use translated replicas of learned feature detectors. This allows them to translate knowledge about good weight values acquired at one position in an image to other positions. This has proven extremely helpful in image interpretation. Even though we are replacing the scalar-output feature detectors of CNNs with vector-output capsules and max-pooling with routing-by-agreement, we would still like to replicate learned knowledge across space. To achieve this, we make all but the last layer of capsules be convolutional. As with CNNs, we make higher-level capsules cover larger regions of the image. Unlike max-pooling however, we do not throw away information about the precise position of the entity within the region. For low level capsules, location information is “place-coded” by which capsule is active. As we ascend the hierarchy, more and more of the positional information is “rate-coded” in the real-valued components of the output vector of a capsule. This shift from place-coding to rate-coding combined with the fact that higher-level capsules represent more complex entities with more degrees of freedom suggests that the dimensionality of capsules should increase as we ascend the hierarchy.

2 How the vector inputs and outputs of a capsule are computed

There are many possible ways to implement the general idea of capsules. The aim of this paper is not to explore this whole space but simply to show that one fairly straightforward implementation works well and that dynamic routing helps.

We want the length of the output vector of a capsule to represent the probability that the entity represented by the capsule is present in the current input. We therefore use a non-linear "squashing" function to ensure that short vectors get shrunk to almost zero length and long vectors get shrunk to a length slightly below $1$. We leave it to discriminative learning to make good use of this non-linearity.

[TABLE]

where ${\bf v}_{j}$ is the vector output of capsule $j$ and ${\mathbf{s}}_{j}$ is its total input.

For all but the first layer of capsules, the total input to a capsule ${\bf s}_{j}$ is a weighted sum over all “prediction vectors” ${\bf\hat{u}}_{j|i}$ from the capsules in the layer below and is produced by multiplying the output ${\bf u}_{i}$ of a capsule in the layer below by a weight matrix ${\bf W}_{ij}$

[TABLE]

where the $c_{ij}$ are coupling coefficients that are determined by the iterative dynamic routing process.

The coupling coefficients between capsule $i$ and all the capsules in the layer above sum to $1$ and are determined by a “routing softmax” whose initial logits $b_{ij}$ are the log prior probabilities that capsule $i$ should be coupled to capsule $j$.

[TABLE]

The log priors can be learned discriminatively at the same time as all the other weights. They depend on the location and type of the two capsules but not on the current input image222For MNIST we found that it was sufficient to set all of these priors to be equal.. The initial coupling coefficients are then iteratively refined by measuring the agreement between the current output ${\bf v}_{j}$ of each capsule, $j$, in the layer above and the prediction ${\bf\hat{u}}_{j|i}$ made by capsule $i$.

The agreement is simply the scalar product $a_{ij}={\bf v}_{j}.{\bf\hat{u}}_{j|i}$. This agreement is treated as if it was a log likelihood and is added to the initial logit, $b_{ij}$ before computing the new values for all the coupling coefficients linking capsule $i$ to higher level capsules.

In convolutional capsule layers, each capsule outputs a local grid of vectors to each type of capsule in the layer above using different transformation matrices for each member of the grid as well as for each type of capsule.

3 Margin loss for digit existence

We are using the length of the instantiation vector to represent the probability that a capsule’s entity exists. We would like the top-level capsule for digit class $k$ to have a long instantiation vector if and only if that digit is present in the image. To allow for multiple digits, we use a separate margin loss, $L_{k}$ for each digit capsule, $k$:

[TABLE]

where $T_{k}=1$ iff a digit of class $k$ is present333We do not allow an image to contain two instances of the same digit class. We address this weakness of capsules in the discussion section. and $m^{+}=0.9$ and $m^{-}=0.1$. The $\lambda$ down-weighting of the loss for absent digit classes stops the initial learning from shrinking the lengths of the activity vectors of all the digit capsules. We use $\lambda=0.5$. The total loss is simply the sum of the losses of all digit capsules.

4 CapsNet architecture

A simple CapsNet architecture is shown in Fig. 1. The architecture is shallow with only two convolutional layers and one fully connected layer. Conv$1$ has $256$, $9\times 9$ convolution kernels with a stride of 1 and ReLU activation. This layer converts pixel intensities to the activities of local feature detectors that are then used as inputs to the primary capsules.

The primary capsules are the lowest level of multi-dimensional entities and, from an inverse graphics perspective, activating the primary capsules corresponds to inverting the rendering process. This is a very different type of computation than piecing instantiated parts together to make familiar wholes, which is what capsules are designed to be good at.

The second layer (PrimaryCapsules) is a convolutional capsule layer with $32$ channels of convolutional $8$D capsules (i.e. each primary capsule contains 8 convolutional units with a $9\times 9$ kernel and a stride of 2). Each primary capsule output sees the outputs of all $256\times 81$ Conv$1$ units whose receptive fields overlap with the location of the center of the capsule. In total PrimaryCapsules has $[32\times 6\times 6]$ capsule outputs (each output is an $8$D vector) and each capsule in the $[6\times 6]$ grid is sharing their weights with each other. One can see PrimaryCapsules as a Convolution layer with Eq. 1 as its block non-linearity. The final Layer (DigitCaps) has one $16$D capsule per digit class and each of these capsules receives input from all the capsules in the layer below.

We have routing only between two consecutive capsule layers (e.g. PrimaryCapsules and DigitCaps). Since Conv$1$ output is $1$D, there is no orientation in its space to agree on. Therefore, no routing is used between Conv$1$ and PrimaryCapsules. All the routing logits ($b_{ij}$) are initialized to zero. Therefore, initially a capsule output (${\bf u}_{i}$) is sent to all parent capsules (${\bf v}_{0}...{\bf v}_{9}$) with equal probability ($c_{ij}$).

Our implementation is in TensorFlow (Abadi et al. (2016)) and we use the Adam optimizer (Kingma and Ba (2014)) with its TensorFlow default parameters, including the exponentially decaying learning rate, to minimize the sum of the margin losses in Eq. 4.

4.1 Reconstruction as a regularization method