Capsule Networks (CapsNets)

Based on the paper: Dynamic Routing Between Capsules, by Sara Sabour, Nicholas Frosst and Geoffrey E. Hinton (NIPS 2017).

Inspired in part from Huadong Liao's implementation: CapsNet-TensorFlow.

Introduction

Watch this video to understand the key ideas behind Capsule Networks:

You may also want to watch this video, which presents the main difficulties in this notebook:

Imports

To support both Python 2 and Python 3:

To plot pretty figures:

We will need NumPy and TensorFlow:

Reproducibility

Let's reset the default graph, in case you re-run this notebook without restarting the kernel:

Let's set the random seeds so that this notebook always produces the same output:

Load MNIST

Yes, I know, it's MNIST again. But hopefully this powerful idea will work as well on larger datasets, time will tell.

Let's look at what these hand-written digit images look like:

And these are the corresponding labels:

Now let's build a Capsule Network to classify these images. Here's the overall architecture, enjoy the ASCII art! ;-) Note: for readability, I left out two arrows: Labels → Mask, and Input Images → Reconstruction Loss.

                            Loss
                              ↑
                    ┌─────────┴─────────┐
      Labels → Margin Loss      Reconstruction Loss
                    ↑                   ↑
                  Length             Decoder
                    ↑                   ↑ 
             Digit Capsules ────Mask────┘
               ↖↑↗ ↖↑↗ ↖↑↗
             Primary Capsules
                    ↑      
               Input Images

We are going to build the graph starting from the bottom layer, and gradually move up, left side first. Let's go!

Input Images

Let's start by creating a placeholder for the input images (28×28 pixels, 1 color channel = grayscale).

Primary Capsules

The first layer will be composed of 32 maps of 6×6 capsules each, where each capsule will output an 8D activation vector:

To compute their outputs, we first apply two regular convolutional layers:

Note: since we used a kernel size of 9 and no padding (for some reason, that's what "valid" means), the image shrunk by 9-1=8 pixels after each convolutional layer (28×28 to 20×20, then 20×20 to 12×12), and since we used a stride of 2 in the second convolutional layer, the image size was divided by 2. This is how we end up with 6×6 feature maps.

Next, we reshape the output to get a bunch of 8D vectors representing the outputs of the primary capsules. The output of conv2 is an array containing 32×8=256 feature maps for each instance, where each feature map is 6×6. So the shape of this output is (batch size, 6, 6, 256). We want to chop the 256 into 32 vectors of 8 dimensions each. We could do this by reshaping to (batch size, 6, 6, 32, 8). However, since this first capsule layer will be fully connected to the next capsule layer, we can simply flatten the 6×6 grids. This means we just need to reshape to (batch size, 6×6×32, 8).

Now we need to squash these vectors. Let's define the squash() function, based on equation (1) from the paper:

$\operatorname{squash}(\mathbf{s}) = \dfrac{\|\mathbf{s}\|^2}{1 + \|\mathbf{s}\|^2} \dfrac{\mathbf{s}}{\|\mathbf{s}\|}$

The squash() function will squash all vectors in the given array, along the given axis (by default, the last axis).

Caution, a nasty bug is waiting to bite you: the derivative of $\|\mathbf{s}\|$ is undefined when $\|\mathbf{s}\|=0$, so we can't just use tf.norm(), or else it will blow up during training: if a vector is zero, the gradients will be nan, so when the optimizer updates the variables, they will also become nan, and from then on you will be stuck in nan land. The solution is to implement the norm manually by computing the square root of the sum of squares plus a tiny epsilon value: $\|\mathbf{s}\| \approx \sqrt{\sum\limits_i{{s_i}^2}\,\,+ \epsilon}$.

Now let's apply this function to get the output $\mathbf{u}_i$ of each primary capsules $i$ :

Great! We have the output of the first capsule layer. It wasn't too hard, was it? However, computing the next layer is where the fun really begins.

Digit Capsules

To compute the output of the digit capsules, we must first compute the predicted output vectors (one for each primary / digit capsule pair). Then we can run the routing by agreement algorithm.

Compute the Predicted Output Vectors

The digit capsule layer contains 10 capsules (one for each digit) of 16 dimensions each:

For each capsule $i$ in the first layer, we want to predict the output of every capsule $j$ in the second layer. For this, we will need a transformation matrix $\mathbf{W}_{i,j}$ (one for each pair of capsules ($i$, $j$)), then we can compute the predicted output $\hat{\mathbf{u}}_{j|i} = \mathbf{W}_{i,j} \, \mathbf{u}_i$ (equation (2)-right in the paper). Since we want to transform an 8D vector into a 16D vector, each transformation matrix $\mathbf{W}_{i,j}$ must have a shape of (16, 8).

To compute $\hat{\mathbf{u}}_{j|i}$ for every pair of capsules ($i$, $j$), we will use a nice feature of the tf.matmul() function: you probably know that it lets you multiply two matrices, but you may not know that it also lets you multiply higher dimensional arrays. It treats the arrays as arrays of matrices, and it performs itemwise matrix multiplication. For example, suppose you have two 4D arrays, each containing a 2×3 grid of matrices. The first contains matrices $\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D}, \mathbf{E}, \mathbf{F}$ and the second contains matrices $\mathbf{G}, \mathbf{H}, \mathbf{I}, \mathbf{J}, \mathbf{K}, \mathbf{L}$. If you multiply these two 4D arrays using the tf.matmul() function, this is what you get:

$ \pmatrix{ \mathbf{A} & \mathbf{B} & \mathbf{C} \\ \mathbf{D} & \mathbf{E} & \mathbf{F} } \times \pmatrix{ \mathbf{G} & \mathbf{H} & \mathbf{I} \\ \mathbf{J} & \mathbf{K} & \mathbf{L} } = \pmatrix{ \mathbf{AG} & \mathbf{BH} & \mathbf{CI} \\ \mathbf{DJ} & \mathbf{EK} & \mathbf{FL} } $

We can apply this function to compute $\hat{\mathbf{u}}_{j|i}$ for every pair of capsules ($i$, $j$) like this (recall that there are 6×6×32=1152 capsules in the first layer, and 10 in the second layer):

$ \pmatrix{ \mathbf{W}{1,1} & \mathbf{W}{1,2} & \cdots & \mathbf{W}{1,10} \ \mathbf{W}{2,1} & \mathbf{W}{2,2} & \cdots & \mathbf{W}{2,10} \ \vdots & \vdots & \ddots & \vdots \ \mathbf{W}{1152,1} & \mathbf{W}{1152,2} & \cdots & \mathbf{W}_{1152,10} } \times \pmatrix{ \mathbf{u}_1 & \mathbf{u}_1 & \cdots & \mathbf{u}_1 \ \mathbf{u}_2 & \mathbf{u}_2 & \cdots & \mathbf{u}2 \ \vdots & \vdots & \ddots & \vdots \ \mathbf{u}{1152} & \mathbf{u}{1152} & \cdots & \mathbf{u}{1152}

}

\pmatrix{ \hat{\mathbf{u}}{1|1} & \hat{\mathbf{u}}{2|1} & \cdots & \hat{\mathbf{u}}{10|1} \ \hat{\mathbf{u}}{1|2} & \hat{\mathbf{u}}{2|2} & \cdots & \hat{\mathbf{u}}{10|2} \ \vdots & \vdots & \ddots & \vdots \ \hat{\mathbf{u}}{1|1152} & \hat{\mathbf{u}}{2|1152} & \cdots & \hat{\mathbf{u}}_{10|1152} } $

The shape of the first array is (1152, 10, 16, 8), and the shape of the second array is (1152, 10, 8, 1). Note that the second array must contain 10 identical copies of the vectors $\mathbf{u}_1$ to $\mathbf{u}_{1152}$. To create this array, we will use the handy tf.tile() function, which lets you create an array containing many copies of a base array, tiled in any way you want.

Oh, wait a second! We forgot one dimension: batch size. Say we feed 50 images to the capsule network, it will make predictions for these 50 images simultaneously. So the shape of the first array must be (50, 1152, 10, 16, 8), and the shape of the second array must be (50, 1152, 10, 8, 1). The first layer capsules actually already output predictions for all 50 images, so the second array will be fine, but for the first array, we will need to use tf.tile() to have 50 copies of the transformation matrices.

