Training capsules as a routing-weighted product of expert neurons
Michael Hauser
cs.NE cs.LG stat.ML


Role of coherence in the plasmonic control of molecular absorption
    Stefano Corni
    December 30, 2023
==================================================================




    Capsules are the multidimensional analogue to scalar neurons in neural networks, and because they are multidimensional, much more complex routing schemes can be used to pass information forward through the network than what can be used in traditional neural networks. This work treats capsules as collections of neurons in a fully connected neural network, where sub-networks connecting capsules are weighted according to the routing coefficients determined by routing by agreement. An energy function is designed to reflect this model, and it follows that capsule networks with dynamic routing can be formulated as a product of expert neurons. By alternating between dynamic routing, which acts to both find subnetworks within the overall network as well as to mix the model distribution, and updating the parameters by the gradient of the contrastive divergence, a bottom-up, unsupervised learning algorithm is constructed for capsule networks with dynamic routing. The model and its training algorithm are qualitatively tested in the generative sense, and is able to produce realistic looking images from standard vision datasets.








§ INTRODUCTION


Capsule networks <cit.> with routing by agreement <cit.> are a substantial departure from standard feedforward neural networks in the sense of the layer wise process by which they transform the data representation. The McCulloch-Pitts model <cit.> takes the forward activation of a neuron to be a weighted sum of the scalar-valued activations of the previous layer neurons, where the weighting is determined by the parameter weights. Some deep learning functions break from this model, such as the softmax and max pooling functions, since to which node the softmax and max pooling functions route data depend not only on the previous layer activations, but on the activations of neurons in the current layer as well. Most feedforward deep learning systems, such as the original convolutional LeNet <cit.>, or its larger counterparts such as AlexNet <cit.>, are such that (i) they transform scalar values from neurons and (ii) the routing functions are either non existent or extremely simple.

Comparatively, capsule networks take the forward activation of a capsule to be vector valued <cit.>. One advantage of a vector valued capsule over a scalar valued neuron is that with a vector one can separate the activation of the capsule from what the capsule is representing. This is done by encoding the activation of the capsule into the magnitude of the vector, and the instantiation parameters for representing the object in the orientation of the vector.

A second advantage of vector valued capsules is that one can develop much more complex routing schemes to pass information forward through the network. With routing by agreement <cit.>, each vector is linearly mapped to another vector by a learned prediction matrix, and a weighted sum of these predicted vectors becomes the vector to be passed through a nonlinear activation, where the weights of the weighted sum are determined online by routing by agreement. This type of forward pass of information, or routing, is substantially different than the parameter-weighted sum of scalar activations of the McCulloch-Pitts model. It is different because the scalar routing weights determined by the routing by agreement is a complex nonlinear weighting of all of the next layer activations, whereas the McCulloch-Pitts model has no dependence on the next layer activations, the softmax function is a linear weighting of the next layer activations and the max pooling function is a simple nonlinear function of the next layer activations and discards a large amount of information.

Because capsule network models with dynamic routing are substantially different from neural network models in the sense described above, it isn't necessarily straight-forward to extend the learning algorithms that were designed for neural networks to capsule networks. Nevertheless, many of the neural network learning algorithms are quite general and can be extended to capsule networks without much modification, which has led to using backpropagation <cit.> and expectation-maximization <cit.> in capsule networks. Similarly, the unsupervised learning method of generative adversarial networks <cit.> has been extended to capsule networks <cit.>, but this is built off of backpropagation <cit.>. 

We build on these efforts by developing an unsupervised learning algorithm for capsule networks with routing by agreement. We do this by considering the capsule network with routing by agreement as a routing-weighted product of expert neurons <cit.>. In this way collections of scalar valued neurons are grouped together to form vector valued capsules, and as capsules we can leverage the routing by agreement algorithm to dynamically route the information.



§ ALGORITHM REVIEW


Because the techniques developed in this paper are based substantially off of the mathematics and intuitions from both product of experts learning as well as routing by agreement, we will briefly review these techniques.



 §.§ Product of experts learning


