ATSO: Asynchronous Teacher-Student Optimization for Semi-Supervised Medical Image Segmentation

We study the problem of medical image segmentation in abdominal CT scans, and make use of 2D-based methods to process volumetric data. Each volume is partitioned into a set of slices along three directions, which are named the coronal, sagittal, and axial views in the context of medical images. Each slice is then sent into a 2D network for segmentation, and the output is stacked into a 3D volume as the final prediction. From this point of view, the key problem is a standard semantic segmentation task and the evaluation protocol is called the Dice-Sørensen coefficient (DSC), with the formula being ${\mathrm{DSC}}={2\times\left|\mathcal{Y}\cap\mathcal{Z}\right|/\left(\left|\mathcal{Y}\right|+\left|\mathcal{Z}\right|\right)}$, where $\mathcal{Y}$ and $\mathcal{Z}$ are the set of target voxels in the prediction and ground-truth, respectively. We choose the 2D segmentation baseline to be RSTN [33], an open-sourced, coarse-to-fine framework that achieved good performance in a few medical image segmentation tasks. Let us denote the network as ${\mathbf{y}}={\mathbf{f}\!\left(\mathbf{x};\bm{\theta}\right)}$, in which $\mathbf{x}$ and $\mathbf{y}$ denote the 2D input and output of the deep network $\mathbf{f}\!\left(\cdot\right)$ parameterized by weights $\bm{\theta}$.

Although RSTN as well as other algorithms [8, 29, 17] achieved satisfying performance under full supervision, the performance can degenerate dramatically in the scenarios of limited training data. For example, a fully-supervised RSTN ($60$ training samples) achieves an average accuracy of over $84\%$ on the NIH dataset, however, if only $10\%$ of training data ($6$ training samples) are preserved, the accuracy quickly drops to nearly $70\%$. This raises an important setting named semi-supervised learning, in which a large portion of training data do not have labels but we need to learn as much knowledge as possible from them. As a formal definition, the training set $\mathcal{T}$ is partitioned into two parts, namely, the supervised (labeled) set $\mathcal{S}$ and the reference (unlabeled) set $\mathcal{R}$. Most often, we have ${\left|\mathcal{S}\right|}\ll{\left|\mathcal{R}\right|}$. Also, there is a testing set, $\mathcal{E}$, which is invisible to the algorithm.

A simple and effective pipeline for semi-supervised learning is named self-learning. It starts with an initial model, denoted by $\mathbb{M}_{0}$, which is trained under supervised learning on $\mathcal{S}$. $\mathbb{M}_{0}$ gets updated for a total of $T$ rounds. In the $t$-th (${t}={1,2,\ldots,T}$ round (a.k.a., generation)), the reference subset, $\mathcal{R}$, is sent into the old model $\mathbb{M}_{t-1}$ (often referred to as the teacher model), and the prediction is named the pseudo label in the current round. The training process of the student model, $\mathbb{M}_{t}$, then follows a regular supervised learning procedure on both $\mathcal{T}$ and $\mathcal{R}$, with the supervision on $\mathcal{R}$ coming from the pseudo labels. The trained student model of the current round is used as the teacher model of the next round and the iteration continues till the end.

We start with training the model using $10\%$ of data labeled and the remaining $90\%$ unlabeled on the NIH dataset. The detailed settings are elaborated in Section 4.1. After $2$ generations, the self-learning pipeline achieves an accuracy of $78.98\%$ on the test set, claiming a significant improvement over $72.62\%$ that is reported by the model trained only on the $10\%$ labeled data. However, this number is still far below $85.08\%$ when the full labeled training set is available. That is to say, self-learning indeed extracts knowledge from the reference set, but the efficiency is far below satisfaction.

To understand the reason for this, we investigate the accuracy of the reference set, which is expected to grow with the learning procedure. However, we find that the accuracy quickly arrives at a plateau. After the $0$th (the initial training stage with only labeled data), $1$st, $2$nd, and $3$rd generations, the reference set accuracy is $71.41\%$, $74.54\%$, $75.44\%$, and $74.72\%$, respectively. Compared to the scenario that $100\%$ training data have been labeled, the accuracy on the same subset of training data is $86.70\%$. In other words, the training procedure has entered a ‘trap’ that the reference set stops at a low accuracy (what is worse, the accuracy starts to drop in the $3$rd generation), but the algorithm does not know because the ground-truth is missing.

