Capsules with Inverted Dot-Product Attention Routing

1 INTRODUCTION . Capsule Networks ( CapsNets ) represent visual features using groups of neurons . Each group ( called a “ capsule ” ) encodes a feature and represents one visual entity . Grouping all the information about one entity into one computational unit makes it easy to incorporate priors such as “ a part can belong to only one whole ” by routing the entire part capsule to its parent whole capsule . Routing is mutually exclusive among parents , which ensures that one part can not belong to multiple parents . Therefore , capsule routing has the potential to produce an interpretable hierarchical parsing of a visual scene . Such a structure is hard to impose in a typical convolutional neural network ( CNN ) . This hierarchical relationship modeling has spurred a lot of interest in designing capsules and their routing algorithms ( Sabour et al. , 2017 ; Hinton et al. , 2018 ; Wang & Liu , 2018 ; Zhang et al. , 2018 ; Li et al. , 2018 ; Rajasegaran et al. , 2019 ; Kosiorek et al. , 2019 ) . In order to do routing , each lower-level capsule votes for the state of each higher-level capsule . The higher-level ( parent ) capsule aggregates the votes , updates its state , and uses the updated state to explain each lower-level capsule . The ones that are well-explained end up routing more towards that parent . This process is repeated , with the vote aggregation step taking into account the extent to which a part is routed to that parent . Therefore , the states of the hidden units and the routing probabilities are inferred in an iterative way , analogous to the M-step and E-step , respectively , of an Expectation-Maximization ( EM ) algorithm . Dynamic Routing ( Sabour et al. , 2017 ) and EMrouting ( Hinton et al. , 2018 ) can both be seen as variants of this scheme that share the basic iterative structure but differ in terms of details , such as their capsule design , how the votes are aggregated , and whether a non-linearity is used . We introduce a novel routing algorithm , which we called Inverted Dot-Product Attention Routing . In our method , the routing procedure resembles an inverted attention mechanism , where dot products are used to measure agreement . Specifically , the higher-level ( parent ) units compete for the attention of the lower-level ( child ) units , instead of the other way around , which is commonly used in attention models . Hence , the routing probability directly depends on the agreement between the parent ’ s pose ( from the previous iteration step ) and the child ’ s vote for the parent ’ s pose ( in the current iteration step ) . We also propose two modifications for our routing procedure – ( 1 ) using Layer Normalization ( Ba et al. , 2016 ) as normalization , and ( 2 ) doing inference of the latent capsule states and routing probabilities jointly across multiple capsule layers ( instead of doing it layer-wise ) . These modifications help scale up the model to more challenging datasets . Our model achieves comparable performance as the state-of-the-art convolutional neural networks ( CNNs ) , but with much fewer parameters , on CIFAR-10 ( 95.14 % test accuracy ) and CIFAR-100 ( 78.02 % test accuracy ) . We also introduce a challenging task to recognize single and multiple overlapping objects simultaneously . To be more precise , we construct the DiverseMultiMNIST dataset that contains both single-digit and overlapping-digits images . With the same number of layers and the same number of neurons per layer , the proposed CapsNet has better convergence than a baseline CNN . Overall , we argue that with the proposed routing mechanism , it is no longer impractical to apply CapsNets on real-world tasks . We will release the source code to reproduce the experiments . 2 CAPSULE NETWORK ARCHITECTURE . An example of our proposed architecture is shown in Figure 1 . The backbone is a standard feedforward convolutional neural network . The features extracted from this network are fed through another convolutional layer . At each spatial location , groups of 16 channels are made to create capsules ( we assume a 16-dimensional pose in a capsule ) . LayerNorm is then applied across the 16 channels to obtain the primary capsules . This is followed by two convolutional capsule layers , and then by two fully-connected capsule layers . In the last capsule layer , each capsule corresponds to a class . These capsules are then used to compute logits that feed into a softmax to computed the classification probabilities . Inference in this network requires a feed-forward pass up to the primary capsules . After this , our proposed routing mechanism ( discussed in the next section ) takes over . In prior work , each capsule has a pose and some way of representing an activation probability . In Dynamic Routing CapsNets ( Sabour et al. , 2017 ) , the pose is represented by a vector and the activation probability is implicitly represented by the norm of the pose . In EM Routing CapsNets ( Hinton et al. , 2018 ) , the pose is represented by a matrix and the activation probability is determined by the EM algorithm . In our work , we consider a matrix-structured pose in a capsule . We denote the capsules in layer L as PL and the i-th capsule in layer L as pLi . The pose p L i ∈ RdL in a vector form and will be reshaped to R √ dL× √ dL when representing it as a matrix , where dL is the number of hidden units grouped together to make capsules in layer L. The activation probability is not explicitly represented . By doing this , we are essentially asking the network to represent the absence of a capsule by some special value of its pose . 3 INVERTED DOT-PRODUCT ATTENTION ROUTING . The proposed routing process consists of two steps . The first step computes the agreement between lower-level capsules and higher-level capsules . The second step updates the pose of the higher-level capsules . Step 1 : Computing Agreement : To determine how capsule j in layer L + 1 ( pL+1j ) agrees with capsule i in layer L ( pLi ) , we first transform the pose p L i to the vote vLij for the pose p L+1 j . This Higher-level Capsules Lower-level Capsules pose … … Agreement by Dot-Product Attention multiplication ∑ LayerNorm ( ⋅ ) routing coefficients Pose UpdateRouting Coefficients as Normalized Agreement agreement weight matrix multiplication reshape ⋅dot product ⊤ Softmax ( ⋅ ) … … … … Figure 2 : Illustration of the Inverted Dot-Product Attention Routing with the pose admitting matrix structure . Procedure 1 Inverted Dot-product Attention Routing algorithm returns updated poses of the capsules in layer L + 1 given poses in layer L and L + 1 and weights between layer L and L + 1 . 1 : procedure INVERTED DOT-PRODUCT ATTENTION ROUTING ( PL , PL+1 , WL ) 2 : for all capsule i in layer L and capsule j in layer ( L + 1 ) : vLij ←WLij · pLi . vote 3 : for all capsule i in layer L and capsule j in layer ( L + 1 ) : aLij ← p L+1 j > · vLij . agreement 4 : for all capsule i in layer L : rLij ← exp ( aLij ) / ∑ j′ exp ( a L ij′ ) . routing coefficient 5 : for all capsule j in layer ( L + 1 ) : pL+1j ← ∑ i r L ijv L ij . pose update 6 : for all capsule j in layer ( L + 1 ) : pL+1j ← LayerNorm ( p L+1 j ) . normalization 7 : return PL+1 transformation is done using a learned transformation matrix WLij : vLij = W L ij · pLi , ( 1 ) where the matrix WLij ∈ RdL+1×dL if the pose has a vector structure and WLij ∈ R √ dL+1× √ dL ( requires dL+1 = dL ) if the pose has a matrix structure . Next , the agreement ( aLij ) is computed by the dot-product similarity between a pose pL+1j and a vote v L ij : aLij = p L+1 j > · vLij . ( 2 ) The pose pL+1j is obtained from the previous iteration of this procedure , and will be set to 0 initially . Step 2 : Computing Poses : The agreement scores aLij are passed through a softmax function to determine the routing probabilities rLij : rLij = exp ( aLij ) ∑ j′ exp ( a L ij′ ) , ( 3 ) where rLij is an inverted attention score representing how higher-level capsules compete for attention of lower-level capsules . Using the routing probabilities , we update the pose pL+1j for capsule j in layer L + 1 from all capsules in layer L : pL+1j = LayerNorm ( ∑ i rLijv L ij ) . ( 4 ) We adopt Layer Normalization ( Ba et al. , 2016 ) as the normalization , which we empirically find it to be able to improve the convergence for routing . The routing algorithm is summarized in Procedure 1 and Figure 2 . 4 INFERENCE AND LEARNING . To explain how inference and learning are performed , we use Figure 1 as an example . Note that the choice of the backbone , the number of capsules layers , the number of capsules per layer , the design of the classifier may vary for different sets of experiments . We leave the discussions of configurations in Sections 5 and 6 , and in the Appendix . Procedure 2 Inference . Inference returns class logits given input images and parameters for the model . Capsule layer 1 denotes the primary capsules layer and layer N denotes the class capsules layer . 1 : procedure INFERENCE ( I ; θ , W1 : N−1 ) / * Pre-Capsules Layers : backbone features extraction * / 2 : F← backbone ( I ; θ ) . backbone feature / * Capsules Layers : initialization * / 3 : P1 ← LayerNorm ( convolution ( F ; θ ) ) . primary capsules 4 : for L in layers 2 to N : PL ← 0s . non-primary capsules / * Capsules Layers ( 1st Iteration ) : sequential routing * / 5 : for L in layers 1 to ( N − 1 ) do 6 : PL+1 ← Routing ( PL , PL+1 ; WL ) . non-primary capsules / * Capsules Layers ( 2nd to tth Iteration ) : concurrent routing * / 7 : for ( t− 1 ) iterations do 8 : for L in layers 1 to ( N − 1 ) : P̄L+1 ← Routing ( PL , PL+1 ; WL ) 9 : for L in layers 2 to N : PL ← P̄L . non-primary capsules / * Post-Capsules Layers : classification * / 10 : for all capsule i in layer N : ŷi ← classifier ( pNi ; θ ) . class logits 11 : return ŷ 4.1 INFERENCE . For ease of exposition , we decompose a CapsNet into pre-capsule , capsule and post-capsule layers . Pre-Capsule Layers : The goal is to obtain a backbone feature F from the input image I . The backbone model can be either a single convolutional layer or ResNet computational blocks ( He et al. , 2016 ) . Capsule Layers : The primary capsules P1 are computed by applying a convolution layer and Layer Normalization to the backbone feature F. The non-primary capsules layers P2 : N are initialized to be zeros 1 . For the first iteration , we perform one step of routing sequentially in each capsule layer . In other words , the primary capsules are used to update their parent convolutional capsules , which are then used to update the next higher-level capsule layer , and so on . After doing this first pass , the rest of the routing iterations are performed concurrently . Specifically , all capsule layers look at their preceding lower-level capsule layer and perform one step of routing simultaneously . This procedure is an example of a parallel-in-time inference method . We call it “ concurrent routing ” as it concurrently performs routing between capsules layers per iteration , leading to better parallelism . Figure 3 illustrates this procedure from routing iteration 2 to t. It is worth noting that , our proposed variant of CapsNet is a weight-tied concurrent routing architecture with Layer Normalization , which Bai et al . ( 2019 ) empirically showed could converge to fixed points . Previous CapsNets ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) used sequential layer-wise iterative routing between the capsules layers . For example , the model first performs routing between layer L − 1 and layer L for a few iterations . Next , the model performs routing between layer L and L + 1 for a few iterations . When unrolled , this sequential iterative routing defines a very deep computational graph with a single path going from the inputs to the outputs . This deep graph could lead to a vanishing gradients problem and limit the depth of a CapsNet that can be trained well , especially if any squashing non-linearities are present . With concurrent routing , the training can be made more stable , since each iteration has a more cumulative effect . Post-Capsule Layers : The goal is to obtain the predicted class logits ŷ from the last capsule layer ( the class capsules ) PN . In our CapsNet , we use a linear classifier for class i in class capsules : ŷi = classifier ( p N i ) . This classifier is shared across all the class capsules . 1 As compared to 0 initialization , we observe that a random initialization leads to similar converged performance but slower convergence speed . We also tried to learn biases for capsules ’ initialization , which results in similar converged performance and same convergence speed . As a summary , we initialize the capsule ’ s value to 0 for simplicity .

Summary: Authors improve upon dynamic routing between capsules by removing the squash function (norm normalization) and apply a layerNorm normalization instead. Furthermore, they experiment with concurrent routing rather than sequential routing (route all caps layers once, then all layers concurrently again and again). This is an interesting development since provides better gradient in conjunction with layerNorm. They report results on Cifar10 and Cifar100 and achieve similar to CNN (resnet) performance.