Coupled Segmentation and Edge Learning via Dynamic Graph Propagation

1 Introduction . Image segmentation and edge detection have been widely studied as important perception problems . The two problems are closely related . In fact , segmentation subsumes edge detection since any segmentation contour makes a closed boundary of a region . The converse is however not true since general edges do not always form closed contours . Nevertheless , edge detection can serve as an auxiliary task to improve segmentation performance since edges provide important pairwise similarity cues for segmentation . Early works tend to focus on the grouping and contrast of pixels from a perceptual similarity perspective . Martin et al . [ 1 ] proposed the Berkeley Segmentation Dataset , a popular benchmark for segmentation and boundary detection that inspired many impactful works in perceptual grouping [ 2–5 ] . The recent surge of deep learning renders powerful representations with learned features using convolutional neural networks ( CNNs ) [ 6 ] . This has led to great advances in both areas [ 7–12 ] , but the two tasks are often considered separately . In light of the status quo , we consider coupled edge and segmentation learning . Our goal is twofold : ( 1 ) Multi-task learning - being able to produce high quality edge detection and segmentation . ( 2 ) Mutual improvement - the two tasks can help each other with non-trivial performance gains . Designing a principled framework is however nontrivial . The key question is how sparse edge signals can be effectively transformed into dense region-level ones to interact with segmentation . To this end , we propose a learnable recurrent message passing layer where semantic edges are considered ∗Equal contribution . Correspondence to Zhiding Yu < zhidingy @ nvidia.com > . †Work partially done during an internship at NVIDIA . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . as explicitly learned gating signals to refine segmentation . An overview of our framework is shown in Fig . 1 . Specifically , the dynamic message passing layer uses affinity gates to select the neighbor for message passing using max-pooling . It conducts message passing sweeps in each of the four directions : left to right , right to left , top to bottom and bottom to top . The message passing is jointly gated by both the affinity and edge gates , therefore allowing edge cues to naturally influence long-range dense predictions . As such , our framework presents a context module that is clean , compact yet powerful . Our technical contributions can be summarized as follows : • We formulate recurrent message passing as dynamic graph propagation ( DGP ) . We show that such a formulation simplifies the required normalization in propagation networks [ 13 ] . It also dynamically finds graph structures that encodes pixel affinities and improves dense prediction . • We propose a double-gate design where message passing is jointly gated by both the affinity and edge gates to refine the segmentation feature map . We show that this design together with the dynamic 1-way connection in DGP better couples segmentation and edge learning . • We obtain state-of-the-art results on joint semantic segmentation and edge detection . We also show that DGP leads to strong zero-shot robustness to natural corruptions with significant improvement over prior methods on Corrupted Cityscapes ( Cityscapes-C ) . Multitasking segmentation and edge learning is desirable for several reasons : 1 ) There are many downstream applications where both are needed , such as occlusion reasoning [ 14 ] , localization [ 15 ] , proposal generation [ 16–19 ] and conditional generation [ 20 ] . 2 ) There are many challenging cases where segmentation quality is poor but edge quality is far more superior , e.g . segmentation tends to be inferior near object boundaries since they are often optimized for IoU rather than precision [ 21 ] . In these cases , edge learning can potentially capture details missed by segmentation . 3 ) Implicitly improved model generalization as a result of the coupled learning [ 22 ] . 2 Related Work . Semantic Segmentation . There is a rich set of prior work in semantic segmentation . Long et al. , [ 7 ] proposed fully convolutional end-to-end training of semantic segmentation and pointed out its connection to recognition [ 6 ] . Chen et al . [ 8 ] introduced atrous convolution and atrous spatial pyramid pooling ( ASPP ) to capture multi-scale image contexts . Another contemporary work is the wide ResNet-38 which explores a relatively shallower but wider backbone [ 23 ] . It has also been shown that context plays an important role in segmentation , including context encoding [ 24 , 25 ] , multi-scale context [ 26–28 ] as well as relational context [ 29 , 30 ] . More recently , there has been a surge of interests in segmentation with Transformers [ 31–34 ] Boundary/Edge Detection . Similar to segmentation , boundary/edge detection have been widely studied as perceptual grouping problems in early literature [ 1–3 ] . Recent methods tend to resort to CNNs . For example , Bertasius et al . proposed a multi-scale deep network architecture for top-down contour detection , where as Xie et al . [ 11 ] further introduced holistically-nested edge detection ( HED ) for end-to-end edge learning . Besides detecting binary edges , Hariharan et al . [ 35 ] proposed the Semantic Boundaries Dataset ( SBD ) which has become a popular benchmark for semantic edge detection . Compared to binary edge detection , semantic edge detection involves the semantic classification of edge pixels in addition to localization which presents additional challenges to existing frameworks . A series of works including HFL [ 36 ] , CASENet [ 12 ] , SEAL [ 37 ] and STEAL [ 38 ] have followed up [ 35 ] and pushed the boundaries . Multi-task Segmentation and Edge Learning . Multi-task segmentation and edge learning remains under-studied but is not entirely new . Edges have been pairwise similarity cues to improve segmentation [ 3 , 39 , 18 ] and superpixels [ 4 ] . It was shown that edges can be transformed into dense regions through the Laplacian Eigenmaps of boundary transformed affinity [ 36 , 5 , 40 ] . Despite being robust , these methods are generally slow and not end-to-end trainable . For more recent CNN based models , end-to-end multi-task learning on a shared backbone is a natural choice [ 41 ] . In addition , proper regularization between segmentation and edge has been shown to improve the performance of both tasks . For example , Takikawa et al . [ 21 ] use softmax with temperature to impose consistency between segmentation and semantic edges . Zhen et al . [ 42 ] correlate semantic segmentation and edge detection tasks with a consistency loss . Our work goes beyond multi-tasking and let the two task more deeply coupled through dynamic graph propagation . Structured Dense Prediction . Segmentation is a dense prediction task where structured information can be useful . Structured prediction models such as Markov random fields ( MRFs ) [ 43 ] , conditional random fields ( CRFs ) [ 44–46 ] and energy minimization [ 47 , 48 ] have widely proved helpful in segmentation problems by imposing contrast-sensitive smoothness . In addition , structured inference can also be unfolded as network layers for end-to-end training [ 49 , 50 ] , therefore combining the advantages from both ends . Recently , there is an increasing trend to directly model the message passing process itself using multi-dimensional RNN [ 51–53 ] , graph RNN [ 54 , 55 ] and spatial propagation network ( SPN ) [ 13 , 56 , 57 ] . A common advantage is that these methods render more context-aware prediction with larger receptive fields while preserving local details similar to CRFs . Their ability to train end-to-end allows more powerful representation of inter-pixel similarities and thus better dense prediction quality . Our work can be broadly categorized into this category . It is worth mentioning that the contrast sensitive smoothness term in CRF and propagation networks [ 52 , 13 ] is an implicit modeling of edge signals by learning to relax smoothness constraints at high contrast areas . There have also been variants of SPN [ 58 ] that take binary edge as gate regularization . However , none of these works have explicitly addressed multi-tasking learning with category-aware edges whereas the proposed DGP framework presents a novel and effective solution to this problem . 3 Multi-task network . As a first step towards coupled segmentation and edge learning , we introduce two novel multi-task networks that are able to perform multi-task learning for both segmentation and edge detection . An overview of their architectures is shown in Fig . 2 . The rest of the paper follows these notations and covers more details for each module . Backbone . A CNN giving a semantic feature map ( denoted as F ) encoding the segmentation information . There could be multiple choices of the backbone networks , ranging from the popular architectures of DeepLabv2 [ 8 ] , DeepLabv3 [ 27 ] to latest state-of-the-arts such as DeepLabv3+ [ 28 ] and Dual-Attention Network [ 30 ] . Adopting powerful backbones will surely benefit the system-level performance , but the major purpose of this work does not completely lie in achieving state-of-the-arts . We are more interested in showing the effectiveness of a proposed framework on standard backbones . To this end , we consider both CASENet [ 12 ] and ResNet-38 [ 23 ] as the standard backbones for benchmarking . Both of them comprehensively covers the feature map resolutions of 1 , 1/2 , 1/4 and 1/8 , with the last resolution being a rule of thumb adopted by many segmentation networks . For both networks , we adopt atrous spatial pyramid pooling layers with dilation rate [ 6 , 12 , 18 , 24 , 30 , 36 ] to capture context information . We set the output channels to be 128 to produce a semantic feature map , fowllowed by either direct segmentation prediction ( baselines ) or the proposed DGP layer . Affinity stream . Convolution layers that aggregate the cross-layer information of the backbone and produces an affinity feature map ( denoted as A ) to encode the affinity of pairwise neighboring pixels . This stream is only used with the presence of recurrent message passing layer ( including its related baselines ) . The stream aims to model the simialrity of pair-wise neighboring pixels and serve as a major input to the gating signal in recurrent message passing layer . Compared to edge stream , the affinity stream seeks to capture coarser-level inter-pixel affinity with slightly deeper convolutions and higher dimensions . The resulting feature map A is a 256 dimensional tensor concatenated by ASPP and the side convolutional feature maps from the affinity stream . Edge stream . Dense skip-layer features with abundant details that are combined with edge classification layer through shared concatenation [ 12 ] to produce a semantic edge map ( denoted as E ) . We design an improved edge stream over CASENet [ 12 ] to better leverage detail information at multiple scales/levels with dense side connections . Unlike CASENet where the bottom blocks only provide 1-channel side features , we densely obtain side features from every sub-block residual modules . In CASENet , this gives side features with a total number of 31 dimensions . Similar rules apply to ResNet-38 , where the side feature has a total number of 17 dimensions We also found that applying sigmoid after side features with resolution 1/8 greatly benefits edge training from two aspects : ( 1 ) It helps to stabilize and prevent gradient explosions from dense skip connections . ( 2 ) It removes the typical aliasing edge issue caused by upsampling/deconvolution , and produces elegant smooth edges . Even though we do not explicitly apply techniques such as edge alignment [ 37 ] in this work , the proposed backbone is able to produce high quality edge predictions under noisy labels . We also notice that it is better to remove sigmoid for side features with higher resolutions . Edge prediction is multi-tasked alongside with ASPP using 1× 1 convolution on top of Res Block5 . This returns the K classes coarse edge predictions . Similar to CASENet , we upsample all side features together with the K-dimensional edge predictions to full image resolution and apply shared concatenation , where side features are repetitively shared and concatenated with each class for K times , followed by a K-way 3× 3 group convolution to produce semantic edges . Dynamic graph propagation . Learnable recurrent message passing layer that takes the above three branches as input to produce a refined semantic feature map ( denoted as F ∗ ) . More details regarding this module will be introduced in the next section .

Summary: This paper proposes a dynamic (i.e. recurrent) graph propagation module to improve the results of a feedforward semantic segmentation algorithm on the Cityscapes dataset. In addition to outputting the typical feature map where each pixel is assigned to one of 19 categories (i.e. the semantic segmentation map), the proposed network produces supervised predictions of semantic _edges_ (the boundaries between classes) and a set of local pixel-pixel affinities. These two novel outputs are used as gates in a recurrent message-passing algorithm, which refines the semantic segmentation map (as well as the edge map.) The authors show that the refined map gives improved semantic segmentation scores over the unadorned baseline and a number of other state-of-the-art algorithms. Finally, they find that the dynamic graph propagation module gives substantial improvements on "corrupted" versions of the Cityscape dataset (at least for some types of corruption.)