From the viewpoint of optimization, this phenomenon can be explained as a local optimum of the self-learning system. Let $\mathbf{x}_{n}$ be a training sample in the reference set, $\mathcal{R}$, $\mathbf{y}_{n}^{\star}$ is the ground-truth label, and $\mathbf{y}_{n}$ is the predicted output by the teacher model. We assume that ${\mathbf{y}_{n}}={\mathbf{y}_{n}^{\star}+\bm{\epsilon}_{n}}$ where $\bm{\epsilon}_{n}$ is the prediction error. When $\mathcal{R}$ is labeled, i.e., $\mathbf{y}_{n}^{\star}$ is known, $\bm{\epsilon}_{n}$ follows a zero-mean distribution because the optimization goal is to minimize $\left|\bm{\epsilon}_{n}\right|$ on the training set. However, in the scenario of self-learning, $\mathbf{y}_{n}^{\star}$ remains unknown and thus $\bm{\epsilon}_{n}$ may follow a non-zero-mean distribution. Since $\mathbf{y}_{n}$ is used as the pseudo label, such noise can persist across in the student model. What makes things worse, each generation of the teacher-student optimization can introduce new noise and the noise can accumulate on the reference set. We use the name of lazy learning to refer to the behavior that the student model is unable to identify and eliminate the noise of the teacher model.

To alleviate lazy learning, the key is to escape from the optimization trap, i.e., breaking up the chain that the noise, $\bm{\epsilon}_{n}$, is propagated from the teacher to the student. This is done from two aspects, both of which focus on weakening the relationship between the teacher and student models.

First, we prevent using the setting of continual learning that initializes the student model from the teacher model and performs fine-tuning, but starts training from the very beginning. To reveal the advantage, we notice that the optimization goal is to minimize $\left|\mathbf{y}_{n}^{\star}-\mathbf{f}\!\left(\mathbf{x}_{n};\bm{\theta}\right)\right|$, but since $\mathbf{y}_{n}^{\star}$ is unknown, we use $\left|\mathbf{y}_{n}-\mathbf{f}\!\left(\mathbf{x}_{n};\bm{\theta}\right)\right|$ to estimate it. The error of this estimation is bounded by $\left|\bm{\epsilon}_{n}\right|$:

That being said, provided that $\mathbf{y}_{n}^{\star}$ and $\mathbf{y}_{n}$ are fixed. Hence, when $\bm{\theta}$ is initialized independently from the teacher model, $\left|\mathbf{y}_{n}-\mathbf{f}\!\left(\mathbf{x}_{n};\bm{\theta}\right)\right|$ is increased and thus the fraction that $\left|\bm{\epsilon}_{n}\right|$ occupies in the loss function is smaller. This effectively prevents the training procedure from being dominated by $\left|\bm{\epsilon}_{n}\right|$.

Second, we prevent using the same set of reference data continuously, i.e., always using the pseudo labels generated by a teacher model to supervise its direct student. To verify this assumption, we first notice that in the aforementioned self-learning procedure, the accuracy gained on the test set is much higher than that on the reference set, e.g., $6.36\%$ test accuracy gain vs. $4.03\%$ reference accuracy gain after $2$ self-learning generations. To make things clearer, we partition the reference set into two parts, and perform self-learning on different combinations of reference data. Results of two hard examples are shown in Figure 1. Interestingly, the segmentation accuracy is significantly improved when the example is not contained in the reference set. This aligns with the observation that the test accuracy is higher than the reference accuracy. Back to the optimization perspective, the noise on the reference set does not transfer to the data that the current generation does not use for self-learning. This inspires us to partition the reference set into two subsets, denoted by ${\mathcal{R}}={\mathcal{R}^{\left(1\right)}+\mathcal{R}^{\left(2\right)}}$. In each generation, we generate the pseudo labels on $\mathcal{R}^{\left(1\right)}$ using the model that was just self-trained on $\mathcal{R}^{\left(2\right)}$, and vice versa. After the last generation asynchronous update of both subsets, we combine the pseudo labels of both subsets into the complete reference set, based on which the final model is trained.

Integrating the above two aspects obtains the asynchronous teacher-student optimization (ATSO) algorithm that is described in Algorithm 1. Of course, there exist other solutions that can alleviate lazy learning, yet our solution is simple and effective. Moreover, ATSO also has the option of dividing the reference set into more subsets, but this can slow down the training procedure. In practice, we find that using two folds of reference data performs well for semi-supervised learning.

