Bractivate: Dendritic Branching in Medical Image Segmentation Neural Architecture Search

1 INTRODUCTION Researchers manually composing neural networks must juggle multiple goals for their architectures . Architectures must make good decisions ; they must be fast , and they should work even with limited computational resources . These goals are challenging to achieve manually , and researchers often spend months attempting to discover the perfect architecture . To overcome these challenges , we turn to the human brain ’ s efficient neural wiring for automated architecture discovery . Neuroscience already underlies core neural network concepts : The perceptron ( Rosenblatt , 1958 ) is directly analogous to a human neuron . One of the brain ’ s fundamental learning mechanisms is dendritic branching ( Greenough & Volkmar , 1973 ) whereby active neurons send out signals for other neurons to form connections , strengthening signals through that neural pathway . This neuroscience concept inspires us to devise Bractivate , a Neural Architecture Search ( NAS ) algorithm for learning new efficient UNet architectures networks , capable of being trained twice as fast as the traditional UNet , and often one to two orders of magnitude lighter in terms of trainable parameters . We apply Bractivate on three medical imaging segmentation problems : cell nuclei , electron microscopy , and chest X-ray lung segmentation . Medical image segmentation is a growing field in Deep Learning Computer Assisted Detection ( CAD ) : it is a powerful component in clinical decision support tools and has applications in retinal fundus image , lung scan , and mammography analysis . Most papers now approach medical image segmentation with the UNet ( Ronneberger et al. , 2015 ) ; the model architecture is straightforward : it has symmetric , hierarchical convolutional blocks , which are components of an initial contracting path and a final expanding path , with an apex bottleneck layer . Between parallel contracting and expanding blocks , the traditional UNet contains skip connections that pass information through concatenation ( Ronneberger et al. , 2015 ) . Traditional UNet skip connections involve feature map aggregation with same-scale convolutional blocks , but recent advances have yielded more complex connections ranging from the UNet++ ( Zhou et al. , 2018 ) to the NasUNet ( Weng et al. , 2019 ) . While the UNet is a powerful tool , it does have many limitations : 1 . The depth necessary for many segmentation tasks is initially unknown , and traditional neural architecture search ( NAS ) struggles to identify the optimal UNet depth . 2 . Researchers often manually choose skip connection locations , leading to potentially missed optimal connections . 3 . Scientists need a NAS algorithm addressing many implementation objectives , including computational time , number of model parameters , and robust segmentation performance . On a broader level , discovering efficient UNet architectures is crucial because it can generate simpler models for applications on mobile devices , which need low latency for online learning . In the Telemedicine age , many medical applications rely on mobile Deep Learning to segment medical images and process raw patient data ( Xu et al. , 2017 ; Vaze et al. , 2020 ) . We address the Medical and Engineering fields ’ need for efficiency with Bractivate , a NAS algorithm to discover lightweight UNet architectures for medical image segmentation tasks . We present the following three primary contributions : 1 . An evolutionary algorithm that non-randomly samples from a distribution of various UNet Model depths and skip connection configurations , with both tensor concatenation and addition operators . 2 . ” Dendritic Branching ” -inspired mutations that , just as in the brain , cause salient UNet blocks to branch to other blocks in the network through dendritic skip connections , creating efficient networks that preserve information signals through the network . 3 . Bractivate generates high-performing models with lower space complexity than the current state-of-the-art . The remainder of the paper is structured as follows : In Section 2 , we discuss prior works , and what gaps in the literature inspire us to propose Bractivate . Then , in Section 3 , we discuss the search algorithm and the dendritic branching mutation . Later , in Section 4 , we implement our algorithm with various experiments ranging from changing the search space depth to an ablation study . We report our quantitative and qualitative results , along with baseline comparisons in Section 5 before concluding in Section 6 . 2 RELATED WORKS . Deep learning algorithms are often restricted to manual model design ( Simonyan & Zisserman , 2014 ; He et al. , 2016 ; Oktay et al. , 2018 ; Ronneberger et al. , 2015 ) . To automate model schemes , NAS is the process of selecting candidate architectures through various search strategies to achieve optimal performance ( Elsken et al. , 2019 ) . Advances in NAS have branched into different areas , including evolutionary algorithms ( Miller et al. , 1989 ; de Garis , 1990 ; Yao , 1993 ; Fogel et al. , 1990 ; Angeline et al. , 1994 ; Real et al. , 2018 ; Yao , 1999 ) and automatic pattern recognition ( Cai et al. , 2018 ; Radosavovic et al. , 2020 ) . While both approaches are merited , the tasks address image classification problems , and although some focus on skip connections , they lack deeper investigation of their optimal configurations . Recent advances in the UNet have led to alternative skip connection implementations , including addition ( Ghamdi et al. , 2020 ) , max out operations ( Estrada et al. , 2019 ; Goodfellow et al. , 2013 ) and multiplication by a gating function ( Oktay et al. , 2018 ) . Ghamdi et al . ( 2020 ) reports these connections ’ improved efficacy over traditional concatenation , as they overcome vanishing gradients and preserve salient features . Auto-DeepLab , which Liu et al . ( 2019 ) present for semantic segmentation , is a graph-based NAS algorithm that addresses changing model depth and connection locations in hierarchical models . Building off this work , Zhou et al . ( 2020 ) propose a similar graph-search algorithm , termed UNet++ , for improved NAS ; the final model incorporates dense skip connections to achieve multi-scale feature aggregation . Although UNet++ successfully addresses the model depth problem , it ignores choosing the skip connection operator and relies on pretraining and pruning to generate skip connection configurations . The Differential Architecture Search ( DARTs ) algorithm by Liu et al . ( 2018 ) continuously relaxes the architecture representation to enable gradient-based optimization . Advancing this algorithm , Chen et al . ( 2019 ) proposes the Progressive Differentiable Architecture Search Algorithm ( PDARTs ) to allow the searched model ’ s depth to grow during the search ; when applied to ImageNet ( Deng et al. , 2009 ) , CIFAR-10 ( Krizhevsky et al. , 2009 ) , or CIFAR-100 ( Krizhevsky et al. , 2009 ) , the total training time is approximately seven hours . Although the DARTS and PDARTs algorithms are specific to image classification and sequential model architecture , they lack applications for segmentation models . Weng et al . ( 2019 ) suggest a NASUNet method with modified DARTs search for medical imaging segmentation ; their approach addresses searching for model parameters in the convolutional blocks to reduce the space complexity found in attention-based ( Oktay et al. , 2018 ; Hy , 2018 ) and recurrent ( Alom et al. , 2018 ; Hy , 2018 ) UNets , yet NASUNet still preserves same-scale concatenation skip connections , overlooking alternate skip connection possibilities across network blocks . Many existing NAS algorithms use modified objective functions for evaluating the searched model performance e.g . NAS-Bench-101 ( Ying et al. , 2019 ) uses the cross-entropy loss , Stochastic Neural Architecture Search ( SNAS ) ( Xie et al. , 2019 ) devises a cost function deemed Memory Access Cost ( MAC ) that incorporates the floating-point operations ( FLOPs ) and number of parameters , and PDARTs ( Chen et al. , 2019 ) employs auxiliary loss ( Szegedy et al. , 2014 ) . To target gaps in the literature related to skip connection search for efficient models , we propose Bractivate , a NAS algorithm inspired by the brain ’ s dendritic branching to facilitate optimal architecture discovery . 3 THE BRACTIVATE NAS ALGORITHM 3.1 DENDRITIC ARBORIZATION Table 1 translates neuroscience into computational terms we use throughout the paper . In the neuroscience field , dendritic branching occurs when stimulating environments cause neurons to form new connections ( Greenough & Volkmar , 1973 ; Greenough et al. , 1985 ) . These neural connections are associated with learning , and even learning-impaired children with fetal alcohol syndrome display lower dendritic branching levels ( Hamilton et al. , 2010 ) compared to their healthy peers . This branching phenomenon parallels Deep Neural Networks : in the brain , dendrites form new connections to the hyperactive soma : the perceptron ’ s activation function is to the biological soma as the incoming connections are to dendrites . Perceptrons can be stacked together to form multi-layer perceptrons ( Rumelhart et al. , 1986 ) , with parallel architecture similar to the brain , and this structure underlies convolutional neural networks ( LeCun et al. , 1995 ) . For the UNet , if we consider layer in the network ’ s blocks to be a neural soma , then we can think about a block ’ s ” activity ” as the mean-absolute value of its layers ’ activations , as shown by Equa- tion 1 . Ab = 1 L L∑ l=0 |Al| ( 1 ) where Ab represents block activation , b ∈ B , Al is the weight of each layer in the block , l , and L is the total number of layers in the block . Knowing the block ’ s location , b , with max ( Ab ) , surrounding blocks then form skip connections around this active node , a process analogous to dendritic branching . We apply this method to target conv and deconv layers , excluding batch normalization layers as they contain static weights and high values that overwhelm the mutation ’ s layer selection . When new connections are formed , across blocks with various tensor dimensions , we overcome spatial dimensional mismatch by resizing the incoming connection tensors to the receiving tensor by bilinear interpolation . 3.2 NAS WITH DENDRITIC BRANCHING MUTATIONS . Sample Randomly Initialized Domain Chosen Model Best Model Genome Efficiency Evaluator Dendritic Branching Mutation for each block , and the skip connection operator type ( concatenation or addition ) . A detailed discussion on the genotype initialization and its micro-architecture is found in Appendix A.1 . When initializing the genotype , we constrain the number of feature maps to grow by 1.5 in the encoder and decrease it by the same factor for each block in the decoder . 3.3 EFFICIENT LOSS . NAS methods often focus on accuracy as the main performance metric ( Real et al. , 2018 ; Radosavovic et al. , 2020 ) , but often lack consideration for discovered model space and time complexity . To address this , we propose an efficient loss function for the NAS evaluation step . Traditional binary cross-entropy is given by Equation 2 . During the NAS mutation and selections , the search process accelerates as ” better ” models have faster training steps . BCL = − 1 m m∑ i=1 ( yi × log ( ŷi ) + ( 1− yi ) × log ( 1− ŷi ) ) ( 2 ) where m is the number of samples , yi is the sample image ’ s true segmentation mask tensor , and ŷi is the model ’ s prediction tensor . We propose an alternate efficient loss function equation , Efficient Loss Scaling ( ELS ) . It uses the number of model parameters , P , and the training time per epoch , T . EFFICIENCY LOSS SCALING We also explore Efficiency penalty scaling where log ( P ) and log ( T ) scale the overall loss function through multiplication , hence : ELS = γ × log ( P ) × log ( T ) ×BCL ( 3 ) In our experiments we set γ = 0.01 . We use Equation 3 in Section 4.4 during model search . A detailed ablation study on how this equation favors efficient networks and outperforms standard BCL can be found in Appendix A.3 .

Summary: The authors propose a neural architecture search (NAS) algorithm inspired by brain physiology. In particular, they propose a NAS algorithm based on neural dendritic branching, and apply it to three different segmentation tasks (namely cell nuclei, electron microscopy, and chest X-ray lung segmentation). The authors share their codes with the scientific community, which is highly appreciated.