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: This paper introduces capsule network for object detection. To solve the issue of capsule network when applied to large-scale detection problems, this paper develops deformable capsules, a new  prediction head SplitCaps, and a dynamic routing algorithm, SE-Routing. Experiments are conducted on COCO where it performs slightly worse than the baselines but arguably predicts less false positives.