Small Lesion Segmentation in Brain MRIs with Subpixel Embedding

A stroke occurs when a lack of blood flow prevents brain tissue from receiving adequate oxygen and nutrients. This condition affects over 795,000 people annually [31]. The severity of the outcome, including disability and paralysis, depends on the location and intensity of the stroke, as well as the time of diagnosis [2, 33]. Preserving cognitive and motor functions, therefore, hinges on localizing stroke lesions quickly and precisely. However, doing so manually requires expert knowledge, is time consuming, and is ultimately subjective [13, 15].

We focus on automatically segmenting ischemic stroke lesions, which account for 87% of all strokes [31], from T1-weighted anatomical magnetic resonance imaging (MRI) brain scans. These lesions are characterized by high variability in location, shape, and size – the latter two are problematic for conventional convolutional neural networks (CNNs) where precision of irregularly shaped lesion boundaries and recall of small lesions are critical measures of success. Due to aggressive spatial downsampling (i.e. max pooling, strided convolutions) customary in CNNs, details of local structures are lost in the process. Yet, the spatial downsampling is necessary for obtaining a global representation of the input while using fixed-size filters with limited receptive fields. The outcome of which are segmentations with ambiguous boundaries between lesion and normal tissues and missed lesions that occupy small number of voxels in the MRIs.

We propose to retain small local structures by learning an embedding that maps the input to high dimensional feature maps of twice the input resolution. Unlike the typical CNN, we do not perform lossy downsampling on this representation; hence, the embedding preserves local structures, but lacks global context. When combined with the standard encoder-decoder e.g. U-Net [22], the embedding complements the encoder-decoder by supplying the decoder with fine-grained detail information to guide segmentation. Our network also outputs at twice the resolution of the input, representing each element in the input with a $2\times 2$ neighborhood of predictions. The final output is obtained by combining the four predictions (akin to an ensemble) as a weighted sum where the contribution of each prediction is learned from the data. Our design not only enables the network to produce robust segmentations but also localize small lesions (Fig. 3).

Our contributions include (i) an embedding function that preserves fine-grained details of the input by mapping it to larger spatial dimensions, (ii) a neural network architecture that leverages the complementary strengths of the proposed embedding and an encoder-decoder to produce predictions at twice the input resolution, and (iii) a learnable downsampler that combines local predictions in an ensemble fashion to yield robust segmentations at the input resolution. Our approach improves the baseline U-Net architecture by $\approx 11.7\%$ and achieves the state of the art on the ATLAS [13, 14] dataset with lower computational burden than the best competing method.

Lesion Segmentation. Early works [6] aggregated classification results for the center pixel of patches sampled from an image. However, [6] lacked global context, so [24] addressed this with multi-stage cascaded hierarchical models. More recent works build upon the U-Net [22], a 2D fully-convolutional network with skip connections and up-convolutions. For example, [16] used a Dual Path Network [5] encoder while [29] leveraged dilated convolutions to inexpensively increase receptive fields. Furthermore, [1] fused the U-net with other high-performing modules, the BConvLSTM [27] and the SENet [10], and [21] introduced X-blocks to the U-Net, leveraging depthwise separable convolutions to reduce computational load. [34] used skip connections between successive encoder resolutions to prevent the loss of features and ConvLSTM [26] modules to maintain localization.

Recent works also leveraged 3D architectural backbones to improve localization. [35] performed 3D convolutions on a subsection of the scan and fused the results with 2D convolutions. [11] proposed an attention gate to combine 2D segmentations along the axial, sagittal, coronal planes into a 3D volume. However, these works use significantly larger memory footprints and 3D convolutions are computationally expensive – limiting the models’ practicality. We note that while conventional architectures perform well globally (i.e. recovering the coarse shape of lesions) they struggle to segment small lesions that blend into the background.

Super-Resolution. There is an abundance of works in natural images super-resolution [7, 8, 25, 28, 32] and a growing number in medical imaging. [23] proposed to map MRI images from low to high-resolution with an overcomplete dictionary. [19] leveraged SRCNN [7] for super-resolving 2D MRI images and fused them to obtain a 3D volume. [20] handled arbitrary scaling factors with a 3D architecture for multi-modal 3D data. However, these works require low and high-resolution image pairs for training and are limited to the super-resolution task while our method does not rely on a larger resolution ground truth. More recently, [30] introduced Kite-Net, an upsampling encoder that outputs a latent at $8\times$ resolution followed by a max-pooling decoder to downsample back to the original resolution. Kite-Net is used in parallel with a U-Net for lesion segmentation. Our approach draws inspiration from super resolution and latent over-representations as methods to retain local structure that are often lost in spatial downsampling. However, unlike [30], we avoid downsampling the latent with pooling (which discards information), and instead employ lossless space-to-depth and depth-to-space [25] operations to retain fine-grained details. Furthermore, we propose to learn a subpixel embedding at $2\times$ the original resolution to guide our segmentation, which uses a much smaller memory footprint than [30]. We show that our approach can capture small lesions that are missed by [21, 22, 30, 34, 35].

