Meta Segmentation Network for Ultra-Resolution Medical Images

1 Introduction

With the rising up of deep learning, semantic segmentation achieves prominent progress. However, the semantic segmentation of ultra-resolution image (URI) is seldom studied, especially in the application of medical diagnosis. Many medical URIs [Tschandl et al., 2018; Aresta et al., 2019] contain more than 4 M pixels per image, and as for whole-slide image (WSI), a special type of medical URIs, its size even exceeds 40000 × 60000 (about 30 M pixels). The URIs with huge size require large computational resources, which some most popular semantic segmentation framework, such as UNet [Ronneberger et al., 2015], PSPNet [Zhao et al., 2017], and DeepLab [Chen et al., 2014;

∗Equal contribution.

†Corresponding author.

Chen et al., 2018], are hard to afford.

There are two common ways to process URIs: image downsampling and sliding patches [Altunbay et al., 2009; Chang et al., 2015]. The former resizes a large image to a suitable size, e.g., 512 × 512, then feeds it into the model, which leads to the great loss of local details, especially for WSIs. The latter crops original image into many small patches, then segments on patch-level, and finally combines the segmentation results of these patches. While these methods can effectively reduce the computational burden, the global information provided by spatial context and neighborhood dependency is almost abandoned, which makes it difficult to obtain accurate segmentation results.

The latest representative patch-based method is the AWMF-CNN [Tokunaga et al., 2019], it is a multi-branch structure to aggregate contextual information from multiple magnification patches that contain target regions and receptive fields in different resolutions and scales. As a popular strategy, multi-branch induced methods have to find a tradeoff between small inputs and multi-scales, thus, they inevitably introduce two challenge problems: Firstly, they usually need a carefully designed fusion mechanism for final result, e.g., the fusion layers consist of many stacked convolutions with relatively more channels or the auxiliary weighting net in [Tokunaga et al., 2019], resulting in a complicated and redundant structure. Fortunately, by developing a meta-learning fusion method, only a simple structure is needed in our method to ensure good results. Secondly, all the branches are independent so as to be trained from scratch separately, increasing overall parameters significantly.

In this paper, we propose a novel multi-branch framework guided by a meta-learning way for ultra-resolution medical image segmentation, namely Meta Segmentation Network (MSN). The negative loss gradient, which contains more detailed information, e.g., the target-specific difference between prediction and label, can be used as very useful information to fast generate the convolutional weights, as it has been confirmed in [Munkhdalai and Yu, 2017]. Moreover, the structures of most meta-learning frameworks are quite simple but highly effective [Finn et al., 2017; Liu et al., 2019], hence the elaborative structural design is unnecessary.

For this purpose, we develop a meta-fusion mechanism, which can elegantly solve the first challenge in multi-branch methods. Specifically, we use the negative gradients of the output layers of branches as the meta-information to train a meta-learner and directly predict the weights of the fusion layer. Our method is superior to the training way of traditional end-to-end BP that needs more iterative steps and more training samples to converge. It is also noticeable that, unlike the elaborately designed fusion structure in AWMF-CNN, i.e., many stacked convolution layers as well as a redundant weighting net, the structure of meta-fusion contains only two convolutions with a few channels and a simple meta-learner.

To avoid learning all the branches from scratch, we further introduce a quite effective weight sharing mechanism. Although the inputs for those branches have different magnifications, they are still in the same domain. So, we believe that the knowledge among these branches can be shared to some extent, i.e., weight sharing. We adopt a special memory mechanism to achieve fast knowledge adaptation between metabranch and non-meta-branches. The meta-branch represents a reference branch containing basic parameters that need to be shared with other branches. Moreover, we experimentally find that direct weight sharing leads to some knowledge gaps, as illustrated in Fig. 1. To bridge these gaps, we use the memory mechanism to store some useful memory (feature) from the meta-branch, then make a memory transformation between meta-branch and non-meta-branches to realize a fast knowledge adaption.

The contributions can be summarized as follows:

• A Meta-fusion mechanism is proposed for a multibranch deep model for URI segmentation by utilizing meta-learning. The weights of fusion layer can be fast generated through the meta-learner, leading to a simple but highly effective model.
• A novel weight sharing mechanism is introduced to realize fast knowledge adaptation, resulting in a significant reduction in the training process and overall parameters.
• The proposed MSN achieves the best performance on two challenging datasets: BACH and ISIC. Especially, our method achieves a significant performance improvement over the latest AWMF-CNN, and the overall parameters are close to that of a single branch. Thus, it is a practical segmentation method both in resource-saving and accuracy.

