Attention Gated Networks: Learning to Leverage Salient Regions in Medical Images

Introduction

Automated medical image analysis has been extensively studied in the medical imaging community due to the fact that manual labelling of large amounts of medical images is a tedious and error-prone task. Accurate and reliable solutions are required to increase clinical work flow efficiency and support decision making through fast and automatic extraction of quantitative measurements.

With the advent of convolutional neural networks (CNNs), near-radiologist level performance can be achieved in automated medical image analysis tasks including classification of Alzheimer’s disease [Sarraf et al., 2017], skin lesions [Esteva et al., 2017, Kawahara and Hamarneh, 2016] and echo-cardiogram views [Madani et al., 2018], lung nodule detection in CT/X-ray [Liao et al., 2017, Zhu et al., 2018] and cardiac MR segmentation [Bai et al., 2017]. An extensive list of applications can be found in [Litjens et al., 2017, Zaharchuk et al., 2018]. High representation power, fast inference, and weight sharing properties have made CNNs the de facto standard for image classification and segmentation.

Methods for existing applications rely heavily on multi-stage, cascaded CNNs when the target organs show large inter-patient variation in terms of shape and size. Cascaded frameworks extract a region of interest (ROI) and make dense predictions on that particular ROI. The application areas include cardiac MRI [Khened et al., 2018], cardiac CT [Payer et al., 2017], abdominal CT [Roth et al., 2017, 2018] segmentation, and lung CT nodule detection [Liao et al., 2017]. However, this approach leads to excessive and redundant use of computational resources and model parameters; for instance, similar low-level features are repeatedly extracted by all models within the cascade.

To address this general problem, we propose a simple and yet effective solution, named attention gates (AGs). By incorporating AGs into standard CNN models, model parameters and intermediate feature maps are expected to be utilised more efficiently while minimising the necessity of cascaded models to solve localisation and classification tasks separately. In more detail, AGs automatically learn to focus on target structures without additional supervision. At test time, these gates generate soft region proposals implicitly on-the-fly and highlight salient features useful for a specific task. In return, the proposed AGs improve model sensitivity and accuracy for global and dense label predictions by suppressing feature activations in irrelevant regions. In this way, the necessity of using an external organ localisation module can be eliminated while maintaining the high prediction accuracy. In addition, they do not introduce significant computational overhead and do not require a large number of model parameters as in the case of multi-model frameworks. CNN models with AGs can be trained from scratch in a standard way similar to the training of fully convolutional network (FCN) models. Similar attention mechanisms have been proposed for natural image classification [Jetley et al., 2018] and captioning [Anderson et al., 2017] to perform adaptive feature pooling, where model predictions are conditioned only on a subset of selected image regions. In this paper, we generalise this design and propose image-grid based gating that allows attention coefficients to be specific to local regions.

We demonstrate the performance of AG in real-time fetal ultrasound scan plane detection and CT pancreas segmentation. The first task is challenging due to low interpretability of the images and localising the object of interest is key to successful classification of the plane. To this end, we incorporate AGs into a variant of a VGG network, termed AG-Sononet, to demonstrate that attention mechanism can automatically localise the object of interest and improve the overall classification performance. The second task of pancreas segmentation is challenging due to low tissue contrast and large variability in organ shape and size. Moreover, we extend a standard U-Net architecture (Attention U-Net). We choose to evaluate our implementation on two commonly used benchmarks: TCIA Pancreas $CT$-$82$ [Roth et al., 2016] and multi-class abdominal $CT$-$150$. The results show that AGs consistently improve prediction accuracy across different datasets and training sizes while achieving state-of-the-art performance without requiring multiple CNN models.

Attention Gates: AGs are commonly used in classification tasks such as in the analysis of citation graphs [Veličković et al., 2017] and natural images [Jetley et al., 2018, Wang et al., 2017a]. Similarly in the context of natural language processing (NLP), such as image captioning [Anderson et al., 2017] and machine translation [Bahdanau et al., 2014, Luong et al., 2015, Shen et al., 2017, Vaswani et al., 2017], there have been several use cases of soft-attention models to efficiently use the given context information. In particular, given a sequence of text and a current word, a task is to extract a next word in a sentence generation or translation. The idea of attention mechanisms is to generate a context vector which assigns weights on the input sequence. Thus, the signal highlights the salient feature of the sequence conditioned on the current word while suppressing the irrelevant counter-parts, making the prediction more contextualised.

