PaperlayerView exact revision on arXiv

Source-linked AI summary

Enhancing MR Image Segmentation with Realistic Adversarial Data Augmentation

Chen Chen, Chen Qin, Cheng Ouyang, Zeju Li, Shuo Wang, Huaqi Qiu, Liang Chen, Giacomo Tarroni, Wenjia Bai, Daniel Rueckert

arXiv:2108.03429v3eess.IVcs.CVcs.LGq-bio.QM

TL;DR

Medical image segmentation training typically requires large amounts of labeled data. AdvChain introduces adversarial data augmentation for segmentation, and reports data efficiency and stronger performance than consistency-regularized methods in low-shot semi-supervised settings across two tasks.

  • Problem

    Neural-network medical image segmentation typically requires large amounts of labeled training data.

    Source: §1 · p. 1

  • Method

    AdvChain is a generic adversarial data augmentation framework for medical image segmentation.

    Source: §1 · p. 2

  • Results

    In low-shot semi-supervised settings on two tasks, AdvChain demonstrates data efficiency and outperforms several strong consistency-regularized methods.

    Source: §7 · p. 19

  • Takeaways & Limitations

    The reported data-efficiency result spans two tasks in challenging low-shot semi-supervised settings.

    Source: §7 · p. 19

    The method still requires expertise to specify transformation magnitudes.

    Source: §6 · p. 18

Abstract

from arXiv · show

The success of neural networks on medical image segmentation tasks typically relies on large labeled datasets for model training. However, acquiring and manually labeling a large medical image set is resource-intensive, expensive, and sometimes impractical due to data sharing and privacy issues. To address this challenge, we propose AdvChain, a generic adversarial data augmentation framework, aiming at improving both the diversity and effectiveness of training data for medical image segmentation tasks. AdvChain augments data with dynamic data augmentation, generating randomly chained photo-metric and geometric transformations to resemble realistic yet challenging imaging variations to expand training data. By jointly optimizing the data augmentation model and a segmentation network during training, challenging examples are generated to enhance network generalizability for the downstream task. The proposed adversarial data augmentation does not rely on generative networks and can be used as a plug-in module in general segmentation networks. It is computationally efficient and applicable for both low-shot supervised and semi-supervised learning. We analyze and evaluate the method on two MR image segmentation tasks: cardiac segmentation and prostate segmentation with limited labeled data. Results show that the proposed approach can alleviate the need for labeled data while improving model generalization ability, indicating its practical value in medical imaging applications.

1. Introduction

Medical segmentation is constrained by scarce labeled data and conventional augmentation that may ignore image content or downstream tasks. AdvChain uses optimized chained transformations and reports improved performance in limited-label MR segmentation settings.

1. Introduction

  • Labeling is time-consuming and expensive, while collecting and sharing large clinical datasets is difficult; consequently, training often uses few labeled images.The paper identifies privacy and ethical issues as barriers to dataset sharing.Source: §1 · p. 1
  • Conventional random transformations may ignore image content and downstream tasks, producing redundant samples that do not necessarily improve generalization.The proposed framework also combines pixel-level and contour-based consistency.Source: §1 · p. 2
  • AdvChain dynamically chains photometric and geometric transformations, optimizing their parameters to reflect image information and current network fragility.The framework targets realistic, challenging variations to improve training-data diversity and regularize the network.Source: §1 · p. 2
  • AdvChain outperforms several strong consistency-regularized methods and RandAugment when training with only 1 or 3 labeled subjects.The reported comparisons indicate improved generalizability when labeled data is limited.Source: §1 · p. 2
  • Experiments on public cardiac and prostate MR datasets show improved segmentation performance over competitive consistency-regularization methods.The paper extends the framework to multiclass segmentation and evaluates it on both datasets.Source: §1 · p. 2

2. Related work

The section reviews data mixing, adversarial augmentation, and consistency regularization for semi-supervised learning. It situates the proposed approach against prior weak, random augmentations, describing complex photometric and geometric transformations that are adversarially optimized to generate challenging images and potentially support robust representations and generalization.

2.1. Data mixing

  • Mixup interpolates random image-label pairs (Zhang et al., 2018), but mixed images can be unrealistic and generated samples remain within the training-data span.Mixup has also been adapted to medical image segmentation.Source: §2.1 · p. 3

