Network-Agnostic Knowledge Transfer for Medical Image Segmentation

1 INTRODUCTION . Deep learning often requires a sufficiently large training dataset , which is expensive to build and not easy to share between users . For example , a big challenge with semantic segmentation of medical images is the limited availability of annotated data ( Litjens et al. , 2017 ) . Due to ethical concerns and confidentiality constraints , medical datasets are not often released with the trained networks . This highlights the need for knowledge transfer between neural networks , wherein the original training dataset does not need to be accessed . On the other hand , according to the black-box metaphor in deep learning based methods , transferring knowledge is difficult between heterogeneous neural networks . To address these limitations , algorithms were proposed to reuse or share the knowledge of neural networks , such as network weight transfer ( Tan et al. , 2018 ) , knowledge distillation ( Hinton et al. , 2015 ) , federated learning ( Yang et al. , 2019 ) , and self-training ( Xie et al. , 2020b ) . Some conventional algorithms directly transfer the weights of standard large models that were trained on natural image datasets for different tasks ( Kang & Gwak , 2019 ; Motamed et al. , 2019 ; Jodeiri et al. , 2019 ; Raghu et al. , 2019 ) . For example , Iglovikov & Shvets ( 2018 ) adopted VGG11 pre-trained on ImageNet as the encoder of U-Net for 2D image segmentation . Similarly , the convolutional 3D ( Tran et al. , 2015 ) , pre-trained on natural video datasets , was used as the encoder of 3D U-Net for the 3D MR ( Magnetic Resonance ) medical image segmentation ( Zeng et al. , 2017 ) . Transferring the network weights generally requires adjustments to be made to the architecture of the receiver model , this in turn , limits the flexibility of the receiver network . Another technique that involves knowledge transfer is federated learning ( Yang et al. , 2019 ) ; it has received attention for its capability to train a large-scale model in a decentralized manner without requiring users ’ data . In general , federated learning approaches adopt the central model to capture the shared knowledge of all users by aggregating their gradients . Due to the difficulties in transferring knowledge between heterogeneous networks , federated learning often requires all devices , including both the central servers and local users , to use the same neural architecture ( Xie et al. , 2020a ) . To our best knowledge , there has been no federated learning system that uses heterogeneous networks . Knowledge distillation is the process of transferring the knowledge of a large neural network or an ensemble of neural networks ( teacher ) to a smaller network ( student ) ( Hinton et al. , 2015 ) . Given a set of trained teacher models , one feeds training data to them and uses their predictions instead of the true labels to train the student model . For effective transfer of knowledge , however , it is essential that a reasonable fraction of the training examples are observable by the student ( Li et al. , 2018 ) or the metadata at each layer is provided ( Lopes et al. , 2017 ) . Yoo et al . ( 2019 ) used a generative network to extract the knowledge of a teacher network , which generated labeled artificial images to train another network . As can be seen , Yoo et al. ’ s method had to train an additional generative network for each teacher network . Different from knowledge distillation , self-training aims to transfer knowledge to a more capable model . The self-training framework ( Scudder , 1965 ) has three main steps : train a teacher model on labeled images ; use the teacher to generate pseudo labels on unlabeled images ; and train a student model on the combination of labeled images and pseudo labeled images . Xie et al . ( 2020b ) proposed self-training with a noisy student for classification , which iterates this process a few times by treating the student as a teacher to relabel the unlabeled data and training a new student . These studies required labeled images for training the student ; moreover , they implicitly required that the pseudolabeled images should be similar in content to the original labeled images . Inspired by self-training , we propose a network-agnostic knowledge transfer algorithm for medical image segmentation . This algorithm transfers the knowledge of a teacher model to a student model by training the student on a transferal dataset whose annotations are generated by the teacher . The algorithm has the following characteristics : the transferal dataset requires no manual annotation and is independent of the teacher-training dataset ; the student does not need to inherit the weights of the teacher , and such , the knowledge transfer can be conducted between heterogeneous neural architectures ; it is straightforward to implement the algorithm with fine-tuning to solve downstream task , especially by using the downstream task dataset as the transferal dataset ; the algorithm is able to transfer the knowledge of an ensemble of models that are trained independently into one model . We conducted extensive experiments on semantic segmentation using five state-of-the-art neural networks and seven datasets . The neural networks include DeepLabv3+ ( Chen et al. , 2018 ) , UNet ( Ronneberger et al. , 2015 ) , AttU-Net ( Oktay et al. , 2018 ) , SDU-Net ( Wang et al. , 2020 ) , and Panoptic-FPN ( Kirillov et al. , 2019 ) . Out of the seven datasets , four public datasets involve breast lesion , nerve structure , skin lesion , and natural image objects , and three internal/in-house datasets involve breast lesion ( a single dataset with two splits ) and thyroid nodule . Experiments showed that the proposed algorithm performed well for knowledge transfer on semantic image segmentation . 2 ALGORITHM . The main goal of the proposed algorithm is to transfer the knowledge of one or more neural networks to an independent one , without access to the original training datasets . This section presents the proposed knowledge transfer algorithm for semantic segmentation in Algorithm 1 and its application with fine-tuning for a downstream task in Algorithm 2 ( A.4 ) . The knowledge transfer algorithm first employs one or more teacher models to generate pseudo masks for the transferal dataset and then trains the student model on the pseudo-annotated transferal dataset . For an image x in the transferal dataset D , we get the pseudo mask by the weighted average of the teacher models ’ outputs : y = ∑ Ti∈T wi · Ti ( x ) , where wi is the weight for the model Ti . The output of a model Ti ( x ) is either a soft mask ( pixel value varies from 0 to 1 ) or a binary mask ( pixel value is 0 or 1 ) . Since the ensemble of teacher models is not our primary focus , we simply set the weights equally for all teacher models . We adopt two constraints to exclude images without salient targets , as shown in Figure 1 . The first constraint is on the target pixel number in the pseudo mask , where target pixel indicates the pixel with a value above 0.5 . We exclude the image ( of size 384 ×384 pixels ) if the target pixel number is less than a threshold of 256 . The second constraint is on the gray-level entropy of the pseudo Algorithm 1 Network-Agnostic Knowledge Transfer for Semantic Segmentation Require : Teacher set T with trained models , randomly initialized student model S Require : Transferal dataset D Require : Target pixel number threshold α , gray-level entropy threshold β 1 : for each datapoint x in D do 2 : Obtain pseudo mask y = ∑ Ti∈T wi · Ti ( x ) 3 : N ← number of target pixels ( value above 0.5 ) in y 4 : E ← gray-level entropy of y 5 : if N < α or E > β then 6 : Exclude x from D 7 : else 8 : Update x as ( x , y ) 9 : end if 10 : end for 11 : Train student model S on pseudo-annotated D 12 : return Trained S mask . A high entropy value implies the target and the background have no obvious difference and vice versa . We exclude images with a gray-level entropy higher than a threshold of 2.5 . We employ the second constraint only in the case that the teacher model outputs a soft mask rather than a binary mask . In this algorithm , the teacher models and student model can be independent of each other and do not rely on each other during the training or testing phase . The transferal dataset is independent of the teacher-training datasets , which are not observable for the student model . The student model gains the knowledge of the teacher by the transferal dataset . As a result , the student model is trained to work well on datasets similar to the teacher-training dataset , while it is not aimed to predict the ground truth of the transferal dataset . In the case that a small dataset Dt with ground truth is available for a downstream segmentation task ( target task ) , it is straightforward to further fine-tune the student model that has been trained on pseudo-annotated transferal dataset , as Algorithm 2 shows . The transferal dataset D is independent of the target task , while it can also be a non-annotated dataset from the target task . 3 RELATED NEURAL NETWORKS . We posit that if the teacher model is ineffective , a more capable student model can easily achieve similar performance ; if the student model , however , is ineffective , the cause can not be easily identified and the reason can be attributed to the knowledge transfer algorithm and the inherent limitations of the student model . We experimented with five neural networks , all of which have different architectures and have been shown to provide the best outcome on various segmentation tasks . DeepLabv3+ : DeepLabv3+ ( Chen et al. , 2018 ) is one of Google ’ s latest and best performing semantic segmentation models . It combines the advantages of spatial pyramid pooling with an encoderdecoder structure . Spatial pyramid pooling captures rich contextual information while the decoder path is able to gradually recover object boundaries . U-Net : U-Net ( Ronneberger et al. , 2015 ) has been the most widely used network for medical image segmentation . U-Net consists of an encoder-decoder structure with the encoder path extracting rich semantic information , and the decoder path recovering resolution and enabling contextual information . AttU-Net : AttU-Net ( Oktay et al. , 2018 ) is a modification of U-Net , where attention gates are used to control the skip connections from the encoder to the decoder . Attention gates make the model focus on target structures by highlighting important features and suppressing irrelevant features . SDU-Net : SDU-Net ( Wang et al. , 2020 ) utilizes parallel atrous convolutions with different dilation rates in each layer , which is effective in capturing features in both small and large receptive fields . SDU-Net has demonstrated better performance using merely 40 percent of the parameters in U-Net . Panoptic-FPN : Panoptic-FPN ( Kirillov et al. , 2019 ) merges semantic segmentation and object detection , such that each pixel is given a class as in semantic segmentation and each object is given a unique ID as in object detection . Panoptic-FPN provides a more rich and complete segmentation .

Summary: The paper proposes to use student-teacher training as a way of knowledge transfer between neural networks with different architectures without access to the source data. Instead the authors propose to use a separate dataset to transfer the knowledge of the teacher network and a potential different dataset for fine-tuning. The paper evaluates their method with various segmentation architectures by pretraining a DeepLab v3+ on an internal breast lesion dataset and testing transfer and fine-tuning using different medical datasets. The authors find that knowledge transfer performs similar to regular transfer learning in most combinations of datasets.