Initial work on attention modelling has explored salient image regions by interpreting gradient of output class scores with respect to the input image. Trainable attention, on the other hand, is enforced by design and categorised as hard- and soft-attention. Hard attention [Mnih et al., 2014], e.g. iterative region proposal and cropping, is often non-differentiable and relies on reinforcement learning for parameter updates, which makes model training more difficult. Ypsilantis and Montana used recursive hard-attention to detect anomalies in chest X-ray scans. Contrarily, soft attention is probabilistic, end-to-end differentiable, and utilises standard back-propagation without need for posterior sampling. For instance, additive soft attention is used in sentence-to-sentence translation [Bahdanau et al., 2014, Shen et al., 2017] and more recently applied to image classification [Jetley et al., 2018, Wang et al., 2017a].

In computer vision, attention mechanisms are applied to a variety of problems, including image classification [Jetley et al., 2018, Wang et al., 2017a, Zhao et al., 2017], segmentation [Ren and Zemel, 2016], action recognition [Liu et al., 2017, Pei et al., 2016, Wang et al., 2017b], image captioning [Lu et al., 2016, Xu et al., 2015], and visual question answering [Nam et al., 2016, Yang et al., 2015]. Hu et al. used channel-wise attention to highlight important feature dimensions, which was the top-performer in the ILSVRC 2017 image classification challenge. Similarly, non-local self attention was used by Wang et al. [2017b] to capture long range dependencies.

In the context of medical image analysis, attention models have been exploited for medical report generation [Zhang et al., 2017a, b] as well as joint image and text classification [Wang et al., 2018]. However, for standard medical image classification, despite often the information to be classified are extremely localised, only a handful of works use attention mechanisms [Guan et al., 2018, Pesce et al., 2017]. In these methods, either bounding box labels are available to guide the attention, or local context is extracted by a hard-attention model (i.e. region proposal followed by hard-cropping).

2D Ultrasound Scan Plane Detection: Fetal ultrasound screening is an important diagnostic protocol to detect abnormal fetal development. During screening examination, multiple anatomically standardised [NHS Screening Programmes, 2015] scan planes are used to obtain biometric measurements as well as identifying abnormalities such as lesions. Ultrasound suffers from low signal-to-noise ratio and image artefacts. As such, diagnostic accuracy and reproducibility is limited and requires a high level of expert knowledge and training. In the past, several approaches were proposed [Chen et al., 2015, Yaqub et al., 2015], however, they are computationally expensive and cannot be deployed for the real-time application. More recently, Baumgartner et al. proposed a CNN architecture called Sononet. It achieves very good performance in real-time plane detection, retrospective frame retrieval (retrieving the most relevant frame) and weakly supervised object localisation. However, it suffers from low recall value in differentiating different planar views of the cardiac chambers, which requires the method to be able to exploit the subtle differences in the local structure and it makes the problem challenging.

Pancreas Segmentation in 3D-CT Images: Early work on pancreas segmentation from abdominal CT used statistical shape models [Cerrolaza et al., 2016, Saito et al., 2016] or multi-atlas techniques [Oda et al., 2017, Wolz et al., 2013]. In particular, atlas approaches benefit from implicit shape constraints enforced by propagation of manual annotations. However, in public benchmarks such as the TCIA dataset [Roth et al., 2016], Dice similarity coefficients (DSC) for atlas-based frameworks are relatively low, ranging from $69.6\%$ to $73.9\%$ [Oda et al., 2017, Wolz et al., 2013]. A classification based framework was proposed by Zografos et al. to remove the dependency of atlas to image registration. Recently, cascaded multi-stage CNN models [Roth et al., 2017, 2018, Zhou et al., 2017] have been proposed to address the problem. Here, an initial coarse-level model (e.g. U-Net or Regression Forest) is used to obtain a ROI and then a cropped ROI is used for segmentation refinement by a second model. Similarly, combinations of 2D-FCN and recurrent neural network (RNN) models are utilised by Cai et al. to exploit dependencies between adjacent axial slices. These approaches achieve state-of-the-art performance in the TCIA benchmark ($81.2\%-82.4\%$ DSC). Without using a cascaded framework, the performance drops between $2.0$ and $4.4$ DSC points. Recently, Yu et al. proposed an iterative two-stage model that recursively updates local and global predictions, and both models are trained end-to-end. Besides standard FCNs, dense connections [Gibson et al., 2017] and sparse convolutions [Heinrich et al., 2018, Heinrich and Oktay, 2017] have been applied to the CT pancreas segmentation problem. Dense connections and sparse kernels reduce computational complexity by requiring less number of non-zero parameters.