Okay, let's start by creating a trainable variable of shape (1, 1152, 10, 16, 8) that will hold all the transformation matrices. The first dimension of size 1 will make this array easy to tile. We initialize this variable randomly using a normal distribution with a standard deviation to 0.1.

Now we can create the first array by repeating W once per instance:

That's it! On to the second array, now. As discussed earlier, we need to create an array of shape (batch size, 1152, 10, 8, 1), containing the output of the first layer capsules, repeated 10 times (once per digit, along the third dimension, which is axis=2). The caps1_output array has a shape of (batch size, 1152, 8), so we first need to expand it twice, to get an array of shape (batch size, 1152, 1, 8, 1), then we can repeat it 10 times along the third dimension:

Let's check the shape of the first array:

Good, and now the second:

Yes! Now, to get all the predicted output vectors $\hat{\mathbf{u}}_{j|i}$, we just need to multiply these two arrays using tf.matmul(), as explained earlier:

Let's check the shape:

Perfect, for each instance in the batch (we don't know the batch size yet, hence the "?") and for each pair of first and second layer capsules (1152×10) we have a 16D predicted output column vector (16×1). We're ready to apply the routing by agreement algorithm!

Routing by agreement

First let's initialize the raw routing weights $b_{i,j}$ to zero:

We will see why we need the last two dimensions of size 1 in a minute.

Round 1

First, let's apply the softmax function to compute the routing weights, $\mathbf{c}_{i} = \operatorname{softmax}(\mathbf{b}_i)$ (equation (3) in the paper):

Now let's compute the weighted sum of all the predicted output vectors for each second-layer capsule, $\mathbf{s}_j = \sum\limits_{i}{c_{i,j}\hat{\mathbf{u}}_{j|i}}$ (equation (2)-left in the paper):

There are a couple important details to note here:

• To perform elementwise matrix multiplication (also called the Hadamard product, noted $\circ$), we use the tf.multiply() function. It requires routing_weights and caps2_predicted to have the same rank, which is why we added two extra dimensions of size 1 to routing_weights, earlier.

• The shape of routing_weights is (batch size, 1152, 10, 1, 1) while the shape of caps2_predicted is (batch size, 1152, 10, 16, 1). Since they don't match on the fourth dimension (1 vs 16), tf.multiply() automatically broadcasts the routing_weights 16 times along that dimension. If you are not familiar with broadcasting, a simple example might help:

  $ \pmatrix{1 & 2 & 3 \\ 4 & 5 & 6} \circ \pmatrix{10 & 100 & 1000} = \pmatrix{1 & 2 & 3 \\ 4 & 5 & 6} \circ \pmatrix{10 & 100 & 1000 \\ 10 & 100 & 1000} = \pmatrix{10 & 200 & 3000 \\ 40 & 500 & 6000} $

And finally, let's apply the squash function to get the outputs of the second layer capsules at the end of the first iteration of the routing by agreement algorithm, $\mathbf{v}_j = \operatorname{squash}(\mathbf{s}_j)$ :

Good! We have ten 16D output vectors for each instance, as expected.

Round 2

First, let's measure how close each predicted vector $\hat{\mathbf{u}}_{j|i}$ is to the actual output vector $\mathbf{v}_j$ by computing their scalar product $\hat{\mathbf{u}}_{j|i} \cdot \mathbf{v}_j$.

• Quick math reminder: if $\vec{a}$ and $\vec{b}$ are two vectors of equal length, and $\mathbf{a}$ and $\mathbf{b}$ are their corresponding column vectors (i.e., matrices with a single column), then $\mathbf{a}^T \mathbf{b}$ (i.e., the matrix multiplication of the transpose of $\mathbf{a}$, and $\mathbf{b}$) is a 1×1 matrix containing the scalar product of the two vectors $\vec{a}\cdot\vec{b}$. In Machine Learning, we generally represent vectors as column vectors, so when we talk about computing the scalar product $\hat{\mathbf{u}}_{j|i} \cdot \mathbf{v}_j$, this actually means computing ${\hat{\mathbf{u}}_{j|i}}^T \mathbf{v}_j$.

Since we need to compute the scalar product $\hat{\mathbf{u}}_{j|i} \cdot \mathbf{v}_j$ for each instance, and for each pair of first and second level capsules $(i, j)$, we will once again take advantage of the fact that tf.matmul() can multiply many matrices simultaneously. This will require playing around with tf.tile() to get all dimensions to match (except for the last 2), just like we did earlier. So let's look at the shape of caps2_predicted, which holds all the predicted output vectors $\hat{\mathbf{u}}_{j|i}$ for each instance and each pair of capsules:

And now let's look at the shape of caps2_output_round_1, which holds 10 outputs vectors of 16D each, for each instance:

To get these shapes to match, we just need to tile the caps2_output_round_1 array 1152 times (once per primary capsule) along the second dimension:

And now we are ready to call tf.matmul() (note that we must tell it to transpose the matrices in the first array, to get ${\hat{\mathbf{u}}_{j|i}}^T$ instead of $\hat{\mathbf{u}}_{j|i}$):

We can now update the raw routing weights $b_{i,j}$ by simply adding the scalar product $\hat{\mathbf{u}}_{j|i} \cdot \mathbf{v}_j$ we just computed: $b_{i,j} \gets b_{i,j} + \hat{\mathbf{u}}_{j|i} \cdot \mathbf{v}_j$ (see Procedure 1, step 7, in the paper).

The rest of round 2 is the same as in round 1:

We could go on for a few more rounds, by repeating exactly the same steps as in round 2, but to keep things short, we will stop here:

Static or Dynamic Loop?

In the code above, we created different operations in the TensorFlow graph for each round of the routing by agreement algorithm. In other words, it's a static loop.

Sure, instead of copy/pasting the code several times, we could have written a for loop in Python, but this would not change the fact that the graph would end up containing different operations for each routing iteration. It's actually okay since we generally want less than 5 routing iterations, so the graph won't grow too big.

However, you may prefer to implement the routing loop within the TensorFlow graph itself rather than using a Python for loop. To do this, you would need to use TensorFlow's tf.while_loop() function. This way, all routing iterations would reuse the same operations in the graph, it would be a dynamic loop.

For example, here is how to build a small loop that computes the sum of squares from 1 to 100:

As you can see, the tf.while_loop() function expects the loop condition and body to be provided via two functions. These functions will be called only once by TensorFlow, during the graph construction phase, not while executing the graph. The tf.while_loop() function stitches together the graph fragments created by condition() and loop_body() with some additional operations to create the loop.

Also note that during training, TensorFlow will automagically handle backpropogation through the loop, so you don't need to worry about that.

Of course, we could have used this one-liner instead! ;-)