Product of experts models <cit.> are in a way similar to mixture models, except one replaces the sum of the densities with a product of densities. The intuition is that mixture models act as logical or units since each component of the mixture model can alone say if the data belongs to the distribution or not, whereas a product of experts is analogous to logical and units, since any of the experts composing the product can exclude the data belonging to the distribution with a low-enough probability.

For models that distinguish between a visible layer and a hidden layer, where connections do not exist within each layer, only across layers, the energy of a certain configuration of visible and hidden binary neurons is defined as follows:


    E(v,h) =  -
    ∑_i,j h_j w_ij v_i -
    ∑_i b_i v_i - 
    ∑_j c_j h_j


With this, the probability for a configuration of visible and hidden units is
p(v,h) = 1/Z e^-E(v,h) where Z is the normalizing partition function. In the product of experts formulation, the density over the visible states, after marginalizing out the binary hidden states h_j, is given by


    p(v) =
        ∑_h
        p(v,h) =
        1/Z∑_h
        e^-E(v,h)


It follows that the gradient of the log likelihood is given by


    ∂log p/∂Θ(v|Θ) =
    - ∑_h p(h|v) ∂ E/∂Θ(v,h) +
    ∑_v p(v) ∑_h p(h|v) ∂ E/∂Θ(v,h)



On the right-hand side of this equation, the first term is the expectation of the gradient of the energy, with the expectation being over the true data density. The second term is the expectation of the gradient of the energy, with the expectation being over the data generated by the product of experts model itself. This then leads to the well known learning rule:


    ∂log p/∂ w_ij(v|Θ) =
        ⟨ v_i h_j ⟩_data -
        ⟨ v_i h_j ⟩_model


Finding the model distribution requires mixing with the Markov chain generated by reconstructing the visible and hidden units to infinity. Instead, contrastive divergence <cit.> minimizes the difference of (i) the KL divergence between the data and the model at infinity with (ii) the one-step Markov mixing and the model at infinity, which means the Markov chain only needs to go through one iteration for each gradient update.



 §.§ Routing by agreement in capsule networks


As mentioned earlier, capsule networks <cit.> with routing by agreement <cit.> are a substantial rethinking in the forward flow of information through the network, and this will be briefly reviewed.

At layer l, given a collection of vector-valued capsules x^(l)_i for i=1,2,…,I and matrix-valued prediction maps W^(l)_ij, the pre-activation, vector-valued predicted capsule j from capsule i is z^(l+1)_j|i=W^(l)_ij· x^(l)_i. We then take a weighted average of all of the predictions made at layer l to yield the final, pre-activation, vector-valued capsule j at layer l+1, i.e.

    z^(l+1)_j=∑_ic^(l)_ij W^(l)_ij· x^(l)_i

where the scalar-valued c^(l)_ij's are determined by routing by agreement and ∑_ic^(l)_ij=1. This is done for all capsules j=1,2,…,J at layer l+1. 

Routing by agreement is an iterative procedure to re-weigh each of the individual predictions to yield a final prediction, of a similar flavor to boosting weak learners to yield a strong learner. If an individual prediction z^(l+1)_j|i agrees well with the collective prediction z^(l+1)_j, then we would like to increase the i to j routing weight c^(l)_ij, whereas if they disagree we would like to decrease that same routing weight. Agreement has been measured as both the inner product <cit.>, as well as the cosine distance <cit.>, between the individual predictions and the squashed collective prediction. In this work we use the cosine distance.

The squashing function we use, from <cit.>, is defined as follows:


    squash(z^(l+1)_j) )=
        ‖ z^(l+1)_j‖^2/1+‖ z^(l+1)_j‖^2z^(l+1)_j/‖ z^(l+1)_j‖


The intention of the squashing function is to scale the magnitude of z^(l+1)_j between 0 and 1 while keeping the orientation the same. We can then interpret capsule j is on if the magnitude of the squashing function is close to 1, and j is off if the magnitude is close to 0:

    P(j=on|x^(l)_1,…,x^(l)_I) = 
        ‖squash(z^(l+1)_j) ‖