2 Contributions

In this paper, we propose a novel soft-attention gating module that can be utilised in CNN based standard image analysis models for dense label predictions. Additionally, we explore the benefit of AGs to medical image analysis, in particular, in the context of image classification and segmentation. The contributions of this work can be summarised as follows:

We take the attention approach proposed by Jetley et al. a step further by proposing grid-based gating that allows attention gates to be more specific to local regions. This improves performance compared to gating based on a global feature vector. Moreover, our approach is not only limited to adaptive pooling [Jetley et al., 2018] but can be also used for dense predictions as in segmentation networks.

We propose one of the first use cases of soft-attention in a feed-forward CNN model applied to a medical imaging task that is end-to-end trainable. The proposed attention gates can replace hard-attention approaches used in image classification [Ypsilantis and Montana, 2017] and external organ localisation models in image segmentation frameworks [Khened et al., 2018, Oda et al., 2017, Roth et al., 2017, 2018]. This also eliminates the need for any bounding box labels and backpropagation-based saliency map generation used by Baumgartner et al. .

For classification, we apply the proposed model to real-time fetal ultrasound scan plane detection and show its superior classification performance over the baseline approach. We show that attention maps can used for fast (weakly-supervised) object localisation, demonstrating that the attended features indeed correlate with the anatomy of interest.

For segmentation, an extension to the standard U-Net model is proposed that provides increased sensitivity without the need of complicated heuristics, while not sacrificing specificity. We demonstrate that accuracy improvements when using U-Net are consistent across different imaging datasets and training sizes.

We demonstrate that the proposed attention mechanism provides fine-scale attention maps that can be visualised, with minimal computational overhead, which helps with interpretability of predictions.

Methodology

CNNs are now the state-of-the-art method for many tasks including classification , localisation and segmentation [Bai et al., 2017, Kamnitsas et al., 2017, 2018, Lee et al., 2015, Litjens et al., 2017, Long et al., 2015, Ronneberger et al., 2015, Roth et al., 2017, 2018, Xie and Tu, 2015, Zaharchuk et al., 2018]. CNNs outperform traditional approaches in medical image analysis while being an order of magnitude faster than, e.g., graph-cut and multi-atlas segmentation techniques [Wolz et al., 2013]. The success of CNNs is attributed to the fact that (I) domain specific image features are learnt using stochastic gradient descent (SGD) optimisation, (II) learnt kernels are shared across all pixels, and (III) image convolution operations exploit the structural information in medical images in an optimal fashion. However, it remains difficult to reduce false-positive predictions for small objects that show large shape variability. In such cases, in order to improve the accuracy, current frameworks [Guan et al., 2018, Khened et al., 2018, Roth et al., 2017, 2018] rely on additional preceding object localisation models to simplify the task into separate localisation and subsequent classification/segmentation steps, or guide the localisation using weak labels [Pesce et al., 2017]. Here, we demonstrate that the same objective can be achieved by integrating attention gates (AGs) in a standard CNN model. This does not require the training of multiple models and a large number of extra model parameters. In contrast to the localisation model in multi-stage CNNs, AGs progressively suppress feature responses in irrelevant background regions without the requirement to crop a ROI between networks.

2 Attention Gate Module

We now introduce Attention Gate (AG), which is a mechanism which can be incorporated in any existing CNN architecture. Let $\mathbf{x}^{l}=\{\mathbf{x}_{i}^{l}\}_{i=1}^{n}$ be the activation map of a chosen layer $l\in\{1,\dots,L\}$, where each $\mathbf{x}_{i}^{l}$ represents the pixel-wise feature vector of length $F_{l}$ (i.e. the number of channels). For each $\mathbf{x}_{i}^{l}$, AG computes coefficients $\alpha^{l}=\{\alpha_{i}^{l}\}_{=1}^{n}$, where $\alpha_{i}^{l}\in$, in order to identify salient image regions and prune feature responses to preserve only the activations relevant to the specific task as shown in Figure 1. The output of AG is $\hat{\mathbf{x}}^{l}=\{\alpha_{i}^{l}\mathbf{x}_{i}^{l}\}_{i=1}^{n}$, where each feature vector is scaled by the corresponding attention coefficient.

