Source-linked AI summary

Few-shot 3D Multi-modal Medical Image Segmentation using Generative Adversarial Learning

Arnab Kumar Mondal, Jose Dolz, Christian Desrosiers

arXiv:1810.12241v1cs.CV

TL;DR

Medical image segmentation needs many annotations, while adversarial semi-supervised segmentation had seen limited use for 3D multi-modal volumes. This paper trains a GAN-based segmentation model with labeled, unlabeled, and generated patches, reporting gains over 3D-UNet with few labels and performance close to full supervision.

  • Problem

    Medical segmentation requires costly annotations, and adversarial semi-supervised methods had been applied only to a limited extent to 3D multi-modal medical images.

    Source: §I · p. 1
  • Method

    The method uses GAN-based training with labeled, unlabeled, and generated images, treating generated samples as a fake class.

    Source: §I · p. 1, B. Reduced supervision with GANs · p. 3
  • Results

    The method significantly outperforms 3D-UNet with very few training samples and achieves accuracy close to full supervision.

    Source: §I · p. 1
  • Takeaways & Limitations

    Results on iSEG-2017 and MRBrainS 2013 indicate the method's potential to reduce the burden of acquiring annotated medical data.

    Source: §V · p. 8

    Feature matching may assign probability mass within the data support and, because it matches only first-order statistics, may collapse to a trivial solution that fails to cover regions between manifolds.

    Source: C. Complement (Bad) generator · p. 4

Abstract

from arXiv · show

We address the problem of segmenting 3D multi-modal medical images in scenarios where very few labeled examples are available for training. Leveraging the recent success of adversarial learning for semi-supervised segmentation, we propose a novel method based on Generative Adversarial Networks (GANs) to train a segmentation model with both labeled and unlabeled images. The proposed method prevents over-fitting by learning to discriminate between true and fake patches obtained by a generator network. Our work extends current adversarial learning approaches, which focus on 2D single-modality images, to the more challenging context of 3D volumes of multiple modalities. The proposed method is evaluated on the problem of segmenting brain MRI from the iSEG-2017 and MRBrainS 2013 datasets. Significant performance improvement is reported, compared to state-of-art segmentation networks trained in a fully-supervised manner. In addition, our work presents a comprehensive analysis of different GAN architectures for semi-supervised segmentation, showing recent techniques like feature matching to yield a higher performance than conventional adversarial training approaches. Our code is publicly available at https://github.com/arnab39/FewShot_GAN-Unet3D

I. INTRODUCTION

The paper targets few-shot segmentation of 3D multi-modal medical images, where annotation is costly and prior adversarial segmentation work had limited application to this setting. It situates its approach among semi-supervised methods and reports strong performance with few labeled samples.

I. INTRODUCTION

  • Medical segmentation remains difficult because scans are noisy and low-contrast, anatomy varies, and obtaining annotations is time-consuming and error-prone.CNNs have improved segmentation but require large annotated datasets.Source: §I · p. 1

A. Semi-supervised learning

  • Semi-supervised learning uses unlabeled samples alongside a limited number of strongly annotated samples; related work also includes self-training and teacher-based methods.Some existing approaches require multiple networks or an explicit matching function.Source: §I · p. 1, A. Semi-supervised learning · p. 2

I. INTRODUCTION

  • The paper addresses the limited use of adversarial semi-supervised learning for 3D multi-modal medical images by applying GANs in this setting.Prior GAN-based semi-supervised segmentation work had focused on 2D natural images.Source: §I · p. 1
  • The proposed GAN approach significantly outperforms 3D-UNet with very few training samples and achieves accuracy close to full supervision.The paper presents this as a first application of GANs to semi-supervised segmentation of 3D multi-modal images.Source: §I · p. 1

II. RELATED WORK

  • The paper also compares GAN architectures and reports higher performance for feature matching than conventional adversarial training.The broader related-work context identifies adversarial learning as a promising approach for semi-supervised segmentation.Source: §I · p. 1, §II · p. 1