2.2. Adversarial data augmentation

  • Prior adversarial augmentation often uses gradient-based additive noise, while medical segmentation studies have used attacks to evaluate robustness, including Paschali et al. (2018).Paschali et al. applied a dense adversary generation attack (Xie et al., 2017) to create pixel-wise noise that degrades brain-image segmentation.Source: §2.2 · p. 3
  • AdvChain dynamically optimizes parameters for chained photometric and geometric transformations, generating realistic, challenging variations rather than using a single fixed transformation.It can optimize the transformations in a chain with one forward pass.Source: §2.2 · p. 3
  • GAN-based composite augmentation requires separate models for photometric and geometric parameters and is computationally intensive and unstable, limiting optimization of randomly chained transformations.AdvChain can optimize the transformations in a chain with one forward pass.Source: §2.2 · p. 3
  • AdvChain also optimizes a stationary velocity field and integrates it to generate diffeomorphic deformations for realistic morphological variation.This approach differs from GAN-based deformation methods that model additive displacement fields.Source: §2.2 · p. 4

2.3. Consistency regularization

  • Consistency regularization supports semi-supervised learning by augmenting unlabeled data and encouraging consistent predictions for original and augmented inputs.The approach assumes predictions for a data point and its augmented or perturbed example should agree.Source: §2.3 · p. 4
  • Prior consistency-regularization methods used weak random augmentations, including Gaussian noise, affine transformations, and adversarial bias fields (Chen et al., 2020).Source: §2.3 · p. 4
  • AdvChain instead models more complex photometric and geometric transformations and adversarially optimizes their parameters to generate more challenging augmented images.Source: §2.3 · p. 4
  • The authors propose that diverse, effective, realistic augmentation can regularize training, strengthen high-level robust representations, and improve generalization.Source: §2.3 · p. 4

3. Methods

The methods frame segmentation training as supervised learning on labeled images combined with consistency regularization on labeled and unlabeled data. The consistency term compares predictions before and after transformation.

3. Methods

  • The motivation is to regularize networks that can overfit when training data is insufficient.The paper describes medical image segmentation as a mapping from image space to label space.Source: §3 · p. 4
  • The network objective combines supervised loss on labeled images with weighted consistency regularization on labeled and unlabeled images.The consistency term measures prediction inconsistency between the original image and its transformed version.Source: §3 · pp. 4–5
  • The consistency term compares fθ(x) with fθ(Tt(x)), where Tt denotes a transformation parameterized by t.This formulation provides the regularization term used alongside supervised loss.Source: §3 · p. 5

3.1. AdvChain: A generic adversarial data augmentation

AdvChain searches over randomly chained transformations for perturbations that increase prediction disagreement, then trains the segmentation network with the resulting examples. Constraints keep transformed images meaningful and realistic.

3.1. AdvChain: A generic adversarial data augmentation

  • The method chains randomly selected transformation functions and optimizes their parameters, enabling pixel-level perturbations and geometric variations.The framework is intended to reflect complex image variations in MR imaging.Source: §3.1 · p. 5
  • AdvChain first optimizes transformation parameters to maximize prediction disagreement, then updates the network using the resulting augmented images and consistency loss.The network objective also includes supervised loss when labels are available.Source: §3.1 · p. 5
  • Parameter constraints define the search space and are essential for keeping augmented images meaningful and realistic.The method uses projected gradient descent to update randomly initialized transformation parameters under constraints.Source: §3.1 · pp. 5–6
  • Normalized gradients update parameters along the augmentation chain to reduce gradient explosion or vanishing when chains are long.The chain rule is used to compute gradients through the transformations.Source: §3.1 · p. 6

3.2. Increasing the data diversity of AdvChain with dynamic transformations

AdvChain randomly chains differentiable transformations in varying orders, then adversarially optimizes their parameters to generate augmented images. This flexible strategy increases data diversity while avoiding the high computational cost of searching for optimal transformation combinations.

