SHAMANN: Shared Memory Augmented Neural Networks

Current state-of-the-art methods for semantic segmentation use deep neural networks to learn the segmentation mask from the input image signal as an imageto-image mapping . While these methods effectively exploit global image context , the learning and computational complexities are high . We propose shared memory augmented neural network actors as a dynamically scalable alternative . Based on a decomposition of the image into a sequence of local patches , we train such actors to sequentially segment each patch . To further increase the robustness and better capture shape priors , an external memory module is shared between different actors , providing an implicit mechanism for image information exchange . Finally , the patch-wise predictions are aggregated to a complete segmentation mask . We demonstrate the benefits of the new paradigm on a challenging lung segmentation problem based on chest X-Ray images , as well as on two synthetic tasks based on the MNIST dataset . On the X-Ray data , our method achieves state-of-the-art accuracy with a significantly reduced model size compared to reference methods . In addition , we reduce the number of failure cases by at least half . 1 INTRODUCTION . In the medical image analysis domain , the automatic parsing of medical images represents a fundamental task that impacts the efficiency of the entire clinical workflow from diagnosis to therapy planning , intervention and follow-up investigations . An essential step in this sense is the semantic segmentation of anatomical structures which supports the radiologist to read and understand the image content . Recent approaches are inspired from the vision domain and rely on fully convolutional neural networks , e.g. , ( Ronneberger et al. , 2015 ; Yang et al. , 2017 ) , to achieve state-of-the-art results on various segmentation problems ( Menze et al. , 2015 ) . Usually , these methods use the entire image to directly predict the complete segmentation mask . While this facilitates the incorporation of valuable global image context , it also increases the complexity of the learning task , requiring the models to capture the complete variability in the shape and structure of different objects . In addition , this strategy does not scale well to ( volumetric ) high resolution data due to memory limitations . In this paper , we propose a new paradigm for semantic medical image segmentation based on a novel neural architecture called shared memory augmented neural network ( SHAMANN ) . Based on a decomposition of the original image into a sequence of image subregions , e.g. , local patches , we define different so called SHAMANN actors which traverse the sequence differently and segment each image subregion . An external memory module enables each actor to capture relations between different subregions within its sequence and increase the robustness of its predictions . In particular , this external module is shared among all actors and serves as a means to implicitly exchange local image context information in order to better capture global image properties , such as shape priors . Finally , the predictions of all actors are fused to obtain a semantic segmentation mask for the original image . An overview of the proposed framework with two SHAMANN actors is given in Figure 1 . The contributions of our work are : ( i ) a reformulation of the semantic segmentation problem as a sequence learning task ( ii ) SHAMANN - a memory efficient and dynamically scalabale alternative to end-to-end fully convolutional segmentation networks , that can also implicitly capture global image properties through a shared external memory module ; and ( iii ) a comprehensive analysis of the method and comparison against state-of-the-art methods on a large chest X-Ray dataset . 2 RELATED WORK . Segmentation . In the fields of computer vision and medical imaging , segmentation is a fundamental task for understanding the semantic content of an image . State-of-the-art results on different segmentation benchmarks ( Cordts et al. , 2015 ; Everingham et al. , 2015 ) , have been achieved by using fully convolutional neural networks ( He et al. , 2017 ; Shelhamer et al. , 2017 ) . However , one limitation of such networks is the use of pooling layers . By down-sampling and increasing the field-of-view , precise localization information is lost . To tackle this issue , two different approaches have been proposed . First , encoder-decoder architectures , e.g. , U-NET ( Ronneberger et al. , 2015 ) , recover the details and spatial dimension using de-convolutions and shortcut connections ( Badrinarayanan et al. , 2017 ; Lin et al. , 2017 ; Yang et al. , 2017 ) . The alternative is to use dilated convolutions to increase the field-of-view without decreasing the spatial dimension ( Chen et al. , 2018 ; Peng et al. , 2017 ; Yu et al. , 2017 ; Yu & Koltun , 2015 ; Zhao et al. , 2017 ) . In this context , graphical models such as conditional random fields ( Lafferty et al. , 2001 ) are used to further improve the results . In the medical context , a standard approach for medical segmentation is multi-atlas label propagation ( MALP ) ( Wang et al. , 2013 ; Zikic et al. , 2013 ) . In MALP , a collection of atlases , i.e. , labeled images , is required . At runtime , one needs to perform expensive non-linear registration operations of each atlas to unseen data to achieve a segmentation . These solutions typically scale poorly and are inefficient . Alternatively , one can address the segmentation problem by using random forests ( Glocker et al. , 2012 ) , providing stronger unary predictions through joint class and shape regression . Milletari et al . ( 2017 ) employed an additional patch-voting scheme to increase the robustness against outliers . Other approaches use linear shape models to incorporate prior information ( SSM ) ( Heimann & Meinzer , 2009 ) . In marginal space deep learning ( Ghesu et al. , 2016 ) , SSMs have been coupled with deep learning to enable the segmentation of anatomical structures . While these methods provide good results and are relatively easy to train , they do not exploit global anatomical information . In addition , the inference is time-consuming , especially for 3D images . Memory networks . Recently , neural networks have been augmented with an external memory module to decouple the memorization capacity from the network parameters , making these methods better suitable for capturing long-range dependencies in sequential data . These networks have been used in the context of classification ( Vinyals et al. , 2016 ) , meta-learning ( Santoro et al. , 2016 ; Sprechmann et al. , 2018 ) , reinforcement learning ( Mnih et al. , 2015 ; Pritzel et al. , 2017 ) , graph problems ( Graves et al. , 2016 ) or question answering ( Graves et al. , 2016 ; Sukhbaatar et al. , 2015 ) , to name a few . Closest to our work are generative methods ( van den Oord et al. , 2016 ) , which model the conditional probability of a pixel based on previous pixels , using LSTMs . In contrast , we propose a sequence learning task for image segmentation and show that the memorization capacity can be improved using a shared external memory . Bahdanau et al . ( 2014 ) and Wang et al . ( 2016 ) proposed a memory-based strategy for the task of machine translation . They use a bidirectional RNN to encode the input and save the concatenation of the outputs of the two units in a memory . After the sequence is processed and saved in the memory , a decoder reads from the memory and outputs the final predictions . In contrast , our proposed method allows information exchange between SHAMANN actors while processing the input sequence thereby enabling each agent to access global image context . To the best of our knowledge , this is the first paper that proposes a method based on memory networks for the task of image segmentation . 3 PROPOSED METHODS . In this section , we present our main contribution , the shared memory augmented neural networks ( SHAMANN ) architecture for semantic segmentation . Our observation is that in a bidirectional setup , information from different directions is not being explicitly exchanged . We hypothesize that by sharing an external memory , our networks can better capture global context , leading to a more accurate segmentation . 3.1 PROBLEM FORMULATION . In the following x and xT will denote a row and column vector respectively , and A a matrix . Following formulations are focused on but not limited to 2D images . Formally , let us consider an input image I : Ω → RHI×WI×C , with Ω ⊂ R2 the image domain ; HI , WI and C denoting the height , width and number of channels of the image signal . The goal of the segmentation task is to assign a label to every pixel/voxel x in the input image , considering a predefined set of K object classes { y1 , . . . , yK } . The segmentation result can be represented as a set of segmentation channels Y : Ω → RHI×WI×K , where the value of a pixel ( x , y ) of a given channel k encodes the probability of observing the class yk . A final segmentation mask can then be obtained by applying a softmax function along the different class-specific channels . In this work , we propose to reformulate the segmentation problem as a sequential learning task . Let us consider a sequence of T patches P = { P0 , . . . , PT } covering the image domain , with Pt : Ω→ RHp×Wp×C , where Hp , Wp are the height and width of the patch . For example , these patches may be extracted using uniform sampling . We propose to learn a function f that maps the sequence of input patches to a sequence of patch segmentation masks as f ( Pt ) Tt=0 = ( Φt ) T t=0 , with f : RT×Hp×Wp×C → RT×Hp×Wp×K . 3.2 ARCHITECTURE OVERVIEW . In this section , we introduce in more detail the components of our model ( as can be seen in Figure 1 ) . The encoder extracts a rich visual representation from the raw patch intensities . We model it as a function ( e.g. , a convolutional network ) , mapping the input to a d-dimensional feature space : E ( Pt ) T t=0 = ( ψt ) T t=0 , with E : RHp×Wp×C → Rd . The actor module , defined as component 2 , learns the sequence of input feature vectors Ψ = { ψ0 , . . . , ψT } and captures distal spatial dependencies . Each actor scans the input sequence Ψ differently , to produce an output sequence HJ = { hJ0 , . . . , hJT } , with hJ ∈ Rd . Here , we use two actors , one scanning the input in the forward direction ( J : = F ) , and the other in the backward direction ( J : = B ) . The patch-level fusion step combines the outputs of the actors as σ ( HF ⊕HB ) = H , with σ : R2×d → Rd and ⊕ the concatenation operator . The mapping σ could be a simple function , e.g. , an average or a concatenation operation . In our work , we propose to explicitly learn how to combine the different outputs using a neural network with a single fully connected layer . The decoder maps the fused outputs of the actors to patch segmentation masks as D ( ht ) Tt=0 = ( Φt ) T t=0 , with D : Rd → RHp×Wp×K . In the final image-level fusion step ( see component 4 ) , all patch segmentation masks Φ = { Φ0 , . . . , ΦT } are aggregated over the full image domain to generate the final segmentation mask Y . For fusion , we propose to use averaging ( Iglesias & Sabuncu , 2015 ) .

Summary: The paper proposes a system of semantic segmentation based on sequential processing of the image in a patch-wise manner with multiple "actors", sharing a common external memory. This approach stands in contrast to the more usual approach of single-shot prediction for the whole image, where encoder-decoder architectures or dilated convolutions are used to capture the global context. The authors then discuss three-variants of this method, out of which two use external memory (Bi-MANN, SHAMANN), and one uses memory shared between actors (SHAMANN). Results are presented on segmentation of lung X-ray data and on MNIST digit completion.