We propose a method to partition a 3D MRI volume $X\in\mathbb{R}^{C\times H\times W}$ into lesion (positive, 1) and normal (negative, 0) classes. Our method takes, as input, a 3D slice of $c$ consecutive 2D images $x\in\mathbb{R}^{c\times H\times W}$ ($c$ is an odd integer) from $X$ and predicts the binary segmentation for the image $\bar{x}\in\mathbb{R}^{1\times H\times W}$, the $\frac{c+1}{2}$-th image of $x$. In other words, $x$ is a sliding window of $c$ images centered at a target image $\bar{x}$. To avoid sampling out of bounds, we perform mean padding of size $\frac{c-1}{2}\times H\times W$ on both sides of $X$ before sampling $x$ (see Sec. 1 of Supp. Mat. for more details). To segment a single image $\bar{x}$, we propose to learn a deep neural network $f_{\omega}$, parameterized by $\omega$, where $f:\mathbb{R}^{c\times H\times W}\mapsto[0,1]^{1\times H\times W}$ is a function that takes the 3D slice $x$ as an input and outputs the sigmoid response $f_{\omega}(x)$, a confidence map corresponding to lesions in $\bar{x}$. To obtain the binary segmentation of $X$, we aggregate our predictions by running $f_{\omega}$ for all $x$ and setting any response greater than a threshold of 0.5 to the lesion class. We note that our method can be extended to multi-class segmentation simply by expanding our output to $[0,1]^{K\times H\times W}$ for $K$ classes, and choosing the class with highest response, i.e. $\arg\max f_{\omega}(\cdot)$, to yield the segmentation.

We demonstrate our method on the Anatomical Tracings of Lesion After Stroke (ATLAS) MRI dataset [13, 14], using the metrics defined in Table 1. ATLAS contains 304 T1-weighted MRI scans of stroke patients with corresponding lesion annotations. The data is collected from 11 research sites worldwide, manually annotated, and post-processed (i.e. smoothing and defacing for privacy), leaving 239 patient scans with 189 2D images ($197\times 233$ resolution) each. Since no official data split is provided by [13], previous works [21, 34, 35] evaluated their methods using $k$-fold cross validation and randomly sampled data splits. However, the value of $k$ and samples within each split varied across works. Due to the lack of consistency, the reported results are not directly comparable. Thus, we propose a training (212 patients) and a held-out testing (27 patients) split to standardize the evaluation protocol for more rigorous comparisons. We provide quantitative comparisons against [21, 22, 30, 34, 35] on the proposed training and testing split in Table 2. We also show qualitative (Fig. 3) and quantitative (Table 3) comparisons on segmenting small lesions using a subset of test set: 490 images containing only lesions smaller than 100 pixels (0.2% of the image). All reported results for previous works are obtained using their training procedures and open-sourced code. We also provide details on our training and testing split in Sec. 2 of Supp. Mat. and further $k$-fold cross validation comparisons in Sec. 3. of Supp. Mat.

Implementation details. Our model is implemented in PyTorch [18] and optimized using Adam [12]. We used an initial learning rate of $3\times 10^{-4}$, decreased it to $1\times 10^{-4}$ after 400 epochs, and to $5\times 10^{-5}$ after 1400 epochs for a total of 1600 epochs. We choose $c=5$ for the number images in the input $x$. During training, $\bar{x}$ and its corresponding $x$ are randomly sampled from $X$. Training takes $\approx 8$ hours on an Nvidia GTX 1080 GPU, and inference takes $\approx 11$ ms per 2D image. For data augmentation, we randomly perform (i) horizontal and vertical flips, (ii) rotation between -30^∘ and 30^∘, and (iii) add zero-mean Gaussian noise with standard deviation of $1\times 10^{-2}$ to training samples. We perform augmentation with a probability of 1 for 1400 epochs and decrease it to 0.5 thereafter so training samples will be closer to the true distribution of the dataset.