3.2. Increasing the data diversity of AdvChain with dynamic transformations

  • Each image can use a different single or composite transformation chain, increasing augmentation diversity at low cost.The paper contrasts this flexibility with the extraordinary computational cost of searching for improved augmentation policies.Source: §3.2 · p. 6
  • Randomly selected transformations are chained in random order, and adversarial training optimizes each transformation’s parameters through the chain rule.The resulting chain augments the image before the segmentation network loss is computed.Source: §3.2 · p. 6
  • Transformation operations are selected with probability p and randomly initialized before k update steps optimize their parameters.The algorithm applies the chain rule to compute gradients for the parameters of each transformation in the chain.Source: §3.2 · p. 6

3.3. Realistic image transformation functions

AdvChain uses differentiable photometric and geometric transformations to model common MR-image variations, with four transformation functions presented as a proof of concept. The framework defines these transformations in 2D, with potential extension to 3D.

3.3. Realistic image transformation functions

  • AdvChain is a lightweight framework for on-the-fly optimization of differentiable photometric and geometric transformations, demonstrated with four functions in 2D.The paper notes that the approach could potentially extend to 3D.Source: §3.3 · pp. 6–7

3.3.1. Image noise model Tnoise

  • Tnoise adds pixel-wise noise to images, with its magnitude bounded by a scalar that controls the noise level.The noise model is described as an additive perturbation.Source: §3.3 · p. 6, §3.3.1 · p. 7

3.3.2. Image intensity transformation with bias field Tbias

  • Tbias multiplies image intensities by a smoothly varying bias field to introduce intensity non-uniformity.The field is constructed by smoothing control points with B-splines and exponentiating the result; zero control points yield the identity field.Source: §3.3.2 · p. 7

3.3.3. Image geometric transformation model Taffine

  • Taffine models spatial variation through translation, rotation, and scaling applied in succession, with parameters restricted to user-defined ranges.The model uses five parameters for translation, rotation, and scaling.Source: §3.3.3 · p. 8

3.3.4. Image deformation model Tmorph

  • Tmorph models intra- and inter-subject shape variation with smooth, invertible diffeomorphic deformations that preserve topology and spatial layout.The deformation is parameterized by a stationary velocity field and smoothed with a Gaussian kernel.Source: §3.3.4 · p. 8

3.4. Consistency loss function R

Consistency regularization compares predictions on original and transformed images. For geometric transformations, perturbed predictions are first mapped back to the original image coordinates, including when transformations are chained.

3.4.1. Consistency loss function for photometric transformations

  • For photometric transformations, consistency loss compares original and perturbed predictions using pixel-wise and contour-based terms.The contour term uses Sobel filters to capture differences at foreground-object boundaries; the pixel-wise term uses mean-squared error.Source: §3.4.1 · p. 9

3.4.2. Consistency loss function for geometric transformations

  • For geometric transformations, the perturbed prediction is mapped back to the original image coordinates before consistency loss is computed.The inverse mappings are computed for affine transformations and for morph deformations by integrating the negative velocity field.Source: §3.4.2 · p. 9

3.4.3. Consistency loss function for a chained transformation

  • For a chain containing photometric and geometric transformations, the loss combines the corresponding consistency rules and reverses geometric changes before comparison.The paper illustrates the chained case with an affine transformation followed by noise.Source: §3.4.3 · p. 9

4. Experiments Settings

The experiments evaluated standard and adversarial augmentation for cardiac and prostate MR segmentation with limited labeled training data, using specified network, optimization, and loss procedures. They applied a shared AdvChain configuration to both tasks, with constrained transformations and a one-step procedure for updating augmentation parameters.

4.1.1. Cardiac MR dataset

  • The study split a dataset of 100 subjects, selected N subjects from the remainder for a labeled pool, evaluated one-shot learning (N=1), and tested labeled-subject counts N=10 and N=25 from that pool.Mean performance was reported.Source: §4.1.1 · p. 10

4.1.2. Prostate MR dataset

  • The 32-subject dataset was divided into two 11-subject subsets; randomly selected N (with N ≤11) subjects came from the first, and the unlabeled set comprised M = 11 subjects.Mean performance was reported.Source: §4.1.2 · p. 10

4.2.1. Default data augmentation

  • The default augmentation pipeline applied random affine transforms, flipping, global brightness and contrast changes, and elastic transforms before rescaling intensities to [0, 1].Affine transforms included scaling, rotation, and translation.Source: §4.2.1 · p. 10