Additionally, since the orientation is decoupled from the magnitude, a single capsule can fire over the entire range of input orientations, as compared to a neuron which can only fire under its one specific oriented input. These design feature allows the orientation to encode the instantiation parameters of the object that the capsule is representing, while the magnitude encodes whether the capsule is on or off. With a scalar valued neuron it is only possible to encode whether the neuron is on or off.

To aid in visualizing the routing process, Figure <ref> was made using the unsupervised learning algorithm developed in this paper. The 144 capsules on the first layer are routed to 10 capsules on the second layer. Each of the first layer capsules are composed of 8 neurons while each of the second layer capsules are composed of 16 neurons. If the capsule has a magnitude close to 0 it is a lighter shade, while if it its magnitude is close to 1 it is a darker shade. Similarly the edge connections are the scalar valued routing weights c^(l)_ij's, where light edges correspond to a weak routing weight while a dark edge corresponds to a strong routing weight. As a neural network each edge would be replaced by a fully connected 8× 16 matrix of edges for each of the matrix valued prediction maps W^(l)_ij.






§ CAPSULE NETWORKS AS A PRODUCT OF EXPERT NEURONS 


This section develops the product of experts formulation of capsule networks. The model treats capsules as collections of neurons, and weights these collections of neurons across layers according to the dynamic routing between capsules algorithm.

Our intuition is that during optimization, alternating between routing by agreement and updating parameters, the model will be able to better learn to leverage the routing by agreement mechanism. This is how capsule networks with routing by agreement are trained with backpropagation, and so we are extending this intuition to product of experts learning. In both cases, the optimization procedure is a two step process.


    
  * With W^(l)_ij's fixed, perform routing by agreement to determine the routing weights c^(l)_ij's
    
  * With c^(l)_ij's fixed, define the energy and update the parameter weights W^(l)_ij W^(l)_ij + Δ W^(l)_ij


For clarity, we sometimes write the individual, vector-valued capsule x^(l)_i and x^(l+1)_j in terms of the scalar-valued neuron components x^(l)_i,m and x^(l+1)_j,n, as x^(l)_i=[x^(l)_i,1;x^(l)_i,2;…;x^(l)_i,M] for an M-dimensional capsule at layer l and x^(l+1)_j=[x^(l+1)_j,1;x^(l+1)_j,2;…;x^(l+1)_j,N] for an N-dimensional capsule at layer l+1. Similarly, we write the entire collection of capsules in layer l as x^(l)=[x^(l)_1,x^(l)_2,…,x^(l)_I] for I-many capsules in layer l, and x^(l+1)=[x^(l+1)_1,x^(l+1)_2,…,x^(l+1)_J] for J-many capsules in layer l+1. We use i and m to index over layer l and j and n to index over layer l+1, and we reserve the dot symbol · for matrix-vector multiplication, as in W · x.



 §.§ Routing-Weighted Product of Expert Neurons



Define the energy between vector-valued capsules x^(l)_i and x^(l+1)_j


    E_ij(x^(l)_i,x^(l+1)_j)=
    -
    x^(l+1)T_j· W^(l)_ij· x^(l)_i



With this, the total energy between layer l and l+1 is the sum over all possible combinations of these energies between capsules, weighted by the scalar routing weights c^(l)_ij's that are determined by routing by agreement.


    E(x^(l),x^(l+1)) = 
        ∑_i,j 
        c^(l)_ij
        E_ij(x^(l)_i,x^(l+1)_j)


Defined in this way, the significance of the routing coefficients c^(l)_ij's act to weigh the energy between capsules i and j. If a capsule is understood as being a collection of neurons, then this is a routing-weighted energy between the i^th collection of neurons at layer l and the j^th collection of neurons at layer l+1, given by E_ij(x^(l)_i,x^(l+1)_j) =  ∑_m,n 
 x^(l+1)_j,n W^(l)_ij,mn x^(l)_i,m.

The intuition is that, between capsules i and j, the routing coefficients will increase the energy of the sub-networks that route to each other, and decrease the energy of the sub-networks that do not route to each other.

Following the scalar product of experts formulation, we have the total probability is


    P(x^(l),x^(l+1))=
        1/Z
        e^-E(x^(l),x^(l+1))