There are two commonly used attention types: multiplicative [Luong et al., 2015] and additive attention [Bahdanau et al., 2014]. The former is faster to compute and more memory-efficient in practice since it can be implemented as a matrix multiplication. However, additive attention is experimentally shown to be performing better for large dimensional input features [Britz et al., 2017]. For this reason, we use the latter to obtain the gating coefficient as can be seen in Figure 2, which is formulated as follows:

We note that AG parameters can be trained with the standard back-propagation updates without a need for sampling based optimisation methods as used in hard-attention [Mnih et al., 2014]. While AG does not require auxiliary loss function to optimise, we found that using deep-supervision [Lee et al., 2015] encourages the intermediate feature-maps to be semantically discriminative at each image scale. This ensures that attention units, at different scales, have an ability to influence the responses to a large range of image foreground content. We therefore prevent dense predictions from being reconstructed from small subsets of gated feature-maps.

In case of where multiple semantic classes are present in the image, one can learn multi-dimensional attention coefficients. This is inspired by the approach of Shen et al. , where multi-dimensional attention coefficients are used to learn sentence embeddings. Thus, each AG learns to focus on a subset of target structures. In case of multi-dimensional AGs, each $\alpha^{l}$ corresponds to a vector and produce $\hat{\bm{x}}^{l}=[\alpha^{l}_{(1)}\odot\bm{x}^{l},\dots,\alpha^{l}_{(m)}\odot\bm{x}^{l}]$ where $\alpha^{l}_{(k)}$ is $k$-th sub AG and $\odot$ is element-wise multiplication operation. In each sub-AG, complementary information is extracted and fused to define the output of skip connection.

2.2 Gating Signal and Grid Attention

As the gating signal $\mathbf{g}$ must encode global information from large spatial context, it is usually obtained from the coarsest scale activation map. For example in classification, one could use the activation map just before the final softmax layer. In the context of medical imaging, however, since most objects of interest are highly localised, flattening may have the disadvantage of losing important spatial context. In fact, in many cases a few max-pooling operations are sufficient to infer the global context without explicitly using the global pooling. Therefore, we propose a grid attention mechanism. The idea is to use the coarse scale feature map before any flattening is done. For example, given an input tensor size of $F_{l}\times H_{x}\times W_{x}$, after $r$ max pooling operations, the tensor size is reduced to $F_{g}\times H_{g}\times W_{g}=F_{g}\times H_{x}/(2^{r})\times W_{y}/(2^{r})$. To generate the attention map, we can either downsample or upsample the coarse grid to match the spatial resolution of $\bm{x}^{l}$. In this way, the attention mechanism has more flexibility in terms of what to focus on a regional basis. For upsampling, we chose to use bilinear upsampling. Note that the upsampling can be replaced by a learnable weight, however, we did not opt for this for the sake of simplicity. For segmentation, one can directly use the coarsest activation map as the gating signal.

2.3 Backward Pass through Attention Gates

Information extracted from coarse scale is used in gating to disambiguate irrelevant and noisy responses in input feature-maps. For instance, in the U-Net architecture, gating is performed on skip connections right before the concatenation to merge only relevant activations. Additionally, AGs filter the neuron activations during the forward pass as well as during the backward pass. Gradients originating from background regions are down weighted during the backward pass. This allows model parameters in shallower layers to be updated mostly based on spatial regions that are relevant to a given task. The update rule for convolution parameters in layer $l-1$ can be formulated as follows:

where the first gradient term on the right-hand side is scaled with $\alpha_{i}^{l}$.

3 Attention Gates for Segmentation

In this work, we build our attention-gated segmentation model on top of a standard 3D U-Net architecture. U-Nets are commonly used for image segmentation tasks because of their good performance and efficient use of GPU memory. The latter advantage is mainly linked to extraction of image features at multiple image scales. Coarse feature-maps capture contextual information and highlight the category and location of foreground objects. Feature-maps extracted at multiple scales are later merged through skip connections to combine coarse- and fine-level dense predictions as shown in Figure 3. The proposed AGs are incorporated into the standard U-Net architecture to highlight salient features that are passed through the skip connections. For AGs, we chose sigmoid activation function for normalisation: $\sigma_{2}(x)=\frac{1}{1+exp(-x)}$. While in image captioning [Anderson et al., 2017] and classification [Jetley et al., 2018] tasks, the softmax activation function is used to normalise the attention coefficients $\sigma_{2}$, however, sequential use of softmax yields sparser activations at the output. For dense prediction task, we empirically observed that sigmoid resulted in better training convergence for the AG parameters.