2 ProposedMethod

2.1 Architecture of MSN

In this section, we introduce the framework of MSN, which is illustrated in Fig.2. It mainly contains two components: the multi-branch structure (Mainbody), and the Meta Fusion Module (Meta-FM). Mainbody is an all-in-one structure which realizes multi-resolution segments. Unlike the general multi-resolution structure which requires the multiple branches with a special resolution per branch, Mainbody only uses one branch to realize the multi-resolution segmentation. Let's name the branches as the high-resolution, middleresolution and low-resolution branches according to different resolution inputs. Mainbody contains three key parts: the meta-branch, the Memory Feature Pool (Mem-FP) and the Memory Recall Module (Mem-RM). Mem-FP stores the meta-features of the meta-branch in Mainbody, while MemRM is deployed in non-meta-branch to complement the distinctive features from the non-meta branch with the metafeatures stored in Mem-FP, named Memory.

Mainbody outputs two preliminary segmentation maps, i.e., S1 and S2 in Fig. 2, which will be fused in a metalearning way to achieve a final segmentation result. In the following, we will detail the operating mechanism of Mainbody, Mem-FP, Mem-RM, and Meta-FM.

Mainbody. The architecture of Mainbody is shown in the middle part of Fig. 2. Let X1, X2 and X3 denote the three types of inputs with the same size having a different resolution corresponding to the three branches (e.g., 16×, 4×, 1×). X3 has the widest receptive field with the lowest resolution, while X1 is the opposite. As shown in Fig. 2, X1 is the upscaled patch centered in X3 signed in a green box and has the high-resolution and X2 has the middle-resolution. Mainbody can process the image patches with three resolutions and output the counterpart segmentation maps {S1, S2, S3}.

Considering the commonness and difference among the knowledge of the multi-resolution segmentation, we treat the low-resolution branch as the meta-branch which shares the weights with the middle-resolution and high-resolution branches because it contains the most information, and use Mem-FP and Mem-RM to adjust the weight learning. In detail, a low-resolution image patch X3 is fed into Mainbody, and passed through the non-gap convolution layers and gap convolution layers in the meta-branch. In gap layers, the obtained feature maps are recorded in Mem-FP. As for the high-resolution image patch X1, after fixing all the layers of the meta-branch, it is passed through Mainbody just like X3 in non-gap convolution layer, when meeting gap convolution layer, Memory is recalled from Mem-FP, and Memory as well as the feature maps output by the current gap layer are fed into the Mem-RM for adjusting the weight learning. So do the middle-resolution image patch X2. Subsequently, we fuse the two outputs S1, S2 of non-meta-branches in MetaFM. In the following, we introduce Mem-FP, Mem-RM,and Meta-FM.

Memory Feature Pool. Mem-FP acts as a storage pool. As shown in Fig. 1, the branches of X1 and X2 have a large gap with X3 (meta-branch) in some layers of CNN. When processing the low-resolution image patch X3, feature maps output by the gap layers in the meta-branch, named Metafeatures, is saved in Mem-FP which will be utilized to compensate other branches. Actually, once X3 is passed through the meta branch, the obtained Meta-features are stored.

MemoryRecallModule. In order tomake themeta-branch adapt to other branches in the weight sharing mechanism, non-meta-branches should "recall" the missing features in Mem-FP at these big gap layers. Therefore, we construct Mem-RM and embed it in the meta-branch as an auxiliary module to recall the memory of the gap layers in the metabranch. Specifically, as shown in Fig. 1, when an image patch is fed into the non-meta-branch, such as X1 or X2, it is passed forward along the meta-branch until it meets the gap layers. At the gap layers, both the pre-saved meta-feature and the counterpart feature from non-meta-branch are fed into MemRM. As shown in Fig.3, there are two input branches: the top branch inputs Meta-feature of X3 (A) and the bottom branch inputs the output feature maps of X1 or X2 (B). In order to align the feature maps between A and B, we crop the target region centered in Meta-feature and upscaled to the same size as B. After that we concatenate them and implement convolutions on them. The process is formulated as:

