Dynamic Data Augmentation via MCTS for Prostate MRI Segmentation

TL;DR

This paper introduces DDAug, an efficient automatic data augmentation method using MCTS to optimize augmentation strategies for prostate MRI segmentation, improving model generalization with minimal computational overhead.

Contribution

The paper proposes a novel hierarchical MCTS-based approach for automatic data augmentation tailored to medical imaging datasets, reducing manual tuning and computational costs.

Findings

DDAug outperforms existing augmentation methods on prostate MRI datasets.

The hierarchical MCTS approach effectively adapts augmentation strategies to different datasets.

The method achieves improved segmentation accuracy with negligible additional computation.

Abstract

Medical image data are often limited due to the expensive acquisition and annotation process. Hence, training a deep-learning model with only raw data can easily lead to overfitting. One solution to this problem is to augment the raw data with various transformations, improving the model's ability to generalize to new data. However, manually configuring a generic augmentation combination and parameters for different datasets is non-trivial due to inconsistent acquisition approaches and data distributions. Therefore, automatic data augmentation is proposed to learn favorable augmentation strategies for different datasets while incurring large GPU overhead. To this end, we present a novel method, called Dynamic Data Augmentation (DDAug), which is efficient and has negligible computation cost. Our DDAug develops a hierarchical tree structure to represent various augmentations and utilizes an efficient Monte-Carlo tree searching algorithm to update, prune, and sample the tree. As a result, the augmentation pipeline can be optimized for each dataset automatically. Experiments on multiple Prostate MRI datasets show that our method outperforms the current state-of-the-art data augmentation strategies.