4 Attention Gates for Classification

For attention-gated classification model, we chose Sononet [Baumgartner et al., 2016] to be our base architecture, which is a variant of VGG network [Simonyan and Zisserman, 2014]. The difference is that Sononet can be decoupled into feature extraction module and adaptation module. In the adaptation module, the number of channels are first reduced to the number of target classes $C$. Subsequently, the spatial information is flattened via channel-wise global average pooling. Finally, a softmax operation is applied to the resulting vector and the entry with maximum activation is selected as the prediction. As the network is constrained to classify based on the reduced vector, the network is forced to extract the most salient features for each class.

Given the attended feature vectors at different scales, we highlight that the aggregation strategy is flexible and that it can be adjusted depending on the target problem. We empirically observed that a combination of deep-supervision [Lee et al., 2015] for each scale followed by fine-tuning using a new FC layer fitted on the concatenated vector gave the best performance.

The simplest is to just fit a fully connected layer on the concatenated vector as mentioned above. However, we noticed that sometimes the network abandons the fine-scale attention mechanisms as it is non-trivial to train. An alternative approach is to fit a separate fully connected (FC) layer at each scale and make separate predictions. The final prediction is then given by either weighted mean or max operations. One can also use deep-supervision [Lee et al., 2015] to force each scale to learn a useful prediction as well as when combined. We empirically observed that first training the network at each scale, then fine-tuning using a new FC layer fitted on the concatenated vector worked the best. In the experimentation, we considered the following variations:

Experiments and Results

The proposed AG model is modular and independent of application type; as such it can be easily adapted for pixel and image level classification tasks. To demonstrate its applicability to image classification and segmentation, we evaluate the proposed attention based FCN models on challenging abdominal CT multi-label segmentation and 2D ultrasound image plane classification problems. In particular, pancreas boundary delineation is a difficult task due to shape-variability and poor tissue contrast, similarly image quality and subject variability introduce challenges in 2D-US image classification. Our models are compared against the standard 3D U-Net and Sononet in terms of model prediction performance, model capacity, computation time, and memory requirements.

In this section, we present the image datasets used in classification and segmentation experiments.

For the experiments, two different CT abdominal datasets are used: (I) 150 abdominal 3D CT scans acquired from patients diagnosed with gastric cancer ($CT$-$150$). In all images, the pancreas, liver, and spleen boundaries were semi-automatically delineated by three trained researchers and manually verified by a clinician. The same dataset is used by Roth et al. to benchmark the U-Net model in pancreas segmentation. (II) The second datasethttps://wiki.cancerimagingarchive.net/display/Public/Pancreas-CT ($CT$-$82$) consists of 82 contrast enhanced 3D CT scans with pancreas manual annotations performed slice-by-slice. This dataset (NIH-TCIA) [Roth et al., 2016] is publicly available and commonly used to benchmark CT pancreas segmentation frameworks. The images from both datasets are downsampled to isotropic $2.00$ mm resolution due to the large image size and hardware memory limitations.

1.2 2D Fetal Ultrasound Image Dataset

Our dataset consisted of 2694 2D ultrasound examinations of volunteers with gestational ages between 18 and 22 weeks. The dataset contains 13 types of standard scan planes and background, complying the standard specified in the UK National Health Service (NHS) fetal anomaly screening programme (FASP) handbook [NHS Screening Programmes, 2015]. The standard scan planes are: Brain (Cb.), Brain (Tv.), Profile, Lips, Abdominal, Kidneys, Femur, Spine (Cor.), Spine (Sag.), 4CH, 3VV, RVOT, LVOT. The dataset further includes large portions of frames which contains anatomies that are not part of the scan plane, labelled as “background”. The details of the image acquisition protocol as well as how scan plane labels are obtained can be found in [Baumgartner et al., 2016]. The data was cropped to central $208\times 272$ to prevent the network from learning the surrounding annotations shown in the ultrasound scan screen.

2 Model Training and Implementation Details

