Trans-Caps: Transformer Capsule Networks with Self-attention Routing

1 INTRODUCTION . Convolutional Neural Networks ( CNNs ) have achieved state-of-the-art performance in many different computer vision tasks ( Krizhevsky et al. , 2012 ; He et al. , 2016 ) . This is achieved by local connectivity and parameter sharing across spatial locations so that useful local features learned in one receptive field can then be detected across the input feature space . While such a mechanism is sufficient to learn relationships between nearby pixels and to detect the existence of objects of interest , CNNs often fail to detect objects presented in radically new viewpoints due to the complex effects of the viewpoint changes on the pixel intensity values . This limitation forces us to train each CNN with a large number of data points which is computationally expensive . Capsule Networks ( CapsNets ) were introduced to explicitly learn a viewpoint invariant representation of the geometry of an object . In CapsNets , each group of neurons ( called a “ capsule ” ) encodes and represents the visual features of a higher-level object in an instantiation parameter vector or matrix ( which we refer to as the pose vector or matrix throughout this paper ) . The lower-level capsules ( which we refer to as part capsules ) estimate the pose of the object parts and hierarchically combine them to predict the pose of the whole object in the next layer . The object-part relationship is viewpoint-invariant , meaning that changes in the viewpoint change the pose of parts and objects in a coordinated way . Therefore , regardless of the viewpoint , we can infer the pose of the whole object from its parts using a set of trainable viewpoint-invariant transformation matrices . Capsule routing mechanisms can therefore learn the underlying spatial relationships between parts and objects . This improves the generalization capabilities of the network due to the underlying linear relationship between the viewpoint changes and the pose matrices . In order to route information between capsules , the part capsules vote for the pose of the higher-level capsules ( which we refer to as object capsules ) . A routing-by-agreement mechanism is employed to aggregate votes ( which has been traditionally accomplished using a recurrent clustering procedure ) effectively computing the contribution of each part to the object pose . While various proposed iterative routing mechanisms ( such as Dynamic ( Sabour et al. , 2017 ) and EM ( Hinton et al. , 2018 ) routing ) have been shown to be effective in the detection of viewpoint variations , their iterative nature increases computational cost . Prior research has additionally shown that these routing mechanisms may fail to properly construct a parse tree between each set of part and object capsules , partly due to the inability of the network to learn routing weights through back-propagation ( Peer et al. , 2018 ) . This ultimately limits the performance of CapsNets in realworld image classification tasks . Additionally , the correct number of routing iterations serves as an additional data-dependent hyper-parameter that needs to be carefully selected ; failing to optimize the number of routing operations can result in increased bias or variance in the model ( Hinton et al. , 2018 ) . This issue is amplified when training networks with multiple capsule layers . In this paper , we introduce a novel routing algorithm called self-attention routing ( SAR ) , which is inspired by the structural resemblance between CapsNets and Transformer networks ( Vaswani et al. , 2017 ) . This mechanism eliminates the need for recursive computations by replacing unsupervised routing procedures with a self-attention module , making the use of CapsNets effective in complex and large-scale image classification tasks . Our algorithm also reduces the risk of under and over-fitting associated with selecting a small and large number of routing iterations , respectively . We compare our proposed routing algorithm to two of the most prominent iterative methods , namely dynamic and EM routing , and the recently published non-iterative self-routing mechanism ( Hahn et al. , 2019 ) . We evaluate performance on several image classification datasets including SVHN , CIFAR-10 , CIFAR-100 , Tiny-ImageNet , ImageNet , and SmallNORB . Our results show that our model outperforms other baseline CapsNets and achieves better classification performance and convergence speed while requiring significantly fewer trainable parameters , fewer computations ( in FLOPs ) , and less memory . Moreover , our experimental result on the SmallNORB dataset with novel viewpoints shows that the proposed model is significantly more robust to changes in the viewpoint and is able to retain its performance under severe viewpoint shifts . All source code will be made publicly available . 2 RELATED WORK . 2.1 CAPSULE NETWORKS . CapsNets were originally introduced in Transforming Autoencoders by Hinton et al . ( 2011 ) ; here they pose computer vision tasks as inverse graphics problems to deal with variations in an object ’ s instantiation parameters . This architecture learns to reconstruct an affine-transformed version of the input image , therefore learning to represent each input as a combination of its parts and their respective characteristics . Sabour et al . ( 2017 ) introduced capsules with Dynamic Routing ( DRCaps ) , which allows the network to learn part-whole relationships through an iterative unsupervised clustering procedure . In DR-Caps , capsules output a pose vector whose length ( norm or magnitude ) implicitly represents the capsule activation . The vector norm should be able to scale depending on the pose values ; representing existence with the vector norm can therefore potentially weaken the representation power of any given capsule layer . Hinton et al . ( 2018 ) proposed capsules with EM routing ( EM-Caps ) , where capsule activations and pose matrices are segregated to fit the votes from part capsules through a mixture of Gaussians . While powerful , capsule network ’ s routing procedures have several fundamental limitations : : 1 ) Iterative routing operations are the bottleneck of CapsNets due to their computational complexity , which limits their widespread applicability in complex , large-scale datasets ( Zhang et al. , 2018 ; Li et al. , 2018 ) . 2 ) The number of routing iterations are hyper-parameters that need to be carefully tuned to prevent under and over fitting ( Hinton et al. , 2018 ) . 3 ) Lin et al . ( 2018 ) showed that even after seven iterative routing operations , the entropy of the coupling coefficient was still large , indicating that part capsules pass information to all object capsules . 4 ) EM-Caps have difficulty converging and have been shown to be numerically unstable , which limits their applicability in complex tasks ( Ahmed & Torresani , 2019 ; Gritzman , 2019 ) . Several studies have proposed non-iterative methods to replace the traditional iterative routing mechanisms in CapsNets . STAR-CAPS ( Ahmed & Torresani , 2019 ) combines an attention gate with a straight-through estimator to make a binary decision to either connect or disconnect the route between each part and object capsule . Tsai et al . ( 2020 ) proposed an inverted dot-product attention routing mechanism ( IDPA-Caps ) which generates the routing coefficients between capsules ; they unroll the iterative routing procedure and perform the iterations concurrently which helps improve parallelization . While powerful , the number of concurrent iterations is a hyper-parameter that needs to be tuned , and the unrolling process creates a very large network that is memory intensive . Inspired by Mixture-of-Experts , Hahn et al . ( 2019 ) introduced a self-routing mechanism . While non-iterative , self-routing attaches stationary routing weights to specific locations which limits its ability to generalize to novel viewpoints . This also increases the required number of trainable parameters , making it impractical for high-dimensional images . 2.2 SELF-ATTENTION . Attention operations bias a network to more informative components of the input in order to improve the discriminative capabilities of the model . This operation has been used to tackle the problem of long-range interactions in sequence modeling and has seen great success across the fields of Natural Language Processing ( NLP ) ( Bahdanau et al. , 2014 ) , genomics ( Zaheer et al. , 2020 ) , speech recognition ( Chorowski et al. , 2015 ) , and computer vision ( Hu et al. , 2018 ; Wang et al. , 2017 ) . Various attention mechanisms have been used to improve CNNs by allowing the network to capture interactions between elements of the encoded feature space , which is difficult for a stand-alone convolutional operation ( Woo et al. , 2018 ; Hu et al. , 2018 ) . Transformer based architectures were introduced by Vaswani et al . ( 2017 ) and utilize self-attention as the primary mechanism for representation learning . Self-attention employs the standard dot product operation to generate attention coefficients that effectively capture the long-range interactions between inputs and outputs . These architectures have outperformed recurrent neural networks in a wide range of tasks and have become the SOTA for representation learning ( Devlin et al. , 2018 ; Radford et al. , 2019 ; Huang et al. , 2018 ) . This concept was then expanded to computer vision applications by treating each output pixel as an element in the self-attention operation , thus allowing a CNN to learn global dependencies between receptive fields ( Bello et al. , 2019 ) . While powerful , this mechanism generates global dependencies for all pixels , making it memory intensive and computationally cumbersome . This issue was later addressed by restricting the scope of each self-attention operation to local patches as opposed to applying self-attention to the global feature space ( Hu et al. , 2019 ; Ramachandran et al. , 2019 ) . The Set Transformer is an encoder-decoder architecture that utilizes a self-attention mechanism to cluster a group of independent inputs by modeling the interactions among the elements of the set ( Lee et al. , 2019 ) . Given that the order of the part capsules does not contribute to the understanding of each object capsule , we took inspiration from the Set Transformer to replace the recurrent routing mechanism with a self-attention-based aggregation of the “ votes ” in a permutation invariant manner . 3 TRANSFORMER CAPSULE NETWORK . 3.1 MODEL ARCHITECTURE OVERVIEW . The Transformer Capsule Network ( Trans-Caps ) is a capsule-based neural network architecture where each capsule represents an encoded pose matrix . Trans-Caps starts with a convolutional backbone , followed by a sequence of capsule layers . Note that the choice of the convolutional backbone , the number of capsule layers and the number of capsules per layer varies for different sets of experiments . Given an input image X ∈ RH×W×D , the role of the convolutional backbone is to extract a set of features F ∈ RH′×W ′×D′ from the input . The backbone can be either a single convolutional layer , a cascade of a few convolutional layers , or a cascade of residual blocks ( He et al. , 2016 ) . Previous studies have shown improved performance for more complex datasets when using a residual backbone ( Tsai et al. , 2020 ; Hahn et al. , 2019 ) . We provide a detailed discussion of the various backbone configurations in the Experiments section . The first capsule layer ( PrimaryCaps ) is a convolutional layer , followed by BN applied to the output backbone features F . The outputs are then reshaped to form the primary capsule pose elements . All layers following this layer are convolutional capsule layers ( ConvCaps ) with SAR performed between the layers to construct the pose of the object capsules . With our non-iterative SAR mechanism , the pose computation at all stages of the network can be performed sequentially in one forward pass , yielding numerical stability and efficiency . The final capsule layer ( ClassCaps ) has as many capsules as the number of classes and predicts the pose of the objects , PLj where j ∈ { 1 , ... , J } . J and L represent the number of object capsules and the total number of layers , respectively . This layer is followed by a linear Classifier which is shared across all class capsules and computes the final class logits as ŷc = Classifier ( PLc ) where P L c is the output of the final ConvCaps layer and c ∈ { 1 , ... , C } in a C class problem .

Summary: The paper proposes to swap the typical routing mechanisms in capsules for a more standard attention mechanism. The attention mechanism is based on computing similarity scores using gaussians instead of dot-products . The authors show that this leads to better downstream performance of more natural tasks while preserving robustness to viewpoint changes, one of the main strengths of capsules.