Joke aside, apart from reducing the graph size, using a dynamic loop instead of a static loop can help reduce how much GPU RAM you use (if you are using a GPU). Indeed, if you set swap_memory=True when calling the tf.while_loop() function, TensorFlow will automatically check GPU RAM usage at each loop iteration, and it will take care of swapping memory between the GPU and the CPU when needed. Since CPU memory is much cheaper and abundant than GPU RAM, this can really make a big difference.

Estimated Class Probabilities (Length)

The lengths of the output vectors represent the class probabilities, so we could just use tf.norm() to compute them, but as we saw when discussing the squash function, it would be risky, so instead let's create our own safe_norm() function:

To predict the class of each instance, we can just select the one with the highest estimated probability. To do this, let's start by finding its index using tf.argmax():

Let's look at the shape of y_proba_argmax:

That's what we wanted: for each instance, we now have the index of the longest output vector. Let's get rid of the last two dimensions by using tf.squeeze() which removes dimensions of size 1. This gives us the capsule network's predicted class for each instance:

Okay, we are now ready to define the training operations, starting with the losses.

Labels

First, we will need a placeholder for the labels:

Margin loss

The paper uses a special margin loss to make it possible to detect two or more different digits in each image:

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

• $T_k$ is equal to 1 if the digit of class $k$ is present, or 0 otherwise.
• In the paper, $m^{+} = 0.9$, $m^{-} = 0.1$ and $\lambda = 0.5$.
• Note that there was an error in the video (at 15:47): the max operations are squared, not the norms. Sorry about that.