The datasets used in this manuscript contain large class imbalance issue that needs to be addressed. For ultrasound dataset, due to the nature of screening process, the background label dominates the dataset. To address this, we used a weighted sampling strategy, where we matched the probability of sampling one of the foreground labels to the probability of sampling a background label. For the segmentation models, the class imbalance problem is tackled using the Sorensen-Dice loss [Drozdzal et al., 2016, Milletari et al., 2016] defined over all semantic classes. Dice loss is experimentally shown to be less sensitive to class imbalance in segmentation tasks.

For both tasks, batch-normalisation, deep-supervision [Lee et al., 2015], and standard data-augmentation techniques (affine transformations, axial flips, random crops) are used in training attention and baseline networks. Intensity values are linearly scaled to obtain a normal distribution $N(0,1)$. For classification models, we empirically found that optimising with Stochastic Gradient Descent with Nesterov momentum ($\rho=0.9$) worked the best. The initial learning rate was set to 0.1, which was subsequently reduced by a factor of 0.1 for every 100 epoch. We also used a warm-start learning rate of 0.01 for the first 5 epochs. For segmentation models, we used Adam with $\alpha=10^{-4},\beta_{1}=0.9,\beta_{2}=0.999$. The batch size for the Sononet models was set to 64. However, for the 3D-CT segmentation models, gradient updates are computed using small batch sizes of $2$ to $4$ samples. For larger segmentation networks, gradient averaging is used over multiple forward and backward passes. This is mainly because we propose a 3D-model to capture sufficient semantic context in contrast to the state-of-the-art CNN segmentation frameworks [Cai et al., 2017, Roth et al., 2018]. Gating parameters are initialised so that attention gates pass through feature vectors at all spatial locations. Moreover, we do not require multiple training stages as in hard-attention based approaches therefore simplifying the training procedure.

The architecture for AG-sononet is shown in Fig. 4. The parameters for AG-Sononet was initialised using a partially trained Sononet. We compare our models with different capacities, with the initial number of features 8, 16 and 32. For U-net and Attention U-net, the initial number of features is set to $F_{1}=8$, which is doubled after every max-pooling operation. Our implementation using PyTorch [Paszke et al., 2017] is publicly availablehttps://github.com/ozan-oktay/Attention-Gated-Networks.

3 3D-CT Abdominal Image Segmentation Results

The proposed Attention U-Net model is benchmarked against the standard U-Net [Ronneberger et al., 2015] on multi-class abdominal CT segmentation. We use $CT$-$150$ dataset for both training ($120$) and testing ($30$). The corresponding Dice scores (DSC) and surface distances (S2S) are given in Table 1. The results on pancreas predictions demonstrate that attention gates (AGs) increase recall values ($p=.005$) by improving the model’s expression power as it relies on AGs to localise foreground pixels. The difference between predictions obtained with these two models are qualitatively compared in Figure 5. In the second experiment, the same models are trained with fewer training images ($30$) to show that the performance improvement is consistent and significant for different sizes of training data ($p=.01$). For both approaches, we observe a performance drop on spleen DSC as the training size is reduced. The drop is less significant with the proposed framework. For kidney segmentation, the models achieve similar accuracy since the tissue contrast is higher.

In Table 1, we also report the number of trainable parameters for both models. We observe that by adding 8% extra capacity to the standard U-Net, the performance can be improved by 2-3% in terms of DSC. For a fair comparison, we also train higher capacity U-Net models and compare against the proposed model with smaller network size. The results shown in Table 2 demonstrate that the addition of AGs contributes more than simply increasing model capacity (uniformly) across all layers of the network ($p=.007$). Therefore, additional capacity should be used for AGs to localise tissues, in cases when AGs are used to reduce the redundancy of training multiple, individual models.

