PMT: Progressive Mean Teacher via Exploring Temporal Consistency for Semi-Supervised Medical Image Segmentation

Ning Gao,Sanping Zhou,Le Wang,Nanning Zheng

2024-09-15

Abstract:Semi-supervised learning has emerged as a widely adopted technique in the field of medical image segmentation. The existing works either focuses on the construction of consistency constraints or the generation of pseudo labels to provide high-quality supervisory signals, whose main challenge mainly comes from how to keep the continuous improvement of model capabilities. In this paper, we propose a simple yet effective semi-supervised learning framework, termed Progressive Mean Teachers (PMT), for medical image segmentation, whose goal is to generate high-fidelity pseudo labels by learning robust and diverse features in the training process. Specifically, our PMT employs a standard mean teacher to penalize the consistency of the current state and utilizes two sets of MT architectures for co-training. The two sets of MT architectures are individually updated for prolonged periods to maintain stable model diversity established through performance gaps generated by iteration differences. Additionally, a difference-driven alignment regularizer is employed to expedite the alignment of lagging models with the representation capabilities of leading models. Furthermore, a simple yet effective pseudo-label filtering algorithm is employed for facile evaluation of models and selection of high-fidelity pseudo-labels outputted when models are operating at high performance for co-training purposes. Experimental results on two datasets with different modalities, i.e., CT and MRI, demonstrate that our method outperforms the state-of-the-art medical image segmentation approaches across various dimensions. The code is available at <a class="link-external link-https" href="https://github.com/Axi404/PMT" rel="external noopener nofollow">this https URL</a>.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

### What problems does this paper attempt to solve?

This paper aims to address the key challenges in semi - supervised medical image segmentation, namely how to continuously improve the model's ability and generate high - quality pseudo - labels with limited labeled data. Specifically, the paper proposes a new semi - supervised learning framework - Progressive Mean Teachers (PMT) - to solve the following problems:

1. **Acquisition of high - quality supervision signals**:
- In medical image segmentation tasks, labeled data is very scarce, so it is necessary to effectively utilize unlabeled data to improve model performance.
- Existing methods mainly focus on consistency constraints and pseudo - label generation, but these methods are insufficient in continuously improving the model's ability.

2. **Maintenance of model diversity and robustness**:
- Existing methods have difficulty maintaining model diversity during the training process, resulting in low - quality pseudo - labels.
- The paper solves this problem by introducing two Mean Teacher architectures for collaborative training and using a Discrepancy Driven Alignment (DDA) regularizer.

3. **Selection of high - fidelity pseudo - labels**:
- To ensure that the pseudo - labels used for training have high fidelity, the paper designs a simple Pseudo Label Filtering (PLF) algorithm, which selects high - quality pseudo - labels by comparing the model performance at different training stages.

### Overview of the solution

The PMT framework proposed in the paper mainly includes the following core components:

- **Mean Teacher (MT)**: Updates parameters through exponential moving average (EMA) to maintain consistency between the teacher network and the student network, thereby enhancing the model's representation ability.
- **Discrepancy Driven Alignment (DDA)**: Quickly aligns the representation ability of the lagging model with that of the leading model by comparing the prediction differences between models at different training stages.
- **Pseudo Label Filtering (PLF)**: Filters out high - fidelity pseudo - labels generated by better - performing models to avoid the negative impact of low - quality pseudo - labels on the training process.

### Experimental results

The experimental results show that the PMT framework outperforms existing methods on datasets of two different modalities, CT and MRI. In particular, experiments on the left atrium (LA) and pancreas (Pancreas - NIH) datasets show that PMT can achieve higher segmentation accuracy with less labeled data.

### Formula summary

1. **Mean Teacher parameter update formula**:
\[
\theta_{\text{teacher}}=\alpha\theta_{\text{teacher}}+(1 - \alpha)\theta_{\text{student}}
\]
where $\theta_{\text{teacher}}$ and $\theta_{\text{student}}$ represent the parameters of the teacher network and the student network respectively, and $\alpha$ is the EMA decay rate.

2. **Pseudo - label sharpening formula**:
\[
Y_U=\frac{P_s^{1/T}}{P_s^{1/T}+(1 - P_s)^{1/T}}
\]
where $P_s$ represents the output of the good student network and $T$ is a hyperparameter.

3. **Loss function**:
\[
L_{\text{total}}=L_s+\lambda_1P(f_i,f_{i + 1})L_U+\lambda_2L_T
\]
where $L_s$ is the supervised learning loss, $L_U$ is the unsupervised learning loss, $L_T$ is the Mean Teacher consistency loss, $\lambda$