We evaluate ATSO on the NIH dataset [21] for pancreas segmentation. It contains $82$ normal CT scans, each of which is a 3D volume of $512\times 512\times L$ voxels, where $L$ is the length of the long axis. We follow the prior work [37, 33] to partition each dataset into four folds and use the first three folds ($62$ cases) as the training data. For semi-supervised learning, we follow the prior work [36, 28] to use a small portion ($10\%$ or $20\%$) of labeled training data and leave the remaining part to the reference set. We report the average DSC over all test cases.

The configuration of the deep networks follows that in the original RSTN paper [33]. In the training stage, we optimize three individual networks for segmentation along with the coronal, sagittal and axial views, respectively. In the inference stage, either on the reference set or the test set, predictions from these three views are fused into the final segmentation by majority voting. Please refer to the original paper for further technical details.

We first investigate semi-supervised segmentation on the NIH dataset with $10\%$ of the training set ($6$ cases) labeled. The naive baseline, by only using the labeled data for training, reports a $72.62\%$ test accuracy which is far lower than the upper-bound, $85.04\%$, when the labels of all training data are available. In what follows, we gradually add the key components of ATSO into the baseline and show how these components push the training procedure towards higher accuracy.

We first compare the options with and without continual learning, which we refer to as the self-learning baseline and synchronous teacher-student optimization (STSO), respectively. The former uses the last snapshot of the teacher model to initialize the student model, while the latter trains the student model from scratch. To save computational costs, we define the scratch model as the snapshot after the first two sections of RSTN, after which a randomly initialized transformation module is inserted that injects sufficient perturbations to break up the effect of continual learning.

Results are summarized in Table 1. The difference between the self-learning baseline and STSO is significant. After two generations, the self-learning baseline achieves a $78.98\%$ accuracy on the testing set and $75.44\%$ on the reference set. Starting from the third generation, these numbers start to drop, demonstrating that lazy learning has obstructed the model from obtaining useful information. The best accuracy of STSO, $79.67\%$, is obtained after $4$ generations. That is to say, simply switching off continual learning leads to a non-trivial $0.69\%$ improvement and, as we shall see in Table 2, surpasses all the competitors.

Next, we study the difference between STSO and ATSO. Results are summarized in Table 1. ATSO improves segmentation accuracy on both the reference and test sets. Detailed results on the reference subsets during iteration are provided in Appendix A. Interestingly, ATSO enjoys faster growth in both numbers: after only two generations, the test accuracy has increased to over $80\%$, claiming a nearly $3\%$ advantage over the corresponding number of STSO. After five generations, ATSO still enjoys a $2\%$ advantage over STSO. That being said, ATSO has a broad range of applications in the scenario of limited computational resource for model training.

More importantly, these results indicate that the lazy learning phenomenon indeed persists during the entire training process. Therefore, by not generating pseudo labels on the reference set that was just used, the algorithm can escape from the optimization trap. We notice that compared to STSO that always refers to the entire reference set, ATSO trains each model using only half of the reference data, which has a natural deficit but still achieves higher accuracy. We expect ATSO to have a larger advantage when the amount of unlabeled data becomes larger.

In Table 2, we compare ATSO against state-of-the-art approaches, and show that ATSO outperforms all of them. In particular, ASTO surpasses [28] by more than $2.5\%$ in both scenarios that $10\%$ and $20\%$ labeled data have been used. Note that [28] is a recently published method which involved uncertainty in multi-view learning – in comparison, our solution is easier and more effective. In addition, being simple and easily implemented, ATSO can be combined with other training strategies, e.g., adversarial training [18] or uncertainty evaluation [28], towards better performance.

Figure 2 shows some typical examples of how segmentation errors are fixed with semi-supervised learning. When the labeled set is small, it is very likely that the labeled training set does cover sufficient situations, causing some failure cases in the reference set. In the self-learning baseline or even STSO, it is relatively difficult for the model to fix these errors during iteration, and the persisted errors can hinder the ability of the student model. In comparison, ATSO offers extra opportunities to jump out of the current distribution and thus get rid of the failure case. Hence, the efficiency of utilizing unlabeled training data is improved.