The proposed architecture is evaluated on the public TCIA CT Pancreas benchmark to compare its performance with state-of-the-art methods. Initially, the models trained on $CT$-$150$ dataset are directly applied to $CT$-$82$ dataset to observe the applicability of the two models on different datasets. The corresponding results (BFT) are given in Table 3. U-Net model outperforms traditional atlas techniques [Wolz et al., 2013] although it was trained on a disjoint dataset. Moreover, the attention model performs consistently better in pancreas segmentation across different datasets. These models are later fine-tuned (AFT) on a subset of TCIA dataset ($61$ train, $21$ test). The output nodes corresponding to spleen and kidney are excluded from the output softmax computation, and the gradient updates are computed only for the background and pancreas labels. The results in Table 3 and 4 show improved performance compared to concatenated multi-model CNN approaches [Cai et al., 2017, Roth et al., 2018, Zhou et al., 2017] due to additional training data and richer semantic information (e.g. spleen labels). Additionally, we trained the two models from scratch (SCR) with $61$ training images randomly selected from the $CT$-$82$ dataset. Similar to the results on $CT$-$150$ dataset, AGs improve the segmentation accuracy and lower the surface distances ($p=.03$) due to increased recall rate of pancreas pixels ($p=.09$).

Results from state-of-the-art CT pancreas segmentation models are summarised in Table 4 for comparison purposes. Since the models are trained on the same training dataset, this comparison gives an insight on how the attention model compares to the relevant literature. It is important to note that, post-processing (e.g. using conditional random field) is not utilised in our framework as the experiments mainly focus on quantification of performance improvement brought by AGs in an isolated setting. Similarly, residual and dense connections can be used as in [Gibson et al., 2017] in conjunction with AGs to improve the segmentation results. In that regard, our 3D Attention U-Net model performs similar to the state-of-the-art, despite the input images are downsampled to lower resolution. More importantly, our approach significantly improves the results compared to single-model based segmentation frameworks (see Table 4). We do not require multiple CNN models to localise and segment object boundaries. Lastly, we performed $5$-fold cross-validation on the $CT$-$82$ dataset using the Attention U-Net for a better comparison, which achieved $81.48\pm 6.23$ DSC for pancreas labels.

4 2D Fetal Ultrasound Image Classification Results

The dataset was split to training ($122,233$), validation ($30,553$) and testing ($38,243$) frames on subject basis. For evaluation, we used macro-averaged precision, recall, F1, overall accuracy, the number of parameters and execution speed, summarised in Table 5.

In general, AG-Sononet improves the results over Sononet at all capacity levels. In particular, AG-Sononet achieves higher precision. AG-Sononet reduces false positive examples because the gating mechanism suppresses background noise and forces the network to make the prediction based on class-specific features. As the capacity of Sononet is increased, the gap between the methods are tightened, but we note that the performance of AG-Sononet is also close to the one of Sononet with double the capacity. In Table 6, we show the class-wise F1, precision and recall values for AG-Sononet-8, where the improvement over Sononet is indicated in brackets. We see that the precision increased by around 5% for kidney, profie and spines. For the most challenging cardiac views, we see on average 3% improvement for 4CH and 3VV ($p<0.05$).

5 Attention Map Analysis

The attention coefficients of the proposed U-Net model, which are obtained from 3D-CT test images, are visualised with respect to training epochs (see Figure 6). We commonly observe that AGs initially have a uniform distribution and pass features at all spatial locations. This is gradually updated and localised towards the targeted organ boundaries. Additionally, at coarser scales AGs provide a rough outline of organs which are gradually refined at finer resolutions. Moreover, by training multiple AGs at each image scale, we observe that each AG learns to focus on a particular subset of organs.

With the proposed architecture, the localisation maps can obtained for almost no additional computational cost. In Figure 7, we show the attention maps of AG-Sononet across different subjects, together the red bounding box annotation generated using the attention maps (see Appendix for the heuristics). We see that the network consistently focuses on the object of interest, consistent with the blue ground truth annotation. We note, however, attention map outlines the discriminant region; in particular, it does not necessarily coincide with the entire object. Nevertheless, as it does not guided backpropagation for localisation (a strategy in [Baumgartner et al., 2016]), attention models are advantageous for the real-time applications.

Finally, in Figure 7, we show the attention maps of AG-Sononet-FT across different subjects, together the bounding box annotation generated using the attention maps (see Appendix for the heuristics). We see that the network consistently focuses on the object of interest, which indicates that the network indeed learnt the most important feature for each class. We note, however, attention map outlines the discriminant region; in particular, it does not necessarily coincide with the entire object. This behaviour makes sense because some part of object will appear in background label (i.e. when the ideal plane is not reached). Qualitatively, however, the bounding boxes well agree with the annotated ground truth. Most crucially, the attention map is obtained for almost no additional computational cost; In comparison, [Baumgartner et al., 2016] requires guided backpropagation for localisation, which limits the localisation speed. This highlights the advantage of attention model for the real-time applications.