where Z is the normalizing partition function. This can be written as a routing-weighted product of expert neurons:


    P(x^(l),x^(l+1))=
    1/ZΠ_j Π_n
    exp(
        ∑_i 
    c^(l)_ij∑_m 
    x^(l+1)_j,n W^(l)_ij,mn x^(l)_i,m)


In vector notation, Equation <ref> is equivalent to


    P(x^(l),x^(l+1))=
        1/ZΠ_j
        exp(
            ∑_i 
        c^(l)_ij
        x^(l+1)T_j· W^(l)_ij· x^(l)_i)


Vector notation has the advantage that it makes clear the capsule organization of the collections of neurons, and how it relates to routing by agreement. We marginalize over the hidden states with P(x^(l))=∑_x^(l+1) P(x^(l),x^(l+1)) to yield the visible state distribution:


    P(x^(l))=
        1/ZΠ_j
        [
        1+
        exp(
            ∑_i 
        c^(l)_ij
        W^(l)_ij· x^(l)_i)]


We can now perform inference for Gibb's sampling as P(x^(l+1)_j,n=1 | x^(l))=
    1/1+
    exp(-
        ∑_i 
    c^(l)_ij∑_m 
    W^(l)_ij,mn x^(l)_i,m), or in vector notation:


    P(x^(l+1)_j=1 | x^(l))=
        1/1+
        exp(-
            ∑_i 
        c^(l)_ij
        W^(l)_ij·
        x^(l)_i)


In the reverse direction the analogous equation is:


    P(x^(l)_i=1 | x^(l+1))=
        1/1+
        exp(-
            ∑_j 
        c^(l)_ij
        W^(l)T_ij·
        x^(l+1)_j)

where W^(l)T_ij is the matrix-transpose of W^(l)_ij, and is not the same thing as W^(l)_ji.

Subtleties regarding Equations <ref> and  <ref> deserve discussion. In routing by agreement, the individual capsule i makes its prediction about capsule j as z^(l+1)_j|i=W^(l)_ij· x^(l)_i, and then the routing-weighted average with the other predictions about j from each of the i's is taken, yielding z^(l+1)_j=∑_i c^(l)_ij z^(l+1)_j|i, where the weighted average is over the routing weights c^(l)_ij's such that ∑_i c^(l)_ij=1.

In this way, Equation <ref> is the sigmoid activation of this routing-weighted capsule prediction, where the sigmoid nonlinearity acts element-wise, and the value of each dimension is independent of the values of the other dimensions. This is in comparison to the squashing nonlinearity used during the routing by agreement, where the values of each dimension are dependent on the values of the other dimensions, to ensure that the vector is only scaled without changing its orientation.

Interpretting Equation <ref> in a similar way suggests that capsule x^(l+1)_j is making its prediction for the layer below as z^(l)_i|j=W^(l)T_ij· x^(l+1)_j, and then keeping the same routing weights c^(l)_ij's to give z^(l)_i=∑_j c^(l)_ij z^(l)_i|j. This is not a proper weighted average since ∑_i c^(l)_ij=1, whereas the weighted sum is over the j's, as opposed to the i's. It is not straightforward to use a different set of routing weights e^(l+1)_ji such that ∑_j e^(l+1)_ji=1, because this would mean having a different energy function for the forwards and backwards passes. 

Since we divide the optimization procedure into two steps these are not issues, just as is the case when using backpropagation with routing by agreement. The first step uses the capsule structure of the network to perform routing by agreement in order to decide how to weigh pairs of capsules i and j. The second step ignores the capsule structure of the network, except for the fact that it weighs the i^th group of neurons in layer l with the j^th group of neurons in layer l+1 by c^(l)_ij. In this way Equations <ref> and <ref> can be viewed as groups of neurons in one layer inferring activations about groups of neurons in another layer.



 §.§ Gradient of the log-likelihood


