Matrix capsules with EM routing

Abstract: A capsule is a group of neurons whose outputs represent different properties of the same entity. Each layer in a capsule network contains many capsules. We describe a version of capsules in which each capsule has a logistic unit to represent the presence of an entity and a 4x4 matrix which could learn to represent the relationship between that entity and the viewer (the pose). A capsule in one layer votes for the pose matrix of many different capsules in the layer above by multiplying its own pose matrix by trainable viewpoint-invariant transformation matrices that could learn to represent part-whole relationships. Each of these votes is weighted by an assignment coefficient. These coefficients are iteratively updated for each image using the Expectation-Maximization algorithm such that the output of each capsule is routed to a capsule in the layer above that receives a cluster of similar votes. The transformation matrices are trained discriminatively by backpropagating through the unrolled iterations of EM between each pair of adjacent capsule layers. On the smallNORB benchmark, capsules reduce the number of test errors by 45\% compared to the state-of-the-art. Capsules also show far more resistance to white box adversarial attacks than our baseline convolutional neural network.

Reviews and author responses:

The objective function in details is:
\sum_c a'_c (-\beta_a) + a'_c ln(a'_c) + (1-a'_c)ln(1-a'_c)+\sum_h cost_{ch} + \sum_i a_i *  r_{ic} * ln(r_{ic})

a'_c is the activation for capsule c in layer L+1 and a_i is the activation probability for capsule i in layer L. The rest of the notations follow paper. 

Plots showing the decay of objective function and the absolute difference between two routing iterations in the above objective function can be found at:
https://imgur.com/a/eeD2X

The paper proposes a novel architecture for capsule networks. Each capsule has a logistic unit representing the presence of an entity plus a 4x4 pose matrix representing the entity/viewer relationship. This new representation comes with a novel iterative routing scheme, based on the EM algorithm.
Evaluated on the SmallNORB dataset, the approach proves to be more accurate than previous work (beating also the recently proposed "routing-by-agreement" approach for capsule networks by Sabour et al.). It also generalizes well to new, unseen viewpoints and proves to be more robust to adversarial examples than traditional CNNs.

Capsule networks have recently gained attention from the community. The paper addresses important shortcomings exhibited by previous work (Sabour et al.), introducing a series of valuable technical novelties.
There are, however, some weaknesses. The proposed routing scheme is quite complex (involving an EM-based step at each layer); it's not fully clear how efficiently it can be performed / how scalable it is. Evaluation is performed on a small dataset for shape recognition; as noted in Sec. 6, the approach will need to be tested on larger, more challenging datasets. Clarity could be improved in some parts of the paper (e.g.: Sec. 1.1 may not be fully clear if the reader is not already familiar with (Sabour et al., 2017); the authors could give a better intuition about what is kept and what is discarded, and why, from that approach. Sec. 2: the sentence "this is incorrect because the transformation matrix..." could be elaborated more. V_{ih} in eq. 1 is defined only a few lines below; perhaps, defining the variables before the equations could improve clarity. Sec. 2.1 could be accompanied by mathematical formulation).
All in all, the paper brings an original contribution and will encourage further research / discussion on an important research question (how to effectively leverage knowledge about the part-whole relationships).

Other notes:
- There are a few typos (e.g. Sec. 1.2 "(Jaderberg et al. (2015)",  Sec. 2 "the the transformation", Sec. 4 "cetral crop" etc.).
- The authors could discuss in more detail why the approach does not show significant improvement on NORB with respect to the state of the art.
- The authors could provide more insights about why capsule gradients are smaller than CNN ones.
- It would be interesting to discuss how the network could potentially be adapted, in the future, to: 1. be more efficient 2. take into account other changes produced by viewpoint changes (pixel intensities, as noted in Sec. 1).
- In Sec, 4, the authors could provide more details about the network training.
- In Procedure 1, for indexing tensors and matrices it might be better to use a comma to separate dimensions (e.g. V_{:,c,:} instead of V_{:c:}).

This paper proposes a new kind of capsules for CNN. The capsule contains a 4x4 pose matrix motivated by 3D geometric transformations describing the relationship between the viewer and the object (parts). An EM-type of algorithm is used to compute the routing.

The authors use the smallNORB dataset as an example. Since the scenes are simulated from different viewer angles, the pose matrix quite fits the motivation. It would be more beneficial to know if this kind of capsules is limited to the motivation or is general. For example, the authors may consider reporting the results of the affNIST dataset where the digits undergo 2D affine transformations (in which case perhaps 3x3 pose matrices are enough?).

Minor: The arguments in line 5 of the procedure RM Routing(a,V) do not match those in line 1 of the procedure E-Step.

Section 2.1 (objective of EM) is unclear. The authors may want to explicitly write down the free energy function.

The section about robustness against adversarial attacks is interesting.

Overall the idea appears to be useful but needs more empirical validation (affNIST, ImageNet, etc).


The paper describes another instantiation of "capsules" which attempt to learn part-whole relationships and the geometric pose transformations between them.  Results are presented on the smallNORB test set obtaining impressive performance.

Although I like very much this overall approach, this particular paper is so opaquely written that it is difficult to understand exactly what was done and how the network works.  It sounds like the main innovation here is using a 4x4 matrix for the pose parameters, and an iterative EM algorithm to find the correspondence between capsules (routing by agreement).  But what exactly the pose matrix represents, and how they get transformed from one layer to the next, is left almost entirely to the reader's imagination.  In addition, how EM factors in, what the probabilities P_ih represent, etc. is not clear.  I think the authors could do a much better job explaining this model, the rationale behind it, and how it works.

Perhaps the most interesting and compelling result is Figure 2, which shows how ambiguity in object class assignment is resolved with each iteration.  This is very intriguing, but it would be great to understand what is going on and how this is happening.

Although the results are impressive, if one can't understand how this was achieved it is hard to know what to make of it.



Thank you for your comments. upon reflection we agree that the paper was confusing and we have taken several steps to reduce the opacity of our work to the reader. To that end we have done the following: 
- We have added section 2 which gives a general and intuitive explanation of the mechanism of capsule networks, paying close attention to how pose matrices get transformed from one layer to the next.
- Having identified the EM objective as another source of confusion, we added an extended appendix in which we provide a gentle and approachable explanation for the free energy view of EM and how our routing algorithm builds upon it. 
- We have also added a paragraph to further explain figure 2 in the experiments section. 
- Finally we have made several changes to the language of the paper, focusing in particular on the notation.  
We believe that the comprehensibility of the paper has thus improved and appreciate your criticism. 

Thank you for your detailed reading of the paper and suggestions!    
As per your comments on the EM routing, we agree that it was not presented as best it could have been, and have added an appendix to present a gentle and thorough introduction to the free energy view of EM and the objective function which our routing operation minimizes. In response to the question about efficiency, we would like to draw your attention to the total number of arithmetic operations required for the routing procedure - each iteration of routing represents fewer arithmetic operations than a single layer feed forward pass, but due to architectural optimization decisions in tensorflow, our current capsule implementation is not as fast as it could be. 

We agree that larger scale testing would ideal, but due to the aforementioned efficiency limitations were not able to include it in this paper. 

In regards to your other comments we have done the following: 
- To increase the clarity of the paper,  we have made several changes to the language used, and improved the mathematical notation.
- We have added section 2 which provides an intuitive explanation of capsules and makes clear when the routing occurs. We feel that improves the readers' ability to engage with the rest of the presented content. We also defined the variables and notation used in the rest of the paper more explicitly.  
- We have expanded on the sentence "this is incorrect because the transformation matrix..." you mentioned which is now in the appendix. 
- We have also made several changes to the nation and language throughout the paper to make it more comprehensible. 
thank you for your feedback, and hope that we have addressed your comments to your satisfaction. 

thank you for the feedback! To address your comments we have done the following: 
- To clarify the EM objective we have added an extended and thorough appendix which presents a gentle and intuitive explanation of the free energy view of EM, and explicit free energy function, and how our routing algorithm makes use of it.
- We believe that the benefit of capsules is not limited to smallNORB and will generalize. As suggested, we replicated the affNIST generalization experiment reported in the previous Capsule paper (Sabour et al. 2017). We found that our EM capsule model (the exact architecture used for smallNORB and MNIST in the paper), when trained to 0.8% test error on expanded MNIST (40x40 pixel MNIST images, created by padding and shifting MNIST), achieved 6.9% test error on affNIST. We trained a baseline CNN (with AlexNet architecture, without pooling) to 0.8% test error and it was only able to achieve 14.1% test error on affNIST. Our capsule model was able to half the test error of a CNN when trained on MNIST and tested on affNIST.  Due to time and space constraints these results are not reported in the paper as it is now. 
- finally we address the minor issue raised in line 5 of the routing procedure. 
we hope this has addressed your concerns, and thank you for your suggestions. 

Author: Hang Yu | Suofei Zhang

Email: hangyu5 at illinois.edu | zhangsuofei at njupt.edu.cn

## Reproduce Method

#### Hyperparameters
smallNORB dataset:
* Samples per epoch: 46800
* Sample dimensions: 96x96x1
* Batch size: 50
* Preprocessing:
    * training:
        1. add random brightness with max delta equals 32 / 255.
        2. add random contrast with lower 0.5 and upper 1.5.
        3. resize into HxW 48x48 with bilinear method.
        4. crop into random HxW 32x32 piece.
        5. apply batch norm to have zero mean and unit variance.
        6. squash the image from 4 so that each entry has value from 0 to 1. This image is to be compared with the reconstructed image.
    * testing:
        1. resize into HxW 48x48 with bilinear method.
        2. crop the center HxW 32x32 piece.
        3. apply batch norm with moving mean and moving variance collected from training data set.

#### Method

1. The so called dynamic routing is in analog to the fully-connected layer in CNN. The so called ConvCaps structure extends dynamic routing into convolutional filter structure. The ConvCaps are implemented similarly as the dynamic routing for the whole feature map. The only difference is to tile the feature map into kernel-wise data and treat different kernels as batches. Then EM routing can be implemented within each batch in the same way as dynamic routing.

2. Different initialization strategies are used for convolutional filters. Linear weights are initialized with Xavier method. Biases are initialized with truncated normal distribution. This configuration provide higher numerical stability of input to EM algorithm.

3. The output of ConvCaps2 layer is processed by em routing with kernel size of 1*1. Then a global average pooling is deployed here to results final Class Capsules. Coordinate Addition is also injected during this stage.

4. Equation 2 in E-step of Procedure 1 from original paper is replaced by products of probabilities directly. All the probabilities are normalized into [0, 10] for higher numerical stability in products. Due to the division in Equation 3, this operation will not impact the final result. Exponent and logarithm are also used here for the same purpose.

5. A common l2 regularization of network parameters is considered in the loss function. Beside this, reconstruction loss and spread loss are implemented as the description in the original paper.

6. Learning rate: starts from 1e-3, then decays exponentially in a rate of 0.8 for every 46800/50 steps, and ends in 1e-5 (applied for all trainings).

7. We use Tensorflow 1.4 API and python programming language.

## Reproduce Result

#### Overview

Experiments on is done by Suofei Zhang. His hardware is:

* cpu：Intel(R) Xeon(R) CPU E5-2680 v4@ 2.40GHz，
* gpu：Tesla P40


**On test accuracy**:

smallNORB dataset test accuracy (our result/proposed result):

* CNN baseline (4.2M): 88.7%(best)/94.8%
* Matrix Cap with EM routing (310K, 2 iteration): 91.8%(best)/98.6%

There are two comments to make:

1. Even though the best of Matrix Cap is over by 3% to the best of CNN baseline, the test curve suggest Matrix Cap fluctuates between roughly 80% to 90% test dataset.
2. We are curious to know the learning curve and test curve that can be generated by the author.

**Training speed**:

1. CNN baseline costs 6m to train 50 epochs on smallNORB dataset. Each batch costs about 0.006s.
2. Matrix Cap costs 15h55m36s to train. Each batch costs about 1.2s.

**Recon image**:

Will come soon.

**routing histogram**:

We have difficulty in understanding how the histogram is calculated.

**AD attack**:

We haven't planned to run AD attack yet.

### Notes

> **Status:**
> According to github commit history, this reproduce project had its init commit on Nov.19th. We started writing this report on Dec.19th. Mainly, it is cost by undedicated code review so that we have to fix bug and run it again, otherwise the project should be able to finish in a week.

> **Current Results on smallNORB:**
- Configuration: A=32, B=8, C=16, D=16, batch_size=50, iteration number of EM routing: 2, with Coordinate Addition, spread loss, batch normalization
- Training loss. Variation of loss is suppressed by batch normalization. However, there still exists a gap between our best results and the reported results in the original paper.

- Test accuracy(current best result is 91.8%)

> **Current Results on MNIST:**
- Configuration: A=32, B=8, C=16, D=16, batch_size=50, iteration number of EM routing: 2, with Coordinate Addition, spread loss, batch normalization, reconstruction loss.

- Test accuracy(current best result is 99.3%, only 10% samples are used in test)

##Reference

[1] [MATRIX CAPSULES WITH EM ROUTING (paper)](https://openreview.net/pdf?id=HJWLfGWRb)

[2] [Matrix-Capsules-EM-Tensorflow (our github repo: code and comments)](https://github.com/www0wwwjs1/Matrix-Capsules-EM-Tensorflow/)


Thanks for all your research effort. It is great to read on this new paradigm and to see it actually working.

One part that is missing in my opinion, or I am very ignorantly glossing over it, is the downsampling from ConvCaps2 (L_final-1) to Class Capsules (L_final).

As mentioned, weights are shared among same entity capsules, so this would result in a one dimensional convolution (because K=1, stride=1), i.e. keeping the two spatial dimensions of the ConvCaps2 layer.
Whereas the other layer transitions indeed keep their spatial information and result in multiple, same-entity capsules spread over the 2 input image dimensions, the final layer only has one capsule for each class for the entire image.
IMHO this therefore requires a downsampling of the ConvCaps2 votes, a la maxpool, averagepool, or some extra dimension added to the EM routing algorithm.

beta_v and beta_a are per capsule type. Therefore, they are vectors for both convolutional capsules and final capsules. For example in terms of the notation in fig.1 beta_a and beta_v for convCaps1 are C dimensional vectors.

Thanks! We will revise the paper in regard to these points.

The dimensionality of the two trained beta parameters is not very clear to me from the paper. Are they shared across all capsules in the same layer (making them scalars) or does each capsule type have their own beta (meaning they are vectors).  I have had a look at the current implementation attempts of the model on GitHub and there the interpretations vary widely as well. Could you please clarify this point?

Minor notes on the algorithm (Procedure 1):
- "V_ich  is an H dimensional vote...": Did you mean V_ic?
- M-Step line 5: Missing quantifier. Like mu and sigma, cost_h is computed for all h

The spread-loss in 3.1 is the square of the WW-hinge-loss for multi-class SVMs, a large-margin loss.

See:

Weston, Jason; Watkins, Chris (1999). "Support Vector Machines for Multi-Class Pattern Recognition" (PDF). European Symposium on Artificial Neural Networks.

and the following paper describes the relations of the different variants of this loss:

http://jmlr.org/papers/v17/11-229.html
In the notation of that paper, it would be the combination of sum-over-others aggregation with relative margin concept and squared hinge loss.

For theoretical considerations, the log-probability should be used, in which case m  = 1 is fine and the last layer would not need to be normalized any more.

The sentence fragment:
   Spatial transformer networks (Jaderberg et al. (2015) seeks
is missing a ), and the subject is plural and not singular.  So it should be:
    Spatial transformer networks (Jaderberg et al. (2015)) seek

1) "V_ih is the product of the the transformation matrix W_ic that is learned discriminatively"
There is part of the sentense missing. Also, I believe this sentence describes "V_i" and not "V_ih". Suggestion:
" ... and V_ih is the value on dimension h of the vote V_i from capsule i to capsule c. V_i is obtained by taking the matrix product of the pose p_i of capsule i and the transformation Matrix W_ic. W_ic is learned discriminatively."

2) From what I understand, the vote V_i is a matrix (since it's obtained by multiplying a 4x4 matrix with a 4x4 matrix), and v_ih is a scalar. I found "V_ih is the value on dimension h of the vote ..." to be missleading. Maybe it should be mentioned that V_i has to be reshaped into a vector first and then its h'th entry is V_ih?

W_{ic} is 4*4 if you flatten the capsule types and grid positions. Therefore i goes over changes in the range of (1, channels * height * width) in this formulation.

However, We share the W_ic between different positions of two capsule types as in a convolutional layer with a kernel size k. Therefore, the total number of trainable parameters between two convolutional capsule layer types is 4*4*k*k and for the whole layer is 4*4*k*k*B*C. Where B is the number of different capsule types in layer bellow and C is the number of different capsule types in the next layer.

Please note that it is 4*4 rather than (4*4)*(4*4). 

As Jianfei has explained, the primary capsule layer is a convolutional layer with 1x1 kernel. It transforms the A channels in the first layer to B*(4x4+1) channels. Then we split the B*(4x4+1) channels into B*(4x4) as the pose matrices for B capsules and B*1 as the activation logits of B capsules in primary layer. Then we apply sigmoid nonlinearity on the activation logits.

In the convolutional capsule layers, what's the dimensionality of  transformation matrix W_{ic}?
Is it still (4*4)->(4*4) which correspond to a 1*1 linear convolutional layer?
or it is (4*4*k*k)->(4*4) which correspond to a k*k linear convolutional layer?

Figure 1 explains that. I guess they use a A*B*(4*4+1) kernel to (linear) transform a 1 width * 1 height * 32 channels patch to 32 capsules, each shape is 4*4+1. Then they reshape the 4*4 part as a matrix and apply a sigmoid on the 1 part. 

I still don't understand the transformation from convolution layer to primary capsule layer? Is it achieved by slicing 4x4*32 patches from the conv layer and then do a linear transformation for each 4x4 matrices?  what is the weight in "The activations of the primary capsules are produced by applying the sigmoid function to weighted sums of the same set of lower layer ReLUs." is it the 4x4 variable? I found it confusing, can you elaborate how this works.

Can you write down what exactly is the objective function in Section 2.1?

They gain a lot by using the meta data at test time. Without using that information (which normally is not available at test time) they get 2.6%. 

The meta data is not used during test time only during training time.

1.5% error rate has previously been reported on small NORB.
https://www.researchgate.net/publication/265335724_Nonlinear_Supervised_Locality_Preserving_Projections_for_Visual_Pattern_Discrimination

Meta-review: Authors present a new multi-layered capsule network architecture, implemented an EM routing procedure, and introduced "Coordinate Addition".  Capsule architectures are gaining interest because of their ability to achieve equivariance of parts, and employ a new form of pooling called "routing" (as opposed to max pooling) which groups parts that make similar predictions of the whole to which they belong, rather than relying on spatial co-locality. New state-of-art performances are being achieved on focused datasets, for which the authors have continued the trend.

Pros:
- New significant improvement to state-of-art performance is obtained on smallNORB, both in comparison to CNN structure as well as the most recent previous implementation of capsule network.

Cons:
- Some concern arose regarding the writing of the paper and the ability to understand the material, which authors have made an effort to address.

Given the general consensus of the reviewers that this work should be accepted, the general applicability of the technology to multiple domains, and the potential impact that improvements to capsule networks may have on an early field, area chair recommends this work be accepted as a poster presentation. 

Decision: accepted-poster-papers