Discussion

In this work, we considered soft-attention mechanism and discussed how to incorporate this idea into segmentation and scan plane detection frameworks to better exploit local structures in CT abdominal and fetal ultrasound images. In particular, we highlighted several aspects: gridded attention mechanisms, a normalisation strategy for the attention map, and aggregation strategies. We empirically observed and reported that using soft-max as the activation function tends to generate a map that is sparsely activated and is overly sensitive to local intensity changes. The latter is problematic as in ultrasound imaging, image quality is often low. In the classification setting, We found that dividing the activations by the sum of the activations helped generate attention map with larger contextual support. As demonstrated in the segmentation framework, Sigmoid function is a good alternative as it only normalises the range and allows more information to flow. However, we found that training is non-trivial due to the gradient saturation problem.

We noted that training the attention-mechanism was slightly more complex than the standard network architecture. In particular, we observed that the strategy employed to aggregate the attention maps at different scales affects both the learning of the attention mechanism itself and hence the performance. Having a loss term defined at each scale ensures that the network learns to attend at each scale. We observed that first training the network at each scale separately, followed by fine-tuning was the most stable approach to get the optimal performance.

There is a vast body of literature in machine learning exploring different gating architectures. For example, highway networks [Greff et al., 2016] make use of residual connections around the gate block to allow better gradient back-propagation and slightly softer attention mechanisms. Although our segmentation experiments with residual connections have not provided any significant performance improvement, future work will focus on this aspect to obtain a better training behaviour.

Lastly, we note that the presented quantitative comparisons between the Attention 3D-Unet and state-of-the-art 2D cascaded models might not be sufficient enough to draw a final conclusion, as the proposed approach takes advantage of rich contextual information in all spatial dimensions. On the other hand, the 2D models utilise the high resolution information present in axial CT planes without any downsampling. We think that with the advent of improved GPU computation power and memory, larger capacity 3D-CT segmentation models can be trained with larger image grids without the need for image downsampling. In this regard, future research will focus more and more on deploying 3D models, and the performance of Attention U-Net can be further enhanced by utilising fine resolution input batches without any additional heuristics.

Conclusion

In this work we proposed a novel and modular attention gate model that can be easily incorporated into existing segmentation and classification architectures. Our approach can eliminate the necessity of applying an external object localisation model by implicitly learning to highlight salient regions in input images. Moreover, in a classification setting, AGs leverage the salient information to perform task adaptive feature pooling operation.

We applied the proposed attention model to standard scan plane detection during fetal ultrasound screening and showed that it improves overall results, especially precision, with much less parameters. This was done by generating the gating signal to pinpoint local as well as global information that is useful for the classification. Similarly, experimental results on CT segmentation task demonstrate that the proposed AGs are highly beneficial for tissue/organ identification and localisation. This is particularly true for variable small size organs such as the pancreas, and similar behaviour is observed in image classification tasks.

Additionally, AGs allow one to generate fine-grained attention map that can be exploited for object localisation. We envisage that the proposed soft-attention module could support explainable deep learning, which is a vital research area for medical imaging analysis.

References

Appendix - Weakly Supervised Object Localisation (WSL)

In [Baumgartner et al., 2016], WSL was performed by exploiting the pixel-level saliency map obtained by guided-backpropagation, followed by ad-hoc procedure to extract bounding boxes. The same heuristics can be applied for the given network, however, owing to the attention map, we can device a much efficient way of performing object localisation. In particular, we generate object location by simply: (1) blur the attention maps, (2) threshold the low activations, (3) perform connected-component analysis, (4) select a component that overlaps at each scale and (5) apply bounding box around the selected components. In this heuristics, backpropagation is not required so it can be executed efficiently. We note, however, attention map outlines salient region used by the network to perform classification; in particular, it does not necessarily agree with the object of interest. This behaviour makes sense because some part of object will appear both in the class as well as background frame until the ideal plane is reached. Therefore, the quantitative result is shown in 7, however, the result is biased. We however define new metric called Relative Correctness, which is defined as 50% of maximum achievable IOU (due to bias). We see that in this metric, the method achieves very high results, indicating that it can detect relevant features of the object of interest in its proximity.

Acknowledgements

We thank the volunteers, radiographers and experts for providing manually annotated datasets, Wellcome Trust IEH Award , NVIDIA for their GPU donations, and Intel.