To optimize the log-likelihood, we take the following gradient:


    ∂/∂ W^(l)_ijlog P(x^(l)) = 
        ∑_x^(l+1)_j 
        P( x^(l+1)_j | x^(l)_i ) 
        ∂ E_ij/∂ W^(l)_ij(x^(l)_i,x^(l+1)_j) 
        
    
        -
        ∑_x^(l)_i 
        P( x^(l)_i )
        ∑_x^(l+1)_j 
        P( x^(l+1)_j | x^(l)_i ) 
        ∂ E_ij/∂ W^(l)_ij(x^(l)_i,x^(l+1)_j)

where we used the fact that ∂ E/∂ W^(l)_ij = ∑_i'j'∂ E_i'j'/∂ W^(l)_ij = ∂ E_ij/∂ W^(l)_ij. The gradient of the energy has a simple form:


    ∂ E_ij/∂ W^(l)_ij(x^(l),x^(l+1)) = 
        c^(l)_ij
        x^(l+1) T_j x^(l)_i


We now have the update rule for the parameter weights:

    ∂/∂ W^(l)_ijlog P(x^(l)) = 
        c^(l)_ij(
        ⟨ x^(l)_i x^(l+1)_j⟩_data
        -
        ⟨ x^(l)_i x^(l+1)_j⟩_model)


We see that without routing by agreement, this update rule would be identical to the RBM update rule. With routing by agreement, this tells us that the parameters are updated proportionally to the routing weights c^(l)_ij's. This means that the routing weights appear in two ways, one in determining the forward/backward sampling of x^(l+1)_j = σ( ∑_i c^(l)_ijW^(l)_ij· x^(l)_i ), and second in proportionally weighing the parameter updates.

As is well known, estimating the model distribution requires sampling the Markov chain to infinity. Instead we minimize the contrastive divergence <cit.>, and so we only run the Markov chain for a single iteration.




§ EXPERIMENTS




This section details the experiments conducted, training the capsule networks in the unsupervised ways described in the preceding section. The network diagram can be seen in Figure <ref>. In all experiments, we use an autoencoding convolutional neural network <cit.>, with dropout <cit.>, to learn the first two layers of filters in an unsupervised way. We denote filter bank sizes by [ height, width, channels-in, channels-out]. The first filter bank is of size [9,9,1,128] for grey-scale or [9,9,3,128] for rgb, followed by a Leaky ReLU activation, while the second filter bank is [9,9,128,128] followed by a sigmoid activation. We use a sigmoid for the second hidden layer activation because we want to map the data between 0 and 1 to match the energy model's sigmoids. For the decoder, first we use a Leaky Relu, and then a sigmoid; this sigmoid is to map the pixel intensities between 0 and 1.

Once the autoencoder has been sufficiently trained, we fix these filter weights and next train the capsule network as a product of expert neurons. The 128-channel hidden layer is reshaped to 6× 6× 128/8=576 capsules, each having 8 dimensions. These are then mapped to 20 capsules, each having 16 dimensions, by x^(l+1)_j = σ( ∑_i c^(l)_ijW^(l)_ij· x^(l)_i ), and in reverse x^(l)_i = σ( ∑_j c^(l)_ijW^(l)T_ij· x^(l+1)_j ) with the c^(l)_ij's fixed from the forward pass of the routing by agreement. All updates use sgd with momentum, learning rate decay, an ℓ_2 regularization on the weights and is implemented in TensorFlow <cit.>. On a single gpu it takes about 30-60 minutes to train, depending on the dataset.



 §.§ Routing-Weighted Product of Expert Neurons


Results generated from the routing-weighted product of expert neurons model of Section <ref> can be seen in Figure <ref>. For each of the 20 hidden layer capsules, four points are randomly sampled from a 16-dimensional, zero mean, unit variance Gaussian, and then fed into a sigmoid function. All hidden layer capsules are set to 0 except for one, which is assigned this random sigmoided 16-dimensional point. The network is then run in reverse, from the hidden layer to the input, where capsules are routed according to z^(l)_i = ∑_j c^(l)_ij W^(l) T_ij· x^(l+1)_j, where the c^(l)_ij's are determined by routing by agreement, but are normalized along the l^th-layer, i.e. ∑_i c^(l)_ij=1. In this manner, this then acts as an unsupervised, generative model, creating the 4 × 20 images of Figure <ref>. It is noted that these 4 × 20 random points were sampled together from the Gaussian, and so these are not cherry picked examples.

