B RACTIVATE : D ENDRITIC B RANCHING IN M EDICAL I MAGE S EGMENTATION N EURAL A RCHITECTURE S EARCH
Anonymous authors
Paper under double-blind review
A BSTRACT
Researchers manually compose most neural networks through painstaking experimentation.
This process is taxing and explores only a limited subset of possible architecture.
Researchers design architectures to address objectives ranging from low space complexity to high accuracy through hours of experimentation.
Neural architecture search (NAS) is a thriving ﬁeld for automatically discovering architectures achieving these same objectives.
Addressing these ever-increasing challenges in computing, we take inspiration from the brain because it has the most efﬁcient neuronal wiring of any complex structure; its physiology inspires us to propose Bractivate, a NAS algorithm inspired by neural dendritic branching.
An evolutionary algorithm that adds new skip connection combinations to the most active blocks in the network, propagating salient information through the network.
We apply our methods to lung x-ray, cell nuclei microscopy, and electron microscopy segmentation tasks to highlight Bractivate’s robustness.
Moreover, our ablation studies emphasize dendritic branching’s necessity: ablating these connections leads to signiﬁcantly lower model performance.
We ﬁnally compare our discovered architecture with other state-of-the-art UNet models, highlighting how efﬁcient skip connections allow Bractivate to achieve comparable results with substantially lower space and time complexity, proving how Bractivate balances efﬁciency with performance.
We invite you to work with our code here: https://tinyurl.com/bractivate.
I NTRODUCTION
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 efﬁcient 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 efﬁcient 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 ﬁeld 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 ﬁnal 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:
The depth necessary for many segmentation tasks is initially unknown, and traditional neural architecture search (NAS) struggles to identify the optimal UNet depth.
Researchers often manually choose skip connection locations, leading to potentially missed optimal connections.
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 efﬁcient 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 ﬁelds’ need for efﬁciency with Bractivate, a NAS algorithm to discover lightweight
UNet architectures for medical image segmentation tasks.
We present the following three primary contributions:
An evolutionary algorithm that non-randomly samples from a distribution of various UNet Model depths and skip connection conﬁgurations, with both tensor concatenation and addition operators.
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 efﬁcient networks that preserve information signals through the network.
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.
R ELATED W ORKS
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 classiﬁcation problems, and although some focus on skip connections, they lack deeper investigation of their optimal conﬁgurations.
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 efﬁcacy 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 ﬁnal 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 conﬁgurations.
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 speciﬁc to image classiﬁcation and sequential model architecture, they lack applications for segmentation models.
Weng et al.
(2019) suggest a
NASUNet method with modiﬁed 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 modiﬁed 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 ﬂoating-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 efﬁcient models, we propose Bractivate, a NAS algorithm inspired by the brain’s dendritic branching to facilitate optimal architecture discovery.
T HE B RACTIVATE NAS A LGORITHM 3. D ENDRITIC A RBORIZATION
[Table]
Table 1 translates neuroscience into computational terms we use throughout the paper.
In the neuroscience ﬁeld, 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-
[Figure]
C onv - T r a n s po s e ( M , M ,
F )
C onv - 1 -( M ,
M , F )
B a t c h N o r m - 1 ( M , M ,
F )
D r op O u t ( M , M , F )
C onv - 2 ( M ,
M , F )
B a t c h N o r m - 2 ( M , M ,
F )
C onv - T r a n s po s e ( M , M ,
F )
C onv - 1 -( M ,
M , F )
B a t c h N o r m - 1 ( M , M ,
F )
D r op O u t ( M , M , F )
C onv - 2 ( M ,
M , F )
B a t c h N o r m - 2 ( M , M ,
F )
C onv - T r a n s po s e ( M , M ,
F )
C onv - 1 -( M ,
M , F )
B a t c h N o r m - 1 ( M , M ,
F )
D r op O u t ( M , M , F )
C onv - 2 ( M ,
M , F )
B a t c h N o r m - 2 ( M , M ,
F )
Concatenation
[Figure]
tion 1. [Equation] where A b represents block activation, b ∈ B , A l 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( A b ) , 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. NAS WITH D ENDRITIC B RANCHING M UTATIONS
[Figure]
We implement our dendritic branching mutations to our evolutionary algorithm based on adding dendritic branching mutations to a randomized search domain, as displayed in Figure 3.
This mutation applies a mutation operator to form new skip connection combinations to the most active block in the network.
This NAS algorithm initializes a random domain with n model genotype.
Each genotype codes for the model’s depth, D .
It also encodes the number of ﬁlters per conv and deconv layer, the degree of skip connections 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. E FFICIENT L OSS
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 efﬁcient 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.
[Equation]
where m is the number of samples, y i is the sample image’s true segmentation mask tensor, and ˆ y i is the model’s prediction tensor.
We propose an alternate efﬁcient loss function equation, Efﬁcient
Loss Scaling ( EL S ).
It uses the number of model parameters, P , and the training time per epoch,
T .
E FFICIENCY L OSS S CALING
We also explore Efﬁciency penalty scaling where log( P ) and log( T ) scale the overall loss function through multiplication, hence:
[Equation]
In our experiments we set γ = 0 .
We use Equation 3 in Section 4.4 during model search.
A detailed ablation study on how this equation favors efﬁcient networks and outperforms standard
BCL can be found in Appendix A.3.
E XPERIMENTS 4. I MPLEMENTATION
For the following experiments, we execute our programs using TensorFlow 2.3.0 and Keras 2.4.3, using Python 3.6.9.
For training the models and inference, we use a virtual machine with 25 GB of
RAM, 150 GB of disk space, and one Tesla v-100 GPU.
T HE D ATA
To validate Bractivate NAS, we work with three datasets:
The Montgomery Pulmonary Chest X-Ray (Jaeger et al., 2014) dataset (N = 138) 2.
The EPFL CVlab Electron Microscopy (Lucchi et al., 2011) dataset (N=330) 3.
The Kaggle Data Science Bowl Cell Nuclei Segmentation Challenge (Caicedo et al., 2019) dataset (N=670)
A detailed discussion of the dataset follows in Section 4.3.
4. P REPROCESSING
We resize all images in the three datasets to be 128 × 128 pixels with bilinear interpolation, and later min-max normalize the pixels, so that pixel intensities are ∈ [0 , 1] .
We apply the same preprocessing to the image masks.
We use a 0.1 validation split percentage on all three datasets to evaluate the optimized model loss function and the Dice coefﬁcient during the search.
4. P ENALIZING T IME AND S PACE C OMPLEXITY
We experiment with using the standard binary BCL loss function in Equation 2, and the EL S function in Equation 3.
We record our results both with the Dice Score, representing the precision and recall harmonic mean.
We note that these loss functions are used only during the search step.
During full model training, we use the BCL loss on the ”best model” and use that for our later experiments in Sections 4.7 and 4.8.
4. D ISCOVERING A RCHITECTURE
For all three datasets, we use the same Bractivate search algorithm based on dendritic branching.
We initialize our search domain as a queue with 20 randomly generate model genotypes and initialized generated models from the genotypes; we then train the models for 50 epochs on the datasets with
Early Stopping.
The Early Stopping has patience of 10 and min ∆ of 0.05.
In a ﬁrst-in-ﬁrst-out (FIFO) fashion, we evaluate each genotype in the queue: if the generated model has EL min , then it becomes the ”Best Model,” and we place this genotype back in the queue.
Suppose it has a
EL > EL min .
In that case, the search algorithm mutates the ”Best Model” genotype with the dendritic branching method described in Figure 2 before replacing the mutated genotype in the search queue.
4. G OING D EEPER
We notice that the original Bractivate search algorithm with a minimum depth of two yields mainly shallow networks, usually with two or three expanding and contracting blocks.
To explore how model depth affects the search algorithm’s results, we constrain the search space such that depth is ∈ [5 , 10] and later ∈
[7 , 10] , and observe how the Dice coefﬁcient and Dice:time ratio change while using the EL S (Equation 3).
4. A BLATING B RANCHING C ONNECTIONS
We hypothesize that active block branching increases signal propagation, reducing time complexity, and improving overall model performance.
Thus, we must prove these new branches are necessary for the model’s success by ablating them and measuring how their absence affects the model performance.
We perform these experiments by training the selected model architecture for 200 epochs on the dataset, and use the Keract (Remy, 2018) library to measure layer activations on the most active layer (the second convolutional layer in the decoder of a D = 5 deep UNet.
For each layer, we calculate the layer’s average activation from Equation 1 and then ablate all dendritic (input) connections of the most active block.
We record the results quantitatively with the Dice coefﬁcient and visualize them by analyzing changes in the activation heat maps.
4. B ASELINE C OMPARISON
We compare Bractivate models to other state-of-the-art methods, including the standard UNet model (Ronneberger et al., 2015),
Wide-UNet (Zhou et al., 2018), UNet++ (Zhou et al., 2020), and attention-based models .
We obtain all model architectures from GitHub repositories made by these models’ authors or by contributors who create implementations in Keras and use them with Xavier weight initialization on the three datasets for comparison.
R ESULTS
We observe clear patterns pointing to how dendritic branching allows for efﬁcient neural network selection through our experimentation.
Figure 4 presents an example of the discovered architecture with the ELS function.
5. D ISCOVERED A RCHITECTURES
[Figure]
Because the mutations add more connections to salient blocks, the most active node has the highest input connections.
Bractivate randomly assigns addition or concatenation connections to the most active block, optimizing the best connecting operator combination.
Figure 4 shows the discovered architecture when the search space is constrained such that D ∈ [5 , 10] .
Al Ghamdi et al .
(Ghamdi et al., 2020) report that the addition operator can perform as well or better than the concatenation operator, so we incorporate both operators as modules into the NAS, with each mutation yielding a new random connecting operator combination to the salient block.
5. D EPTH AND T IME -S PACE C OMPLEXITY
[Figure]
In Figure 5, we observe a negative correlation between the training time per epoch and the Dice coefﬁcient, indicating that
Bractivate favors shallower over deeper networks, as the skip connections branching towards active blocks carry the necessary information for the segmentation task.
This slope varies between datasets; a simple problem like segmenting twocomponent lungs confers an overall larger
Dice: time ratio than the cell nuclei and electron microscopy tasks.
These problems are challenging because they have more than two connected components, yielding lower Dice: time ratios.
More importantly, our methods determine that given the early stopping constraint, shallower models ( D = 2 ) with fewer parameters have comparable performance to deeper
models. When UNet models are too deep ( D ∈ [7 , 10] ), the input signal may be lost in transformations, leading to lower performance; shallower models preserve this signal
5. S KIP C ONNECTION A BLATION S TUDY
Our ablation study most strongly conﬁrms our hypothesis that dendritic branching to active blocks signiﬁcantly improves segmentation performance: Figure 6 examines the saliency maps produced by the model architecture in Figure 4 before and after ablating connections to the most active block.
Before ablation, the most active block’s saliency maps show encoded information that signiﬁcantly inﬂuences the decoded feature maps during deconvolution.
After ablation, the saliency maps for the EM and nuclei segmentation tasks lack accurate late-stage saliency maps.
When the salient block has conﬁgured dendritic branches from neighboring blocks,
[Figure] a t [Figure] d
S [Figure] r c h e d
A b l a t e d
S ea r c h e d
[Figure]
l a t e d
[Figure]
c [Figure] e d
L ung s
E M
[Figure]
u c l e i
[Figure]
the output signal is highly accurate.
However, when these vital encodings in the Decode 2 block lack input from neighboring blocks, the output signal is degraded.
This degradation is especially true for the EM and nuclei segmentation tasks.
The EM and nuclei segmentation tasks contain more than two connected components; removing dendrites to salient blocks prevents valuable information from neighboring blocks to travel to the most salient block, degrading the signal through the last four blocks in the network.
The model’s
Dice score is signiﬁcantly lower in the ablated architecture than in the intact Bractivate-selected model.
The added information from these dendritic skip connections, explicitly targeting a salient block in the model, generates more accurate saliency maps, helping the model learn faster.
Before ablation, activations are more salient during the decoding phase than post-ablation, where saliency concentrates in the encoder.
This observation may be because removing connections towards an active block forces surrounding layers to compensate by increasing their activations.
5. B ASELINE C OMPARISON
[Table]
Figure 7 highlights how Bractivate achieves comparable performance to larger models when initialized with Xavier initialization (Glorot & Bengio, 2010).
Table 2 highlights how Bractivate is signiﬁcantly smaller than many of the other state-of-the-art models: it exchanges high spatial com- plexity for more skip connections, as these branches allow information to propagate through salient blocks in the network.
For domain-speciﬁc tasks, high parameters reduce the signal: noise ratio in the network; simpler models like Bractivate rely on powerful skip connections, analogous to dendrites, to carry most of the signal.
Because these connections consist of simple concatenation or addition operators, they greatly reduce the number of trainable parameters, preventing overﬁtting; this speaks to Bractivate’s comparable–or better–Dice scores as compared to the baseline models.
[Figure]
C ONCLUSION
Throughout this paper, we highlight how dendritic branching in the brain inspires efﬁcient skip connections in Deep Learning models.
With our focus on segmentation, we present Bractivate as a method for identifying skip connection conﬁgurations to elevate the traditional UNet.
During the search, Bractivate mutates the architecture so that the most salient blocks in the network branch out their ”dendrites” to other network blocks.
By replacing the oldest model in the search space with the new mutated architecture, we accelerate the search rate.
The ablation study strongly supports our hypothesis that dendritic branching is necessary for efﬁcient model discovery; when we ablate dendritic connections to the most salient block, the Dice Score decreases.
Before and after ablation, the saliency maps reveal stark contrasts, with the ablated activation maps lacking apparent features for the ﬁnal segmentation layer in the UNet’s decoder.
We ﬁnally weigh our methods with other baselines, highlighting how smaller networks can perform segmentation tasks well given limited pretraining data.
Overall, we present how optimally conﬁgured skip connections, inspired by the brain, yield robust signal streaming paths through a lightweight network.
Our algorithm is an asset to many mobile medical computing technologies that rely on low latency and high computational efﬁciency
A A PPENDIX A. G ENOME M ICRO -A RCHITECTURE
When designing our search space, we formulate genotypes that code for model architectures.
Following common patters in convolutional networks, and the UNet Ronneberger et al.
, we ﬁrst impose the following constraints on our search space:
• The network blocks must be symmetrical.
This means that the number of blocks both in the network encoder and decoder are identical, with mirror internal layer conﬁgurations (types of layers, numbers of ﬁlters, and number of layers in the block) •
The network must be hierarchical.
When designing models for medical image segmentation, we rely on hierarchical backbones for both the encoder and decoder, as reﬂected in Figure 4.
• We constrain skip connection directionality.
In the network, skip connections only occur from earlier to later layers in the background.
Figure 8 shows the standard micro-architecture for the contracting and expanding parts of the network.
We also note that while the layer arrangements are constant, the number of ﬁlters, n , for each block is initially random.
However, each integer value of ﬁlter numbers is scaled by a factor of 1.
for each subsequent block, as Figure 8 highlights
[Figure]
A. GPU R UN -T IME
Overall, the search algorithm had GPU run-times, as shown in Table 3.
We note that these results are reported after running the search algorithm on a Tesla-v100.
The reported values are an average of three independent trials.
Oftentimes, the run time was dependent on the dataset size.
Because the
Cell Nuclei dataset had the highest number of sample images, it took longer to train on as compared to the smaller Lung dataset.
[Table]
[Figure]
We also examine the effect of ablating the efﬁciency loss’ parameter and time penalty terms on the overall model selection.
Through our investigation, we ﬁnd that the efﬁciency loss does help the model select smaller models, that can perform at the level of larger models selected by the BCL loss function.
Figure 9 highlights this trend for the Lung Dataset.
The results are averaged over three trial runs.
We see that removing the penalty terms for high space and time complexity still yields highperforming models.
However, these models are larger, and computationally costly.