B. Adversarial learning

Prior adversarial segmentation methods use unlabeled data in different ways, while few-shot segmentation methods often rely on pretraining. This paper adapts GAN-based learning to 3D medical volumes using labeled, unlabeled, and generated patches.

B. Adversarial learning

  • Adversarial segmentation can use unlabeled images to align output distributions, but related approaches face risks from incorrect pseudo-segmentations or the need to choose a confidence threshold.The cited GAN semi-supervised segmentation work focused on 2D natural images, while 3D volumes add computational and over-fitting challenges.Source: B. Adversarial learning · p. 2

C. Few-shot learning

  • Few-shot segmentation methods commonly refine a model pretrained on a similar task, whereas this approach reports competitive performance with as few as two samples without pretraining.The authors note that the source-to-target gap can make pretraining on a similar task challenging for medical images.Source: C. Few-shot learning · pp. 2–3

III. METHODS

  • The method trains with labeled, unlabeled, and GAN-generated images, assigning generated examples an additional fake class to constrain predictions on real unlabeled data.For 3D images, training uses patches to manage memory and computation, and the segmentation network is fully convolutional to handle arbitrary test-image sizes.Source: §III · p. 3

A. Generative Adversarial Networks

The paper builds on the GAN setup of a generator and discriminator, then adapts segmentation training to use unlabeled and generated samples. Its formulation adds a fake class to the segmentation model.

A. Generative Adversarial Networks

  • A standard GAN trains a generator to map noise to synthetic examples and a discriminator to distinguish those examples from real data.The networks are trained in a two-player min-max game.Source: A. Generative Adversarial Networks · p. 3

B. Reduced supervision with GANs

  • The segmentation model is extended to include unlabeled and generated data, with a (K+1)th output representing the probability that an input voxel is fake.The unlabeled-data objective constrains predictions to the K classes of real data.Source: B. Reduced supervision with GANs · p. 3
  • The method retains the standard labeled-image segmentation loss and adds unlabeled and fake-image terms to train the discriminator and generator.The supplied passage introduces these loss functions but does not specify their full formulas.Source: B. Reduced supervision with GANs · p. 3

1) Discriminator loss:

The discriminator objective combines contributions from labeled, unlabeled, and generated images. The formulation also reparameterizes the K+1 output model into K effective outputs.

1) Discriminator loss:

  • The discriminator loss sums labeled, unlabeled, and fake-image terms, with the labeled-image loss remaining the standard segmentation loss.The authors assume similar numbers of labeled, unlabeled, and fake images so each image type has equal importance.Source: 1) Discriminator loss · pp. 3–4
  • For unlabeled images, the objective favors predictions among the real classes, while generated images are assigned the fake class at every pixel.The fake-image loss is defined for generated inputs.Source: 1) Discriminator loss · p. 4
  • The K+1-output formulation is over-parameterized because a shared function can be subtracted from all logits without changing softmax outputs; the paper uses K effective outputs instead.The modified losses use normalized logits after reparameterization.Source: 1) Discriminator loss · p. 4

2) Generator loss:

The generator uses feature matching rather than maximizing the fake-sample loss, which can be unstable and perform poorly in semi-supervised learning. Feature matching aligns expected intermediate-layer features.

2) Generator loss:

  • Feature matching replaces maximizing Lfake because that objective can cause instability and poor performance in semi-supervised learning.The generator instead matches the expected value of features from an intermediate layer.Source: 2) Generator loss · p. 4
  • The feature representation f(x) uses activations from the second-last layer of the model’s encoding path.Source: 2) Generator loss · p. 4
  • Using the encoder’s second-last layer gave slightly higher performance in preliminary experiments than using its last layer.Source: 2) Generator loss · p. 4

C. Complement (Bad) generator

The complement generator is intended to produce samples outside the true data manifold. The method addresses potential feature-matching collapse by increasing generated-distribution entropy, and pairs the generator with an encoder trained end to end.

