Shape-Tailored Deep Neural Networks Using PDEs for Segmentation

1 INTRODUCTION . Convolutional neural networks ( CNNs ) have been used extensively for segmentation problems in computer vision He et al . ( 2017 ) ; He et al . ( 2016 ) ; Chen et al . ( 2017 ) ; Xie & Tu ( 2015 ) . CNNs provide a framework for learning descriptors that are able to discriminate different textured or semantic regions within images . Much progress has been made in segmentation with CNNs but results are still far from human performance . Also , significant engineering must be performed to adapt CNNs to segmentation problems . A basic component in the architecture for segmentation problems involves labeling or grouping dense descriptors returned by a backbone CNN . A difficulty in grouping these descriptors arises , especially near the boundaries of segmentation regions , as CNN descriptors aggregate data from fixed shape ( square neighborhoods ) at each pixel and may thus aggregate data from different regions . This makes grouping these descriptors into a unique region difficult , which often results in errors in the grouping . In segmentation problems ( e.g. , semantic segmentation ) , current methods attempt to mitigate these errors by adding post-processing layers that aim to group simultaneously the ( coarse-scale ) descriptors from the CNN backbone and the fine-level pixel data . However , the errors introduced might not always be fixed . A more natural approach to avoid this problem is to consider the coarse and fine structure together , avoiding aggregation across boundaries , to prevent errors at the outset . To avoid such errors , one could design descriptors that aggregate data only within boundaries . To this end , Khan et al . ( 2015 ) introduced “ shape-tailored ” descriptors that aggregate data within a region of interest , and used these descriptors for segmentation . However , these descriptors are hand-crafted and do not perform on-par with learned approaches . Khan & Sundaramoorthi ( 2018 ) introduced learned shape-tailored descriptors by learning a neural network operating on the input channel dimension of input hand-crafted shape-tailored descriptors for segmentation . However , these networks , though deep in the channel dimension , did not filter data spatially within layers . Since an advantage of CNNs comes from exploiting spatial filtering at each depth of the network , in this work , we design shape-tailored networks that are deep and perform shape-tailored filtering in space at each layer using solutions of the Poisson PDE . This results in shape-tailored networks that provide more discriminative descriptors than a single shape-tailored kernel . This extension requires development of techniques to back-propagate through PDEs , which we derive in this work . Our contributions are specifically : 1 . We construct and show how to train ST-DNN , deep networks that perform shape-tailored spatial filtering via the Poisson PDE at each depth so as to generalize a CNN to arbitrarily shaped regions . 2 . We show analytically and empirically that ST-DNNs are covariant to translations and rotations as they inherit this property from the Poisson PDE . In segmentation , covariance ( a.k.a. , equivariance ) to translation and rotation is a desired property : if a segment in an image is found , then the corresponding segment should be found in the translated / rotated image ( or object ) . This property is not generally present with existing CNN-based segmentation methods even when trained with augmented translated and rotated images Azulay & Weiss ( 2019 ) , and requires special consideration . 3 . We show analytically and empirically that ST-DNNs are robust to domain deformations . These result from viewpoint change or object articulation , and so they should not affect the descriptor . 4 . To demonstrate ST-DNN and the properties above , we validate them on the task of segmentation , an important problem in low-level vision Malik & Perona ( 1990 ) ; Arbelaez et al . ( 2011b ) . Because of properties of the PDE , ST-DNN also have desirable generalization properties . This is because : a ) The robustness and covariance properties are built into our descriptors and do not need to be learned from data , b ) The PDE solutions , generalizations of Gabor-like filters Olshausen & Field ( 1996 ) ; Zador ( 2019 ) , have natural image structure inherent in their solutions and so this does not need to be learned from data , and c ) Our networks have fewer parameters compared to existing networks in segmentation . This is because the PDE solutions form a basis and only linear combinations of a few basis elements are needed to learn discriminative descriptors for segmentation . In contrast , CNNs spend a lot of parameters to learn this structure . 1.1 RELATED WORK . Traditional approaches to segmentation rely on hand-crafted features , e.g. , through a filter bank Haralick & Shapiro ( 1985 ) . These features are ambiguous near the boundaries of objects . In Khan et al . ( 2015 ) hand-crafted descriptors that aggregate data within object boundaries are constructed to avoid this , but lack sufficient capacity to capture the diversity of textures or be invariant to nuisances . Deep-learning based approaches have showed state-of-the-art results in edge-based methods Xie & Tu ( 2017 ) ; He et al . ( 2019 ) ; Deng et al . ( 2018 ) . Watershed is applied on edge-maps to obtain the segmentation . The main drawback of these methods is it is often difficult to form segmentations due to extraneous or faint edges , particularly when `` textons '' in textures are large . CNNs have been applied to compute descriptors for semantic segmentation , where pixels in an image are classified into certain semantic object classes Li et al . ( 2019 ) ; Huang et al . ( 2019 ) ; Du et al . ( 2019 ) ; Pang et al . ( 2019 ) ; Zhu et al . ( 2019 ) ; Liu et al . ( 2019 ) . Usually these classes are limited to a few object classes and do not tackle general textures , where the number of classes may be far greater , and thus such approaches are not directly applicable to texture segmentation . But semantic segmentation approaches may eventually benefit from our methodology as descriptors aggregating data only within objects or regions are also relevant to these problems . A learned shape-tailored descriptor Khan & Sundaramoorthi ( 2018 ) is constructed with a Siamese network on hand-crafted shape-tailored descriptors . However , Khan & Sundaramoorthi ( 2018 ) only does shape-tailored filtering in pre-processing as layering these requires new methods to train . We further examine covariance and robustness , not examined in Khan & Sundaramoorthi ( 2018 ) . Covariance to rotation in CNNs has been examined in recent works , e.g. , Weiler et al . ( 2018 ) ; Yin et al . ( 2019 ) ; Anderson et al . ( 2019 ) . They , however , are not shape-tailored so do not aggregate data only within shaped regions . Lack of robustness to deformation ( and translation ) in CNNs is examined in Azulay & Weiss ( 2019 ) and theoretically in Bietti & Mairal ( 2017 ) . Sifre & Mallat ( 2013 ) constructs deformation robust descriptors inspired by CNNs , but are hand-crafted . 2 CONSTRUCTION OF SHAPE-TAILORED DNN AND PROPERTIES . In this section , we design a deep neural network that outputs descriptors at each pixel within an arbitrary shaped region of interest and aggregates data only from within the region . We want the descriptors to be discriminative of different texture , yet robust to nuisances within the region ( e.g. , local photometric and geometric variability ) to be useful for segmentation . Our construction uses a Poisson PDE , which naturally smooths data only within a region of interest . Smoothing naturally yields robustness to geometric nuisances ( domain deformations ) . By taking linear combinations of derivatives of the output of the PDE , we can approximate the effect of general convolutional kernels but avoid mixing data across the boundary of region of interest . ST-DNN is also covariant to translations and rotations , inheriting it from the Poisson equation , which leads to the segmentation algorithm being covariant to such transformations . 2.1 SHAPE-TAILORED DNN DESCRIPTORS THROUGH POISSON PDE . Shape-tailored Smoothing via Poisson PDE : To construct a shape-tailored deep network , we first smooth the input to a layer using the Poisson PDE so as to aggregate data only within the region of interest , similar to what is done in Khan et al . ( 2015 ) for just the first layer . Let R ⊂ Ω ⊂ R2 be the region of interest , where Ω is the domain of the input image I : Ω→ Rk and k is the number of input channels to the layer . Let u : R → RM ( M is the number of output channels ) be the result of the smoothing ; the components u of u solve the PDE within R : { u ( x ) − α∆u ( x ) = I ( x ) x ∈ R ∇u ( x ) ·N = 0 x ∈ ∂R , ( 1 ) where I is a channel of I , ∂R is the boundary of R , N is normal to ∂R , α is the scale of smoothing and ∆/∇ are the Laplacian and the gradient respectively . It can be shown that the smoothing can be written in the form u ( x ) = ∫ R K ( x , y ) I ( y ) dy where K ( . , . ) is the Green ’ s function of the PDE , a smoothing kernel , which further shows that the PDE aggregates data only within R. Shape-tailored Deep Network : We can now generalize the operation of convolution tailored to the region of interest by taking linear combinations of partial derivatives of the output of the PDE equation 1 . This is motivated by the fact that in R = R2 , linear combinations of derivatives of Gaussians can approximate any kernel arbitrarily well . Gaussian filters are the solution of the heat equation , and the PDE equation 1 relates to the heat equation , i.e. , equation 1 is the steady state solution of a heat equation . Thus , linear combinations of derivatives of equation 1 generalize convolution to an arbitrary shape R ; in experiments , a few first order directional derivatives are sufficient for our segmentation tasks ( see Section 5 for details ) . A layer of the ST-DNN takes such linear combinations and rectifies it as follows : fi ( x ) = r ◦ Li ◦ T [ I ] ( x ) , ( 2 ) where I : R→ Rk is the input to the layer , T is an operator that outputs derivatives of the solution of the Poisson PDE equation 1 , Li ( y ) = wiy+bi is a point-wise linear function ( i.e. , a 1×1 convolution applied to combine different channels ) , r is the rectified linear function , and i indexes the layer of the network . Notice that since r and Li are pointwise operations , they preserve the property of T that it aggregates data only within the region R. We now compose layers to construct a ST-DNN as follows : F [ I ] ( x ) = s ◦ fm ◦ fm−1 ◦ fm−2 ◦ .... f0 ◦ I ( x ) , ( 3 ) where F [ I ] ( x ) is the output of the ST-DNN , f0 , ... , fm are the m+ 1 layers of the network , I is the input image , and s represents the soft-max operation ( to bound the output values ) . ST-DNN does not have a pooling layer because the PDE already aggregates data from a neighborhood by smoothing ; further , the lack of reduction in spatial dimension allows for more accurate shape estimation in our subsequent segmentation , and avoids the need for up-sampling layers . We will show that we can retain efficiency in training and inference .

Summary: The paper suggests replacing convolutional layers with ST-DNNs. ST-DNNs, in contrast to conv. layers, can natively support non-rectangular neighborhoods (regions). Similarly to conv. layers, ST-DNNs can be stacked to increase expressivity. ST-DNNs themselves correspond to solutions to the Poisson PDE. The paper describes how these layers can be trained and used for prediction. Results are reported on a texture dataset.