4.2.2. Training details

  • The network-independent method used a 2D U-net trained with Adam at batch size 20 for 1,000 epochs (learning rate=10−3) to obtain a pretrained segmentation network, which was then fine-tuned.Source: §4.2.2 · p. 10
  • The supervised loss used soft Dice, and BG, PZ, and CZ class weights were 0.01:0.66:0.33; PZ received higher weight because it has significantly fewer pixels.Source: §4.2.2 · pp. 10–11
  • Constraints for noise, bias field, rotation, and translation were directly taken from Chen et al. (2020) and RandAugment, while velocity-magnitude constraint was set by visual inspection.Source: §4.2.2 · p. 11
  • Transformation parameters were initialized by random sampling within specified ranges, then updated with projected gradient descent to remain in the search space.Noise and velocity fields were normalized and rescaled to meet magnitude constraints; bias fields were clipped, and affine parameters were HardTanh-activated and rescaled.Source: §4.2.2 · p. 11
  • The search used a step size of α_i = 1 and one update step (k = 1), following Miyato et al. (2018).The authors chose this setup for simplicity and training efficiency.Source: §4.2.2 · p. 11
  • The shared configuration yielded substantial improvements in both cardiac and prostate segmentation tasks.The authors used the same configuration to test generality across the two applications.Source: §4.2.2 · p. 11

5. Results

The results evaluate AdvChain on cardiac and prostate segmentation across supervised and semi-supervised settings, comparing augmentation methods, chain configurations, and consistency-loss functions. They report Dice-score outcomes and examine how performance varies with the learning setting and transformation choices.

5.1. Comparison study

  • Replacing AdvChain with RandAugment lowers average Dice on cardiac segmentation; on prostate, it does not significantly outperform AdvChain.Source: §5.1 · p. 12
  • On cardiac and prostate segmentation, AdvChain achieves the highest performance; on prostate, it exceeds the U-net upperbound trained on labeled images from both sets, with average Dice 0.7325 versus 0.7288.Source: §5.1 · pp. 12–13, Table 2 · p. 13
  • AdvChain supports diffeomorphic transformations that model realistic morphological variation, whereas RandAugment and other compared methods have limited geometric variation.AdvChain uses the segmentation model’s predictions to identify and deform local structures of interest.Source: §5.1 · p. 13
  • The authors attribute AdvChain’s efficacy to increased data diversity and effectiveness from dynamic adversarial transformations, plus consistency regularization with a contour-based loss.Source: §5.1 · p. 14

5.2.1. E

  • Individual augmentations with adversarial training consistently outperform their random counterparts in semi-supervised experiments using three labeled subjects.Source: Table 3 · p. 14, §5.2.1 · p. 14
  • Morphological transformations outperformed the other three individual data augmentations, highlighting the value of local anatomical variation.Source: §5.2.1 · p. 14
  • AdvChain achieved the highest Dice scores on both cardiac and prostate segmentation, while random-based composite augmentation did not always outperform individual augmentations.On cardiac segmentation, random-based composite augmentation scored 0.7802 versus 0.7857 for random bias-field augmentation.Source: §5.2.1 · p. 14

5.2.2. The generality of AdvChain with di

  • Adversarially optimized augmentations perturbed network predictions more effectively than random initialization and focused on deforming local target structures.Source: §5.2.2 · p. 15
  • AdvChain consistently improved segmentation performance over RandChain with maximum chain lengths from 1 to 4.Source: §5.2.2 · p. 14, Figure 6 · p. 15
  • AdvChain also improved segmentation performance with transformations chained in different orders.Source: §5.2.2 · pp. 14–15, Figure 7 · p. 16
  • The better transformation order varied by task: one ordering performed slightly better for cardiac segmentation, while another did so for prostate segmentation.Source: §5.2.2 · p. 16
  • The best maximum chain length also differed by task—1 for cardiac and 2 for prostate—and depended on the training-set selection.Source: §5.2.2 · p. 16

5.2.3. Effect of different consistency loss functions

  • The consistency-loss experiment compared proposed regularization R with RDKL, RDMSE, and RDKL+Contour; Table 4 reports average Dice and foreground-class AVG.Source: §5.2.3 · p. 16
  • RDMS E+Contour outperformed the other three consistency losses in most cases and achieved the highest average Dice scores on both tasks; contour loss generally improved paired standalone losses.The comparisons were conducted on the cardiac and prostate segmentation tasks.Source: §5.2.3 · p. 16

