[Paper Review] Automatic Data Augmentation via Deep Reinforcement Learning for Effective Kidney Tumor Segmentation

Tiexin Qin, Ziyuan Wang+4 more|arXiv (Cornell University)|Feb 22, 2020

Advanced Neural Network Applications11 references4 citations

TL;DR

This paper proposes a deep reinforcement learning (DRL)-based automatic data augmentation method for kidney tumor segmentation in CT scans, where an agent learns optimal augmentation policies end-to-end by maximizing Dice score improvements on a validation set. The method outperforms conventional and GAN-based augmentation techniques, achieving a 75.8% mIoU and 84.0% DSC, demonstrating superior boundary localization and robustness on limited data.

ABSTRACT

Conventional data augmentation realized by performing simple pre-processing operations (\eg, rotation, crop, \etc) has been validated for its advantage in enhancing the performance for medical image segmentation. However, the data generated by these conventional augmentation methods are random and sometimes harmful to the subsequent segmentation. In this paper, we developed a novel automatic learning-based data augmentation method for medical image segmentation which models the augmentation task as a trial-and-error procedure using deep reinforcement learning (DRL). In our method, we innovatively combine the data augmentation module and the subsequent segmentation module in an end-to-end training manner with a consistent loss. Specifically, the best sequential combination of different basic operations is automatically learned by directly maximizing the performance improvement (\ie, Dice ratio) on the available validation set. We extensively evaluated our method on CT kidney tumor segmentation which validated the promising results of our method.

Motivation & Objective

• To address performance degradation in medical image segmentation due to limited annotated training data.
• To overcome the limitations of conventional data augmentation, which relies on fixed, random operations and lacks feedback from segmentation performance.
• To develop an end-to-end learning framework that jointly optimizes data augmentation and segmentation by using reinforcement learning to discover image-specific augmentation policies.
• To directly maximize segmentation performance (Dice score) during augmentation policy learning, ensuring that generated data improves downstream task accuracy.
• To validate the method on a real-world CT kidney tumor segmentation dataset and compare it against state-of-the-art augmentation baselines.

Proposed method

• The method models data augmentation as a Markov decision process (MDP), where the agent selects image transformations (e.g., flip, crop, rotate) to improve segmentation performance.
• A Dueling Deep Q-Network (DQN) is used to learn the optimal sequence of augmentation operations, with the reward function defined as the improvement in Dice score on the validation set.
• The segmentation model is used as a substitute policy to provide feedback during DRL training, enabling end-to-end joint optimization of augmentation and segmentation.
• The agent interacts with the environment (CT images and masks), selecting actions to transform the input image, and receives rewards based on the resulting segmentation quality.
• The final augmented dataset is formed by mixing original and DRL-generated augmented images, doubling the training set size.
• The framework is trained using Adam optimizer with shared hyperparameters across all methods for fair comparison.

Experimental results

Research questions

• RQ1 Can a reinforcement learning agent learn a data-specific sequence of augmentations that maximizes segmentation performance on limited medical imaging data?
• RQ2 Does end-to-end joint training of data augmentation and segmentation lead to better performance than separate training of augmentation and segmentation modules?
• RQ3 How does the proposed DRL-based augmentation method compare to conventional data augmentation and GAN-based data generation in kidney tumor segmentation?
• RQ4 What types of augmentation operations (e.g., flip, crop) are most frequently selected by the agent, and do they correlate with improved segmentation of tumor boundaries?
• RQ5 Can the method generalize to real-world clinical datasets with limited annotations and achieve state-of-the-art performance?

Key findings

• The proposed method achieved a 75.8% mean Intersection over Union (mIoU), significantly outperforming conventional data augmentation (74.3%) and other baselines.
• The method achieved a Dice Similarity Coefficient (DSC) of 84.0%, surpassing the next best method (83.2% with traditional augmentation) and showing the highest sensitivity (94.0%) and specificity (PPV: 80.2%).
• The centroid distance (CD) was reduced to 8.86 mm, the Hausdorff distance (HD) to 8.52 mm, and the average surface distance (ASD) to 4.01 mm, indicating superior boundary localization accuracy.
• The agent frequently selected horizontal flipping and cropping operations, which enhanced tumor visibility and improved generalization on small or irregularly shaped tumors.
• Visual results confirmed that the method produced more accurate tumor boundary predictions, especially in challenging regions with low contrast or irregular shapes.
• The ablation study confirmed that end-to-end joint learning with feedback from the segmentation module led to better policy learning than isolated augmentation or fixed operation sets.