TLDR: A new curriculum learning method called Progressive Growing of Patch Size (PGPS) improves 3D medical image segmentation by gradually increasing image patch sizes during training. It offers an “Efficiency Mode” that drastically cuts training time (to 44%) with comparable performance, and a “Performance Mode” that significantly boosts segmentation accuracy (1.26% Dice Score gain) while also reducing training time (to 89%). PGPS is especially effective for imbalanced tasks and works across various AI architectures, making AI model development faster, more accurate, and resource-efficient.
A new approach to training artificial intelligence models for medical image segmentation, called Progressive Growing of Patch Size (PGPS), promises to make these crucial tools both faster and more accurate. Developed by a team of researchers including Stefan M. Fischer and Julia A. Schnabel, this method introduces a curriculum learning strategy that changes how AI models process 3D medical images during training.
Traditionally, AI models for 3D medical image segmentation are trained using a constant patch size, meaning they always look at image sections of the same size. However, the PGPS method challenges this by starting with smaller image patches and gradually increasing their size as training progresses. This mimics how humans learn, starting with simpler concepts before moving to more complex ones.
The core idea behind starting with smaller patches is to improve the “class balance” early in the training. In medical images, especially for tasks like lesion segmentation, there’s often a huge imbalance between the foreground (the area of interest, like a tumor) and the background (healthy tissue). Smaller patches are more likely to contain a higher proportion of foreground pixels, giving the AI model a better chance to learn about these important, but rare, features from the start. As the patches grow larger, the model gains a broader context of the image, which is essential for accurate segmentation.
The researchers developed two main modes for PGPS: an “Efficiency Mode” and a “Performance Mode.” The Efficiency Mode focuses on significantly reducing training time and computational costs while maintaining segmentation performance comparable to traditional methods. It achieves this by keeping the batch size constant, meaning fewer computations are needed during the early stages with smaller patches. This mode can cut training time by more than half, making AI development more resource-friendly.
The Performance Mode, on the other hand, aims for the highest possible segmentation accuracy. It dynamically increases the batch size when using smaller patches to fully utilize the available GPU memory. This strategy not only improves the mean segmentation performance but also reduces the variability in results across different training runs, leading to more reliable models. In experiments across 15 diverse 3D medical image segmentation tasks, the Performance Mode consistently outperformed constant patch size training, achieving statistically significant improvements in Dice Score (a common metric for segmentation accuracy) while also reducing training time by about 11%.
The benefits of PGPS are particularly noticeable in highly imbalanced segmentation tasks, such as identifying small lesions, where foreground voxels are rare. The improved class balance in early training phases helps the model learn these challenging tasks more effectively. Furthermore, the method has been shown to be broadly applicable, improving performance across various AI architectures, including the popular UNet, UNETR, and SwinUNETR models, demonstrating its versatility beyond just convolutional neural networks.
Also Read:
- Medformer: A Flexible Framework for Medical Imaging AI
- A New Deep Learning Model for Precise Medical Image Segmentation
This research highlights that the way data is sampled and presented to AI models during training is a critical, yet often overlooked, area for improvement. By simply adapting the sampling strategy, PGPS offers a straightforward yet powerful way to accelerate training convergence and enhance the final segmentation performance in 3D medical imaging. This could lead to faster development of more accurate AI tools for diagnostics and treatment planning, ultimately benefiting patient care. For more details, you can read the full research paper here.
