Trace-back along capsules and its application on semantic segmentation  		

In this paper , we propose a capsule-based neural network model to solve the semantic segmentation problem . By taking advantage of the extractable part-whole dependencies available in capsule layers , we derive the probabilities of the class labels for individual capsules through a recursive , layer-by-layer procedure . We model this procedure as a traceback pipeline and take it as a central piece to build an end-to-end segmentation network . Under the proposed framework , image-level class labels and object boundaries are jointly sought in an explicit manner , which poses a significant advantage over the state-of-the-art fully convolutional network ( FCN ) solutions . Experiments conducted on modified MNIST and neuroimages demonstrate that our model considerably enhance the segmentation performance compared to the leading FCN variant . 1 INTRODUCTION . An effective segmentation solution should have a well-equipped mechanism to capture both semantic ( i.e. , what ) and location ( i.e. , where ) information . The fully convolutional network ( FCN ) ( Long et al. , 2015 ) and its variants ( Ronneberger et al. , 2015 ; Noh et al. , 2015 ; Badrinarayanan et al. , 2017 ) constitute a popular class of solutions for this task , producing state-of-the-art results in a variety of applications . FCN and its variants ( FCNs ) are commonly constructed with an encoder-decoder architecture . In the encoding path , input images are processed through a number of “ convolution + pooling ” layers to generate high-level latent features , which are then progressively upsampled in the decoder to reconstruct the target pixel labels . The feature maps produced in higher ( coarser ) layers and those in lower ( finer ) layers contain complementary information : the former is richer in semantics , while the latter carries more spatial details that define class boundaries . Originated from and constructed upon convolutional neural networks ( CNNs ) ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ) , FCNs ’ encoders inherit some common drawbacks of CNNs , one of which is the lack of an internal mechanism in achieving viewpoint-invariant recognition . Traditional CNNs , as well as FCNs , rely on convolution operations to capture various visual patterns , and utilize poolings to enable multi-scale processing of the input images . Rotation invariance , however , is not readily available in both models . As a result , more data samples or additional network setups ( Cohen & Welling , 2016 ; Cohen et al. , 2018 ) would be required for objects from different viewpoints to be correctly recognized . The absence of explicit part-whole relationships among objects imposes another limitation for FCNs – without such a mechanism , the rich semantic information residing in the higher layers and the precise boundary information in the lower layers can only be integrated in an implicit manner ( Zhang et al. , 2018 ) . Capsule nets ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) , operating on a different paradigm , can provide a remedy . Capsule nets are built on capsules , each of which is a group of neurons representing one instance of a visual entity , i.e. , an object or one of its parts ( Hinton et al. , 2011 ) . Capsules output both activation probabilities of their presence and the instantiation parameters that describe their properties , such as pose , deformation and texture , relative to a viewer ( Hinton et al. , 2011 ) . During inference propagation , the principle of coincidence filtering is employed to activate higher-level capsules and set up part-whole relationships among capsule entities . Such part-whole hierarchy equips capsule nets with a solid foundation for viewpoint-invariant recognition , which can be implemented through dynamic routing ( Sabour et al. , 2017 ) or EM routing ( Hinton et al. , 2018 ) . The same hierarchy , if properly embedded into a segmentation network , would provide a well-grounded platform to specify contextual constraints and enforce label consistency . With this thought , we develop a capsule-based semantic segmentation solution in this paper . Our approach treats capsule nets as probabilistic graphical models capable of inferring probabilistic dependences among visual entities , through which part-whole relationships can be explicitly constructed . As a concrete implementation , we propose a new operation sequence , which we call traceback pipeline , to capture such part-whole information through a recursive procedure to derive the class memberships for individual pixels . We term our model Tr-CapsNet . The contributions of our Tr-CapsNet can be summarized as : 1 . In Tr-CapsNet , the class labels for individual spatial coordinates within each capsule layer are analytically derived . The traceback pipeline in our model , taking advantage of the graphical properties of capsule nets , is mathematically rigorous . To the best of our knowledge , this is the first work to explore a capsule traceback approach for image segmentation . In addition , probability maps at each capsule layer are readily available , which makes it convenient to conduct feature visualization and layer interpretation . 2 . In parallel with segmentation , Tr-CapsNet carries out explicit class recognition at the same time . Such explicitness poses a powerful practical advantage over FCNs . 3 . The traceback pipeline is designed under a general context , making it applicable to many other potential tasks , including object localization and detection , action localization and network interpretation . 2 BACKGROUND . A capsule ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) is aimed to represent an instance of a visual entity , and its outputs include two parts : 1 ) the probability of the existence of this instance and 2 ) instantiation parameters , which carries certain visual properties of the instance . Built on capsules , capsule nets are designed to overcome the drawbacks of CNNs , in particular , the inability of handling viewpoint changes and the absence of part-whole relationships among visual entities . Unlike CNNs , where image patterns are captured through convolution ( correlation ) operations , capsule nets rely on high-dimensional coincidence filtering ( HDCF ) to detect objects that attract concentrated votes from their parts/children . The HDCF mechanism is carried out based on the computation of two network parameter sets : weights and coupling coefficients , of the connected capsules in consecutive layers ( Fig . 1 ) . In CNNs , the weights between neurons represent convolution filters , which are to be learned in training . In capsule nets , each weight matrix W , however , represents a linear transformation that would map the parts of an object into a cluster for the same whole . Similar to CNN , weight matrices in capsule nets are globally learned in training through labeled samples . As viewpoint changes do not alter the relative orientation between parts and whole , maintaining the same W would be sufficient to handle inputs from different angles , leading a viewpoint-invariant system . The coupling coefficients cij are assignment probabilities between a pair of whole-part capsules . Different from the weights , cij is dynamically determined solely in inference time . The cijs between each capsule i and all its potential parent capsules are summed to 1 , and thus have a similar effect of average pooling in CNNs . While the two versions of capsule nets ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) follow the same HDCF general principle for object recognition , they differ in several aspects , primarily the setup of the instantiation parameter u and the routing process to estimate cij . In Sabour et al . ( 2017 ) , u is set as a one-dimensional vector , whose length represents the probability of the capsule ’ s presence . The routing algorithm , dynamic routing , updates cij indirectly by the scalar product of votes and outputs of possible parents . Meanwhile , Hinton et al . ( 2018 ) uses a 4× 4 matrix to represent instantiation parameters of a capsule , and formulates the routing procedure as a mixture of Gaussians problem , which is then solved with a modified Expectation-Maximization algorithm . In this work , we take advantage of the part-whole relationships available in capsule nets to produce the labels for the pixels at the input image space . To facilitate the derivation of our solution , we would first summarize the related concepts under a probabilistic framework . Capsules can be categorized into different types ( e.g . cat , bicycle , or sofa ) . Each capsule layer l is associated with a set of capsule types T l = { tl1 , tl2 , ... , tlm , ... } . At each spatial coordinate ( position ) of certain layer l , there exists exactly one capsule of every type in T l. The probability of a capsule i existing/activated in the network is denoted as P ( i ) . If capsule i on layer l is within the receptive field of the capsule j of layer l + 1 , we call j a possible parent of i . The assignment probability of i being a part of j is denoted as P ( j|i ) = cij . The assignment probabilities between capsule i and the capsules co-located with j are forced to meet the normalization requirements of ∑ j∈T l+1 cij = 1 ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) . Three types of capsule layers were introduced in ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) , which will also be used as building blocks in our Tr-CapsNet . Therefore , we briefly describe them as follows , • Primary capsule layer is the first capsules layer , where features from previous convolution layer are processed and transitioned into capsules via convolution filtering . • Convolutional capsule layer ( s ) function similarly to CNN convolutional layers in many aspects . However , they take capsules as inputs and utilize certain routing algorithm to infer outputs , which are also capsules . • Class capsule layer L is a degenerated layer with one capsule for each predefined class label Ck ∈ C = { C1 , C2 , ... , Ck , ... } . Each capsule in the previous layer is fully connected to the capsules in this layer . The capsule nets are designed for object classification . In this work , we aim to take advantage of the notion and structure of capsule nets to design a semantic segmentation solution , hoping the viewpoint invariance and network interoperability mechanism can lead to improved segmentation performance . We present the description of our model in next section . 3 THE ARCHITECTURE OF OUR TR-CAPSNET . The general structure of our Tr-CapsNet , which consists of three components , is shown in Fig . 2 . It starts with a feature extraction module that capture the discriminative features of the input data to be fed into the later modules . For this purpose , a sequence of convolution layers as in CNNs would be a good choice ( Sabour et al. , 2017 ; Hinton et al. , 2018 ) . A capsule & traceback module comes next . It consists of a sequence of capsule layers and a traceback pipeline , followed by a convolution layer . The lineup of the capsule layers , shown as orange and brown arrows in Fig . 2 , can be a combination of one primary capsule layer , optionally one or more convolutional capsule layers , and a class capsule layer at the end . The traceback pipeline , shown as red arrows in Fig . 2 , is the major innovation of this paper . It produces class maps of the same size as the primary capsule layer , which are then taken as inputs to a convolution layer . This pipeline is specifically designed for segmentation purpose , and its details will be presented in next subsection . The third component of our model , a upsampling module , restores the original resolution by upsampling the label map computed in the capsule & trackback module . This module is implemented based on the deconvolution scheme in Long et al . ( 2015 ) in this paper , but it can also be implemented through atrous convolution ( Chen et al. , 2018 ) , unpooling ( Badrinarayanan et al. , 2017 ) , or simply bilinear interpolation . It should be noted that the feature extraction and upsampling layers can be regarded as somewhat symmetric in terms dimensionality , where the latter can be optional if no dimensionality reduction occurs in feature extraction layers .

Summary: This paper proposes a traceback layer for capsule networks to do semantic segmentation. Comparing to previous works that use capsule networks for semantic segmentation, this paper makes explicit use of part-whole relationship in the capsule layers. Experiments are done on modified MNIST and Hippocampus dataset. Results demonstrate encouraging improvements over U-Net. The writing could be tremendously improved if some background of the capsule networks is included.