The pancreatic tumor subset of the MSD dataset (http://medicaldecathlon.com/) has $281$ abnormal CT scans. Each sample contains the annotation of the pancreas and the tumor. Since we train the base model on the NIH dataset (normal pancreas), we only evaluate the pancreas segmentation accuracy on the MSD data. Transferring from NIH to MSD (the NIH training data are labeled while the MSD training data are not) is a challenging task since the distributions are very different between these two datasets: the scanners are different, and the MSD data contain abnormality (pancreatic tumor) while the NIH data do not. In this scenario, the key is to extract information from the unlabeled training set that has a close distribution to the test set. Since the resolution on the long axis varies significantly, we normalize the inter-slice distance along the long axis during training and testing, but rescale the final output to the original size for a fair comparison. A global DSC (combining all test cases into a single volume) is computed.

Directly applying the pre-trained model for segmentation reports an average accuracy of $69.95\%$ on the test set. After four generations, ATSO reports an average accuracy of $76.71\%$ (an over $6.76\%$ gain over the direct transfer baseline) which is higher than that of STSO ($75.48\%$) and the self-learning baseline ($74.78\%$) after the same number of generations. Note that $76.71\%$ is even higher than the accuracy ($75.49\%$) obtained from using $62$ labeled MSD data for training. Similar to the NIH experiments, we also obtain more accurate pseudo labels. After four generations, ATSO achieves a $71.67\%$ accuracy on the reference set, but both the self-learning baseline and STSO reports around $70\%$. This aligns with our expectation: ATSO has the ability of transferring knowledge from a labeled dataset to another unlabeled, even when the data distributions differ considerably. Detailed results of transfer segmentation experiments from NIH to MSD are provided in Appendix B.

The second transfer learning scenario is defined between two natural image segmentation datasets. We use two popular datasets for autonomous driving, i.e., training the model on Cityscapes [4] and transferring it to Mapillary (https://www.mapillary.com/)¹¹1As a side comment, we also investigate the setting of semi-supervised learning within the Cityscapes dataset, and provide the results in Appendix C. The conclusions are similar to that reported in the NIH experiments.. The labeled training set contains $2\rm{,}975$ finely annotated images from Cityscapes and the reference and test sets are from Mapillary, which have $18\rm{,}000$ and $2\rm{,}000$ images, respectively. Note that Cityscapes has $19$ semantic classes while Mapillary has $66$. We train a model to infer $19$ classes and evaluate it on the reduced ground-truth on Mapillary by map each of the $66$ classes to one of the $19$ classes. The direct transfer method reports a mIOU of $24.70\%$ on the Mapillary test set that is dramatically lower than the number ($77.86\%$) on the Cityscapes test set, revealing a significant distribution gap between these two datasets.

We directly apply STSO and ATSO in this scenario, and obtain mIOU values of $28.11\%$ and $28.26\%$ (see Table 3), respectively. ATSO does not show advantages over STSO. This is mainly because $19$-class segmentation is more challenging, so the noise in the pseudo labels may have been magnified to overwhelm the difference between STSO and ATSO. As an example, for the objects with a small ground-truth area (e.g., a train), the teacher model (not seeing the ground-truth) can easily ignore the entire object so that the pseudo label will guide the student models to make the same critical mistake. To alleviate this instability, we propose a solution that weakens the pseudo labels to get rid of the trouble. We further perform a pre-defined mapping to reduce the number of classes from $19$ to $5$ (i.e., rode, vehicle, person, vegetation, others) so as to reduce the risk of missing some class completely. In this reduced, $5$-class segmentation task, ATSO reports a mIOU of $77.11\%$ which is significantly higher than the mIOU produced by STSO, $69.94\%$.

We then use the $5$-class pseudo label, generated by the final generation of either STSO or ATSO, to guide the target network. During the training procedure, the target network generates $19$-class segmentation and then reduces it to $5$ classes for computing the loss function with respect to the pseudo labels. With only one training generation, the target model produces mIOU of $29.76\%$ and $34.36\%$ using the pseudo labels from STSO and ATSO, respectively, i.e., better pseudo labels lead to higher mIOU. Interestingly, when we continue fine-tuning the latter model (a $34.36\%$ mIOU) with $19$-class pseudo labels, its performance is further boosted to $37.23\%$, claiming an over $12\%$ gain beyond the direct transfer baseline ($24.70\%$). Class-wise IOU numbers are detailed in Table 3. These results provide new insights to apply semi-supervised segmentation to the challenging datasets, in particular, with difficult semantic classes that can be easily missed. More details and results of the transfer experiments can be found in Appendix D.