where ˆB is the final output of Mem-RM, f(·) is the nonlinear transformation function, and cat(·, ·), up(·) and crop(·) are the operations of concatenation, upsampling and cropping, respectively.

Meta Fusion Module. The final step of our framework is the fusion of different branches. Since the branches of X1 and X2 have already captured the memory of X3, we only need to consider the fusion of the branches of X1 and X2. One of the most common way is to use an elaboratively designed structure that might contain dozens of convolution layers to perform feature fusion, and then conduct common optimizer, e.g., SGD [Bottou, 2012], to adjust the parameters of these convolutional layers. This process may take many iterations for optimization to achieve convergence.

To pursue a simple but highly effective fusion structure, we propose using a specific target provided by an auxiliary metalearner. It is well known that, the negative gradient, which is used in SGD to determine the direction of optimization, contains detail information that measures the difference between prediction and ground truth. Draw lessons from the theory of negative gradient, we construct Meta-FM to predict the weights of these convolutional layers directly. The structure of Meta-FM is shown in Fig. 2, and Meta-FM receive the negative gradients (meta-information) of the output layers of two branches, and output the predicted weights through two fully connected layers (FC). Meta-FM can be formulated as:

where W1 and W2 are the parameters of two fusion convolutions, respectively. Note that, W1 and W2 should be reshaped to the weight matrixes because the output of FC is a vector. f(·) is a nonlinear function which contains the structure of FC-Relu-FC. σ is the gradient vector of the output layers of two branches, which is formulated as:

where L(·, ·) is the loss function, S1 is the segmentation prediction of the branch of X1, and Y1 is the ground truth. S′2 = up(crop(S2)), since the resolution of X2 is lower than X1, we crop the target region from S2 and upsample it to the size of S1. Wo1 and Wo2 are the weights of the output layers of the branch of X1 and X2, respectively. The operation v(·) reshapes the gradient matrix to a column vector.

2.2 Loss Function

We use the cross entropy as the loss function of our model, which can be formulated as:

where P is the predicted segmentation maps, Y denotes the counterpart ground truths, N is the total number of samples, and j is the j-th pixel of Pi. This loss function will be used in multiple segmentation results of our model, i.e., S1, S2, S3 as well as the final fused result S, to train our model.

2.3 Training

We adopt a 3-step training scheme to train MSN. Step 1. we train the meta-branch in Mainbody to obtain the meta parameters which will be shared with other branches. Step 2. the Mem-RM is trained for the non-meta-branches to fix knowledge gaps. Step 3. Meta-FM is learned to fuse the multiresolution segmentation results. We divide the training data into two parts: a training set and sub-training set. The training set is used for the first step, while the sub-training set is involved in the last two steps. We initialize all layers similar to [He et al., 2015].

Training Meta-branch. The low-resolution image patch X3 is fed into the branch and obtain the segmentation map S3, and then the weights of this branch are updated with the loss function L(S3, Y3) formulated in Eq. (4).

Training Mem-RM. After Step 1, we have the meta parameters and fix them. Next, we train Mem-RM on subtraining set to alleviate the influence of gap layers w.r.t the meta-branch. We input X1 or X2 to the fixed meta-branch with their specific Mem-RMs and get the counterpart segmentation results S1 and S2. Then we use the loss function L(S1, Y1) and L(S2, Y2) to update each Mem-RM which is specific to the branch of X1 or X2.

TrainingMeta-FM. Firstly, we fix the trainedmeta-branch and Mem-RM. Then we obtain the segmentation maps S1, S2 by using the same way of Step 2. After some operation as mentioned before, e.g., cropping and concatenation, finally we feed the processed S1 and S2 to the fusion layers whose weights are generated by Meta-FM, and obtain the fusion result S. Since the reshape operation on the weights vector before padding into fusion layer is differentiable, we thus tune the parameters of Meta-FM in few epochs by minimizing the loss function L(S, Y1) on sub-training set.

3 Experiments

In this section, we evaluate our method on two ultraresolution medical datasets: BACH and ISIC. We take two criteria for evaluation: the mean Intersection over Union (mIoU) and the amount of model parameters.

3.1 Datasets

