Deformable Capsules for Object Detection

1 INTRODUCTION . Capsule networks promise many potential benefits over convolutional neural networks ( CNNs ) . These include practical benefits , such as requiring less data for training or better handling unbalanced class distributions ( Jiménez-Sánchez et al. , 2018 ) , and important theoretical benefits , such as buildingin stronger internal representations of objects ( Punjabi et al. , 2020 ) , and modeling the agreement between those intermediate representations which combine to form final object representations ( e.g . part-whole relationships ) ( Kosiorek et al. , 2019 ; Sabour et al. , 2017 ) . Although these benefits might not be seen in the performance metrics ( e.g . average precision ) on standard benchmark computer vision datasets , they are important for real-world applications . As an example , it was found by Alcorn et al . ( 2019 ) that CNNs fail to recognize 97 % of their pose space , while capsule networks have been shown to be far more robust to pose variations of objects ( Hinton et al. , 2018 ) ; further , real-world datasets are not often as extensive and cleanly distributed as ImageNet or MS COCO . These benefits are achieved in capsule networks by storing richer vector ( or matrix ) representations of features , rather than the simple scalars of CNNs , and dynamically choosing how to route that information through the network . The instantiation parameters for a feature are stored in these capsule vectors and contain information ( e.g . pose , deformation , hue , texture ) useful for constructing the object being modeled . Early studies have shown strong evidence that these vectors do in fact capture important local and global variations across objects ’ feature components ( or parts ) within a class ( Punjabi et al. , 2020 ; Sabour et al. , 2017 ) . Inside their networks , capsules dynamically route their information , seeking to maximize the agreement between these vector feature representations and the higher-level feature vectors they are attempting to form . Despite their potential benefits , many have remained unconvinced about the general applicability of capsule networks to large-scale computer vision tasks . To date , no capsule-based study has achieved classification performance comparable to a CNN on datasets such as ImageNet , instead relegated to smaller datasets such as MNIST or CIFAR . Worse still , to the best of our knowledge , no capsule network has shown successful results in object detection , a very important problem in computer vision , robotics , and medical imaging . Now , the argument can be made that standard benchmark datasets such as ImageNet or MS COCO likely contain majority of objects in that 3 % range of usual poses , and thus CNNs will appear to perform extremely well when measured in terms of accuracy , stripping capsule networks of one of their largest advantages . However , until capsule networks can perform on-par with CNNs on these typical object poses , few will care about the benefits of stronger internal representations and better generalization to unseen poses . Summary of Our Contributions : ( 1 ) We propose the first ever capsule-based object detection framework in the literature . Our network is a one-stage ( single-shot ) architecture , where objects are both localized and classified using capsules , and can perform on-par with the state-of-the-art CNNs on a large-scale dataset ( MS COCO ) . ( 2 ) We address the geometric constraint of convolutional capsules ( and locally-constrained routing ) by introducing deformable capsules , where parent capsules learn to adaptively sample child capsules , effectively eliminating rigid spatial restrictions while remaining memory efficient . ( 3 ) We design a new capsule-based prediction head structure , SplitCaps , which reformulates the projections of an objects ’ instantiation parameters , presence , and class , eliminating the previous dimensional increase of capsules by the number of classes . This crucial addition enables the training of capsule networks on large-scale computer vision datasets for the first time in the literature . ( 4 ) To route information across SplitCaps ’ unique structure , we introduce a novel Squeezeand-Excitation inspired dynamic routing algorithm , SE-Routing , which seeks to maximize agreement between child capsule projections , without the need of iterative loops 2 DEFORMABLE CAPSULES : FIXING LOCALLY-CONSTRAINED DYNAMIC ROUTING . The capsule network architecture proposed by Sabour et al . ( 2017 ) acted on global information , where digit capsules represented the pose and presence of digits in an image regardless of spatial location . The information from all children in the previous layer was sent to every parent in the following layer , weighted via the routing coefficients found in a cosine similarity routing algorithm . While this proved to be a highly-effective strategy , it was also computationally expensive , limiting its use to only small-scale datasets . Recent works attempted to scale up capsule networks to larger problems such as biomedical image segmentation ( LaLonde & Bagci , 2018 ) or action detection in video ( Duarte et al. , 2018 ) by using convolutional capsules and locally-constraining the routing algorithm . Although efficient solutions were presented in those studies , the representation power of capsule networks was fundamentally limited due to imposing local constraints . This is because convolutions , by design , have a fixed geometric structure ( Dai et al. , 2017 ) , and such a geometric constraint significantly inhibits capsules ’ ability to model part-whole relationships , relying on parts of objects to fall within a fixed local grid . Therefore , it is unreasonable to expect a capsule to effectively represent the pose and deformations of an object when the information related to the parts of that object are locked into a fixed spatial relationship . In this study , we propose to effectively solve this aforementioned problem by introducing a method that balances efficiency with the ability for a capsule to represent any pose and deformation of an object ( i.e . where child capsules can be found in different spatial relationships to one another for the same parent ) . In such a formulation , global information is not explicitly required , but it does require parent capsules to have more flexibility over which child capsules they draw information from . Our proposed solution is deformable capsules . The idea behind the proposed algorithm is simple : if parent capsules are supposed to capture common deformations of the objects they represent within their vectors , then the choice of which children to aggregate information from must be handled in a deformable manner as well . Deformable capsules allow parents to adaptively gather projections from a non-spatially-fixed set of children , and thus effectively and efficiently model objects ’ poses . To achieve this overall goal , we follow the same efficient convolutional capsule paradigm , where projection vectors are formed via a convolution operation with a kernel centered on the parent capsules ’ spatial location , but now we learn an additional set of weights for each parent capsule . These learnable weights are the same shape as each parent ’ s kernel , and represent the offset values for the spatial sampling of child capsules for that parent . Based on the child capsule representation vectors in the previous layer , these weights learn which children a parent capsule should adaptively sample from for a given input image . Dynamic routing then determines how to weight the information coming from each of these children based on their agreement for each projected parent . Let us give a concrete example to better illustrate the parameter savings of this technique . Given a H = 128 by W = 128 grid of child capsules with ci = 32 capsule types of ai = 8 atoms each , being routed to a set of cj = 10 parent capsule types of aj = 16 atoms each , the fully-connected capsules of Sabour et al . ( 2017 ) would requireH×W×ci×ai×cj×aj ⇒ 128×128×32×8×10×16 ≈ 671M parameters for this layer alone ( assuming the goal is classification , with detection requiring a multiplicative increase by the detection grid size ) . Instead , using our proposed deformable capsules with a k2 = 52 kernel , we only require 2× k× k× ai × cj × aj ⇒ 2× 5× 5× 8× 10× 16 ≈ 64K parameters . Convolutional capsules with locally-constrained routing require 32K parameters , not needing the additional spatial offsets kernel , but as mentioned above , they are fundamentally limited in the poses and deformations that they can represent . In our experiments , we found deformable capsules to converge faster and to much higher performance than convolutional capsules . 3 OBJECTS AS CAPSULES : SplitCaps WITH SE-Routing We propose a novel one-stage ( single-shot ) capsule network architecture for object detection , called DeformCaps , where objects are detected , classified , and modeled with capsules . Our overall network architecture , shown in Fig . 1 , is built upon CenterNet by Zhou et al . ( 2019a ) who proposed to represent objects in images as scalar point values located at the center of their bounding boxes . The authors then regress the remaining characteristics of the object ( e.g . height , width , depth ) for each center-point detected . In our work , we follow the same center-point detection paradigm , but represent our objects with capsule vectors instead . Since several recent studies have found utilizing a CNN backbone before forming capsule types to be beneficial to overall performance ( Duarte et al. , 2018 ; Kosiorek et al. , 2019 ; Tsai et al. , 2020 ) , we adopt the preferred backbone of CenterNet , DLA-34 ( Yu et al. , 2018 ) , as this gave the best trade-off between speed and accuracy . Features extracted from the backbone are sent to our capsule object detection head and a bounding box regression head . In order to perform object detection using capsules , we introduce a new capsule structure , called SplitCaps , composed of class-agnostic and class-presence capsules , and a novel routing algorithm , called SE-Routing . 3.1 SPLITCAPS : CLASS-AGNOSTIC CAPSULES AND CLASS-PRESENCE CAPSULES . As discussed in Section 2 , the original CapsNet by Sabour et al . ( 2017 ) was extremely expensive in computation and our proposed deformable capsules is a best possible solution to balance non-rigid deformations of objects while remaining memory efficient . However , there is a more significant memory hurdle capsule networks must overcome when scaling up to large-scale datasets , such as MS COCO . In their current implementation , capsule networks represent each class with its own parent capsule vector . On small scale classification datasets ( and using deformable routing ) , this is not an issue ; it amounts to 2× k× k× ai × cj × aj parameters and N × ci × cj × aj × 4 bytes to store the intermediate representations to be routed for the parents , where cj is usually around 10 classes to represent . Let us suppose we have 5× 5 kernels with 32 input capsule types of 8 atoms per capsule , 10 output capsule types of 16 atoms per capsule , and a batch size of 32 . In total , we would have 2× 5× 5× 8× 10× 16 = 64K parameters and 32× 32× 10× 16× 4 ≈ 655 KB . When we scale this up to object detection , we now need to store representations for every possible object location , and for MS COCO we need to represent 80 possible classes . This gives us 2× k × k×ai× cj ×aj ⇒ 2×5×5×8×80×16 = 512K parameters and N ×H×W × ci× cj ×aj ×4 bytes⇒ 32× 128× 128× 32× 80× 16× 4 ≈ 86 GB for the intermediate representations , where we assume the output grid of detections is 128 × 128 with a single detection ( i.e . bounding box ) predicted per class per location . The problem is not any better for large-scale classification datasets such as ImageNet either , where we lose the grid of predictions but grow to 1000 classes , which would require 2× 5× 5× 8× 1000× 16 = 6.4M parameters and 32× 32× 1000× 16× 4 ≈ 66 GB for the intermediate representations . Clearly , with most GPU memories limited to 12–24 GB , a solution is needed to be found for capsule networks to scale up to larger-scale computer vision tasks . To overcome this issue , we propose a new type of capsule architecture , SplitCaps , to more efficiently scale capsule networks to large-scale computer vision tasks . SplitCaps contains two parent capsule types , each with a different number of atoms per capsule , for each location of the detection grid . As before , the idea is to balance efficiency with the ability to learn powerful representations . Towards this goal , SplitCaps proposes to divide up between its two parent capsules the tasks of ( i ) learning the instantiation parameters necessary to model the possible variations of an object and ( ii ) predicting which classes of objects are present in a given input . The first capsule type we refer is our classagnostic object instantiation capsules , and the second we refer is class presence capsules . Class-agnostic object instantiation capsules : The purpose of these capsules is similar to those in previous works : model the possible variations ( in pose , deformation , texture , etc . ) of objects within a vector of instantiation parameters , the span of which should cover all possible variations for that object at test ( hence why capsules are better than CNNs at generalizing to unseen poses ) . While previous capsule networks did this in a class-wise manner , we argue such a formulation is not required and possibly it is redundant . Many variations ( e.g . rotation , skew , stroke thickness ) may be class-independent , and thus to model these variations class-wise would require repetition across each capsule type . Instead , we propose to model all classes within a single capsule type ( i.e . class-agnostic ) . In this way , while class-dependent variations would each require their own dimensions of the vector , class-independent variations can each be modeled along a single dimension for all possible objects . Since it is reasonable to assume there will be at least some class-specific variations , we increase the default capsule vector dimension from 16 to 64 to accommodate for possible class-specific instantiation parameters . During training , if an object is present at a given spatial location , the 64-dimensional capsule vector for that location is fed to a reconstruction regularization sub-network to construct the mask of that object as similar to the reconstruction regularization used by Sabour et al . ( 2017 ) for classification . This sub-network is a relatively small and fast addition : a set of three ReLU-activated 1× 1 convolutional layers with 256 filters each , followed by a final sigmoid-activated 1 × 1 convolution with N = n2 = 282 = 784 filters , before reshaping outputs to n× n. Since objects ’ scales vary dramatically , we scale normalize all objects ’ ground-truth masks to be 28 × 28 ( by following He et al . ( 2017 ) ) . Supervised training is conducted by computing the Dice loss ( Milletari et al. , 2016 ) between the predicted reconstruction , r , and the object ’ s mask , m Lr = 2 ∑N i rimi∑N i r 2 i + ∑N i m 2 i , ( 1 ) where Lr is used to provide a regularization signal to the instantiation parameters being learned . Class presence capsules : The class presence capsules attempt to model which classes of objects are present in the input at each spatial location , if any . We accomplish this by setting the atoms per capsule to the number of classes being represented ( i.e . 80 for MS COCO ) . Just as Hinton et al . ( 2018 ) separately modeled pose ( with a matrix ) and activation ( with a scalar ) , this 80-dimensional vector can be viewed as a class-dependent set of activation values . The activation values are then passed through a sigmoid function and thresholded ; if one or more activation values are above the threshold , an object is determined to be at that spatial location with the strongest activated dimension determining the class label . In order to produce a smooth loss function during training , we create a ground-truth heatmap by fitting a Gaussian distribution rather than a single point , to the center-point of each object ’ s bounding box , with variance proportional to the size of the box following Zhou et al . ( 2019a ) . More specifically , we create a heatmap H ∈ [ 0 , 1 ] Xd ×Yd ×K containing each down-scaled ground truth center-point p̃ = ( px d , py d ) for class k ∈ K using a Gaussian kernelHxyk = exp ( − ( x−p̃x ) 2+ ( y−p̃y ) 2 2σ2p ) , where d is the amount of downsampling in the network and σp is an object-size-adaptive standard deviation ( Law & Deng , 2018 ) . In the case of overlapping Gaussians , we take the element-wise maximum . To handle the large class imbalance between objects and background in our heatmaps , we use a penalty-reduced pixel-wise logistic regression with a focal loss ( Lin et al. , 2017 ) : Lh = −1 P ∑ xyk { ( 1− Ĥxyk ) α log ( Ĥxyk ) if Hxyk = 1 , ( 1−Hxyk ) β ( Ĥxyk ) α log ( 1− Ĥxyk ) , otherwise ; ( 2 ) where α , β are hyper-parameters of the focal loss , P is the number of center-points in the input , used to normalize all positive focal loss instances to 1 ( Zhou et al. , 2019a ) . We use α = 2 and β = 4 in all our experiments , following Law & Deng ( 2018 ) . At test , to efficiently retrieve the object ’ s exact center , we run a 3× 3 max-pooling over the thresholded spatial map . To predict the height and width of the bounding boxes of objects and recover the x , y offsets needed to map back to the upscaled image dimensions , we follow the same formulation as Zhou et al . ( 2019a ) and pass the backbone features through a 3× 3 convolutional layer with 256 feature maps , then a 1× 1 convolutional layer with 2 feature maps . These layers predict the local offset , Ô ∈ RWd ×Hd ×2 , and size prediction , Ŝ ∈ RWd ×Hd ×2 , for each center-point and are supervised by Lo = 1 P ∑ p ∣∣∣Ôp̃ − ( p d − p̃ ) ∣∣∣ and Ls = 1 P ∑ p ∣∣∣Ŝp − ( x2 − x1 , y2 − y1 ) ∣∣∣ , ( 3 ) respectively . Our final objective function is thus defined as L = Lh + λrLr + λsLs + λoLo . We keep λs = 0.1 and λo = 1 as done in Zhou et al . ( 2019a ) , and set λr = 0.1 initially , then step up to λr = 2.0 at the half-way in training .

Summary: The paper proposes to use the capsules to perform object detection on COCO. Capsules, while showing promises, are usually too expensive for tasks beyond MNIST and Cifar. The authors propose three key improvements in DeformCaps, SplitCaps and SE-Routing to improve the efficiency and therefore allow capsules to be applied on larger tasks such as object detection. The authors claim novelties in: