The Power of Proxy Data and Proxy Networks for Hyper-Parameter Optimization in Medical Image Segmentation

TL;DR

This paper introduces proxy data and proxy networks to efficiently estimate hyper-parameters in medical image segmentation, significantly reducing computational costs while maintaining accuracy across CT and MR datasets.

Contribution

The paper proposes novel proxy data and proxy network methodologies to accelerate hyper-parameter optimization in medical image segmentation tasks.

Findings

Proxy data achieves higher correlation with full data training on external validation.

Proxy networks show high correlation with full networks on validation Dice scores.

Using proxy data and networks speeds up AutoML hyper-parameter search by up to 4.4x.

Abstract

Deep learning models for medical image segmentation are primarily data-driven. Models trained with more data lead to improved performance and generalizability. However, training is a computationally expensive process because multiple hyper-parameters need to be tested to find the optimal setting for best performance. In this work, we focus on accelerating the estimation of hyper-parameters by proposing two novel methodologies: proxy data and proxy networks. Both can be useful for estimating hyper-parameters more efficiently. We test the proposed techniques on CT and MR imaging modalities using well-known public datasets. In both cases using one dataset for building proxy data and another data source for external evaluation. For CT, the approach is tested on spleen segmentation with two datasets. The first dataset is from the medical segmentation decathlon (MSD), where the proxy data is constructed, the secondary dataset is utilized as an external validation dataset. Similarly, for MR, the approach is evaluated on prostate segmentation where the first dataset is from MSD and the second dataset is PROSTATEx. First, we show higher correlation to using full data for training when testing on the external validation set using smaller proxy data than a random selection of the proxy data. Second, we show that a high correlation exists for proxy networks when compared with the full network on validation Dice score. Third, we show that the proposed approach of utilizing a proxy network can speed up an AutoML framework for hyper-parameter search by 3.3x, and by 4.4x if proxy data and proxy network are utilized together.