BACH [Aresta et al., 2019] is composed of Hematoxylin and Eosin (H&E) stained breast histology microscopy and wholeslide images (WSI). There are 10 WSIs, with an average size of 42113 × 62625 pixels (about 3000 M pixels), included in BACH. These WSIs are stored in a multi-resolution pyramid structure, i.e., 1×, 4× and 16×. Four classes are presented in BACH: normal, benign, in situ, and invasive carcinoma. We randomly split 10 WSIs into 7, 1, 2 images for the training set, the sub-training set and the test set, respectively.

ISIC [Tschandl et al., 2018; Codella et al., 2018] is an ultra-resolution medical dataset for pigmented skin lesions, which total contains 2596 images. Its average resolution is up to 9 M, while the highest resolution is up to 6748 × 4499. The dense annotations contain two classes: lesion, normal. We randomly divide the dataset into training, sub-training and testing sets with 2077, 360 and 157 images.

3.2 Implementation Details

In our model, we use BiSeNet [Yu et al., 2018] as the backbone, i.e., the CNN structure in Mainbody which contains three branches: the high-resolution branch, the middleresolution branch and the low-resolution branch. We feed the patch with the size of 256 × 256 into MSN. We firstly crop out X1 from left to right in the image without overlapping except the last patch in each row. Then we align the center of the target area to crop out X2 and X3, if the cropping patch exceeds the boundary, then 0 is padded. As for BACH, we use the professional tool "OpenSlide" [Goode et al., 2013] to read the multi-resolution pyramid in WSI, where the resolutions of the input patches fed into the three branches are 16×, 4×, and 1×, respectively. We finally get training, subtraining and test set with 9520, 2379 and 75603 patches for each resolution. As for ISIC, we set three resolutions as 4×, 2× and 1×, where the original image is considered as the highest resolution. Then we crop out X1, X2 and X3 in these three-resolution images, respectively. The number of patches in each resolution for training, sub-training and test set is 21471, 3001 and 52166, respectively. We train the metabranch for 30 epochs, and tune the non-meta-branches as well as Meta-FM for only 10 epochs, with the batch size of 32. The optimizer Adam [Kingma and Ba, 2014] is utilized with an initial learning rate 0.0001 to update the parameters of the network. The whole model is trained in PyTorch [Ketkar, 2017] with a single 1080Ti GPU.

3.3 Comparisons with State-of-the-art Methods

Result on BACH Dataset. We compare our method with five state-of-the-art methods: UNet [Ronneberger et al., 2015], PSPNet [Zhao et al., 2017], BiSeNet [Yu et al., 2018], DeepLab-V3+ [Chen et al., 2018] and AWMF-CNN [Tokunaga et al., 2019], where the first four methods are representative general semantic segmentation frameworks, and the last one is the latest powerful multi-branch method for processing medical URIs. Because the first four methods are not the multi-branch structure, the fusion results are not available, which are denoted by "-" in Table 1. All methods have publicly provided code except AWMF-CNN, thus we reproduce it using the same backbone. We train all competitors by using Eq. (4) on the training set. For the first four methods, we train each model with a specific resolution for 30 epochs. For AWMF-CNN, we adopt two training ways: 1) Similar to the original way in AWMF-CNN, firstly we pretrain its three branches for 10 epochs, then train the fusion parts. After that, we alternately train the multi-resolution branches and fusion part for 20 epochs. 2) we only train its fusion part for 30 epochs with the fixed trained branches, we denote it as AWMF-CNN (fixed). For the convenience of expression, we use different marks in the superscript to denote the different settings: (1) "† ": use BiSeNet as backbone; (2) "‡ ": use UNet as backbone; (3) "∗ ": similar to other comparison methods.

As shown in Table 1, we observe that our method achieves the best results. Note that, with the help of our special weight sharing mechanism, we improve the result significantly for the non-meta-branches by almost 10% mIoU compared with the counterpart branches of BiSeNet and UNet, respectively (For example, the branch of X1 of MSN† achieves 46.4% mIoU, while the one of BiSeNet only gets 28.6%). Meanwhile, the result can be further boosted with our meta-fusion.

It is obviously observed that our method already obtained the best results by only training on the small sub-training set, e.g., MSN† , while other compared methods are trained on the training set. When we also train on the training set, e.g., MSN†∗ , we can get better performance. Therefore, it can be concluded that MSN is more flexible in data requirements.

The amount of parameters of MSN is almost the same as that of a single network (see MSN† vs BiSeNet and MSN‡ vs UNet), and is much smaller than AWMF-CNN, thus our model has lower complexity and is more practical.