C. Complement (Bad) generator

  • A complement, or “bad,” generator aims to generate samples outside the true data manifold, since a good generator can make unlabeled and fake images indistinguishable.Source: C. Complement (Bad) generator · p. 4
  • Feature matching can place probability mass within the data support or collapse to mean unlabeled features, leaving gaps between supports underdetermined for the discriminator.Source: C. Complement (Bad) generator · p. 4
  • The method proposes increasing the generated distribution’s entropy to address feature-matching collapse, while noting that the first identified problem is less likely for complex multi-modal 3D patches.Source: C. Complement (Bad) generator · p. 4
  • The modified generator loss is optimized using a variational upper bound on negative entropy, with q defined as a bounded-variance diagonal Gaussian.Source: C. Complement (Bad) generator · p. 4
  • The fake-image generator is paired with an encoder that maps generated images back to noise vectors, and all components train jointly end to end.Source: C. Complement (Bad) generator · p. 4

IV. EXPERIMENTS

The experiments assess few-shot segmentation of infant and adult brain tissue from multi-modal 3D MRI, including whether performance with one or two training subjects can approach full supervision. The evaluation uses overlap and surface-distance metrics.

A. Materials

  • The experiments assess whether GAN-based segmentation can approach full-supervision performance with only one or two training subjects, and examine the impact of GAN techniques.Source: A. Materials · p. 5

1) Dataset:

  • The evaluation covers infant iSEG-2017 T1- and T2-weighted MRI with three tissue classes and 10 labeled subjects, plus adult MRBrainS data with T1, T1-IR, and FLAIR scans and five labeled training subjects.The datasets also provide unlabeled scans: 13 iSEG subjects and 15 MRBrainS subjects.Source: 1) Dataset · p. 5

2) Evaluation:

  • Segmentation is assessed with Dice similarity coefficient for spatial overlap and Average Symmetric Surface Distance for boundary distances.DSC ranges from 0 to 1, with 1 indicating perfect overlap.Source: 2) Evaluation · p. 5

B. Implementation Details

The implementation adapts 3D U-Net for adversarial semi-supervised training and uses a volumetric generator with an encoder that estimates the noise distribution for an image.

B. Implementation Details

  • The 3D U-Net replaces batch normalization with weight normalization, which the authors say helped stabilize training and improve performance.Source: 1) Architecture · p. 5
  • Leaky ReLUs replace ReLUs, and average pooling replaces max pooling to avoid sparse gradients that hamper GAN training.Source: 1) Architecture · p. 5
  • A volume generator creates 3D patches, while a three-layer 3D CNN encoder estimates the mean and standard deviation of the noise vector used to generate an image.Batch normalization was used in both generator and encoder based on preliminary experiments.Source: 1) Architecture · p. 5

2) Training:

The training setup uses very few labeled brain-MRI examples alongside unlabeled data, and the FM GAN is evaluated against U-Net as labeled training data increases. Across these settings, GAN-based methods improve on basic U-Net, with feature matching performing best among the tested GAN variants.

2) Training:

  • For iSEG-2017, training used 1 or 2 labeled subjects, 1 for validation, and 7 fixed test subjects; the 13 unlabeled subjects also trained the GAN.For MRBrainS 2013, the split was 1 training, 1 validation, and 3 testing examples, with 15 unlabeled subjects used for GAN training.Source: 2) Training · p. 5
  • Training used bias-corrected, intensity-normalized scans divided into overlapping 32×32×32 patches, without other data augmentation.The patch step size was 8 voxels per dimension; training and validation example combinations were varied across three splits while test examples remained fixed.Source: 2) Training · p. 5, C. Few-shot learning with FM GAN · p. 6

