Progressive Growing of Patch Size: Curriculum Learning for Accelerated and Improved Medical Image Segmentation

TL;DR

This work introduces Progressive Growing of Patch Size (PGPS), an input-length curriculum for 3D patch-based medical image segmentation that starts training with small patches to improve class balance and context, then progressively increases patch size. Implemented in nnU-Net, PGPS offers two modes: Efficiency (low runtime, similar Dice) and Performance (higher Dice, reduced runtime). Across 15 diverse datasets and multiple backbones (UNet, UNETR, SwinUNETR), PGPS-Performance delivers statistically significant Dice improvements and faster convergence, with substantial reductions in FLOPs and training time compared to constant patch size training. The approach generalizes across architectures, enhances robustness, and outperforms the Progressive Resolution curriculum for dense segmentation tasks, making PGPS a practical default for 3D medical image segmentation. The findings highlight the importance of training-time data sampling strategies and suggest avenues for further optimization via differentiable sampling and hyperparameter tuning.

Abstract

In this work, we introduce Progressive Growing of Patch Size, an automatic curriculum learning approach for 3D medical image segmentation. Our approach progressively increases the patch size during model training, resulting in an improved class balance for smaller patch sizes and accelerated convergence of the training process. We evaluate our curriculum approach in two settings: a resource-efficient mode and a performance mode, both regarding Dice score performance and computational costs across 15 diverse and popular 3D medical image segmentation tasks. The resource-efficient mode matches the Dice score performance of the conventional constant patch size sampling baseline with a notable reduction in training time to only 44%. The performance mode improves upon constant patch size segmentation results, achieving a statistically significant relative mean performance gain of 1.28% in Dice Score. Remarkably, across all 15 tasks, our proposed performance mode manages to surpass the constant patch size baseline in Dice Score performance, while simultaneously reducing training time to only 89%. The benefits are particularly pronounced for highly imbalanced tasks such as lesion segmentation tasks. Rigorous experiments demonstrate that our performance mode not only improves mean segmentation performance but also reduces performance variance, yielding more trustworthy model comparison. Furthermore, our findings reveal that the proposed curriculum sampling is not tied to a specific architecture but represents a broadly applicable strategy that consistently boosts performance across diverse segmentation models, including UNet, UNETR, and SwinUNETR. In summary, we show that this simple yet elegant transformation on input data substantially improves both Dice Score performance and training runtime, while being compatible across diverse segmentation backbones.

Figures (7)

• Figure 1: Overview of the proposed Progressive Growing of Patch Size curriculum, illustrated for lung cancer segmentation (cancer regions highlighted with yellow bounding boxes). Training begins with the minimal patch size ($H_{min}, W_{min}$) and progressively increases the patch dimensions stepwise until the final maximal patch size ($H_{max}, W_{max}$) is reached. Smaller patch sizes provide a better class balance, which decreases as the patch size grows. During inference, the maximum patch size is used to capture maximal global context. Figure is adapted from fischer2024progressive.
• Figure 2: Input tensors for the two proposed modes of the Progressive Growing of Patch Size (PGPS) curriculum, illustrated for a multi-organ segmentation task with example foreground patches (top) and corresponding label maps (bottom). In both modes, the patch size is progressively increased during training. In PGPS-Efficiency (green), the batch size remains constant, resulting on average in smaller input tensors. In PGPS-Performance (orange), the available GPU budget is fully utilized by increasing the batch size.
• Figure 3: Segmentation performance convergence of CPS, PGPS-Performance, and PGPS-Efficiency across training iterations. Dice scores are tracked across 15 segmentation tasks for different training lengths (1%, 10%, 25%, 50%, and 100% of nnU-Net's default total training iterations). On average, PGPS-Performance exhibits the fastest convergence, while PGPS-Efficiency converges more slowly, due to the smaller input tensors that result in superior training speed. The final performance is, on average, best for PGPS-Performance, followed by PGPS-Efficiency and CPS. The average convergence over all 15 datasets (right) is computed by normalizing each task by its maximum performance and then averaging across all tasks.
• Figure 4: Training characteristics of PGPS curricula for the BTCV dataset with 14 classes. Left: PGPS curricula improve foreground-background voxel balance compared to CPS due to on average smaller patch sizes. Center: PGPS-Performance achieves the highest average number of unique classes per batch/iteration due to larger batch sizes, while PGPS-Efficiency has the lowest. Right: CPS maintains a constant input tensor size; PGPS-Efficiency increases tensor size exponentially, while PGPS-Performance exhibits occasional size drops due to the discrete nature of patch and batch sizes.
• Figure 5: Segmentation performance for repeated training of CPS, PGPS-Efficiency, and PGPS-Performance across different training lengths (1%, 10%, and 100% of nnU-Net’s total training iterations). Each experiment is repeated five times on the same data split with different random seeds. Variability decreases with longer training, while mean performance increases. The highly class-imbalanced MSD Lung Tumor dataset exhibits higher variance than the multi-class BTCV dataset. PGPS-Performance shows improved convergence and less performance variability between repeated runs. PGPS-Efficiency results in slower convergence than CPS and PGPS-Performance and has the highest performance variability. PGPS-Performance outperforms both other strategies in 100% of the 125 possible combinations for MSD Lung Tumor, while for BTCV, PGPS-Performance outperforms both other strategies in 58.4% of all combinations, as well as CPS and 60% of combinations.