Result on ISIC Dataset. The comparison results on ISIC are shown in Table 2. For fast implementation without loss of generality, we compare MSN with the latest three methods: DeepLab-V3+, BiSeNet and AWMF-CNN. MSN also achieves the best result with a comparable amount of parameters compared to other methods.

Visualization. Finally, we visualize the results on BACH and ISIC. Due to space limitation, we directly compare our method (BiSeNet as the backbone) with BiSeNet that trains three branches separately. The results are illustrated in Fig. 4. Obviously, with the special weight sharing mechanism, the non-meta-branches of MSN significantly outperform all the branches of BiSeNet. More importantly, with our metafusion mechanism, some details can be further refined, which makes the final result more complete.

3.4 Ablation Study

Effectiveness ofWeight SharingMechanism. The special weight sharing mechanism can not only significantly reduce the number of parameters of the multi-branch model, but also realize the knowledge transfer between the branches. Thus the results of the non-meta-branches can be promoted on the basis of the meta-branch, which has been verified in Section 3.3. To further verify its effectiveness, we design the following experiments:

Firstly, we compare four methods: (1) Meta-branch: we only use the trained meta-branch to obtain the results of all resolution inputs, without fixing the gap layers. (2) Multibranch: all branches in this structure are separately trained from scratch. (3) MSN† and (4) MSN†∗ . The backbone of all methods is BiSeNet. For a fair comparison, we also conduct our meta-fusion mechanism on the first two compared methods. The results are shown in Table 4. It is observed that our non-meta-branches outperform other methods significantly, it shows that our weight sharing mechanism can effectively eliminate the gaps between meta-branch and other branches, then improve the performance by leveraging existing knowledge. And the fusion results are based on branches' results, therefore, the improvement of branches is also conducive to the improvement of the final performance.

Secondly, we compare the convergency of the non-metabranches of MSN and the ones of Multi-branch (training on training set). The curves are illustrated in Fig. 5. It shows that our method not only performs on segmentation better, but converges faster. And when we use the training set rather than the sub-training set for training (MSN†∗ ), we can obtain better convergence performance.

Thirdly, to explore the impact of 'gap layers', we attempt to only add Mem-RM at the 'non-gap layers', i.e., the layers whose index in black font in Fig. 1. As expected, the result of this approach drops a lot, which shows that our effort to fix the gaps between meta-branch and other branches is reasonable.

Effectiveness of Meta-fusion. To verify the effectiveness of the meta-fusion mechanism, we compare three methods: (1) w/o Meta: the same fusion structure as ours without the meta-fusion mechanism, i.e., two 3 × 3 convolution layers in Fig. 2, and we train it end-to-end from scratch. (2) AWMFCNN: the fusion mechanism in AWMF-CNN, which introduces a heavy weighting net for branches weighting, then uses some convolution layers training from scratch for fusion. (3) MSN: our meta-fusion mechanism. For a fair comparison, we fix the results of three resolutions, which come from our trained three branches, and then we train the fusion part of all comparison methods on the sub-training set.

The comparison results on BACH and ISIC are shown in Table 5. It is observed that our method outperforms w/o Meta significantly. Although we have more parameters than it, it is almost negligible due to the small order of magnitude. The performance of AWMF-CNN is a little lower than ours, but its fusion structure is more complicated than ours, resulting in a sharp increase in the number of parameters.

We further illustrate the train trend of meta-fusion mechanism and w/o Meta's on BACH and ISIC in Fig. 6. As expected, our method achieves an extremely better convergency. e.g., it converges in almost 1 epoch.

4 Conclusions

In this work, we propose MSN for the effective segmentation of medical URIs. A novel meta-fusion module with a very simple but effective structure is introduced for branches fusion through a meta-learning way. Moreover, MSN achieves a lightweight multi-branch structure with the help of our particular weight sharing mechanism. The experimental results on BACH and ISIC demonstrate that our method achieves the best comprehensive performance.

Acknowledgments

This work is supported by the National Natural Science Foundation of China under Grant 61876161, Grant 61772524, Grant U1065252 and partly by the Beijing Municipal Natural Science Foundation under Grant 4182067, and partly by the Fundamental Research Funds for the Central Universities associated with Shanghai Key Laboratory of Trustworthy Computing.