ATLAS test set. Table 2 shows that our approach outperforms competing methods [21, 22, 30, 34, 35] across all evaluation metrics. Specifically, we beat the best performing method X-Net [21] by an average of $\approx 10.4$% with a 72.3% reduction in training memory and a 57.5% runtime reduction during inference. Our approach also uses a smaller memory footprint, containing only $\approx 5.3$M parameters, compred to $\approx 15$M in [21]. Another key comparison is with KiU-Net, which learns a representation at $8\times$ the original input spatial resolution. Unlike us, KiU-Net [30] uses max pooling layers, which discards information, to reduce the size of their high resolution representation to the original ($1\times$) resolution. Whereas, we maintain the $2\times$ resolution of our embedding until the output layer, which yields subpixel predictions that are aggregated by our learnable downsampler to the $1\times$ resolution. Admittedly, this comes at the cost of runtime – our method requires 2.145s per patient and KiU-Net [30] requires 1.05s. However, we outperform [30] by an average of $33.7$% across all metrics and reduce test time memory by half. We show qualitative comparisons in row 1 of Fig. 3 where the segmentation produced by our approach better captures irregularly shaped lesions than those predicted by competing methods.

Small lesion segmentation. Here, we consider the task of segmenting lesions that occupy fewer than 100 pixels or 0.2% of the image. Due to the challenging nature of the task, we observe an expected drop in performance across all methods (trained on the proposed split) when segmenting small lesions (Table 3), as compared to doing so for all lesion sizes (Table 2). However, we still outperform all competing methods – by even larger margins than on the full test set. This shows that competing methods, while able to localize large and medium sized lesions, actually perform poorly on small lesions. With the exception of precision, where we tie for second with X-Net [21], we rank first in all other metrics. We note that while CLCI-Net [34] has the highest precision, it also achieved second lowest recall, meaning that it misses many small lesions, which is critical to clinical prognosis and thus patient recovery. This is also reflected in DSC and IOU where we outperform [34] by 72% and 51%, respectively. Qualitatively, rows 2 and 3 in Fig. 3 show that our method successfully localized small lesions that [21, 22, 30, 34, 35] missed entirely.

Ablation studies. Table 4 shows the effect of each of our contributions to architectural design. Row 1 shows that our baseline, a U-Net [22] based encoder-decoder, performs significantly worse by $11.7\%$ than the proposed approach because it lacks fine local details from SPG and uses bilinear downsampling instead of a learnable downsampler (LD). Including LD alone, but not SPG (row 2) provides no improvement as the network only learns a coarse global representation, but is still missing details lost during spatial downsampling.

In row 3, we show that solely increasing parameters (i.e. adding ResNet blocks [9] to the baseline) brings no improvement, which suggests that the performance boost is not a result of a larger network. In fact, SPG and the learnable downsampler marginally increase the model size as they only combine for 27K parameters. Rows 4 and 5 show that using hand-crafted $2\times$ resolution images (from bilinear, nearest neighbor upsampling) does provide some gain. In these experiments, we replace SPG with different interpolation methods and the higher resolution images undergo $3\times 3$ convolutions before being passed as skip connections to the decoder. However, because the $2\times$ representation is not learned, as it is with SPG, the result is still $\approx 6\%$ worse than our full model. Our learnable downsampler (LD) contributes $4.4\%$ to our performance (row 6) as removing LD and replacing it with bilinear interpolation smooths lesion boundaries, resulting in loss of details. Finally, we justify the use of cross entropy for our loss function; row 7 demonstrates that minimizing a soft Dice loss, as in [34, 35], results in worse performance. The best performance is achieved with our full model using SPG and LD, and minimizing cross entropy (row 8).

We propose SPiN, a network architecture that learns a spatially increasing embedding that, when used as guidance for an encoder-decoder network, helps ensure that small structures are not lost through spatial downsampling in the encoder. We note that our embedding does not create extra spatial information (data processing inequality), but serves as a means for better characterization of local regions for the downstream segmentation task. While we outperform existing works and improve on small lesion segmentation, we do cost more memory and compute than the baseline. However, the extra cost is within reason (1 GB of memory for training and $\approx$ 0.7s in runtime) and does not limit applicability. Despite the improved segmentation performance, we would like to address that there is still room for improvement, especially with small lesions. The highest recall of 0.347 (Table 3) achieved by our model is admittedly low compared to recall metrics on the full dataset, implying that many small lesions still pass undetected. We note that this is one of the first works to study subpixel architectures in lesion segmentation, and we hope that our optimistic results will motivate further exploration in this direction.

Acknowledgements. This work was supported by NIH-NEI 5R01EY029689 and R01EY030595, ARO W911NF-17-1-0304, NRF-2017R1A2B4006023 and IITP-2021-0-01341, AI Graduate School (CAU) in Korea.