Abstract
Deep learning usually requires large amounts of labeled training data, but annotating data is costly and tedious. The framework of semi-supervised learning provides the means to use both labeled data and arbitrary amounts of unlabeled data for training. Recently, semi-supervised deep learning has been intensively studied for standard CNN architectures. However, Fully Convolutional Networks (FCNs) set the state-of-the-art for many image segmentation tasks. To the best of our knowledge, there is no existing semi-supervised learning method for such FCNs yet. We lift the concept of auxiliary manifold embedding for semi-supervised learning to FCNs with the help of Random Feature Embedding. In our experiments on the challenging task of MS Lesion Segmentation, we leverage the proposed framework for the purpose of domain adaptation and report substantial improvements over the baseline model.
C. Baur and S. Albarqouni contributed equally towards this work.
You have full access to this open access chapter, Download conference paper PDF
Similar content being viewed by others
1 Introduction
In order to train deep neural networks, usually huge amounts of labeled data are necessary. In the medical field, however, labeled data is scarce as manual annotation is time-consuming and tedious. At the same time, when training models using a limited amount of labeled data, there is no guarantee that these models will generalize well on unseen data that is distributed slightly different. A prominent example in this context is Multiple Sclerosis (MS) lesion segmentation in MR images, which suffers from both a lack of ground-truth and distribution-shift across images from different devices [2]. However, vast amounts of unlabeled data can often be provided comparably easy. Semi-supervised learning provides the means to leverage both a limited amount of labeled and arbitrary amounts of unlabeled data for training deep networks [11]. In recent years, various frameworks for semi-supervised deep learning have been proposed: In 2012, Weston et al. [11] presented a framework for shallow networks based on auxiliary manifold embedding. Using an additional embedding loss function attached to hidden layers and graph adjacency among input samples, they forced the feature representations of neighboring labeled and unlabeled samples to become more similar, leading to improved generalization. In 2013, Lee et al. [4] also reported improved generalization when fine-tuning a model from predictions on unlabeled data using an entropy regularization loss. More recently, in 2015, Rasmus et al. [8] introduced the ladder network architecture for semi-supervised deep learning. And lately, Yang et al. [12] presented a framework also based on graph embeddings, with both transductive and inductive variants for shallow neural networks.
All of these methods show promising results, but they are tailored to classic CNN architectures and often only examined on small computer vision datasets. In challenging problems such as biomedical image segmentation, Fully Convolutional Networks (FCNs) [5] are preferable as they are efficient and show the ability to learn context [3, 6]. As far as we know, there is no existing semi-supervised learning method for such FCNs yet. In this paper, we lift the concept of auxiliary manifold embedding to FCNs with a, to the best of our knowledge, novel strategy called “Random Feature Embedding”. Subsequently, we successfully perform semi-supervised fine-tuning of FCNs for domain adaptation with our proposed embedding technique on the challenging task of MS lesion segmentation.
Illustration of the semi-supervised deep learning framework. (1) Labeled data is used to optimize the network for a primary objective \(\mathcal {L}_{P}\). (2) The training batches are further augmented with unlabeled samples and their prior maps, which altogether contribute to the embedding loss \(\mathcal {L}_{E}\). Note: Unlabeled data does not influence \(\mathcal {L}_{P}\).
2 Methodology
Semi-supervised learning is based on the concept of training a model, \(f(\cdot )\), using both labeled and unlabeled data \(\mathbf {X_L}\) and \(\mathbf {X_U}\), respectively. In our framework for FCNs, training data has the form of \(\mathcal {\mathbf {D}}=\{\mathbf {X}, \mathbf {Y}\}\), where \(\mathbf {X} = \{\mathbf {X_L} \cup \mathbf {X_U}\} =\{\mathbf {x_1},..., \mathbf {x_{N_L}}, \mathbf {x_{N_{L+1}}},..., \mathbf {x_{N_{L+U}}}\} \in \mathbb {R}^{H\times W \times D \times N_{L+U}}\) are D-channel images, and \(\mathbf {Y} = \{\mathbf {y_1},...,\mathbf {y_{N_L}}\} \in \mathbb {R}^{H\times W \times 1 \times N_L}\) are the corresponding label maps which are only available for the labeled data. Since we deal with FCNs, both images and label maps have the same dimensions, \(H \times W\), allowing a distinct label per pixel.
2.1 Auxiliary Manifold Embedding
In our framework, to model \(f(\cdot )\), we employ a modified version of the U-Net architecture [9] that processes images of arbitrary sizes and outputs a label map of the same size as the input (Fig. 1). We train the network to minimize the primary objective \(\mathcal {L}_{P}\), i.e. the Dice-Loss [6], for our segmentation task (see Fig. 1) from labeled data only (Fig. 1, step 1). Simultaneously, to leverage the unlabeled data, we employ an auxiliary manifold embedding loss \(\mathcal {L}_{E}\) on the latent feature representations \(h(\cdot )\) of both \(\mathbf {X_L}\) and \(\mathbf {X_U}\) to minimize the discrepancy between similar inputs in the latent space. Thereby, similarity among \(h(\cdot )\) of unlabeled data is given by prior knowledge (Fig. 1, step 2). The overall objective function can be written using Lagrangian multipliers as:
where \(\lambda _l\) is the regularization parameter associated with the embedding loss \(E_l\) at hidden layer l. Typically, this objective [11] aims at minimizing the distance among latent representations of similar \(h^l(\mathbf {x_i})\) and \(h^l(\mathbf {x_j})\) of neighboring data samples \(\mathbf {x_i}\) and \(\mathbf {x_j}\), and otherwise tries to push them apart if their distance is within a margin m:
Thereby, \(\mathbf {A} \in \mathbb {R}^{n_E \times n_E}\) is an adjacency matrix between all embedding samples \(n_E\) within a training batch, and \(d(\cdot , \cdot ) \in \mathbb {R}^1\) is an arbitrary distance metric measuring the distance between two latent representations. Unlike the typical \(\ell _2\)-norm distance employed in [11], we opt for the angular cosine distance (ACD) \(1 - \frac{h^l(\mathbf {x_i})^Th^l(\mathbf {x_j})}{||h^l(\mathbf {x_i}) ||_2 ||h^l(\mathbf {x_j}) ||_2}\) for two reasons; first, it is naturally bounded between [0, 1], hence limits the searching area for the marginal distance parameter m, and second, it shows superior performance on high-dimensional feature representations in deep architectures [7].
The definition of \(\mathbf {A}\) is left to the user and can be arbitrary. We define \(a_{ij}\) to be 1 if the respective embeddings share the same label (labeled data) or prior (unlabeled data), and set it to 0 otherwise. In our real experiments, the prior is obtained via template matching with NCC, similar to [1], acting as a noisy surrogate for the label maps.
2.2 Random Feature Embedding
In standard CNNs, the embedding loss can be directly attached to the fully connected layers, where a latent feature vector \(h(\mathbf {x_i})\) represents a single input image \(\mathbf {x_i}\). In the case of FCNs, however, an input of arbitrary size can produce a multi-channel feature map of arbitrary size. Since FCNs make predictions at the pixel level, meaningful embeddings can be obtained per pixel along the channel dimension (Fig. 2). However, sampling and comparing all \(h(\cdot )\) of all pixels in large images is computationally infeasible: a \(W' \times H' \times D \times N_b\) feature map tensor for a batch of size \(N_b\) will yield \(n_E = W' \cdot H' \cdot N_b\) and therefore \(n_E^2\) comparisons. This quickly becomes intractable to compute. Instead, we suggest to do Random Feature Embedding (RFE): we randomly sample a limited number \(n_{E}\) of pixels from the feature maps of the current batch according to some sampling strategy discussed in the next section to limit the number of comparisons. The loss remains valid since we propagate back only the gradients of selected pixels.
Sampling Strategy. Ideally, the distribution of randomly sampled embeddings should mimic the one of all embeddings (Fig. 2), while at the same time paying attention to the class distribution such that unwanted bias is not introduced to the model. Therefore, we investigate the following sampling strategies:
-
50/50 RFE: For each class in the prior the same amount of embeddings is randomly extracted to represent \(h(\cdot )\) from different classes equally. For unbalanced classes, this might lead to oversampling of one class.
-
Distribution-Aware RFE: Embeddings are sampled from the given training batch according to the ratio of negative and positive classes in the prior to preserve the actual class distribution. When classes are unbalanced and \(n_E\) is too small, this might lead to undersampling of one class.
-
80/20 RFE: As a trade-off, embeddings can be randomly sampled from a predefined ratio of 80% background and 20% foreground pixels.
3 Experiments and Results
Our experiments on MS lesion segmentation are motivated by the fact that existing automatic segmentation methods often fail to generalize well to MRI data acquired with different devices [2]. In this context, we leverage our semi-supervised learning framework for domain adaptation, i.e. we try to improve generalization of a baseline model by fine-tuning it with unlabeled data from other domains. Therefore, using an optimal prior for the adjacency matrix \(\mathbf {A}\), we first assess different sampling strategies, the impact of different numbers of embeddings as well as different distance measures for RFE. In succession, we utilize the most promising sampling strategy and distance measure together with a real prior for domain adaptation.
Dataset. Our MRI MS Lesion data is a combination of the publicly available MSSEGFootnote 1 training dataset and the non-publicly available in-house MSKRI dataset (c.f. Table 1). For all patients there are co-registered T1, T2 and FLAIR volumes as well as a corresponding ground-truth segmentation volume. Patients are grouped into domains A, B, C and D based on the scanner they have been acquired with. For training & fine-tuning, we randomly crop \(128\times 128\times 3\) sized patches and label maps around lesions from corresponding T1, T2 and FLAIR axial slices, and randomly split them into training & validation sets with a ratio of 7:3. Actual testing is performed on full slices of the testing volumes.
Implementation. Our framework is built on top of MatConvNet [10]. All models were trained in batches of 12 and the learning rate of the primary objective fixed to 10e–6. For embedding with \(\ell _2\), we set \(\lambda = 0.01\) and the margin parameter \(m = 1000\) (empirically chosen), for embedding with ACD we use \(\lambda = 1\) and \(m = 1\), such that \(h(\cdot )\) from different classes become orthogonal.
3.1 Baseline Models
In order to measure the impact of our auxiliary manifold embedding, we first train so called lower bound and upper bound baseline models in a totally supervised fashion from labeled data only. Per domain, we thereby approx. crop 6000 patches from the respective three training patients. We train the Lower Bound Model \(A_{L}\) from domain A training data for 50 epochs to obtain a model which produces decent segmentations on data from domain A, but does not generalize well to other domains. Further, we train so-called Upper Bound Models by taking \(A_L\) after 35 epochs and fine-tuning it until epoch 50 using mixed, labeled training data from domain A and \(d \in [B,C,D]\). We obtain three different models which should ideally be able to segment MS lesion volumes from domain A & B, A & C or A & D respectively.
3.2 Semi-supervised Embedding
For the purpose of semi-supervised deep learning, we now assume there is labeled data from domain A, and unlabeled/barely labeled MRI data from the other domains \(d \in [B,C,D]\). All of the models trained in the following experiments originally build upon epoch 35 of the lower bound model A\(_L\). One embedding loss is attached to the second-last convolutional layer of the network (Fig. 1). This choice is due to the fact that this particular layer produces feature maps at the same resolution as the input images, thus there is no need to downsample the prior required for embedding and thus no risk involved in losing heavily underrepresented, small lesion pixels (comprising less than 1% of all pixels).
Proof of Concept. We first assume a perfect prior, i.e. we use the label maps of all data we consider as unlabeled for embedding, and concentrate on investigating the best choice of sampling strategy, number of embeddings \(n_E\) and distance metric. We fine-tune models for target domain B using (i) 50/50, Distribution-Aware and 80/20 RFE with (ii) the \(\ell _2\)-norm and the ACD as a distance metric and (iii) different number of embeddings \(n_E \in \{20, 100, 200, 500, 1000, 2000\}\), yielding a total of 36 different models. For fine-tuning in these proof-of-concept experiments, we use only a subset of 200 images from domain A and target domain B each, rather than the full training set. Our results show that, as we ramp up \(n_E\), the distribution of randomly sampled embeddings more and more resembles the full distribution (Fig. 3(b)), which renders the random sampling generally valid. Moreover, with increasing \(n_E\), we notice consistent segmentation improvements on target domain B when using ACD as a distance metric (Fig. 3(a)). In repeated experiments, we always obtain similar results. The improvements are most pronounced with the 80/20 sampling strategy. However, the \(\ell _2\)-norm performs poorly and seems unstable as the number of embeddings \(n_E\) increases. We believe this is because the \(\ell _2\)-norm penalizes the magnitudes of the vectors being compared, whereas ACD is scale-invariant and only constrains their direction.
(a) Comparing the lower bound model, the semi-supervised models fine-tuned with NCC & ACD, and the upper bound models on the respective target domain; (b) The semi-supervised model for domain C (middle) produces much fewer FP than the lower bound model (left image) and is only slightly inferior to the respective upper bound model (right image).
Real Prior. Motivated by previous results, we now train models with a real, noisy prior using ACD and 80/20 RFE. For a target domain \(d \in [B,C,D]\), we assume that one out of the three training patients is labeled and use it to compute a prior for the other two. Before fine-tuning, the first labeled FLAIR training volume \(V_1\) of d is selected and 30 different \(5 \times 5 \times 5\) voxel sized 3D templates are randomly extracted around MS lesions. Then, we perform NCC on the remaining volumes \(V_2\) and \(V_3\) of the current domain d. The choice of 3D template matching is to ensure consistency among neighboring MR slices. For thresholding the template matching output, the same matching is applied to \(V_1\) itself and the threshold which maximizes the Dice-Coefficient between the ground-truth labels and the geometric mean of the responses on \(V_1\) is computed. Using this noisy prior, we fine-tune models for domain B, C and D using approx. 4000 training patches from volumes \(V_2\) and \(V_3\) and all labeled training patches from domain A. We set \(n_E = 100\) to obtain an overall number of embeddings similar to the proof of concept experiment (\(n_E = 2000\)), but with lower computational cost. The models show consistent improvements over the lower bound model A\(_L\) for all target domains (see Fig. 4(a)). Visual inspection (Fig. 4(b)) reveals that the semi-supervised embedding seems to dramatically reduce the number of false positives (FP). Interestingly, it detects some lesions (encircled) where the upper bound model fails, but it fails to spot very small lesions. This is probably because embeddings of smaller lesions are more unlikely to be sampled.
4 Discussion and Conclusion
In summary, we presented the concept of auxiliary manifold embedding and successfully integrated it into a semi-supervised deep learning framework for FCNs. At the heart of this framework is Random Feature Embedding, a simple method for sampling feature representations which serve as a statistic for non-linear embedding. Our experiments on MS lesion segmentation revealed that the method can improve generalization capabilities of existing models when using ACD as a distance metric. Yet, there is a lot of room for follow-up investigations. For instance, in future work, the effect of attaching multiple embedding objectives at different layers of an FCN could be investigated. In general, the proposed method should be applied to other problems as well to reveal its full potential.
References
Bermúdez-Chacón, R., Becker, C., Salzmann, M., Fua, P.: Scalable unsupervised domain adaptation for electron microscopy. In: Ourselin, S., Joskowicz, L., Sabuncu, M.R., Unal, G., Wells, W. (eds.) MICCAI 2016. LNCS, vol. 9901, pp. 326–334. Springer, Cham (2016). doi:10.1007/978-3-319-46723-8_38
García-Lorenzo, D., Francis, S., Narayanan, S., Arnold, D.L., Collins, D.L.: Review of automatic segmentation methods of multiple sclerosis white matter lesions on conventional magnetic resonance imaging. Med. Image Anal. 17(1), 1–18 (2013)
Kamnitsas, K., Ledig, C., Newcombe, V.F., Simpson, J.P., Kane, A.D., Menon, D.K., Rueckert, D., Glocker, B.: Efficient multi-scale 3D CNN with fully connected CRF for accurate brain lesion segmentation. Med. Image Anal. 36, 61–78 (2017)
Lee, D.H.: Pseudo-label: the simple and efficient semi-supervised learning method for deep neural networks. In: Workshop on Challenges in Representation Learning, ICML, vol. 3 (2013)
Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation. In: CVPR, pp. 3431–3440 (2015)
Milletari, F., Navab, N., Ahmadi, S.A.: V-net: fully convolutional neural networks for volumetric medical image segmentation. In: 2016 Fourth International Conference on 3D Vision (3DV), pp. 565–571. IEEE (2016)
Nair, V., Hinton, G.E.: Rectified linear units improve restricted Boltzmann machines. In: ICML, pp. 807–814 (2010)
Rasmus, A., Berglund, M., Honkala, M., Valpola, H., Raiko, T.: Semi-supervised learning with ladder networks. In: NIPS, pp. 3546–3554 (2015)
Ronneberger, O., Fischer, P., Brox, T.: U-net: convolutional networks for biomedical image segmentation. In: Navab, N., Hornegger, J., Wells, W.M., Frangi, A.F. (eds.) MICCAI 2015. LNCS, vol. 9351, pp. 234–241. Springer, Cham (2015). doi:10.1007/978-3-319-24574-4_28
Vedaldi, A., Lenc, K.: Matconvnet - convolutional neural networks for matlab. In: Proceeding of the ACM International Conference on Multimedia (2015)
Weston, J., Ratle, F., Mobahi, H., Collobert, R.: Deep learning via semi-supervised embedding. In: Montavon, G., Orr, G.B., Müller, K.-R. (eds.) Neural Networks: Tricks of the Trade. LNCS, vol. 7700, pp. 639–655. Springer, Heidelberg (2012). doi:10.1007/978-3-642-35289-8_34
Yang, Z., Cohen, W., Salakhudinov, R.: Revisiting semi-supervised learning with graph embeddings. In: ICML, pp. 40–48 (2016)
Acknowledgements
We thank our clinical partners, in particular Dr. med. Paul Eichinger and Dr. med. Benedikt Wiestler, from the Neuroradiology Department of Klinikum Rechts der Isar for providing us with their MRI MS Lesion dataset. Further, we want to thank Rohde & Schwarz GmbH & Co KG for funding the project.
Author information
Authors and Affiliations
Corresponding author
Editor information
Editors and Affiliations
Rights and permissions
Copyright information
© 2017 Springer International Publishing AG
About this paper
Cite this paper
Baur, C., Albarqouni, S., Navab, N. (2017). Semi-supervised Deep Learning for Fully Convolutional Networks. In: Descoteaux, M., Maier-Hein, L., Franz, A., Jannin, P., Collins, D., Duchesne, S. (eds) Medical Image Computing and Computer Assisted Intervention − MICCAI 2017. MICCAI 2017. Lecture Notes in Computer Science(), vol 10435. Springer, Cham. https://doi.org/10.1007/978-3-319-66179-7_36
Download citation
DOI: https://doi.org/10.1007/978-3-319-66179-7_36
Published:
Publisher Name: Springer, Cham
Print ISBN: 978-3-319-66178-0
Online ISBN: 978-3-319-66179-7
eBook Packages: Computer ScienceComputer Science (R0)