Since y will contain the digit classes, from 0 to 9, to get $T_k$ for every instance and every class, we can just use the tf.one_hot() function:

A small example should make it clear what this does:

Now let's compute the norm of the output vector for each output capsule and each instance. First, let's verify the shape of caps2_output:

The 16D output vectors are in the second to last dimension, so let's use the safe_norm() function with axis=-2:

Now let's compute $\max(0, m^{+} - \|\mathbf{v}_k\|)^2$, and reshape the result to get a simple matrix of shape (batch size, 10):

Next let's compute $\max(0, \|\mathbf{v}_k\| - m^{-})^2$ and reshape it:

We are ready to compute the loss for each instance and each digit:

Now we can sum the digit losses for each instance ($L_0 + L_1 + \cdots + L_9$), and compute the mean over all instances. This gives us the final margin loss:

Reconstruction

Now let's add a decoder network on top of the capsule network. It is a regular 3-layer fully connected neural network which will learn to reconstruct the input images based on the output of the capsule network. This will force the capsule network to preserve all the information required to reconstruct the digits, across the whole network. This constraint regularizes the model: it reduces the risk of overfitting the training set, and it helps generalize to new digits.

Mask

The paper mentions that during training, instead of sending all the outputs of the capsule network to the decoder network, we must send only the output vector of the capsule that corresponds to the target digit. All the other output vectors must be masked out. At inference time, we must mask all output vectors except for the longest one, i.e., the one that corresponds to the predicted digit. You can see this in the paper's figure 2 (at 18:15 in the video): all output vectors are masked out, except for the reconstruction target's output vector.

We need a placeholder to tell TensorFlow whether we want to mask the output vectors based on the labels (True) or on the predictions (False, the default):

Now let's use tf.cond() to define the reconstruction targets as the labels y if mask_with_labels is True, or y_pred otherwise.

Note that the tf.cond() function expects the if-True and if-False tensors to be passed via functions: these functions will be called just once during the graph construction phase (not during the execution phase), similar to tf.while_loop(). This allows TensorFlow to add the necessary operations to handle the conditional evaluation of the if-True or if-False tensors. However, in our case, the tensors y and y_pred are already created by the time we call tf.cond(), so unfortunately TensorFlow will consider both y and y_pred to be dependencies of the reconstruction_targets tensor. The reconstruction_targets tensor will end up with the correct value, but:

1. whenever we evaluate a tensor that depends on reconstruction_targets, the y_pred tensor will be evaluated (even if mask_with_layers is True). This is not a big deal because computing y_pred adds no computing overhead during training, since we need it anyway to compute the margin loss. And during testing, if we are doing classification, we won't need reconstructions, so reconstruction_targets won't be evaluated at all.
2. we will always need to feed a value for the y placeholder (even if mask_with_layers is False). This is a bit annoying, but we can pass an empty array, because TensorFlow won't use it anyway (it just does not know it yet when it checks for dependencies).

Now that we have the reconstruction targets, let's create the reconstruction mask. It should be equal to 1.0 for the target class, and 0.0 for the other classes, for each instance. For this we can just use the tf.one_hot() function:

Let's check the shape of reconstruction_mask:

Let's compare this to the shape of caps2_output:

Mmh, its shape is (batch size, 1, 10, 16, 1). We want to multiply it by the reconstruction_mask, but the shape of the reconstruction_mask is (batch size, 10). We must reshape it to (batch size, 1, 10, 1, 1) to make multiplication possible:

At last! We can apply the mask:

One last reshape operation to flatten the decoder's inputs:

This gives us an array of shape (batch size, 160):

Decoder

Now let's build the decoder. It's quite simple: two dense (fully connected) ReLU layers followed by a dense output sigmoid layer:

Reconstruction Loss

Now let's compute the reconstruction loss. It is just the squared difference between the input image and the reconstructed image:

Final Loss

