PaperlayerView exact revision on arXiv

Source-linked AI summary

Semi-Supervised Medical Image Segmentation via Cross Teaching between CNN and Transformer

Xiangde Luo, Minhao Hu, Tao Song, Guotai Wang, Shaoting Zhang

arXiv:2112.04894v2eess.IVcs.CV

TL;DR

Limited annotations make medical image segmentation difficult despite strong fully supervised CNN and Transformer performance. The paper introduces cross teaching, using each network’s unlabeled predictions to supervise the other, and reports superior performance to eight existing semi-supervised methods on a public benchmark.

  • Problem

    Fully supervised CNN and Transformer segmentation methods generally require costly pixel- or voxel-level expert annotations, while good performance remains challenging with limited annotations.

  • Method

    The framework trains a CNN and Transformer on labeled data with ground truth and uses each network’s unlabeled predictions as pseudo labels to supervise the other.

  • Results

    The proposed method outperforms eight existing semi-supervised learning methods on a public benchmark.

  • Takeaways & Limitations

    Cross teaching introduces Transformers into semi-supervised medical image segmentation while using complementary CNN and Transformer learning properties during training.

Abstract

from arXiv · show

Recently, deep learning with Convolutional Neural Networks (CNNs) and Transformers has shown encouraging results in fully supervised medical image segmentation. However, it is still challenging for them to achieve good performance with limited annotations for training. In this work, we present a very simple yet efficient framework for semi-supervised medical image segmentation by introducing the cross teaching between CNN and Transformer. Specifically, we simplify the classical deep co-training from consistency regularization to cross teaching, where the prediction of a network is used as the pseudo label to supervise the other network directly end-to-end. Considering the difference in learning paradigm between CNN and Transformer, we introduce the Cross Teaching between CNN and Transformer rather than just using CNNs. Experiments on a public benchmark show that our method outperforms eight existing semi-supervised learning methods just with a simpler framework. Notably, this work may be the first attempt to combine CNN and transformer for semi-supervised medical image segmentation and achieve promising results on a public benchmark. The code will be released at: https://github.com/HiLab-git/SSL4MIS.

1. Introduction

Limited expert labeling makes fully supervised medical image segmentation costly, motivating semi-supervised methods that combine labeled and unlabeled data. The proposed framework cross-teaches between CNNs and Transformers to exploit their different learning paradigms.

  • Pixel- or voxel-level expert labeling is expensive and time-consuming, creating demand for semi-supervised medical image segmentation.
  • Existing semi-supervised medical image analysis methods largely combine labeled and unlabeled data to train CNN models.
  • Cross Teaching between CNN and Transformer uses predictions from each network on unlabeled images as pseudo labels for supervising the other network.
  • The framework leverages CNNs for local information and Transformers for long-range relations, allowing their different learning paradigms to complement each other during training.
  • The authors report that the approach outperforms eight existing semi-supervised methods on a public benchmark.

2. Method

The method trains a CNN and Transformer together on labeled and unlabeled images through bidirectional cross teaching, using each network’s predictions as pseudo labels for the other. Its objective combines supervised cross-entropy and Dice losses with an unlabeled cross-teaching loss weighted by a time-dependent schedule.

  • Cross teaching: The framework feeds labeled and unlabeled images through a CNN and Transformer, supervising labeled predictions with ground truth and unlabeled predictions through cross teaching.The two networks are trained jointly, with predictions from each network used to update the other on unlabeled data.
  • Cross teaching: CNN and Transformer predictions differ because CNNs use local convolution while Transformers model long-range self-attention.The method treats these learning paradigms and their output properties as complementary sources for pseudo-label supervision.
  • Cross teaching: Pseudo labels are generated from one network’s predictions to supervise the other, with no gradient back-propagation through the pseudo-label predictions.The cross-teaching loss is applied bidirectionally between the CNN and Transformer streams.
  • Cross teaching: Unlike consistency regularization, cross teaching uses a bidirectional loss without explicitly constraining the two predictions to become similar.The Transformer serves as a complementary training network rather than producing the final predictions.
  • Overall objective function: The supervised loss combines cross-entropy and Dice losses, while the overall training objective adds the unlabeled cross-teaching loss.The total objective is a joint supervised and unsupervised loss, with λ controlling the cross-teaching contribution.

3. Experiments

Experiments use the ACDC cardiac cine-MR benchmark and compare the proposed method with multiple semi-supervised baselines. The study also includes ablations of architectures, supervision strategies, and cross-teaching loss functions.

  • Dataset: The ACDC dataset contains 200 annotated short-axis cardiac cine-MR images from 100 patients, with masks for the LV, myocardium, and RV.The data are divided into 140 training images from 70 patients and 60 validation images from 30 patients.
  • Network architectures: The proposed method uses UNet as the CNN segmentation network and Swin-UNet as the Transformer segmentation network.Both are U-shaped architectures based on different learning paradigms.
  • Ablation study: Table 1 ablates Transformer-only segmentation, network architectures, cross teaching versus consistency regularization, and alternative cross-teaching loss functions.Results are reported with 7 labeled cases for RV, myocardium, and LV, with predictions marked by architecture.
  • Comparisons: The evaluation compares the method with eight previous semi-supervised methods, including Mean Teacher, Entropy Minimization, CPS, and CCT.The comparison covers several consistency, adversarial, and pseudo-supervision approaches.