C. Few-shot learning with FM GAN

  • DSC gains over basic U-Net were 5–8% WM, 13–23% GM, and 1–9% CSF with 1 label, and 3–12%, 10–18%, and 1–6%, respectively, with 2.Mean DSC and ASD were evaluated on iSEG-2017 test examples; the GAN used 13 unlabeled examples during training.Source: C. Few-shot learning with FM GAN · p. 6, Table I · p. 6
  • Feature matching without an entropy term yielded the best performance across tissue classes and test cases, improving DSC over bad-GAN by 5% WM, 4% GM, and 3% CSF with 1 label.With 2 labeled examples, the respective improvements were 6%, 4%, and 4%.Source: C. Few-shot learning with FM GAN · p. 6
  • Across 1–5 labeled images, FM GAN achieved higher or equal DSC to U-Net, and their accuracy was comparable with 5 labeled images.This comparison used one validation example, four fixed test examples, and 13 unlabeled examples for FM GAN.Source: C. Few-shot learning with FM GAN · p. 6, Figure 2 · p. 6
  • With 1 or 2 labeled examples, FM GAN produced better-looking brain-tissue segmentations than basic U-Net; with 5, its output looked similar to U-Net’s.Figure 3 compares segmentations for two test subjects at different labeled-example counts.Source: C. Few-shot learning with FM GAN · pp. 6–7, Figure 3 · p. 7

D. Detailed analysis of GAN models

The analysis examines how unlabeled and generated data relate to generalization and compares the tested GAN architectures. FM GAN generalizes better than U-Net in the reported validation analysis, while the authors link useful fake patches to discriminator learning.

D. Detailed analysis of GAN models

  • U-Net quickly overfit with one labeled example, whereas FM GAN learned from unlabeled and generated data and generalized better on validation data.Figure 4 plots U-Net cross-entropy and validation error alongside FM GAN unlabeled-image and fake-image losses.Source: D. Detailed analysis of GAN models · p. 7, Figure 4 · p. 7
  • Unlike U-Net, FM GAN assigned fake-class probability near 1 to all voxels of a fake patch.The visual shows one 2D slice of a 3D patch; Table II reports mean unlabeled and fake losses on test-image patches.Source: Table II · p. 7, D. Detailed analysis of GAN models · p. 8, Figure 5 · p. 8
  • FM GAN outperformed bad-GAN, which adds an entropy term for more diverse generated examples; the authors report that adding the VI loss did not improve FM GAN.They suggest poor sample quality, aggravated by increasing entropy, as one possible explanation.Source: D. Detailed analysis of GAN models · p. 8
  • FM GAN’s feature-matching loss converged quickly and remained below bad-GAN’s, while the authors say generated samples should differ from true unlabeled images yet remain close enough to support discriminator learning.They describe bad-GAN samples as more likely to fall far from the true distribution, potentially weakening the discriminator’s boundary.Source: D. Detailed analysis of GAN models · p. 8

E. Validation on MR Brains dataset

On MRBrains, the proposed technique outperforms standard U-Net with one training example. The conclusion reports comparable-to-full-supervision performance with few samples and finds feature matching outperforms bad-GAN on 3D multi-modal brain MRI.

E. Validation on MR Brains dataset

  • DSC improvements over standard U-Net were 13.6% for WM, 7.5% for GM, and 34% for CSF; ASD reductions were 46% for WM, 37% for GM, and 38% for CSF.These validation experiments used one MRBrains training example.Source: E. Validation on MR Brains dataset · p. 8

V. CONCLUSION

  • With few training samples, the method can achieve performance comparable to full supervision by using unlabeled data to discriminate true from generated fake patches and prevent over-fitting.The method can enhance segmentation networks in low-data settings where they fail to produce good outputs.Source: §V · p. 8
  • Feature matching GAN performed better than bad-GAN for segmenting 3D multi-modal brain MRI.Adding the variational-inference term did not improve feature matching GAN performance in this setting.Source: §V · p. 8
  • The method can extend to other 3D multi-modal segmentation tasks using any state-of-the-art segmentation network as discriminator.The authors also report its potential to reduce the burden of acquiring annotated medical data.Source: §V · p. 8
Loading 1810.12241v1…