The final loss is the sum of the margin loss and the reconstruction loss (scaled down by a factor of 0.0005 to ensure the margin loss dominates training):

Final Touches

Accuracy

To measure our model's accuracy, we need to count the number of instances that are properly classified. For this, we can simply compare y and y_pred, convert the boolean value to a float32 (0.0 for False, 1.0 for True), and compute the mean over all the instances:

Training Operations

The paper mentions that the authors used the Adam optimizer with TensorFlow's default parameters:

Init and Saver

And let's add the usual variable initializer, as well as a Saver:

And... we're done with the construction phase! Please take a moment to celebrate. :)

Training

Training our capsule network is pretty standard. For simplicity, we won't do any fancy hyperparameter tuning, dropout or anything, we will just run the training operation over and over again, displaying the loss, and at the end of each epoch, measure the accuracy on the validation set, display it, and save the model if the validation loss is the lowest seen found so far (this is a basic way to implement early stopping, without actually stopping). Hopefully the code should be self-explanatory, but here are a few details to note:

• if a checkpoint file exists, it will be restored (this makes it possible to interrupt training, then restart it later from the last checkpoint),
• we must not forget to feed mask_with_labels=True during training,
• during testing, we let mask_with_labels default to False (but we still feed the labels since they are required to compute the accuracy),
• the images loaded via mnist.train.next_batch() are represented as float32 arrays of shape [784], but the input placeholder X expects a float32 array of shape [28, 28, 1], so we must reshape the images before we feed them to our model,
• we evaluate the model's loss and accuracy on the full validation set (5,000 instances). To view progress and support systems that don't have a lot of RAM, the code evaluates the loss and accuracy on one batch at a time, and computes the mean loss and mean accuracy at the end.

Warning: if you don't have a GPU, training will take a very long time (at least a few hours). With a GPU, it should take just a few minutes per epoch (e.g., 6 minutes on an NVidia GeForce GTX 1080Ti).

Training is finished, we reached over 99.4% accuracy on the validation set after just 5 epochs, things are looking good. Now let's evaluate the model on the test set.

Evaluation

We reach 99.53% accuracy on the test set. Pretty nice. :)

Predictions

Now let's make some predictions! We first fix a few images from the test set, then we start a session, restore the trained model, evaluate caps2_output to get the capsule network's output vectors, decoder_output to get the reconstructions, and y_pred to get the class predictions:

Note: we feed y with an empty array, but TensorFlow will not use it, as explained earlier.

And now let's plot the images and their labels, followed by the corresponding reconstructions and predictions:

The predictions are all correct, and the reconstructions look great. Hurray!

Interpreting the Output Vectors

Let's tweak the output vectors to see what their pose parameters represent.

First, let's check the shape of the cap2_output_value NumPy array:

Let's create a function that will tweak each of the 16 pose parameters (dimensions) in all output vectors. Each tweaked output vector will be identical to the original output vector, except that one of its pose parameters will be incremented by a value varying from -0.5 to 0.5. By default there will be 11 steps (-0.5, -0.4, ..., +0.4, +0.5). This function will return an array of shape (tweaked pose parameters=16, steps=11, batch size=5, 1, 10, 16, 1):

Let's compute all the tweaked output vectors and reshape the result to (parameters×steps×instances, 1, 10, 16, 1) so we can feed the array to the decoder:

Now let's feed these tweaked output vectors to the decoder and get the reconstructions it produces:

Let's reshape the decoder's output so we can easily iterate on the output dimension, the tweak steps, and the instances:

Lastly, let's plot all the reconstructions, for the first 3 output dimensions, for each tweaking step (column) and each digit (row):

Conclusion

I tried to make the code in this notebook as flat and linear as possible, to make it easier to follow, but of course in practice you would want to wrap the code in nice reusable functions and classes. For example, you could try implementing your own PrimaryCapsuleLayer, and DenseRoutingCapsuleLayer classes, with parameters for the number of capsules, the number of routing iterations, whether to use a dynamic loop or a static loop, and so on. For an example a modular implementation of Capsule Networks based on TensorFlow, take a look at the CapsNet-TensorFlow project.

That's all for today, I hope you enjoyed this notebook!