4. Results

Experiments on ACDC show that cross teaching between CNN and Transformer outperforms alternative semi-supervised configurations and eight existing methods, with qualitative gains under 3- and 7-case supervision.

  • Ablation study: Cross teaching between CNN and Transformer achieves better results than other tested network combinations, while also outperforming consistency regularization.The ablation attributes this to complementary learning paradigms and reports that Dice loss improves performance over cross-entropy loss.
  • Comparison with baselines and existing works: The proposed method outperforms eight existing semi-supervised methods by a large margin using either 3 or 7 labeled cases.All methods use the same UNet backbone for final inference, without auxiliary modules or ensemble strategies.
  • Comparison with baselines and existing works: The Transformer branch performs similarly to the CNN branch, while their ensemble outperforms both at higher computation cost.The Transformer has 27.12M parameters versus 1.81M for the CNN, so CNN outputs are used for fair comparison with existing methods.
  • Comparison with baselines and existing works: Visual comparisons report more plausible segmentations with fewer false-positive and missing-segmentation regions than CCT and CPS.The visualization comparison covers validation images using 3 labeled cases in the first two rows and 7 labeled cases in the last two rows.

5. Conclusion

The paper introduces Transformer-based semi-supervised medical image segmentation through cross teaching between CNN and Transformer. On an open benchmark, the method outperforms eight existing semi-supervised learning methods.

  • 5. Conclusion: Cross teaching uses CNN and Transformer networks to exploit unlabeled data in semi-supervised medical image segmentation.The approach is inspired by co-teaching and cross pseudo supervision.
  • 5. Conclusion: The method assumes CNNs capture local features efficiently and Transformers model long-range relations, allowing their properties to complement each other during training.The conclusion identifies this complementarity as the basis for combining the two architectures.
  • 5. Conclusion: Experimental results on an open benchmark show that the proposed method outperforms eight existing semi-supervised learning methods.The authors propose future work to further reduce annotation costs, especially for dense-annotation-based multi-organ segmentation.

A.1. Dataset

The ACDC dataset contains paired end-diastolic and end-systolic cine-MRI scans from 100 patients, split by patient into training and validation sets.

  • A.1. Dataset: ACDC consists of 200 short-axis cine-MRI scans from 100 patients, with two annotated scans per patient.The scans correspond to end-diastolic and end-systolic phases.
  • A.1. Dataset: The patient-level split assigns 70 patients to training and 30 patients to validation.This yields 140 training scans and 60 validation scans, covering both cardiac phases.

A.2. Training details

The experiments use UNet and Swin-UNet implementations trained with stochastic gradient descent and a polynomial learning-rate schedule.

  • A.2. Training details: UNet and Swin-UNet serve as the CNN and Transformer segmentation networks, respectively.Their PyTorch implementations were borrowed from the PyMIC and Swin-UNet projects.
  • A.2. Training details: All models use stochastic gradient descent with a polynomial learning-rate strategy and an initial learning rate of 0.01.Training runs for 30k total iterations on a GTX1080TI GPU using PyTorch 1.8.1.

B.1. Computational-cost

The computational-cost comparison measures forwarding passes, total training time, and per-case inference time across methods. The proposed method trains longer, while inference costs remain similar when using the same segmentation network.

  • Computational-cost measures: The comparison uses FTimes, TTimes, and ITimes to assess forwarding passes, total training time, and per-case inference time.Table 3 reports these measures using the same software and hardware settings, based on 7 labeled cases.
  • Training cost: The proposed method requires more training time than existing methods because Swin-UNet has more parameters.
  • Inference cost: Inference costs are very similar across methods when the same segmentation network is used.Within the proposed method, the transformer branch is slower than the CNN branch, and the ensemble is slowest because it runs both models simultaneously.

B.2. Performance of CNN branch and Transformer branch

The CNN and Transformer branches achieve similar performance in the evaluated 3- and 7-labeled-case settings, although the Transformer branch has higher inference cost. Their ensemble prediction is defined through an argmax operation.

  • Branch performance: The CNN and Transformer branches achieve very similar performance in the 3- and 7-labeled-case settings.
  • Branch performance: The Transformer branch requires more inference cost than the CNN branch despite their similar performance.
  • Ensemble prediction: The ensemble prediction is defined using an argmax operation, with the result reported in the last row of Table 4.
Loading 2112.04894v2…