5.2.4. Effect of number of labeled images

  • AdvChain improved performance over standard training across all labeled-data settings, with the clearest gain in the one-shot setting and a narrowing gap as labeled images increased.The comparison used the same unlabeled dataset (M = 25).Source: §5.2.4 · p. 17

5.3. Supervised learning with extremely low data settings

  • With no unlabeled data (M = 0), AdvChain significantly outperformed RandChain by a large margin on both tasks.The results indicate that AdvChain can enhance training even when training data is limited.Source: §5.3 · p. 17

6. Discussion

AdvChain applies realistic photometric and geometric transformations to MR segmentation and is reported to outperform competitive baselines in low-shot supervised and semi-supervised settings. The discussion highlights its data efficiency and identifies manual transformation constraints as a limitation.

6. Discussion

  • Physics-based photometric and geometric transformations generate realistic local intensity and morphological variations, encouraging local smoothness in segmentation predictions.The discussion connects this regularization to photometric invariance and geometric equivariance.Source: §6 · pp. 17–18
  • On cardiac and prostate MR segmentation, AdvChain reduced annotation effort and outperformed competitive baselines in low-shot supervised and semi-supervised settings.The paper reports results for both segmentation tasks and both learning settings.Source: §6 · p. 18
  • AdvChain uses a small set of explainable, controllable parameters and can operate with only one or three labeled subjects, even without unlabeled subjects.The authors describe the augmentation as a plug-in module for segmentation pipelines.Source: §6 · p. 18
  • The method still requires expertise to set parameter-magnitude constraints that preserve the naturalism of augmented images.Automatically finding augmentation operations, magnitudes, ordering, and chain length remains an active research area; the authors suggest combining AdvChain with Auto DA.Source: §6 · pp. 18–19
  • The authors suggest exploring stronger segmentation architectures, class-imbalance methods, and improved pseudo-label selection for future extensions.The discussion notes that medical segmentation can involve class imbalance and blurry anatomical contours.Source: §6 · p. 19

7. Conclusion

AdvChain jointly optimizes dynamic data augmentation and a segmentation network, using challenging photometric and geometric variations to improve data diversity, effectiveness, and model generalization. On two segmentation tasks, it performs strongly in low-shot semi-supervised settings, while the conclusion also highlights diffeomorphic deformations, contour-based consistency loss, and the framework’s flexibility.

7. Conclusion

  • The framework jointly optimizes dynamic data augmentation and a segmentation network, using challenging photometric and geometric variations to improve diversity and effectiveness.The conclusion describes this as leveraging labeled and unlabeled data for improved generalization.Source: §7 · p. 19
  • The authors highlight adversarial diffeomorphic deformations for improved data diversity and effectiveness, a component not explored in prior consistency-regularization methods.The conclusion also highlights contour-based consistency loss as a more comprehensive inconsistency measure for network training.Source: §7 · p. 19
  • The generic plug-in framework is flexible enough to incorporate additional image transformations and has potential for other modalities and data-driven medical imaging applications.The conclusion leaves broader applications, including image registration and reconstruction, for future work.Source: §7 · p. 19
  • Across two tasks, AdvChain demonstrates data efficiency with four transformation types and outperforms RandAugment in most cases.The paper also reports that it outperforms the state-of-the-art composite data augmentation method in most cases.Source: §7 · p. 19
  • AdvChain outperforms several strong consistency-regularized methods in challenging low-shot semi-supervised settings.The conclusion reports these comparisons across different scenarios.Source: §7 · p. 19

Appendix

The appendix examines how the best maximum chain length varies across training-set selections and segmentation tasks. The variation persists in cardiac segmentation when the labeled set grows from 3 to 25 subjects.

Appendix

  • Boxplots compare average Dice scores across chain lengths 1–4, with yellow stars marking optima that vary by task and labeled-set selection.The variation also persists for cardiac segmentation as the labeled set increases from 3 to 25 subjects.Source: Figure A1 · p. 23
Loading 2108.03429v3…