Each capsule seems to learn specific objects in this unsupervised setting. For example on MNIST, capsules seems to learn digits of different widths, angles and thicknesses. Nearly all of the capsules seem to have dimensions that correspond to the thickness of the digit. Interestingly, other dimensions are sometimes seen to extend strokes in different ways. For example, one of the capsules extends 9's to 8's by adding a stroke in the bottom left, as well as extends 9's to 4's by eliminating the stroke on the top.





§ CONCLUSIONS


This work develops a formulation of capsule networks as a weighted product of expert neurons, where the energy connecting two capsules is weighted proportionally by the routing coefficients from the routing by agreement mechanism. This contributes to the effort of further developing capsule networks by developing a new unsupervised training algorithm specifically designed for capsule networks that can build capsule networks from the ground up in an unsupervised way. Initial results on MNIST, Fashion-MNIST and CIFAR suggest that this learning procedure retains the attractive attributes of capsule networks, namely their ability to represent an object's instantiation parameters in the orientation of the vector, in the developed unsupervised setting. 







10

hinton2011transforming
Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang.
Transforming auto-encoders.
In International Conference on Artificial Neural Networks,
  pages 44–51. Springer, 2011.

sabour2017dynamic
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton.
Dynamic routing between capsules.
In Advances in neural information processing systems, pages
  3856–3866, 2017.

mcculloch1943logical
Warren S McCulloch and Walter Pitts.
A logical calculus of the ideas immanent in nervous activity.
The bulletin of mathematical biophysics, 5(4):115–133, 1943.

lecun1998gradient
Yann LeCun, Léon Bottou, Yoshua Bengio, Patrick Haffner, et al.
Gradient-based learning applied to document recognition.
Proceedings of the IEEE, 86(11):2278–2324, 1998.

krizhevsky2012imagenet
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.
Imagenet classification with deep convolutional neural networks.
In Advances in neural information processing systems, pages
  1097–1105, 2012.

rumelhart1985learning
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.
Learning internal representations by error propagation.
Technical report, California Univ San Diego La Jolla Inst for
  Cognitive Science, 1985.

hinton2018matrix
Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst.
Matrix capsules with em routing.
2018.

goodfellow2014generative
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley,
  Sherjil Ozair, Aaron Courville, and Yoshua Bengio.
Generative adversarial nets.
In Advances in neural information processing systems, pages
  2672–2680, 2014.

jaiswal2018capsulegan
Ayush Jaiswal, Wael AbdAlmageed, Yue Wu, and Premkumar Natarajan.
Capsulegan: Generative adversarial capsule network.
In Proceedings of the European Conference on Computer Vision
  (ECCV), pages 0–0, 2018.

hinton2002training
Geoffrey E Hinton.
Training products of experts by minimizing contrastive divergence.
Neural computation, 14(8):1771–1800, 2002.

welling2005exponential
Max Welling, Michal Rosen-Zvi, and Geoffrey E Hinton.
Exponential family harmoniums with an application to information
  retrieval.
In Advances in neural information processing systems, pages
  1481–1488, 2005.

fischer2012introduction
Asja Fischer and Christian Igel.
An introduction to restricted boltzmann machines.
In iberoamerican congress on pattern recognition, pages 14–36.
  Springer, 2012.

zeiler2010deconvolutional
Matthew D Zeiler, Dilip Krishnan, Graham W Taylor, and Robert Fergus.
Deconvolutional networks.
In Cvpr, volume 10, page 7, 2010.

srivastava2014dropout
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan
  Salakhutdinov.
Dropout: a simple way to prevent neural networks from overfitting.
The Journal of Machine Learning Research, 15(1):1929–1958,
  2014.

abadi2016tensorflow
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey
  Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.
Tensorflow: A system for large-scale machine learning.
In 12th {USENIX} Symposium on Operating Systems Design and
  Implementation ({OSDI} 16), pages 265–283, 2016.



